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)