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')