diff --git a/docs/docs/pypaimon/python-api.mdx b/docs/docs/pypaimon/python-api.mdx index c1887c537c75..910a40dd0187 100644 --- a/docs/docs/pypaimon/python-api.mdx +++ b/docs/docs/pypaimon/python-api.mdx @@ -471,6 +471,14 @@ 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. 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 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/datasource/ray_datasource.py b/paimon-python/pypaimon/read/datasource/ray_datasource.py index 7a2f72614312..b11b93a756c1 100644 --- a/paimon-python/pypaimon/read/datasource/ray_datasource.py +++ b/paimon-python/pypaimon/read/datasource/ray_datasource.py @@ -161,16 +161,22 @@ 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: + 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 6e641b50963a..b5738b058de7 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -16,7 +16,9 @@ # under the License. import os +import queue import threading +import time from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Dict, Iterator, List, Optional @@ -38,6 +40,69 @@ 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 _ClosableBatchIterator: + """Iterator which closes its source when exhausted, failed, or released.""" + + def __init__(self, batch_iterator, cancel=None): + self._batch_iterator = batch_iterator + self._cancel = cancel + self._close_lock = threading.Lock() + self._closed = False + + def __iter__(self): + return self + + def __next__(self): + try: + return next(self._batch_iterator) + except StopIteration: + self.close() + raise + except BaseException: + try: + self.close() + except BaseException: + pass + raise + + 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 + if error is not None: + raise error + + def __del__(self): + try: + self.close() + except BaseException: + pass + + +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: @@ -82,6 +147,10 @@ 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_LOOKAHEAD_FACTOR = 2 + _PIPELINE_QUEUE_TIMEOUT_SECONDS = 0.1 + _PIPELINE_SHUTDOWN_TIMEOUT_SECONDS = 0.2 def __init__( self, @@ -137,14 +206,192 @@ 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. 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) - batch_iterator = self._arrow_batch_generator(splits, schema, effective_bp) - return pyarrow.ipc.RecordBatchReader.from_batches(schema, batch_iterator) + if self.limit is not None and self.limit <= 0: + 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 _create_arrow_batch_reader( + schema, batch_iterator, cancel=cancel.set) + else: + batch_iterator = self._arrow_batch_generator( + splits, schema, effective_bp) + return _create_arrow_batch_reader(schema, batch_iterator) + + def _pipelined_arrow_batch_generator( + self, + splits: List[Split], + schema: pyarrow.Schema, + blob_parallelism: int, + workers: int, + cancel: Optional[threading.Event] = None, + ) -> 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 = {} + 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(): + 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: + publish_error(exception) + finally: + try: + batch_iterator.close() + except BaseException as exception: + if not cancel.is_set(): + publish_error(exception) + if not cancel.is_set(): + put_item(output_queue, end) + + def worker(): + 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( + 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 + task_queue.put((splits[index], output_queue)) + + remaining = self.limit + lookahead = min( + len(splits), workers * self._PIPELINE_LOOKAHEAD_FACTOR) + try: + for index in range(lookahead): + schedule(index) + + for index in range(len(splits)): + output_queue = queues.pop(index) + while True: + 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 + 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 + lookahead) + raise_first_error() + finally: + cancel.set() + 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: @@ -207,7 +454,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/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 ea0c3d8f61d1..5660f2918321 100644 --- a/paimon-python/pypaimon/tests/reader_parallel_test.py +++ b/paimon-python/pypaimon/tests/reader_parallel_test.py @@ -19,14 +19,17 @@ import shutil import tempfile import threading +import time import types import unittest 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): @@ -120,6 +123,295 @@ def worker(): self.assertTrue(rr.exhausted()) +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): + 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_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: + 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) + + 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) + self.assertIsInstance(reader, pa.ipc.RecordBatchReader) + 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_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() + + 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_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() + + 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() + close_reader = getattr(reader, 'close', None) + if close_reader is not None: + close_reader() + + pipeline.assert_not_called() + serial.assert_not_called() + + class ParallelReaderAppendOnlyTest(unittest.TestCase): """Append-only multi-partition table — parallel must match serial exactly.""" @@ -225,6 +517,61 @@ 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)) + 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) + 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 +604,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 +654,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 +804,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)