From 66b0aa5b3160de78684f268c2a3508789535d09f Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 22 Jul 2026 20:59:00 +0800 Subject: [PATCH 1/3] [python] Pipeline split reads in Arrow batch reader --- docs/docs/pypaimon/python-api.mdx | 5 + .../pypaimon/common/options/core_options.py | 4 +- paimon-python/pypaimon/read/table_read.py | 108 ++++++++++++- .../pypaimon/tests/reader_parallel_test.py | 145 +++++++++++++++++- 4 files changed, 255 insertions(+), 7 deletions(-) diff --git a/docs/docs/pypaimon/python-api.mdx b/docs/docs/pypaimon/python-api.mdx index c1887c537c75..05fe388efb38 100644 --- a/docs/docs/pypaimon/python-api.mdx +++ b/docs/docs/pypaimon/python-api.mdx @@ -471,6 +471,11 @@ for batch in table_read.to_arrow_batch_reader(splits): # f1: ["a","b","c"] ``` +Multiple splits are read concurrently according to `parallelism` or +`read.parallelism` while batches are streamed in split order. Pass +`parallelism=1` to force serial reading. The reader buffers at most two batches +per in-flight split. + ### Read Python Iterator You can read the data row by row into a native Python iterator. diff --git a/paimon-python/pypaimon/common/options/core_options.py b/paimon-python/pypaimon/common/options/core_options.py index 3b890cc0d62d..b9e16ba5d711 100644 --- a/paimon-python/pypaimon/common/options/core_options.py +++ b/paimon-python/pypaimon/common/options/core_options.py @@ -919,8 +919,8 @@ class CoreOptions: "When unset (the default), reads auto-scale to " "min(number of splits, CPU count). Set to 1 to force serial " "reads, or to a specific value >= 1 to cap the thread pool that " - "reads splits concurrently and assembles the result in input " - "order. Has no effect when fewer than 2 splits are passed.") + "reads splits concurrently while preserving input split order. " + "Has no effect when fewer than 2 splits are passed.") ) ADD_COLUMN_BEFORE_PARTITION: ConfigOption[bool] = ( diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index 6e641b50963a..2a8ca8463a52 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -16,6 +16,7 @@ # under the License. import os +import queue import threading from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Dict, Iterator, List, Optional @@ -82,6 +83,8 @@ class TableRead: # workers (P) each spin up blob_parallelism (B) blob threads, so peak # connections ~= P*B. Shrink per-split B to keep the product bounded. _MAX_TOTAL_BLOB_WORKERS = 64 + _PIPELINE_QUEUE_BATCHES = 2 + _PIPELINE_QUEUE_TIMEOUT_SECONDS = 0.1 def __init__( self, @@ -137,15 +140,108 @@ def _record_generator(): return _record_generator() - def to_arrow_batch_reader(self, splits: List[Split], - blob_parallelism: Optional[int] = None) -> pyarrow.ipc.RecordBatchReader: + def to_arrow_batch_reader( + self, + splits: List[Split], + blob_parallelism: Optional[int] = None, + parallelism: Optional[int] = None, + ) -> pyarrow.ipc.RecordBatchReader: + """Stream Arrow batches while reading multiple splits concurrently. + + Batches preserve input split order. At most two batches per in-flight + split are buffered, keeping memory bounded while overlapping split I/O. + """ effective_bp = self._resolve_blob_parallelism(blob_parallelism) + effective = self._resolve_parallelism(parallelism, len(splits)) schema = PyarrowFieldParser.from_paimon_schema(self.read_type) if self.include_row_kind: schema = self._add_row_kind_to_schema(schema) - batch_iterator = self._arrow_batch_generator(splits, schema, effective_bp) + if self._should_run_parallel(splits, effective): + workers = min(effective, len(splits)) + effective_bp = self._cap_blob_parallelism(workers, effective_bp) + batch_iterator = self._pipelined_arrow_batch_generator( + splits, schema, effective_bp, workers) + else: + batch_iterator = self._arrow_batch_generator( + splits, schema, effective_bp) return pyarrow.ipc.RecordBatchReader.from_batches(schema, batch_iterator) + def _pipelined_arrow_batch_generator( + self, + splits: List[Split], + schema: pyarrow.Schema, + blob_parallelism: int, + workers: int, + ) -> Iterator[pyarrow.RecordBatch]: + cancel = threading.Event() + end = object() + queues = {} + + def put_item(output_queue, item) -> bool: + while not cancel.is_set(): + try: + output_queue.put( + item, timeout=self._PIPELINE_QUEUE_TIMEOUT_SECONDS) + return True + except queue.Full: + continue + return False + + def read_split(split, output_queue): + batch_iterator = self._arrow_batch_generator( + [split], schema, blob_parallelism) + try: + for batch in batch_iterator: + if not put_item(output_queue, batch): + return + except BaseException as exception: + put_item(output_queue, exception) + finally: + try: + batch_iterator.close() + except BaseException as exception: + put_item(output_queue, exception) + finally: + put_item(output_queue, end) + + executor = ThreadPoolExecutor( + max_workers=workers, + thread_name_prefix="pypaimon-stream-read", + ) + + def schedule(index: int): + if index >= len(splits): + return + output_queue = queue.Queue(maxsize=self._PIPELINE_QUEUE_BATCHES) + queues[index] = output_queue + executor.submit(read_split, splits[index], output_queue) + + remaining = self.limit + try: + for index in range(workers): + schedule(index) + + for index in range(len(splits)): + output_queue = queues.pop(index) + while True: + item = output_queue.get() + if item is end: + break + if isinstance(item, BaseException): + raise item + batch = item + if remaining is not None and batch.num_rows > remaining: + batch = batch.slice(0, remaining) + yield batch + if remaining is not None: + remaining -= batch.num_rows + if remaining <= 0: + return + schedule(index + workers) + finally: + cancel.set() + executor.shutdown(wait=True) + @staticmethod def _add_row_kind_to_schema(schema: pyarrow.Schema) -> pyarrow.Schema: """Add _row_kind column to the schema as the first column.""" @@ -207,7 +303,11 @@ def to_arrow( if self._should_run_parallel(splits, effective): return self._to_arrow_parallel(splits, schema, effective, effective_bp) - batch_reader = self.to_arrow_batch_reader(splits, blob_parallelism=effective_bp) + batch_reader = self.to_arrow_batch_reader( + splits, + blob_parallelism=effective_bp, + parallelism=effective, + ) table_list = [] for batch in iter(batch_reader.read_next_batch, None): diff --git a/paimon-python/pypaimon/tests/reader_parallel_test.py b/paimon-python/pypaimon/tests/reader_parallel_test.py index ea0c3d8f61d1..f8260b9b4b10 100644 --- a/paimon-python/pypaimon/tests/reader_parallel_test.py +++ b/paimon-python/pypaimon/tests/reader_parallel_test.py @@ -120,6 +120,80 @@ def worker(): self.assertTrue(rr.exhausted()) +class PipelinedBatchGeneratorTest(unittest.TestCase): + + def setUp(self): + self.read = TableRead.__new__(TableRead) + self.read.limit = None + self.schema = pa.schema([('value', pa.int64())]) + + def _batch(self, value): + return pa.RecordBatch.from_arrays( + [pa.array([value], type=pa.int64())], schema=self.schema) + + def test_reads_concurrently_and_emits_in_split_order(self): + second_started = threading.Event() + + def generate(splits, schema, blob_parallelism): + value = splits[0] + if value == 0 and not second_started.wait(5): + raise TimeoutError('second split did not start concurrently') + if value == 1: + second_started.set() + yield self._batch(value) + + with mock.patch.object( + self.read, '_arrow_batch_generator', side_effect=generate + ): + batches = list(self.read._pipelined_arrow_batch_generator( + [0, 1, 2], self.schema, 1, 2)) + + self.assertEqual( + [0, 1, 2], + [batch.column(0)[0].as_py() for batch in batches], + ) + + def test_propagates_worker_error(self): + def generate(splits, schema, blob_parallelism): + if splits[0] == 0: + raise RuntimeError('simulated streaming failure') + yield self._batch(splits[0]) + + with mock.patch.object( + self.read, '_arrow_batch_generator', side_effect=generate + ): + with self.assertRaisesRegex( + RuntimeError, 'simulated streaming failure' + ): + list(self.read._pipelined_arrow_batch_generator( + [0, 1], self.schema, 1, 2)) + + def test_close_cancels_workers(self): + started = threading.Barrier(2) + closed = [] + closed_lock = threading.Lock() + + def generate(splits, schema, blob_parallelism): + value = splits[0] + try: + started.wait(5) + while True: + yield self._batch(value) + finally: + with closed_lock: + closed.append(value) + + with mock.patch.object( + self.read, '_arrow_batch_generator', side_effect=generate + ): + batches = self.read._pipelined_arrow_batch_generator( + [0, 1], self.schema, 1, 2) + next(batches) + batches.close() + + self.assertCountEqual([0, 1], closed) + + class ParallelReaderAppendOnlyTest(unittest.TestCase): """Append-only multi-partition table — parallel must match serial exactly.""" @@ -225,6 +299,57 @@ def test_parallel_via_table_option_matches_serial(self): .sort_values('user_id').reset_index(drop=True) self.assertTrue(serial_df.equals(parallel_df)) + def test_parallel_batch_reader_matches_serial(self): + rb = self.table.new_read_builder() + splits = self._scan_splits(rb) + read = rb.new_read() + serial = pa.Table.from_batches( + read.to_arrow_batch_reader(splits, parallelism=1)) + parallel = pa.Table.from_batches( + read.to_arrow_batch_reader(splits, parallelism=4)) + self.assertEqual(serial, parallel) + + def test_parallel_batch_reader_uses_table_option(self): + rb = self.table_opt_4.new_read_builder() + splits = self._scan_splits(rb) + read = rb.new_read() + with mock.patch.object( + read, + '_pipelined_arrow_batch_generator', + wraps=read._pipelined_arrow_batch_generator, + ) as pipeline: + result = pa.Table.from_batches( + read.to_arrow_batch_reader(splits)) + pipeline.assert_called_once() + self.assertEqual(self.expected_rows, result.num_rows) + + def test_default_batch_reader_auto_uses_pipeline_when_multicore(self): + if (os.cpu_count() or 1) < 2: + self.skipTest('single-core runner: auto resolves to serial') + rb = self.table.new_read_builder() + splits = self._scan_splits(rb) + read = rb.new_read() + with mock.patch.object( + read, + '_pipelined_arrow_batch_generator', + wraps=read._pipelined_arrow_batch_generator, + ) as pipeline: + result = pa.Table.from_batches( + read.to_arrow_batch_reader(splits)) + pipeline.assert_called_once() + self.assertEqual(self.expected_rows, result.num_rows) + + def test_parallel_batch_reader_limit_matches_serial(self): + rb = self.table.new_read_builder().with_limit(600) + splits = self._scan_splits(rb) + read = rb.new_read() + serial = pa.Table.from_batches( + read.to_arrow_batch_reader(splits, parallelism=1)) + parallel = pa.Table.from_batches( + read.to_arrow_batch_reader(splits, parallelism=4)) + self.assertEqual(serial, parallel) + self.assertEqual(600, parallel.num_rows) + def test_default_none_runs_auto_parallel(self): # No runtime arg and no table option => auto. With >= 2 splits on a # multi-core box this takes the parallel fan-out path; the result @@ -257,11 +382,13 @@ def test_default_none_auto_takes_parallel_path_when_multicore(self): def test_method_arg_overrides_option_to_serial(self): # option=4 but caller passes 1: should disable parallelism. read = self.table_opt_4.new_read_builder().new_read() - with mock.patch.object(read, '_to_arrow_parallel') as patched: + with mock.patch.object(read, '_to_arrow_parallel') as patched, \ + mock.patch.object(read, '_pipelined_arrow_batch_generator') as pipeline: patched.side_effect = AssertionError( "_to_arrow_parallel should not be called when arg=1 overrides option") splits = self._scan_splits(self.table_opt_4.new_read_builder()) read.to_arrow(splits, parallelism=1) + pipeline.assert_not_called() def test_method_arg_overrides_option_to_parallel(self): # option=1 (forces serial) but caller passes 4: should enable parallelism. @@ -305,6 +432,8 @@ def test_invalid_method_arg_raises(self): self.assertNotIn("read.parallelism", str(ctx.exception)) with self.assertRaises(ValueError): read.to_pandas(splits, parallelism=-1) + with self.assertRaises(ValueError): + read.to_arrow_batch_reader(splits, parallelism=0) def test_invalid_option_value_raises(self): # Build a fresh table with an invalid option value. @@ -453,6 +582,20 @@ def test_parallel_merge_matches_serial(self): self.assertEqual(updated.iloc[0].behavior, 'v2-updated') self.assertEqual(updated.iloc[0].item_id, 9001) + def test_parallel_batch_reader_merge_matches_serial(self): + rb = self.table.new_read_builder() + splits = rb.new_scan().plan().splits() + read = rb.new_read() + serial = pa.Table.from_batches( + read.to_arrow_batch_reader(splits, parallelism=1)) + parallel = pa.Table.from_batches( + read.to_arrow_batch_reader(splits, parallelism=4)) + serial = serial.to_pandas().sort_values( + ['dt', 'user_id']).reset_index(drop=True) + parallel = parallel.to_pandas().sort_values( + ['dt', 'user_id']).reset_index(drop=True) + self.assertTrue(serial.equals(parallel)) + def test_parallel_with_limit_pk(self): limit = 12 rb = self.table.new_read_builder().with_limit(limit) From 5088b80c954b8efd07c8197ff1241d7f4a816458 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 23 Jul 2026 00:11:34 +0800 Subject: [PATCH 2/3] [python] Make pipelined split reads cancellable --- docs/docs/pypaimon/python-api.mdx | 4 +- .../read/datasource/ray_datasource.py | 22 ++- paimon-python/pypaimon/read/table_read.py | 178 ++++++++++++++++-- paimon-python/pypaimon/tests/ray_data_test.py | 39 ++++ .../pypaimon/tests/reader_parallel_test.py | 111 +++++++++++ 5 files changed, 325 insertions(+), 29 deletions(-) diff --git a/docs/docs/pypaimon/python-api.mdx b/docs/docs/pypaimon/python-api.mdx index 05fe388efb38..bf0d07805942 100644 --- a/docs/docs/pypaimon/python-api.mdx +++ b/docs/docs/pypaimon/python-api.mdx @@ -473,8 +473,8 @@ for batch in table_read.to_arrow_batch_reader(splits): Multiple splits are read concurrently according to `parallelism` or `read.parallelism` while batches are streamed in split order. Pass -`parallelism=1` to force serial reading. The reader buffers at most two batches -per in-flight split. +`parallelism=1` to force serial reading. Each in-flight split can retain two +queued batches plus one batch waiting to enter the queue. ### Read Python Iterator diff --git a/paimon-python/pypaimon/read/datasource/ray_datasource.py b/paimon-python/pypaimon/read/datasource/ray_datasource.py index 7a2f72614312..c098a055e0c5 100644 --- a/paimon-python/pypaimon/read/datasource/ray_datasource.py +++ b/paimon-python/pypaimon/read/datasource/ray_datasource.py @@ -161,16 +161,20 @@ def _get_read_task( table, predicate, read_type, limit=limit, nested_name_paths=nested_name_paths) - batch_reader = worker_table_read.to_arrow_batch_reader(splits) + batch_reader = worker_table_read.to_arrow_batch_reader( + splits, parallelism=1) has_data = False - for batch in iter(batch_reader.read_next_batch, None): - if batch.num_rows == 0: - continue - has_data = True - table = pyarrow.Table.from_batches([batch]) - if table.schema != schema: - table = table.cast(schema) - yield table + try: + for batch in iter(batch_reader.read_next_batch, None): + if batch.num_rows == 0: + continue + has_data = True + table = pyarrow.Table.from_batches([batch]) + if table.schema != schema: + table = table.cast(schema) + yield table + finally: + batch_reader.close() if not has_data: yield pyarrow.Table.from_arrays( diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index 2a8ca8463a52..cf1193fca1c7 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -18,6 +18,7 @@ import os import queue import threading +import time from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Dict, Iterator, List, Optional @@ -41,6 +42,85 @@ ROW_KIND_COLUMN = "_row_kind" +class _ClosableArrowBatchReader: + """RecordBatchReader-compatible wrapper which closes its Python iterator.""" + + def __init__(self, schema, batch_iterator, cancel=None): + self._batch_iterator = batch_iterator + self._cancel = cancel + self._reader = pyarrow.ipc.RecordBatchReader.from_batches( + schema, batch_iterator) + self._close_lock = threading.Lock() + self._closed = False + + @property + def schema(self): + return self._reader.schema + + def read_next_batch(self): + try: + return self._reader.read_next_batch() + except StopIteration: + self.close() + raise + except BaseException: + try: + self.close() + except BaseException: + pass + raise + + def read_all(self): + try: + return self._reader.read_all() + finally: + self.close() + + def read_pandas(self, **options): + try: + return self._reader.read_pandas(**options) + finally: + self.close() + + def close(self): + with self._close_lock: + if self._closed: + return + self._closed = True + + error = None + if self._cancel is not None: + self._cancel() + close_iterator = getattr(self._batch_iterator, "close", None) + if close_iterator is not None: + try: + close_iterator() + except BaseException as exception: + error = exception + try: + self._reader.close() + except BaseException as exception: + if error is None: + error = exception + if error is not None: + raise error + + def __iter__(self): + return self + + def __next__(self): + return self.read_next_batch() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.close() + + def __getattr__(self, name): + return getattr(self._reader, name) + + class _RemainingRows: """Thread-safe remaining-rows counter for parallel reads. @@ -85,6 +165,7 @@ class TableRead: _MAX_TOTAL_BLOB_WORKERS = 64 _PIPELINE_QUEUE_BATCHES = 2 _PIPELINE_QUEUE_TIMEOUT_SECONDS = 0.1 + _PIPELINE_SHUTDOWN_TIMEOUT_SECONDS = 0.2 def __init__( self, @@ -145,26 +226,32 @@ def to_arrow_batch_reader( splits: List[Split], blob_parallelism: Optional[int] = None, parallelism: Optional[int] = None, - ) -> pyarrow.ipc.RecordBatchReader: + ) -> _ClosableArrowBatchReader: """Stream Arrow batches while reading multiple splits concurrently. - Batches preserve input split order. At most two batches per in-flight - split are buffered, keeping memory bounded while overlapping split I/O. + Batches preserve input split order. Each in-flight split can retain + two queued batches plus one producer-held batch, keeping memory bounded + while overlapping split I/O. """ effective_bp = self._resolve_blob_parallelism(blob_parallelism) effective = self._resolve_parallelism(parallelism, len(splits)) schema = PyarrowFieldParser.from_paimon_schema(self.read_type) if self.include_row_kind: schema = self._add_row_kind_to_schema(schema) + if self.limit is not None and self.limit <= 0: + return _ClosableArrowBatchReader(schema, iter(())) if self._should_run_parallel(splits, effective): workers = min(effective, len(splits)) effective_bp = self._cap_blob_parallelism(workers, effective_bp) + cancel = threading.Event() batch_iterator = self._pipelined_arrow_batch_generator( - splits, schema, effective_bp, workers) + splits, schema, effective_bp, workers, cancel) + return _ClosableArrowBatchReader( + schema, batch_iterator, cancel=cancel.set) else: batch_iterator = self._arrow_batch_generator( splits, schema, effective_bp) - return pyarrow.ipc.RecordBatchReader.from_batches(schema, batch_iterator) + return _ClosableArrowBatchReader(schema, batch_iterator) def _pipelined_arrow_batch_generator( self, @@ -172,10 +259,31 @@ def _pipelined_arrow_batch_generator( schema: pyarrow.Schema, blob_parallelism: int, workers: int, + cancel: Optional[threading.Event] = None, ) -> Iterator[pyarrow.RecordBatch]: - cancel = threading.Event() + if self.limit is not None and self.limit <= 0: + return + + cancel = cancel or threading.Event() end = object() + stop = object() queues = {} + task_queue = queue.Queue() + first_error = [] + error_lock = threading.Lock() + + def publish_error(exception): + with error_lock: + if first_error: + return + first_error.append(exception) + cancel.set() + + def raise_first_error(): + with error_lock: + exception = first_error[0] if first_error else None + if exception is not None: + raise exception def put_item(output_queue, item) -> bool: while not cancel.is_set(): @@ -195,26 +303,44 @@ def read_split(split, output_queue): if not put_item(output_queue, batch): return except BaseException as exception: - put_item(output_queue, exception) + publish_error(exception) finally: try: batch_iterator.close() except BaseException as exception: - put_item(output_queue, exception) - finally: + if not cancel.is_set(): + publish_error(exception) + if not cancel.is_set(): put_item(output_queue, end) - executor = ThreadPoolExecutor( - max_workers=workers, - thread_name_prefix="pypaimon-stream-read", - ) + def worker(): + while not cancel.is_set(): + try: + task = task_queue.get( + timeout=self._PIPELINE_QUEUE_TIMEOUT_SECONDS) + except queue.Empty: + continue + if task is stop: + return + read_split(*task) + + threads = [ + threading.Thread( + target=worker, + name="pypaimon-stream-read-%d" % index, + daemon=True, + ) + for index in range(workers) + ] + for thread in threads: + thread.start() def schedule(index: int): if index >= len(splits): return output_queue = queue.Queue(maxsize=self._PIPELINE_QUEUE_BATCHES) queues[index] = output_queue - executor.submit(read_split, splits[index], output_queue) + task_queue.put((splits[index], output_queue)) remaining = self.limit try: @@ -224,11 +350,17 @@ def schedule(index: int): for index in range(len(splits)): output_queue = queues.pop(index) while True: - item = output_queue.get() + raise_first_error() + if cancel.is_set(): + return + try: + item = output_queue.get( + timeout=self._PIPELINE_QUEUE_TIMEOUT_SECONDS) + except queue.Empty: + continue + raise_first_error() if item is end: break - if isinstance(item, BaseException): - raise item batch = item if remaining is not None and batch.num_rows > remaining: batch = batch.slice(0, remaining) @@ -238,9 +370,19 @@ def schedule(index: int): if remaining <= 0: return schedule(index + workers) + raise_first_error() finally: cancel.set() - executor.shutdown(wait=True) + for _ in threads: + task_queue.put(stop) + deadline = ( + time.monotonic() + self._PIPELINE_SHUTDOWN_TIMEOUT_SECONDS + ) + for thread in threads: + remaining_time = deadline - time.monotonic() + if remaining_time <= 0: + break + thread.join(remaining_time) @staticmethod def _add_row_kind_to_schema(schema: pyarrow.Schema) -> pyarrow.Schema: diff --git a/paimon-python/pypaimon/tests/ray_data_test.py b/paimon-python/pypaimon/tests/ray_data_test.py index 6d00d13ff956..7b0d444acc45 100644 --- a/paimon-python/pypaimon/tests/ray_data_test.py +++ b/paimon-python/pypaimon/tests/ray_data_test.py @@ -19,6 +19,7 @@ import tempfile import unittest import shutil +from unittest import mock import pyarrow as pa import pyarrow.types as pa_types @@ -26,9 +27,47 @@ from pypaimon import CatalogFactory, Schema from pypaimon.common.options.core_options import CoreOptions +from pypaimon.read.datasource.ray_datasource import RayDatasource from pypaimon.schema.data_types import PyarrowFieldParser +class RayDatasourceUnitTest(unittest.TestCase): + + def test_ray_task_disables_inner_split_parallelism(self): + schema = pa.schema([('value', pa.int64())]) + batch = pa.record_batch([pa.array([1])], schema=schema) + split = mock.Mock() + split.file_size = 1 + split.row_count = 1 + split.merged_row_count.return_value = None + split.file_paths = [] + table = mock.Mock() + table.is_primary_key_table = False + provider = mock.Mock() + provider.table.return_value = table + provider.predicate.return_value = None + provider.read_type.return_value = [] + provider.nested_name_paths.return_value = None + provider.splits.return_value = [split] + provider.limit.return_value = None + reader = mock.Mock() + reader.read_next_batch.side_effect = [batch, StopIteration] + + with mock.patch( + 'pypaimon.read.datasource.ray_datasource.' + 'PyarrowFieldParser.from_paimon_schema', + return_value=schema, + ), mock.patch('pypaimon.read.table_read.TableRead') as table_read: + table_read.return_value.to_arrow_batch_reader.return_value = reader + task = RayDatasource(provider).get_read_tasks(1)[0] + result = list(task()) + + self.assertEqual([pa.Table.from_batches([batch])], result) + table_read.return_value.to_arrow_batch_reader.assert_called_once_with( + [split], parallelism=1) + reader.close.assert_called_once() + + class RayDataTest(unittest.TestCase): """Tests for Ray Data integration with PyPaimon.""" diff --git a/paimon-python/pypaimon/tests/reader_parallel_test.py b/paimon-python/pypaimon/tests/reader_parallel_test.py index f8260b9b4b10..001edcfd9dfc 100644 --- a/paimon-python/pypaimon/tests/reader_parallel_test.py +++ b/paimon-python/pypaimon/tests/reader_parallel_test.py @@ -19,6 +19,7 @@ import shutil import tempfile import threading +import time import types import unittest from unittest import mock @@ -125,6 +126,9 @@ class PipelinedBatchGeneratorTest(unittest.TestCase): def setUp(self): self.read = TableRead.__new__(TableRead) self.read.limit = None + self.read._read_parallelism = None + self.read.read_type = [] + self.read.include_row_kind = False self.schema = pa.schema([('value', pa.int64())]) def _batch(self, value): @@ -193,6 +197,113 @@ def generate(splits, schema, blob_parallelism): self.assertCountEqual([0, 1], closed) + def test_public_reader_close_cancels_workers(self): + started = threading.Barrier(2) + closed = [] + closed_lock = threading.Lock() + + def generate(splits, schema, blob_parallelism): + value = splits[0] + try: + started.wait(5) + while True: + yield self._batch(value) + finally: + with closed_lock: + closed.append(value) + + with mock.patch( + 'pypaimon.read.table_read.PyarrowFieldParser.from_paimon_schema', + return_value=self.schema, + ), mock.patch.object( + self.read, '_arrow_batch_generator', side_effect=generate + ): + reader = self.read.to_arrow_batch_reader( + [0, 1], parallelism=2) + reader.read_next_batch() + reader.close() + + deadline = time.monotonic() + 1 + while len(closed) < 2 and time.monotonic() < deadline: + time.sleep(0.01) + self.assertCountEqual([0, 1], closed) + + def test_close_does_not_wait_for_blocked_split(self): + blocked = threading.Event() + release = threading.Event() + + def generate(splits, schema, blob_parallelism): + value = splits[0] + yield self._batch(value) + if value == 0: + blocked.set() + release.wait(5) + + try: + with mock.patch( + 'pypaimon.read.table_read.PyarrowFieldParser.from_paimon_schema', + return_value=self.schema, + ), mock.patch.object( + self.read, '_arrow_batch_generator', side_effect=generate + ): + reader = self.read.to_arrow_batch_reader( + [0, 1], parallelism=2) + reader.read_next_batch() + self.assertTrue(blocked.wait(5)) + started = time.monotonic() + reader.close() + self.assertLess(time.monotonic() - started, 1) + finally: + release.set() + + def test_later_error_bypasses_blocked_earlier_split(self): + first_started = threading.Event() + release = threading.Event() + + def generate(splits, schema, blob_parallelism): + value = splits[0] + if value == 0: + first_started.set() + release.wait(5) + yield self._batch(value) + else: + if not first_started.wait(5): + raise TimeoutError('first split did not start') + raise RuntimeError('later split failed') + + started = time.monotonic() + try: + with mock.patch.object( + self.read, '_arrow_batch_generator', side_effect=generate + ): + with self.assertRaisesRegex( + RuntimeError, 'later split failed' + ): + list(self.read._pipelined_arrow_batch_generator( + [0, 1], self.schema, 1, 2)) + finally: + release.set() + self.assertLess(time.monotonic() - started, 1) + + def test_zero_limit_does_not_start_reading(self): + self.read.limit = 0 + with mock.patch( + 'pypaimon.read.table_read.PyarrowFieldParser.from_paimon_schema', + return_value=self.schema, + ), mock.patch.object( + self.read, '_pipelined_arrow_batch_generator' + ) as pipeline, mock.patch.object( + self.read, '_arrow_batch_generator' + ) as serial: + reader = self.read.to_arrow_batch_reader( + [0, 1], parallelism=2) + with self.assertRaises(StopIteration): + reader.read_next_batch() + reader.close() + + pipeline.assert_not_called() + serial.assert_not_called() + class ParallelReaderAppendOnlyTest(unittest.TestCase): """Append-only multi-partition table — parallel must match serial exactly.""" From e98ba5919b238fc9c429ac8ff1947439ef8224f4 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 23 Jul 2026 11:54:58 +0800 Subject: [PATCH 3/3] [python] Harden pipelined split reader --- docs/docs/pypaimon/python-api.mdx | 7 +- .../read/datasource/ray_datasource.py | 4 +- paimon-python/pypaimon/read/table_read.py | 121 ++++++++++-------- .../pypaimon/tests/reader_parallel_test.py | 119 ++++++++++++++++- 4 files changed, 188 insertions(+), 63 deletions(-) diff --git a/docs/docs/pypaimon/python-api.mdx b/docs/docs/pypaimon/python-api.mdx index bf0d07805942..910a40dd0187 100644 --- a/docs/docs/pypaimon/python-api.mdx +++ b/docs/docs/pypaimon/python-api.mdx @@ -473,8 +473,11 @@ for batch in table_read.to_arrow_batch_reader(splits): Multiple splits are read concurrently according to `parallelism` or `read.parallelism` while batches are streamed in split order. Pass -`parallelism=1` to force serial reading. Each in-flight split can retain two -queued batches plus one batch waiting to enter the queue. +`parallelism=1` to force serial reading. Streaming reads use at most 16 split +workers process-wide and keep a two-window lookahead. Each lookahead split can +retain two queued batches, while each active worker can hold one additional +batch. Reads with a `LIMIT` remain serial so speculative later-split failures +cannot change the result. ### Read Python Iterator diff --git a/paimon-python/pypaimon/read/datasource/ray_datasource.py b/paimon-python/pypaimon/read/datasource/ray_datasource.py index c098a055e0c5..b11b93a756c1 100644 --- a/paimon-python/pypaimon/read/datasource/ray_datasource.py +++ b/paimon-python/pypaimon/read/datasource/ray_datasource.py @@ -174,7 +174,9 @@ def _get_read_task( table = table.cast(schema) yield table finally: - batch_reader.close() + close_reader = getattr(batch_reader, "close", None) + if close_reader is not None: + close_reader() if not has_data: yield pyarrow.Table.from_arrays( diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index cf1193fca1c7..b5738b058de7 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -40,26 +40,26 @@ from pypaimon.table.row.offset_row import OffsetRow ROW_KIND_COLUMN = "_row_kind" +_MAX_PIPELINE_SPLIT_WORKERS = 16 +_PIPELINE_WORKER_SLOTS = threading.BoundedSemaphore( + _MAX_PIPELINE_SPLIT_WORKERS) -class _ClosableArrowBatchReader: - """RecordBatchReader-compatible wrapper which closes its Python iterator.""" +class _ClosableBatchIterator: + """Iterator which closes its source when exhausted, failed, or released.""" - def __init__(self, schema, batch_iterator, cancel=None): + def __init__(self, batch_iterator, cancel=None): self._batch_iterator = batch_iterator self._cancel = cancel - self._reader = pyarrow.ipc.RecordBatchReader.from_batches( - schema, batch_iterator) self._close_lock = threading.Lock() self._closed = False - @property - def schema(self): - return self._reader.schema + def __iter__(self): + return self - def read_next_batch(self): + def __next__(self): try: - return self._reader.read_next_batch() + return next(self._batch_iterator) except StopIteration: self.close() raise @@ -70,18 +70,6 @@ def read_next_batch(self): pass raise - def read_all(self): - try: - return self._reader.read_all() - finally: - self.close() - - def read_pandas(self, **options): - try: - return self._reader.read_pandas(**options) - finally: - self.close() - def close(self): with self._close_lock: if self._closed: @@ -97,28 +85,24 @@ def close(self): close_iterator() except BaseException as exception: error = exception - try: - self._reader.close() - except BaseException as exception: - if error is None: - error = exception if error is not None: raise error - def __iter__(self): - return self - - def __next__(self): - return self.read_next_batch() - - def __enter__(self): - return self + def __del__(self): + try: + self.close() + except BaseException: + pass - def __exit__(self, exc_type, exc_value, traceback): - self.close() - def __getattr__(self, name): - return getattr(self._reader, name) +def _create_arrow_batch_reader(schema, batch_iterator, cancel=None): + managed_iterator = _ClosableBatchIterator(batch_iterator, cancel) + source_reader = pyarrow.ipc.RecordBatchReader.from_batches( + schema, managed_iterator) + from_stream = getattr(pyarrow.ipc.RecordBatchReader, "from_stream", None) + if from_stream is None: + return source_reader + return from_stream(source_reader) class _RemainingRows: @@ -164,6 +148,7 @@ class TableRead: # connections ~= P*B. Shrink per-split B to keep the product bounded. _MAX_TOTAL_BLOB_WORKERS = 64 _PIPELINE_QUEUE_BATCHES = 2 + _PIPELINE_LOOKAHEAD_FACTOR = 2 _PIPELINE_QUEUE_TIMEOUT_SECONDS = 0.1 _PIPELINE_SHUTDOWN_TIMEOUT_SECONDS = 0.2 @@ -226,7 +211,7 @@ def to_arrow_batch_reader( splits: List[Split], blob_parallelism: Optional[int] = None, parallelism: Optional[int] = None, - ) -> _ClosableArrowBatchReader: + ) -> pyarrow.ipc.RecordBatchReader: """Stream Arrow batches while reading multiple splits concurrently. Batches preserve input split order. Each in-flight split can retain @@ -239,19 +224,20 @@ def to_arrow_batch_reader( if self.include_row_kind: schema = self._add_row_kind_to_schema(schema) if self.limit is not None and self.limit <= 0: - return _ClosableArrowBatchReader(schema, iter(())) - if self._should_run_parallel(splits, effective): - workers = min(effective, len(splits)) + return _create_arrow_batch_reader(schema, iter(())) + if self.limit is None and self._should_run_parallel(splits, effective): + workers = min( + effective, len(splits), _MAX_PIPELINE_SPLIT_WORKERS) effective_bp = self._cap_blob_parallelism(workers, effective_bp) cancel = threading.Event() batch_iterator = self._pipelined_arrow_batch_generator( splits, schema, effective_bp, workers, cancel) - return _ClosableArrowBatchReader( + return _create_arrow_batch_reader( schema, batch_iterator, cancel=cancel.set) else: batch_iterator = self._arrow_batch_generator( splits, schema, effective_bp) - return _ClosableArrowBatchReader(schema, batch_iterator) + return _create_arrow_batch_reader(schema, batch_iterator) def _pipelined_arrow_batch_generator( self, @@ -263,8 +249,26 @@ def _pipelined_arrow_batch_generator( ) -> Iterator[pyarrow.RecordBatch]: if self.limit is not None and self.limit <= 0: return + if self.limit is not None: + yield from self._arrow_batch_generator( + splits, schema, blob_parallelism) + return cancel = cancel or threading.Event() + worker_slots = _PIPELINE_WORKER_SLOTS + acquired_workers = 0 + for _ in range(workers): + if not worker_slots.acquire(blocking=False): + break + acquired_workers += 1 + if acquired_workers < 2: + for _ in range(acquired_workers): + worker_slots.release() + yield from self._arrow_batch_generator( + splits, schema, blob_parallelism) + return + workers = acquired_workers + end = object() stop = object() queues = {} @@ -314,15 +318,18 @@ def read_split(split, output_queue): put_item(output_queue, end) def worker(): - while not cancel.is_set(): - try: - task = task_queue.get( - timeout=self._PIPELINE_QUEUE_TIMEOUT_SECONDS) - except queue.Empty: - continue - if task is stop: - return - read_split(*task) + try: + while not cancel.is_set(): + try: + task = task_queue.get( + timeout=self._PIPELINE_QUEUE_TIMEOUT_SECONDS) + except queue.Empty: + continue + if task is stop: + return + read_split(*task) + finally: + worker_slots.release() threads = [ threading.Thread( @@ -343,8 +350,10 @@ def schedule(index: int): task_queue.put((splits[index], output_queue)) remaining = self.limit + lookahead = min( + len(splits), workers * self._PIPELINE_LOOKAHEAD_FACTOR) try: - for index in range(workers): + for index in range(lookahead): schedule(index) for index in range(len(splits)): @@ -369,7 +378,7 @@ def schedule(index: int): remaining -= batch.num_rows if remaining <= 0: return - schedule(index + workers) + schedule(index + lookahead) raise_first_error() finally: cancel.set() diff --git a/paimon-python/pypaimon/tests/reader_parallel_test.py b/paimon-python/pypaimon/tests/reader_parallel_test.py index 001edcfd9dfc..5660f2918321 100644 --- a/paimon-python/pypaimon/tests/reader_parallel_test.py +++ b/paimon-python/pypaimon/tests/reader_parallel_test.py @@ -25,9 +25,11 @@ from unittest import mock import pyarrow as pa +import pyarrow.dataset as ds from pypaimon import CatalogFactory, Schema -from pypaimon.read.table_read import TableRead, _RemainingRows +from pypaimon.read.table_read import ( + TableRead, _MAX_PIPELINE_SPLIT_WORKERS, _RemainingRows) class ResolveParallelismTest(unittest.TestCase): @@ -157,6 +159,28 @@ def generate(splits, schema, blob_parallelism): [batch.column(0)[0].as_py() for batch in batches], ) + def test_lookahead_keeps_workers_busy_behind_slow_first_split(self): + third_started = threading.Event() + + def generate(splits, schema, blob_parallelism): + value = splits[0] + if value == 0 and not third_started.wait(5): + raise TimeoutError('lookahead split did not start') + if value == 2: + third_started.set() + yield self._batch(value) + + with mock.patch.object( + self.read, '_arrow_batch_generator', side_effect=generate + ): + batches = list(self.read._pipelined_arrow_batch_generator( + [0, 1, 2, 3], self.schema, 1, 2)) + + self.assertEqual( + [0, 1, 2, 3], + [batch.column(0)[0].as_py() for batch in batches], + ) + def test_propagates_worker_error(self): def generate(splits, schema, blob_parallelism): if splits[0] == 0: @@ -220,6 +244,7 @@ def generate(splits, schema, blob_parallelism): ): reader = self.read.to_arrow_batch_reader( [0, 1], parallelism=2) + self.assertIsInstance(reader, pa.ipc.RecordBatchReader) reader.read_next_batch() reader.close() @@ -228,6 +253,22 @@ def generate(splits, schema, blob_parallelism): time.sleep(0.01) self.assertCountEqual([0, 1], closed) + def test_public_reader_works_with_scanner_from_batches(self): + def generate(splits, schema, blob_parallelism): + yield self._batch(splits[0]) + + with mock.patch( + 'pypaimon.read.table_read.PyarrowFieldParser.from_paimon_schema', + return_value=self.schema, + ), mock.patch.object( + self.read, '_arrow_batch_generator', side_effect=generate + ): + reader = self.read.to_arrow_batch_reader( + [0, 1], parallelism=2) + result = ds.Scanner.from_batches(reader).to_table() + + self.assertEqual([0, 1], result.column(0).to_pylist()) + def test_close_does_not_wait_for_blocked_split(self): blocked = threading.Event() release = threading.Event() @@ -256,6 +297,70 @@ def generate(splits, schema, blob_parallelism): finally: release.set() + def test_blocked_readers_do_not_exceed_global_worker_limit(self): + release = threading.Event() + + def generate(splits, schema, blob_parallelism): + yield self._batch(splits[0]) + release.wait(5) + + worker_slots = threading.BoundedSemaphore(2) + readers = [] + try: + with mock.patch( + 'pypaimon.read.table_read._PIPELINE_WORKER_SLOTS', + worker_slots, + ), mock.patch( + 'pypaimon.read.table_read.PyarrowFieldParser.from_paimon_schema', + return_value=self.schema, + ), mock.patch.object( + self.read, '_arrow_batch_generator', side_effect=generate + ): + for _ in range(3): + reader = self.read.to_arrow_batch_reader( + [0, 1], parallelism=2) + reader.read_next_batch() + reader.close() + readers.append(reader) + + active = [ + thread for thread in threading.enumerate() + if thread.name.startswith('pypaimon-stream-read-') + ] + self.assertLessEqual(len(active), 2) + finally: + release.set() + deadline = time.monotonic() + 1 + while (any(thread.name.startswith('pypaimon-stream-read-') + for thread in threading.enumerate()) + and time.monotonic() < deadline): + time.sleep(0.01) + + def test_public_reader_caps_pipeline_workers(self): + splits = list(range(_MAX_PIPELINE_SPLIT_WORKERS + 10)) + + def generate(splits, schema, blob_parallelism): + yield self._batch(splits[0]) + + with mock.patch( + 'pypaimon.read.table_read.PyarrowFieldParser.from_paimon_schema', + return_value=self.schema, + ), mock.patch.object( + self.read, + '_pipelined_arrow_batch_generator', + wraps=self.read._pipelined_arrow_batch_generator, + ) as pipeline, mock.patch.object( + self.read, '_arrow_batch_generator', side_effect=generate + ): + reader = self.read.to_arrow_batch_reader( + splits, parallelism=len(splits)) + reader.read_all() + + self.assertEqual( + _MAX_PIPELINE_SPLIT_WORKERS, + pipeline.call_args.args[3], + ) + def test_later_error_bypasses_blocked_earlier_split(self): first_started = threading.Event() release = threading.Event() @@ -299,7 +404,9 @@ def test_zero_limit_does_not_start_reading(self): [0, 1], parallelism=2) with self.assertRaises(StopIteration): reader.read_next_batch() - reader.close() + close_reader = getattr(reader, 'close', None) + if close_reader is not None: + close_reader() pipeline.assert_not_called() serial.assert_not_called() @@ -456,8 +563,12 @@ def test_parallel_batch_reader_limit_matches_serial(self): read = rb.new_read() serial = pa.Table.from_batches( read.to_arrow_batch_reader(splits, parallelism=1)) - parallel = pa.Table.from_batches( - read.to_arrow_batch_reader(splits, parallelism=4)) + with mock.patch.object( + read, '_pipelined_arrow_batch_generator' + ) as pipeline: + parallel = pa.Table.from_batches( + read.to_arrow_batch_reader(splits, parallelism=4)) + pipeline.assert_not_called() self.assertEqual(serial, parallel) self.assertEqual(600, parallel.num_rows)