wehub-resource-sync
161 lines
6.3 KiB
Python
161 lines
6.3 KiB
Python
from unittest.mock import MagicMock
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data._internal.execution.interfaces import ExecutionOptions
|
|
from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
|
|
from ray.data._internal.execution.operators.limit_operator import LimitOperator
|
|
from ray.data._internal.execution.streaming_executor import StreamingExecutor
|
|
from ray.data._internal.execution.util import make_ref_bundles
|
|
from ray.data._internal.logical.optimizers import get_execution_plan
|
|
from ray.data.context import DataContext
|
|
from ray.data.tests.conftest import noop_counter
|
|
from ray.data.tests.util import run_op_tasks_sync
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
|
|
def test_limit_estimated_num_output_bundles(ray_start_regular_shared):
|
|
# Test limit operator estimation
|
|
input_op = InputDataBuffer(
|
|
DataContext.get_current(), make_ref_bundles([[i, i] for i in range(100)])
|
|
)
|
|
op = LimitOperator(100, input_op, DataContext.get_current())
|
|
|
|
while input_op.has_next():
|
|
op.add_input(input_op.get_next(), 0)
|
|
run_op_tasks_sync(op)
|
|
assert op._estimated_num_output_bundles == 50
|
|
|
|
op.all_inputs_done()
|
|
|
|
# 2 rows per bundle, 100 / 2 = 50 blocks output
|
|
assert op._estimated_num_output_bundles == 50
|
|
|
|
# Test limit operator estimation where: limit > # of rows
|
|
input_op = InputDataBuffer(
|
|
DataContext.get_current(), make_ref_bundles([[i, i] for i in range(100)])
|
|
)
|
|
op = LimitOperator(300, input_op, DataContext.get_current())
|
|
|
|
while input_op.has_next():
|
|
op.add_input(input_op.get_next(), 0)
|
|
run_op_tasks_sync(op)
|
|
assert op._estimated_num_output_bundles == 100
|
|
|
|
op.all_inputs_done()
|
|
|
|
# all blocks are outputted
|
|
assert op._estimated_num_output_bundles == 100
|
|
|
|
|
|
def test_limit_operator(ray_start_regular_shared):
|
|
"""Test basic functionalities of LimitOperator."""
|
|
num_refs = 3
|
|
num_rows_per_block = 3
|
|
total_rows = num_refs * num_rows_per_block
|
|
# Test limits with different values, from 0 to more than input size.
|
|
limits = list(range(0, total_rows + 2))
|
|
for limit in limits:
|
|
refs = make_ref_bundles([[i] * num_rows_per_block for i in range(num_refs)])
|
|
input_op = InputDataBuffer(DataContext.get_current(), refs)
|
|
limit_op = LimitOperator(limit, input_op, DataContext.get_current())
|
|
counter = noop_counter()
|
|
input_op.start(ExecutionOptions(), counter)
|
|
limit_op.start(ExecutionOptions(), counter)
|
|
limit_op.mark_execution_finished = MagicMock(
|
|
wraps=limit_op.mark_execution_finished
|
|
)
|
|
if limit == 0:
|
|
# If the limit is 0, the operator should be completed immediately.
|
|
assert limit_op.has_completed()
|
|
assert limit_op._limit_reached()
|
|
cur_rows = 0
|
|
loop_count = 0
|
|
while input_op.has_next() and not limit_op._limit_reached():
|
|
loop_count += 1
|
|
assert not limit_op.has_completed(), limit
|
|
assert not limit_op.has_execution_finished(), limit
|
|
limit_op.add_input(input_op.get_next(), 0)
|
|
while limit_op.has_next():
|
|
# Drain the outputs. So the limit operator
|
|
# will be completed when the limit is reached.
|
|
limit_op.get_next()
|
|
cur_rows += num_rows_per_block
|
|
if cur_rows >= limit:
|
|
assert limit_op.mark_execution_finished.call_count == 1, limit
|
|
assert limit_op.has_completed(), limit
|
|
assert limit_op._limit_reached(), limit
|
|
assert limit_op.has_execution_finished(), limit
|
|
else:
|
|
assert limit_op.mark_execution_finished.call_count == 0, limit
|
|
assert not limit_op.has_completed(), limit
|
|
assert not limit_op._limit_reached(), limit
|
|
assert not limit_op.has_execution_finished(), limit
|
|
limit_op.mark_execution_finished()
|
|
# After inputs done, the number of output bundles
|
|
# should be the same as the number of `add_input`s.
|
|
assert limit_op.num_outputs_total() == loop_count, limit
|
|
assert limit_op.has_completed(), limit
|
|
|
|
|
|
def test_limit_operator_memory_leak_fix(ray_start_regular_shared, tmp_path):
|
|
"""Test that LimitOperator properly drains upstream output queues.
|
|
|
|
This test verifies the memory leak fix by directly using StreamingExecutor
|
|
to access the actual topology and check queued blocks after execution.
|
|
"""
|
|
for i in range(100):
|
|
data = [{"id": i * 5 + j, "value": f"row_{i * 5 + j}"} for j in range(5)]
|
|
table = pa.Table.from_pydict(
|
|
{"id": [row["id"] for row in data], "value": [row["value"] for row in data]}
|
|
)
|
|
parquet_file = tmp_path / f"test_data_{i}.parquet"
|
|
pq.write_table(table, str(parquet_file))
|
|
|
|
parquet_files = [str(tmp_path / f"test_data_{i}.parquet") for i in range(100)]
|
|
|
|
ds = (
|
|
ray.data.read_parquet(parquet_files, override_num_blocks=100)
|
|
.limit(5)
|
|
.map(lambda x: x)
|
|
)
|
|
|
|
physical_plan, _ = get_execution_plan(ds._logical_plan)
|
|
|
|
# Use StreamingExecutor directly to have access to the actual topology
|
|
executor = StreamingExecutor(DataContext.get_current())
|
|
output_iterator = executor.execute(physical_plan.dag)
|
|
|
|
# Collect all results and count rows
|
|
total_rows = 0
|
|
for bundle in output_iterator:
|
|
for block_ref in bundle.block_refs:
|
|
block = ray.get(block_ref)
|
|
total_rows += block.num_rows
|
|
assert (
|
|
total_rows == 5
|
|
), f"Expected exactly 5 rows after limit(5), but got {total_rows}"
|
|
|
|
# Find the parquet read operator's OpState. Covers both the V1
|
|
# ``ReadParquet`` op name and the V2 ``ReadFilesParquet{V2,}`` name
|
|
# under the ``DataContext.use_datasource_v2`` path.
|
|
topology = executor._topology
|
|
read_parquet_op_state = None
|
|
for op, op_state in topology.items():
|
|
if "Parquet" in op.name and ("Read" in op.name or "ReadFiles" in op.name):
|
|
read_parquet_op_state = op_state
|
|
break
|
|
|
|
# Check the output queue size
|
|
output_queue_size = len(read_parquet_op_state.output_queue)
|
|
assert output_queue_size == 0, f"Expected 0 items, but got {output_queue_size}."
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|