wehub-resource-sync
658 lines
22 KiB
Python
658 lines
22 KiB
Python
import sys
|
||
from typing import Any, Dict, List
|
||
|
||
import pandas as pd
|
||
import pytest
|
||
|
||
import ray
|
||
from ray.data import Dataset
|
||
from ray.data._internal.logical.interfaces import LogicalOperator, Plan
|
||
from ray.data._internal.logical.operators import Download, Limit
|
||
from ray.data._internal.logical.rules.limit_pushdown import LimitPushdownRule
|
||
from ray.data._internal.util import rows_same
|
||
from ray.data.block import BlockMetadata
|
||
from ray.data.datasource import Datasource
|
||
from ray.data.datasource.datasource import ReadTask
|
||
from ray.data.tests.conftest import * # noqa
|
||
from ray.tests.conftest import * # noqa
|
||
|
||
|
||
def _check_valid_plan_and_result(
|
||
ds: Dataset,
|
||
expected_plan: Plan,
|
||
expected_result: List[Dict[str, Any]],
|
||
expected_physical_plan_ops=None,
|
||
check_ordering=True,
|
||
):
|
||
actual_result = ds.take_all()
|
||
if check_ordering:
|
||
assert actual_result == expected_result
|
||
else:
|
||
assert rows_same(pd.DataFrame(actual_result), pd.DataFrame(expected_result))
|
||
assert ds._logical_plan.dag.dag_str == expected_plan
|
||
|
||
expected_physical_plan_ops = expected_physical_plan_ops or []
|
||
for op in expected_physical_plan_ops:
|
||
assert op in ds.stats(), f"Operator {op} not found: {ds.stats()}"
|
||
|
||
|
||
class _DummyLogicalOperator(LogicalOperator):
|
||
def __init__(self, input_dependencies, name=None):
|
||
object.__setattr__(self, "_input_dependencies", input_dependencies)
|
||
if name is not None:
|
||
object.__setattr__(self, "_name", name)
|
||
|
||
|
||
def test_limit_pushdown_recreates_frozen_download():
|
||
input_op = _DummyLogicalOperator(input_dependencies=[], name="DummyInput")
|
||
download_op = Download(
|
||
uri_column_names=["uri"],
|
||
output_bytes_column_names=["bytes"],
|
||
input_dependencies=[input_op],
|
||
)
|
||
limit_op = Limit(1, input_dependencies=[download_op])
|
||
|
||
result = LimitPushdownRule()._push_limit_down(limit_op)
|
||
|
||
assert isinstance(result, Download)
|
||
assert isinstance(result.input_dependencies[0], Limit)
|
||
assert result.input_dependencies[0].limit == 1
|
||
assert result.input_dependencies[0].input_dependencies[0] is input_op
|
||
|
||
|
||
def test_limit_pushdown_basic_limit_fusion(ray_start_regular_shared_2_cpus):
|
||
"""Test basic Limit -> Limit fusion."""
|
||
# Use override_num_blocks=1 for deterministic row ordering.
|
||
ds = ray.data.range(100, override_num_blocks=1).limit(5).limit(100)
|
||
_check_valid_plan_and_result(
|
||
ds,
|
||
"Read[ReadRange] -> Limit[limit=5]",
|
||
[{"id": i} for i in range(5)],
|
||
check_ordering=False,
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_limit_fusion_reversed(ray_start_regular_shared_2_cpus):
|
||
"""Test Limit fusion with reversed order."""
|
||
# Use override_num_blocks=1 for deterministic row ordering.
|
||
ds = ray.data.range(100, override_num_blocks=1).limit(100).limit(5)
|
||
_check_valid_plan_and_result(
|
||
ds,
|
||
"Read[ReadRange] -> Limit[limit=5]",
|
||
[{"id": i} for i in range(5)],
|
||
check_ordering=False,
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_multiple_limit_fusion(ray_start_regular_shared_2_cpus):
|
||
"""Test multiple Limit operations fusion."""
|
||
# Use override_num_blocks=1 for deterministic row ordering.
|
||
ds = (
|
||
ray.data.range(100, override_num_blocks=1)
|
||
.limit(50)
|
||
.limit(80)
|
||
.limit(5)
|
||
.limit(20)
|
||
)
|
||
_check_valid_plan_and_result(
|
||
ds,
|
||
"Read[ReadRange] -> Limit[limit=5]",
|
||
[{"id": i} for i in range(5)],
|
||
check_ordering=False,
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_through_maprows(ray_start_regular_shared_2_cpus):
|
||
"""Test that Limit pushes through MapRows operations."""
|
||
|
||
def f1(x):
|
||
return x
|
||
|
||
ds = ray.data.range(100, override_num_blocks=100).map(f1).limit(1)
|
||
_check_valid_plan_and_result(
|
||
ds,
|
||
"Read[ReadRange] -> Limit[limit=1] -> MapRows[Map(f1)]",
|
||
[{"id": 0}],
|
||
check_ordering=False,
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_through_mapbatches(ray_start_regular_shared_2_cpus):
|
||
"""Test that Limit pushes through MapBatches operations."""
|
||
|
||
def f2(x):
|
||
return x
|
||
|
||
ds = (
|
||
ray.data.range(100, override_num_blocks=100)
|
||
.map_batches(f2, udf_modifying_row_count=False)
|
||
.limit(1)
|
||
)
|
||
_check_valid_plan_and_result(
|
||
ds,
|
||
"Read[ReadRange] -> Limit[limit=1] -> MapBatches[MapBatches(f2)]",
|
||
[{"id": 0}],
|
||
check_ordering=False,
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_stops_at_filter(ray_start_regular_shared_2_cpus):
|
||
"""Test that Limit does NOT push through Filter operations (conservative)."""
|
||
ds = (
|
||
ray.data.range(100, override_num_blocks=100)
|
||
.filter(lambda x: x["id"] < 50)
|
||
.limit(1)
|
||
)
|
||
_check_valid_plan_and_result(
|
||
ds,
|
||
"Read[ReadRange] -> Filter[Filter(<lambda>)] -> Limit[limit=1]",
|
||
[{"id": 0}],
|
||
check_ordering=False,
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_through_project(ray_start_regular_shared_2_cpus):
|
||
"""Test that Limit pushes through Project operations."""
|
||
ds = ray.data.range(100, override_num_blocks=100).select_columns(["id"]).limit(5)
|
||
_check_valid_plan_and_result(
|
||
ds,
|
||
"Read[ReadRange] -> Limit[limit=5] -> Project[Project]",
|
||
[{"id": i} for i in range(5)],
|
||
check_ordering=False,
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_stops_at_sort(ray_start_regular_shared_2_cpus):
|
||
"""Test that Limit stops at Sort operations (AllToAll)."""
|
||
ds = ray.data.range(100).sort("id").limit(5)
|
||
_check_valid_plan_and_result(
|
||
ds,
|
||
"Read[ReadRange] -> Sort[Sort] -> Limit[limit=5]",
|
||
[{"id": i} for i in range(5)],
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_complex_interweaved_operations(ray_start_regular_shared_2_cpus):
|
||
"""Test Limit pushdown with complex interweaved operations."""
|
||
|
||
def f1(x):
|
||
return x
|
||
|
||
def f2(x):
|
||
return x
|
||
|
||
ds = ray.data.range(100).sort("id").map(f1).limit(20).sort("id").map(f2).limit(5)
|
||
_check_valid_plan_and_result(
|
||
ds,
|
||
"Read[ReadRange] -> Sort[Sort] -> Limit[limit=20] -> MapRows[Map(f1)] -> "
|
||
"Sort[Sort] -> Limit[limit=5] -> MapRows[Map(f2)]",
|
||
[{"id": i} for i in range(5)],
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_between_two_map_operators(ray_start_regular_shared_2_cpus):
|
||
"""Test Limit pushdown between two Map operators."""
|
||
|
||
def f1(x):
|
||
return x
|
||
|
||
def f2(x):
|
||
return x
|
||
|
||
ds = ray.data.range(100, override_num_blocks=100).map(f1).limit(1).map(f2)
|
||
_check_valid_plan_and_result(
|
||
ds,
|
||
"Read[ReadRange] -> Limit[limit=1] -> MapRows[Map(f1)] -> MapRows[Map(f2)]",
|
||
[{"id": 0}],
|
||
check_ordering=False,
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_correctness(ray_start_regular_shared_2_cpus):
|
||
"""Test that limit pushdown produces correct results in various scenarios."""
|
||
|
||
# Test 1: Simple project + limit
|
||
ds = ray.data.range(100).select_columns(["id"]).limit(10)
|
||
result = ds.take_all()
|
||
expected = [{"id": i} for i in range(10)]
|
||
assert result == expected
|
||
|
||
# Test 2: Multiple operations + limit (with MapRows pushdown)
|
||
ds = (
|
||
ray.data.range(100)
|
||
.map(lambda x: {"id": x["id"], "squared": x["id"] ** 2})
|
||
.select_columns(["id"])
|
||
.limit(5)
|
||
)
|
||
result = ds.take_all()
|
||
expected = [{"id": i} for i in range(5)]
|
||
assert result == expected
|
||
|
||
# Test 3: MapRows operations should get limit pushed (safe)
|
||
ds = ray.data.range(100).map(lambda x: {"id": x["id"] * 2}).limit(5)
|
||
result = ds.take_all()
|
||
expected = [{"id": i * 2} for i in range(5)]
|
||
assert result == expected
|
||
|
||
# Test 4: MapBatches operations should not get limit pushed
|
||
ds = ray.data.range(100).map_batches(lambda batch: {"id": batch["id"] * 2}).limit(5)
|
||
result = ds.take_all()
|
||
expected = [{"id": i * 2} for i in range(5)]
|
||
assert result == expected
|
||
|
||
# Test 5: Filter operations should not get limit pushed (conservative)
|
||
ds = ray.data.range(100).filter(lambda x: x["id"] % 2 == 0).limit(3)
|
||
result = ds.take_all()
|
||
expected = [{"id": i} for i in [0, 2, 4]]
|
||
assert result == expected
|
||
|
||
# Test 6: Complex chain with both safe operations (should all get limit pushed)
|
||
ds = (
|
||
ray.data.range(100)
|
||
.select_columns(["id"]) # Project - could be safe if it was the immediate input
|
||
.map(lambda x: {"id": x["id"] + 1}) # MapRows - NOT safe, stops pushdown
|
||
.limit(3)
|
||
)
|
||
result = ds.take_all()
|
||
expected = [{"id": i + 1} for i in range(3)]
|
||
assert result == expected
|
||
|
||
# The plan should show all operations after the limit
|
||
plan_str = ds._logical_plan.dag.dag_str
|
||
assert (
|
||
"Read[ReadRange] -> Limit[limit=3] -> Project[Project] -> MapRows[Map(<lambda>)]"
|
||
== plan_str
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_scan_efficiency(ray_start_regular_shared_2_cpus):
|
||
"""Test that limit pushdown scans fewer rows from the data source."""
|
||
|
||
@ray.remote
|
||
class Counter:
|
||
def __init__(self):
|
||
self.value = 0
|
||
|
||
def increment(self, amount=1):
|
||
self.value += amount
|
||
return self.value
|
||
|
||
def get(self):
|
||
return self.value
|
||
|
||
def reset(self):
|
||
self.value = 0
|
||
|
||
# Create a custom datasource that tracks how many rows it produces
|
||
class CountingDatasource(Datasource):
|
||
def __init__(self):
|
||
self.counter = Counter.remote()
|
||
|
||
def prepare_read(self, parallelism, n_per_block=10):
|
||
def read_fn(block_idx):
|
||
# Each block produces n_per_block rows
|
||
ray.get(self.counter.increment.remote(n_per_block))
|
||
return [
|
||
pd.DataFrame(
|
||
{
|
||
"id": range(
|
||
block_idx * n_per_block, (block_idx + 1) * n_per_block
|
||
)
|
||
}
|
||
)
|
||
]
|
||
|
||
return [
|
||
ReadTask(
|
||
lambda i=i: read_fn(i),
|
||
BlockMetadata(
|
||
num_rows=n_per_block,
|
||
size_bytes=n_per_block * 8, # rough estimate
|
||
input_files=None,
|
||
exec_stats=None,
|
||
),
|
||
)
|
||
for i in range(parallelism)
|
||
]
|
||
|
||
def get_rows_produced(self):
|
||
return ray.get(self.counter.get.remote())
|
||
|
||
# Test 1: Project + Limit should scan fewer rows due to pushdown
|
||
source = CountingDatasource()
|
||
ds = ray.data.read_datasource(source, override_num_blocks=20, n_per_block=10)
|
||
ds = ds.select_columns(["id"]).limit(5)
|
||
result = ds.take_all()
|
||
|
||
# Should get correct results
|
||
assert len(result) == 5
|
||
assert result == [{"id": i} for i in range(5)]
|
||
|
||
# Should have scanned significantly fewer than all 200 rows (20 blocks * 10 rows)
|
||
# Due to pushdown, we should scan much less
|
||
rows_produced_1 = source.get_rows_produced()
|
||
assert rows_produced_1 < 200 # Should be much less than total
|
||
|
||
# Test 2: MapRows + Limit should also scan fewer rows due to pushdown
|
||
source2 = CountingDatasource()
|
||
ds2 = ray.data.read_datasource(source2, override_num_blocks=20, n_per_block=10)
|
||
ds2 = ds2.map(lambda x: x).limit(5)
|
||
result2 = ds2.take_all()
|
||
|
||
# Should get correct results
|
||
assert len(result2) == 5
|
||
assert result2 == [{"id": i} for i in range(5)]
|
||
|
||
# Should also scan fewer than total due to pushdown
|
||
rows_produced_2 = source2.get_rows_produced()
|
||
assert rows_produced_2 < 200
|
||
|
||
# Both should be efficient with pushdown
|
||
assert rows_produced_1 < 100 # Should be much less than total
|
||
assert rows_produced_2 < 100 # Should be much less than total
|
||
|
||
# Test 3: Filter + Limit should scan fewer due to early termination, but not pushdown
|
||
source3 = CountingDatasource()
|
||
ds3 = ray.data.read_datasource(source3, override_num_blocks=20, n_per_block=10)
|
||
ds3 = ds3.filter(lambda x: x["id"] % 2 == 0).limit(3)
|
||
result3 = ds3.take_all()
|
||
|
||
# Should get correct results
|
||
assert len(result3) == 3
|
||
assert result3 == [{"id": i} for i in [0, 2, 4]]
|
||
|
||
# Should still scan fewer than total due to early termination
|
||
rows_produced_3 = source3.get_rows_produced()
|
||
assert rows_produced_3 < 200
|
||
|
||
|
||
def test_limit_pushdown_union(ray_start_regular_shared_2_cpus):
|
||
"""Test limit pushdown behavior with Union operations."""
|
||
|
||
# Create two datasets and union with limit
|
||
ds1 = ray.data.range(100, override_num_blocks=10)
|
||
ds2 = ray.data.range(200, override_num_blocks=10)
|
||
ds = ds1.union(ds2).limit(5)
|
||
|
||
expected_plan = "Read[ReadRange] -> Limit[limit=5], Read[ReadRange] -> Limit[limit=5] -> Union[Union] -> Limit[limit=5]"
|
||
_check_valid_plan_and_result(
|
||
ds, expected_plan, [{"id": i} for i in range(5)], check_ordering=False
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_union_with_maprows(ray_start_regular_shared_2_cpus):
|
||
"""Limit after Union + MapRows: limit should be pushed before the MapRows
|
||
and inside each Union branch."""
|
||
ds1 = ray.data.range(100, override_num_blocks=10)
|
||
ds2 = ray.data.range(200, override_num_blocks=10)
|
||
ds = ds1.union(ds2).map(lambda x: x).limit(5)
|
||
|
||
expected_plan = (
|
||
"Read[ReadRange] -> Limit[limit=5], "
|
||
"Read[ReadRange] -> Limit[limit=5] -> Union[Union] -> "
|
||
"Limit[limit=5] -> MapRows[Map(<lambda>)]"
|
||
)
|
||
_check_valid_plan_and_result(
|
||
ds, expected_plan, [{"id": i} for i in range(5)], check_ordering=False
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_union_with_sort(ray_start_regular_shared_2_cpus):
|
||
"""Limit after Union + Sort: limit must NOT push through the Sort."""
|
||
ds1 = ray.data.range(100, override_num_blocks=4)
|
||
ds2 = ray.data.range(50, override_num_blocks=4).map(
|
||
lambda x: {"id": x["id"] + 1000}
|
||
)
|
||
ds = ds1.union(ds2).sort("id").limit(5)
|
||
|
||
expected_plan = (
|
||
"Read[ReadRange], "
|
||
"Read[ReadRange] -> MapRows[Map(<lambda>)] -> "
|
||
"Union[Union] -> Sort[Sort] -> Limit[limit=5]"
|
||
)
|
||
_check_valid_plan_and_result(ds, expected_plan, [{"id": i} for i in range(5)])
|
||
|
||
|
||
def test_limit_pushdown_multiple_unions(ray_start_regular_shared_2_cpus):
|
||
"""Outer limit over nested unions should create a branch-local limit
|
||
for every leaf plus the global one."""
|
||
ds = (
|
||
ray.data.range(100)
|
||
.union(ray.data.range(100, override_num_blocks=5))
|
||
.union(ray.data.range(50))
|
||
.limit(5)
|
||
)
|
||
|
||
expected_plan = (
|
||
"Read[ReadRange] -> Limit[limit=5], "
|
||
"Read[ReadRange] -> Limit[limit=5] -> Union[Union] -> Limit[limit=5], "
|
||
"Read[ReadRange] -> Limit[limit=5] -> Union[Union] -> Limit[limit=5]"
|
||
)
|
||
_check_valid_plan_and_result(
|
||
ds, expected_plan, [{"id": i} for i in range(5)], check_ordering=False
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_union_with_groupby(ray_start_regular_shared_2_cpus):
|
||
"""Limit after Union + Aggregate: limit should stay after Aggregate."""
|
||
ds1 = ray.data.range(100)
|
||
ds2 = ray.data.range(100).map(lambda x: {"id": x["id"] + 1000})
|
||
ds = ds1.union(ds2).groupby("id").count().limit(5)
|
||
# Result should contain 5 distinct ids with count == 1.
|
||
res = ds.take_all()
|
||
# Plan suffix check (no branch limits past Aggregate).
|
||
assert ds._logical_plan.dag.dag_str.endswith(
|
||
"Union[Union] -> Aggregate[Aggregate] -> Limit[limit=5]"
|
||
)
|
||
assert len(res) == 5 and all(r["count()"] == 1 for r in res)
|
||
|
||
|
||
def test_limit_pushdown_complex_chain(ray_start_regular_shared_2_cpus):
|
||
"""
|
||
Complex end-to-end case:
|
||
1. Two branches each with a branch-local Limit pushed to Read.
|
||
• left : Project
|
||
• right : MapRows
|
||
2. Union of the two branches.
|
||
3. Global Aggregate (groupby/count).
|
||
4. Sort (descending id) – pushes stop here.
|
||
5. Final Limit.
|
||
Verifies both plan rewrite and result correctness.
|
||
"""
|
||
# ── left branch ────────────────────────────────────────────────
|
||
left = ray.data.range(50).select_columns(["id"]).limit(10)
|
||
|
||
# ── right branch ───────────────────────────────────────────────
|
||
right = ray.data.range(50).map(lambda x: {"id": x["id"] + 1000}).limit(10)
|
||
|
||
# ── union → aggregate → sort → limit ──────────────────────────
|
||
ds = left.union(right).groupby("id").count().sort("id", descending=True).limit(3)
|
||
|
||
# Expected logical-plan string.
|
||
expected_plan = (
|
||
"Read[ReadRange] -> Limit[limit=10] -> Project[Project], "
|
||
"Read[ReadRange] -> Limit[limit=10] -> MapRows[Map(<lambda>)] "
|
||
"-> Union[Union] -> Aggregate[Aggregate] -> Sort[Sort] -> Limit[limit=3]"
|
||
)
|
||
|
||
# Top-3 ids are the three largest (1009, 1008, 1007) with count()==1.
|
||
expected_result = [
|
||
{"id": 1009, "count()": 1},
|
||
{"id": 1008, "count()": 1},
|
||
{"id": 1007, "count()": 1},
|
||
]
|
||
|
||
_check_valid_plan_and_result(ds, expected_plan, expected_result)
|
||
|
||
|
||
def test_limit_pushdown_union_maps_projects(ray_start_regular_shared_2_cpus):
|
||
r"""
|
||
Read -> MapBatches -> MapRows -> Project
|
||
\ /
|
||
-------- Union ------------- → Limit
|
||
The limit should be pushed in front of each branch
|
||
(past MapRows, Project) while the original
|
||
global Limit is preserved after the Union.
|
||
"""
|
||
# Left branch.
|
||
left = (
|
||
ray.data.range(30)
|
||
.map_batches(lambda b: b, udf_modifying_row_count=False)
|
||
.map(lambda r: {"id": r["id"]})
|
||
.select_columns(["id"])
|
||
)
|
||
|
||
# Right branch with shifted ids.
|
||
right = (
|
||
ray.data.range(30)
|
||
.map_batches(lambda b: b, udf_modifying_row_count=False)
|
||
.map(lambda r: {"id": r["id"] + 100})
|
||
.select_columns(["id"])
|
||
)
|
||
|
||
ds = left.union(right).limit(3)
|
||
|
||
expected_plan = (
|
||
"Read[ReadRange] -> "
|
||
"Limit[limit=3] -> MapBatches[MapBatches(<lambda>)] -> MapRows[Map(<lambda>)] -> "
|
||
"Project[Project], "
|
||
"Read[ReadRange] -> "
|
||
"Limit[limit=3] -> MapBatches[MapBatches(<lambda>)] -> MapRows[Map(<lambda>)] -> "
|
||
"Project[Project] -> Union[Union] -> Limit[limit=3]"
|
||
)
|
||
|
||
expected_result = [{"id": i} for i in range(3)] # First 3 rows from left branch.
|
||
|
||
_check_valid_plan_and_result(
|
||
ds, expected_plan, expected_result, check_ordering=False
|
||
)
|
||
|
||
|
||
def test_limit_pushdown_map_per_block_limit_applied(ray_start_regular_shared_2_cpus):
|
||
"""Test that per-block limits are actually applied during map execution."""
|
||
|
||
# Create a global counter using Ray
|
||
@ray.remote
|
||
class Counter:
|
||
def __init__(self):
|
||
self.value = 0
|
||
|
||
def increment(self):
|
||
self.value += 1
|
||
return self.value
|
||
|
||
def get(self):
|
||
return self.value
|
||
|
||
counter = Counter.remote()
|
||
|
||
def track_processing(row):
|
||
# Record that this row was processed
|
||
ray.get(counter.increment.remote())
|
||
return row
|
||
|
||
# Create dataset with limit pushed through map
|
||
ds = ray.data.range(1000, override_num_blocks=10).map(track_processing).limit(50)
|
||
|
||
# Execute and get results
|
||
result = ds.take_all()
|
||
|
||
# Verify correct results
|
||
expected = [{"id": i} for i in range(50)]
|
||
assert result == expected
|
||
|
||
# Check how many rows were actually processed
|
||
processed_count = ray.get(counter.get.remote())
|
||
|
||
# With per-block limits, we should process fewer rows than the total dataset
|
||
# but at least the number we need for the final result
|
||
assert (
|
||
processed_count >= 50
|
||
), f"Expected at least 50 rows processed, got {processed_count}"
|
||
assert (
|
||
processed_count < 1000
|
||
), f"Expected fewer than 1000 rows processed, got {processed_count}"
|
||
|
||
print(f"Processed {processed_count} rows to get {len(result)} results")
|
||
|
||
|
||
def test_limit_pushdown_preserves_map_behavior(ray_start_regular_shared_2_cpus):
|
||
"""Test that adding per-block limits doesn't change the logical result."""
|
||
|
||
def add_one(row):
|
||
row["id"] += 1
|
||
return row
|
||
|
||
# Compare with and without limit pushdown
|
||
ds_with_limit = ray.data.range(100).map(add_one).limit(10)
|
||
ds_without_limit = ray.data.range(100).limit(10).map(add_one)
|
||
|
||
result_with = ds_with_limit.take_all()
|
||
result_without = ds_without_limit.take_all()
|
||
|
||
# Results should be identical
|
||
assert result_with == result_without
|
||
|
||
# Both should have the expected transformation applied
|
||
expected = [{"id": i + 1} for i in range(10)]
|
||
assert result_with == expected
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"udf_modifying_row_count,expected_plan",
|
||
[
|
||
(
|
||
False,
|
||
"Read[ReadRange] -> Limit[limit=10] -> MapBatches[MapBatches(<lambda>)]",
|
||
),
|
||
(
|
||
True,
|
||
"Read[ReadRange] -> MapBatches[MapBatches(<lambda>)] -> Limit[limit=10]",
|
||
),
|
||
],
|
||
)
|
||
def test_limit_pushdown_udf_modifying_row_count_with_map_batches(
|
||
ray_start_regular_shared_2_cpus,
|
||
udf_modifying_row_count,
|
||
expected_plan,
|
||
):
|
||
"""Test that limit pushdown preserves the row count with map batches."""
|
||
ds = (
|
||
ray.data.range(100)
|
||
.map_batches(lambda x: x, udf_modifying_row_count=udf_modifying_row_count)
|
||
.limit(10)
|
||
)
|
||
_check_valid_plan_and_result(
|
||
ds,
|
||
expected_plan,
|
||
[{"id": i} for i in range(10)],
|
||
)
|
||
|
||
|
||
def test_does_not_pushdown_limit_past_map_batches_by_default(
|
||
ray_start_regular_shared_2_cpus,
|
||
):
|
||
def duplicate_id(batch):
|
||
yield {"data": list(batch["id"]) * 2}
|
||
|
||
# If the optimizer incorrectly pushes the limit past the map operator, then the
|
||
# returned count is 2.
|
||
num_rows = ray.data.range(1).map_batches(duplicate_id).limit(1).count()
|
||
assert num_rows == 1, num_rows
|
||
|
||
|
||
def test_does_not_pushdown_limit_past_map_groups_by_default(
|
||
ray_start_regular_shared_2_cpus,
|
||
):
|
||
def duplicate_id(batch):
|
||
yield {"data": list(batch["id"]) * 2}
|
||
|
||
# If the optimizer incorrectly pushes the limit past the map operator, then the
|
||
# returned count is 2.
|
||
num_rows = ray.data.range(1).groupby("id").map_groups(duplicate_id).limit(1).count()
|
||
assert num_rows == 1, num_rows
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(pytest.main(["-v", __file__]))
|