def test_bad_function(self):
# type: () -> None
context = EvaluationContext(self.small_table_context)
context.add_table_from_node(TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
with self.assertRaisesRegexp(NotImplementedError, 'NOT_A_FUNCTION not implemented'):
FunctionCall.create('not_a_function', [], EMPTY_NODE).evaluate(context)
def test_exists_reference_outer(self):
table_context = DatasetTableContext({
'my_project': {
'my_dataset': {
'my_table':
TypedDataFrame(pd.DataFrame([[1], [4]], columns=['a']),
types=[BQScalarType.INTEGER]),
'my_table2':
TypedDataFrame(pd.DataFrame([[4], [2]], columns=['b']),
types=[BQScalarType.INTEGER]),
}
}
})
select_query = "select a from `my_project.my_dataset.my_table` where " \
"my_table.a = my_table2.b"
select_node, leftover = apply_rule(select_rule, tokenize(select_query))
self.assertFalse(leftover)
exists = Exists(select_node)
context = EvaluationContext(table_context)
context.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table2')),
EMPTY_NODE)
dataframe = exists.evaluate(context)
self.assertEqual(list(dataframe.series), [True, False])
def test_evaluation_context_path(self, path):
# type: (Tuple[str, ...]) -> None
ec = EvaluationContext(self.table_context)
ec.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table2')),
EMPTY_NODE)
self.assertEqual(ec.get_canonical_path(path), ('my_table2', 'b'))
def test_value_eval(self):
# type: () -> None
# A constant is repeated for each row in the context table.
value = Value(12345, BQScalarType.INTEGER)
context = EvaluationContext(self.small_table_context)
context.add_table_from_node(TableReference(('my_project', 'my_dataset', 'my_table')), 'foo')
typed_series = value.evaluate(context)
assert isinstance(typed_series, TypedSeries)
self.assertEqual(list(typed_series.series), [12345, 12345])
def test_functions(self, function_name, args, expected_result):
# type: (str, List[EvaluatableNode], List[PythonType]) -> None
context = EvaluationContext(self.small_table_context)
context.add_table_from_node(TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
result = FunctionCall.create(function_name, args, EMPTY_NODE).evaluate(context)
assert isinstance(result, TypedSeries)
self.assertEqual(
[result.type_.convert(elt) for elt in result.series],
expected_result)
def test_field(self):
# type: () -> None
field = Field(('a',))
context = EvaluationContext(self.small_table_context)
context.add_table_from_node(TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
typed_series = field.evaluate(context)
assert isinstance(typed_series, TypedSeries)
self.assertEqual(list(typed_series.series), [1, 2])
self.assertEqual(typed_series.series.name, 'a')
def test_evaluation_context_path_error(self, path, error):
# type: (Tuple[str, ...], str) -> None
ec = EvaluationContext(self.table_context)
ec.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
ec.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table2')),
EMPTY_NODE)
with self.assertRaisesRegexp(ValueError, error):
ec.get_canonical_path(path)
def test_selector(self):
# type: () -> None
selector = Selector(Field(('a',)), 'field_alias')
context = EvaluationContext(self.small_table_context)
context.add_table_from_node(TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
typed_series = selector.evaluate(context)
assert isinstance(typed_series, TypedSeries)
self.assertEqual(list(typed_series.series), [1, 2])
self.assertEqual(list(typed_series.dataframe), ['field_alias'])
self.assertEqual(typed_series.types, [BQScalarType.INTEGER])
def test_in_check(self, direction, result):
# type: (str, List[bool]) -> None
expression = Field(('a',))
elements = (Value(1, type_=BQScalarType.INTEGER), Value(3, type_=BQScalarType.INTEGER))
in_check = InCheck(expression, direction, elements)
context = EvaluationContext(self.small_table_context)
context.add_table_from_node(TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
typed_series = in_check.evaluate(context)
assert isinstance(typed_series, TypedSeries)
self.assertEqual(list(typed_series.series), result)
def test_if(self):
condition = BinaryExpression(Field(('a',)), '>', Value(1, BQScalarType.INTEGER))
then = Value('yes', BQScalarType.STRING)
else_ = Value('no', BQScalarType.STRING)
# IF a > 1 THEN "yes" ELSE "no"
if_expression = If(condition, then, else_)
context = EvaluationContext(self.small_table_context)
context.add_table_from_node(TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
typed_series = if_expression.evaluate(context)
assert isinstance(typed_series, TypedSeries)
self.assertEqual(list(typed_series.series), ['no', 'yes'])
def test_evaluation_context(self):
# type: () -> None
ec = EvaluationContext(self.table_context)
ec.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
self.assertEqual(ec.table_ids, set(['my_table']))
self.assertEqual(ec.column_to_table_ids, {'a': ['my_table']})
self.assertEqual(ec.canonical_column_to_type,
{'my_table.a': BQScalarType.INTEGER})
self.assertEqual(ec.table.to_list_of_lists(), [[1], [2]])
self.assertEqual(list(ec.table.dataframe), ['my_table.a'])
self.assertEqual(ec.table.types, [BQScalarType.INTEGER])
def test_context_lookup_success(self):
# type: () -> None
ec = EvaluationContext(self.table_context)
ec.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
ec.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table2')),
EMPTY_NODE)
path = ('b', )
result = ec.lookup(path)
self.assertEqual(list(result.series), [2, 4])
self.assertEqual(result.type_, BQScalarType.INTEGER)
def test_exists(self, select_query, result):
# type: (str, List[bool]) -> None
subquery_node, leftover = apply_rule(query_expression, tokenize(select_query))
assert isinstance(subquery_node, QueryExpression)
self.assertFalse(leftover)
exists = Exists(subquery_node)
context = EvaluationContext(self.small_table_context)
context.add_table_from_node(TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
typed_series = exists.evaluate(context)
assert isinstance(typed_series, TypedSeries)
self.assertEqual(list(typed_series.series), result)
def test_selector_group_by_success(self):
# type: () -> None
selector = Selector(Field(('c',)), EMPTY_NODE)
selector.position = 1
context = EvaluationContext(self.large_table_context)
context.add_table_from_node(TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
context.exclude_aggregation = True
updated_selector, = context.do_group_by([selector], [Field(('my_table', 'c'))])
typed_series = updated_selector.evaluate(context)
assert isinstance(typed_series, TypedSeries)
self.assertEqual(list(typed_series.series), [3])
def test_context_lookup_type_error(self):
# type: () -> None
ec = EvaluationContext(self.table_context)
ec.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table2')),
EMPTY_NODE)
path = ('my_table2', 'b')
error = (
r"path \('my_table2', 'b'\) \(canonicalized to key 'my_table2.b'\) "
r"not present in type dict; columns available: \['my_table2.a'\]")
# Delete the type mapping for this column to test the error
del ec.canonical_column_to_type['my_table2.b']
with self.assertRaisesRegexp(KeyError, error):
ec.lookup(path)
def test_context_lookup_key_error(self):
# type: () -> None
ec = EvaluationContext(self.table_context)
ec.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
path = ('my_table', 'c')
error = (
r"path \('my_table', 'c'\) \(canonicalized to key 'my_table.c'\) "
r"not present in table; columns available: \['my_table.a'\]")
# Fake in an entry to the column -> table mapping to test the error
ec.column_to_table_ids['c'] = ['my_table']
with self.assertRaisesRegexp(KeyError, error):
ec.lookup(path)
def test_null_check(self, direction, result):
# type: (str, List[bool]) -> None
table_context = DatasetTableContext({
'my_project': {
'my_dataset': {
'my_table': TypedDataFrame(
pd.DataFrame([[1, None], [2, 3]], columns=['a', 'b']),
types=[BQScalarType.INTEGER, BQScalarType.INTEGER]
)
}
}
})
context = EvaluationContext(table_context)
context.add_table_from_node(TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
expression = Field(('b',))
null_check = NullCheck(expression, direction)
typed_series = null_check.evaluate(context)
assert isinstance(typed_series, TypedSeries)
self.assertEqual(list(typed_series.series), result)
def test_subcontext_lookup_error_already_has_subcontext(self):
# type: () -> None
ec = EvaluationContext(self.table_context)
ec.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
ec.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table2')),
EMPTY_NODE)
subcontext1 = EvaluationContext(self.table_context)
subcontext1.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table3')),
EMPTY_NODE)
ec.add_subcontext(subcontext1)
subcontext2 = EvaluationContext(self.table_context)
subcontext2.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table4')),
EMPTY_NODE)
with self.assertRaisesRegexp(ValueError,
'Context already has subcontext'):
ec.add_subcontext(subcontext2)
def test_subcontext_lookup_success(self, path, expected_result):
# type: (Tuple[str, ...], List[int]) -> None
ec = EvaluationContext(self.table_context)
ec.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
ec.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table2')),
EMPTY_NODE)
subcontext = EvaluationContext(self.table_context)
subcontext.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table3')),
EMPTY_NODE)
ec.add_subcontext(subcontext)
result = ec.lookup(path)
self.assertEqual(list(result.series), expected_result)
self.assertEqual(result.type_, BQScalarType.INTEGER)
"(BQScalarType.STRING, BQScalarType.INTEGER)"),
)
@unpack
def test_case_error(
self,
comparand, # type: Union[_EmptyNode, EvaluatableNode]
whens, # type: List[Tuple[AbstractSyntaxTreeNode, EvaluatableNode]]
else_, # type: EvaluatableNode
error # type: str
):
# type: (...) -> None
case = Case(comparand, whens, else_)
context = EvaluationContext(self.small_table_context)
context.add_table_from_node(
TableReference(('my_project', 'my_dataset', 'my_table')),
EMPTY_NODE)
with self.assertRaisesRegexp(ValueError, escape(error)):
case.evaluate(context)
@data(
dict(value=Value(1, BQScalarType.INTEGER),
cast_type='STRING',
result='1'),
dict(value=Value(1, BQScalarType.INTEGER),
cast_type='FLOAT',
result=1.0),
dict(value=Value(1, BQScalarType.INTEGER),
cast_type='BOOLEAN',
result=True),
dict(value=Value(1.0, BQScalarType.FLOAT),