def test_with_clause_error(self, query_expression, error):
# type: (str, str) -> None
query_expression_node, leftover = query_expression_rule(
tokenize(query_expression))
self.assertFalse(leftover)
assert isinstance(query_expression_node, QueryExpression)
with self.assertRaisesRegexp(ValueError, error):
query_expression_node.get_dataframe(self.table_context)
def test_with_clause(self, query_expression, expected_result):
# type: (str, List[List[int]]) -> None
query_expression_node, leftover = query_expression_rule(
tokenize(query_expression))
self.assertFalse(leftover)
assert isinstance(query_expression_node, QueryExpression)
dataframe, _ = query_expression_node.get_dataframe(self.table_context)
self.assertEqual(dataframe.to_list_of_lists(), expected_result)
def test_unnest_error(self, query, expected_error):
# type: (str, str) -> None
node, leftover = query_expression_rule(tokenize(query))
self.assertFalse(leftover)
assert isinstance(node, QueryExpression)
with self.assertRaisesRegexp(ValueError, expected_error):
node.get_dataframe(TableContext())
def test_unnest(self, query, result, result_columns):
# type: (str, List[List[Any]], List[str]) -> None
node, leftover = query_expression_rule(tokenize(query))
self.assertFalse(leftover)
assert isinstance(node, QueryExpression)
dataframe, _ = node.get_dataframe(TableContext())
self.assertEqual(dataframe.to_list_of_lists(), result)
self.assertEqual(list(dataframe.dataframe.columns), result_columns)
def test_query_expression_set_operation_error(self, query_expression,
error):
table_context = DatasetTableContext({
'my_project': {
'my_dataset': {
'my_table':
TypedDataFrame(
pd.DataFrame([[1, 2, 3]], columns=['a', 'b', 'c']), [
BQScalarType.INTEGER, BQScalarType.INTEGER,
BQScalarType.INTEGER
])
}
}
})
query_expression_node, leftover = query_expression_rule(
tokenize(query_expression))
self.assertFalse(leftover)
with self.assertRaisesRegexp(ValueError, error):
query_expression_node.get_dataframe(table_context)
def test_query_expression_set_operation(self, query_expression,
expected_result):
table_context = DatasetTableContext({
'my_project': {
'my_dataset': {
'my_table':
TypedDataFrame(
pd.DataFrame([[1, 2, 3]], columns=['a', 'b', 'c']), [
BQScalarType.INTEGER, BQScalarType.INTEGER,
BQScalarType.INTEGER
])
}
}
})
query_expression_node, leftover = query_expression_rule(
tokenize(query_expression))
dataframe, unused_table_name = query_expression_node.get_dataframe(
table_context)
self.assertFalse(leftover)
self.assertEqual(dataframe.to_list_of_lists(), expected_result)
def test_query_expression_syntax(self, query):
query_expression = ' '.join(query)
query_expression_node, leftover = query_expression_rule(
tokenize(query_expression))
self.assertFalse(leftover)
