def test_bottom_up_until_type_break():
s = symbol('s', 'var * {name: string, amount: int}')
data = np.array([('Alice', 100), ('Bob', 200), ('Charlie', 300)],
dtype=[('name', 'S7'), ('amount', 'i4')])
e = (s.amount + 1).distinct()
expr, scope = bottom_up_until_type_break(e, {s: data})
amount = symbol('amount', 'var * real', token=1)
assert expr.isidentical(amount)
assert len(scope) == 1
assert amount in scope
assert (scope[amount] == np.array([101, 201, 301], dtype='i4')).all()
# This computation has a type change midstream, so we stop and get the
# unfinished computation.
e = s.amount.sum() + 1
expr, scope = bottom_up_until_type_break(e, {s: data})
amount_sum = symbol('amount_sum', 'int')
assert expr.isidentical(amount_sum + 1)
assert len(scope) == 1
assert amount_sum in scope
assert scope[amount_sum] == 600
# ensure that we work on binops with one child
x = symbol('x', 'real')
expr, scope = bottom_up_until_type_break(x + x, {x: 1})
assert len(scope) == 1
x2 = list(scope.keys())[0]
assert isinstance(x2, Symbol)
assert isinstance(expr, Symbol)
assert scope[x2] == 2
def test_bottom_up_until_type_break():
s = symbol('s', 'var * {name: string, amount: int}')
data = np.array([('Alice', 100), ('Bob', 200), ('Charlie', 300)],
dtype=[('name', 'S7'), ('amount', 'i4')])
e = (s.amount + 1).distinct()
expr, scope = bottom_up_until_type_break(e, {s: data})
amount = symbol('amount', 'var * int64', token=expr._token)
assert expr.isidentical(amount)
assert len(scope) == 1
assert amount in scope
assert (scope[amount] == np.array([101, 201, 301], dtype='i4')).all()
# This computation has a type change midstream, so we stop and get the
# unfinished computation.
e = s.amount.sum() + 1
expr, scope = bottom_up_until_type_break(e, {s: data})
amount_sum = symbol('amount_sum', 'int64', token=expr.lhs._token)
assert expr.isidentical(amount_sum + 1)
assert len(scope) == 1
assert amount_sum in scope
assert scope[amount_sum] == 600
# ensure that we work on binops with one child
x = symbol('x', 'real')
expr, scope = bottom_up_until_type_break(x + x, {x: 1})
assert len(scope) == 1
x2 = list(scope.keys())[0]
assert isinstance(x2, Symbol)
assert isinstance(expr, Symbol)
assert scope[x2] == 2
