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__]))