def test_table_broadcast():
t = symbol('t', 'var * {x: int, y: int, z: int}')
expr = t.distinct()
expr = (2 * expr.x + expr.y + 1).distinct()
expected = t.distinct()
expected = broadcast(2 * expected.x + expected.y + 1, [expected]).distinct()
assert broadcast_collect(expr).isidentical(expected)
expr = (t.x + t.y).sum()
result = broadcast_collect(expr)
expected = broadcast(t.x + t.y, [t]).sum()
assert result.isidentical(expected)
def test_table_broadcast():
t = Symbol('t', 'var * {x: int, y: int, z: int}')
expr = t.distinct()
expr = (2 * expr.x + expr.y + 1).distinct()
expected = t.distinct()
expected = broadcast(2 * expected.x + expected.y + 1, [t]).distinct()
assert broadcast_collect(expr).isidentical(expected)
expr = (t.x + t.y).sum()
result = broadcast_collect(expr)
expected = broadcast(t.x + t.y, [t]).sum()
assert result.isidentical(expected)
def test_optimize_broadcast():
expr = (t.distinct().x + 1).distinct()
expected = broadcast(t.distinct().x + 1, [t.distinct()]).distinct()
result = broadcast_collect(expr, broadcastable=(Field, Arithmetic), want_to_broadcast=(Field, Arithmetic))
assert result.isidentical(expected)
def test_broadcast_collect_doesnt_collect_scalars():
expr = xx + yy * a
assert broadcast_collect(
expr,
broadcastable=Arithmetic,
want_to_broadcast=Arithmetic,
).isidentical(expr)
def test_optimize_broadcast():
expr = (t.distinct().x + 1).distinct()
expected = broadcast(t.distinct().x + 1, [t.distinct()]).distinct()
result = broadcast_collect(expr,
Broadcastable=(Field, Arithmetic),
WantToBroadcast=(Field, Arithmetic))
assert result.isidentical(expected)
def test_broadcast_collect():
t = symbol('t', 'var * {x: int, y: int, z: int, when: datetime}')
expr = t.distinct()
expr = expr.x + 2 * expr.y
expr = expr.distinct()
result = broadcast_collect(expr)
expected = t.distinct()
expected = broadcast(expected.x + 2 * expected.y, [expected])
expected = expected.distinct()
assert result.isidentical(expected)
def test_broadcast_collect():
t = symbol('t', 'var * {x: int, y: int, z: int, when: datetime}')
expr = t.distinct()
expr = expr.x + 2 * expr.y
expr = expr.distinct()
result = broadcast_collect(expr)
expected = t.distinct()
expected = broadcast(expected.x + 2 * expected.y, [expected])
expected = expected.distinct()
assert result.isidentical(expected)
def test_broadcast_doesnt_affect_scalars():
t = symbol('t', '{x: int, y: int, z: int}')
expr = (2 * t.x + t.y + 1)
assert broadcast_collect(expr).isidentical(expr)
def test_broadcast_doesnt_affect_scalars():
t = Symbol('t', '{x: int, y: int, z: int}')
expr = (2 * t.x + t.y + 1)
assert broadcast_collect(expr).isidentical(expr)
def test_broadcast_collect_doesnt_collect_scalars():
expr = xx + yy * a
assert broadcast_collect(expr, Broadcastable=Arithmetic,
WantToBroadcast=Arithmetic).isidentical(expr)
def test_broadcast_doesnt_affect_scalars():
t = symbol("t", "{x: int, y: int, z: int}")
expr = 2 * t.x + t.y + 1
assert broadcast_collect(expr).isidentical(expr)
