wehub-resource-sync
79 lines
2.7 KiB
Python
79 lines
2.7 KiB
Python
import pyarrow as pa
|
|
import pytest
|
|
|
|
import ray
|
|
|
|
|
|
def test_union_basic(ray_start_10_cpus_shared):
|
|
ds1 = ray.data.from_items([{"id": 0}])
|
|
ds2 = ray.data.from_items([{"id": 1}])
|
|
ds = ds1.union(ds2)
|
|
assert sorted(row["id"] for row in ds.take_all()) == [0, 1]
|
|
|
|
|
|
def test_union_schema(ray_start_10_cpus_shared):
|
|
ds = ray.data.range(1).union(ray.data.range(1))
|
|
assert ds.schema().names == ["id"]
|
|
assert ds.schema().types == [pa.int64()]
|
|
|
|
|
|
def test_union_repr(ray_start_10_cpus_shared):
|
|
ds = ray.data.range(1).union(ray.data.range(1))
|
|
assert repr(ds) == (
|
|
"shape: (?, 1)\n"
|
|
"╭───────╮\n"
|
|
"│ id │\n"
|
|
"│ --- │\n"
|
|
"│ int64 │\n"
|
|
"╰───────╯\n"
|
|
"(Dataset isn't materialized)"
|
|
)
|
|
|
|
|
|
def test_union_with_preserve_order(ray_start_10_cpus_shared, restore_data_context):
|
|
# Test for https://github.com/ray-project/ray/issues/41524
|
|
ctx = ray.data.DataContext.get_current()
|
|
ctx.execution_options.preserve_order = True
|
|
|
|
# To reproduce the bug, you need three or more input datasets, and the first dataset
|
|
# can't finish first. To emulate this behavior, the test adds a `map` to the first
|
|
# dataset and materializes the other two.
|
|
ds1 = ray.data.from_items([{"id": 0}]).map(lambda x: x)
|
|
ds2 = ray.data.from_items([{"id": 1}]).materialize()
|
|
ds3 = ray.data.from_items([{"id": 2}]).materialize()
|
|
ds = ds1.union(ds2, ds3)
|
|
|
|
assert [row["id"] for row in ds.take_all()] == [0, 1, 2]
|
|
|
|
|
|
def test_union_with_filter(ray_start_10_cpus_shared):
|
|
"""Test that filters are pushed through union to both branches."""
|
|
from ray.data._internal.logical.optimizers import LogicalOptimizer
|
|
from ray.data.expressions import col
|
|
|
|
ds1 = ray.data.from_items([{"id": 0}, {"id": 1}, {"id": 2}])
|
|
ds2 = ray.data.from_items([{"id": 3}, {"id": 4}, {"id": 5}])
|
|
ds = ds1.union(ds2).filter(expr=col("id") > 2)
|
|
|
|
# Verify the filter was pushed through the union
|
|
optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
|
|
actual_plan_str = optimized_plan.dag.dag_str
|
|
|
|
# After optimization, filter should be pushed to both union branches
|
|
# So we should see: Filter(Read), Filter(Read) -> Union
|
|
# Not: Read, Read -> Union -> Filter
|
|
assert "Union" in actual_plan_str
|
|
assert "Filter" in actual_plan_str
|
|
# Ensure Filter is before Union (pushed down), not after
|
|
assert actual_plan_str.index("Filter") < actual_plan_str.index("Union")
|
|
|
|
# Verify correctness
|
|
result = sorted(row["id"] for row in ds.take_all())
|
|
assert result == [3, 4, 5]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|