def test_select(self): # type: () -> None from_ = DataSource((TableReference( ('my_project', 'my_dataset', 'my_table')), EMPTY_NODE), []) selector = StarSelector(EMPTY_NODE, EMPTY_NODE, EMPTY_NODE) select = Select(EMPTY_NODE, [selector], from_, EMPTY_NODE, EMPTY_NODE, EMPTY_NODE) dataframe, table_name = select.get_dataframe(self.table_context) self.assertEqual(table_name, None) self.assertEqual(dataframe.to_list_of_lists(), [[1], [2], [3]]) self.assertEqual(list(dataframe.dataframe), ['a']) self.assertEqual(dataframe.types, [BQScalarType.INTEGER])
def test_select_where(self, where): # type: (EvaluatableNode) -> None where_table_context = DatasetTableContext({ 'my_project': { 'my_dataset': { 'my_table': TypedDataFrame( pd.DataFrame([[1, 2], [3, 4]], columns=['a', 'b']), types=[BQScalarType.INTEGER, BQScalarType.INTEGER]) } } }) fields = [Selector(Field(('a', )), EMPTY_NODE)] from_ = DataSource((TableReference( ('my_project', 'my_dataset', 'my_table')), EMPTY_NODE), []) select = Select(EMPTY_NODE, fields, from_, where, EMPTY_NODE, EMPTY_NODE) dataframe, table_name = select.get_dataframe(where_table_context) self.assertEqual(dataframe.to_list_of_lists(), [[3]])
def test_query_expression_limit(self): # type: () -> None from_ = DataSource((TableReference( ('my_project', 'my_dataset', 'my_table')), EMPTY_NODE), []) selector = StarSelector(EMPTY_NODE, EMPTY_NODE, EMPTY_NODE) select = Select(EMPTY_NODE, [selector], from_, EMPTY_NODE, EMPTY_NODE, EMPTY_NODE) limit = Value(1, BQScalarType.INTEGER) offset = Value(1, BQScalarType.INTEGER) qe = QueryExpression(EMPTY_NODE, select, EMPTY_NODE, (limit, offset)) dataframe, table_name = qe.get_dataframe(self.table_context) self.assertEqual(dataframe.to_list_of_lists(), [[2]])
def test_marker_syntax_tree_node(self): # type: () -> None self.assertEqual(Select.literal(), 'SELECT')