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_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)