This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 54d8596ce7 [python] Improve parallel native read scheduling (#10003)
54d8596ce7 is described below
commit 54d8596ce7f1539f3e3c0125debc622b0e7a52b6
Author: Jingsong Lee <[email protected]>
AuthorDate: Sun Sep 20 16:44:57 2026 +0800
[python] Improve parallel native read scheduling (#10003)
---
paimon-python/README.md | 19 +--
paimon-python/pypaimon/read/table_read.py | 108 ++++++++++++++--
paimon-python/pypaimon/tests/native_read_test.py | 150 ++++++++++++++++++++++-
3 files changed, 254 insertions(+), 23 deletions(-)
diff --git a/paimon-python/README.md b/paimon-python/README.md
index c8ee07f1d3..6aaffefb64 100644
--- a/paimon-python/README.md
+++ b/paimon-python/README.md
@@ -136,14 +136,17 @@ rows = builder.new_read().to_arrow(plan.splits())
Native reads return PyArrow batches through the Arrow C Data interface. They
currently require untouched splits produced by the native planner and top-level
-projection. Query authorization, nested projection, row-kind output, and
-explicit blob parallelism retain the Python reader. For materialized
-`to_arrow()` reads, the effective split parallelism (the method argument,
-`read.parallelism`, or the automatic default) runs independent Rust readers
-over contiguous split groups. Streaming `to_arrow_batch_reader()` keeps one
-lazy Rust reader. A missing reader capability, unsupported route, or
-native-reader construction failure falls back to Python; I/O and data errors
-raised after streaming starts surface to the caller.
+projection. Query authorization, nested projection, and row-kind output retain
+the Python reader. For both materialized
+`to_arrow()` and streaming `to_arrow_batch_reader()` reads, the effective split
+parallelism (the method argument, `read.parallelism`, or the automatic default)
+runs independent Rust readers. Splits stay in input order but contiguous groups
+are balanced by physical file bytes to reduce worker skew. An unfiltered row
+limit also caps reader fan-out to avoid speculative work.
+Streaming buffers at most one batch per reader and closing the batch reader
+interrupts native reads that are still in flight. A missing reader capability,
+unsupported route, or native-reader construction failure falls back to Python;
+I/O and data errors raised after streaming starts surface to the caller.
With Rust main's `Table.from_resolved_schema()` binding, filesystem and JDBC
catalog
tables preserve the Python table's resolved schema and complete effective
diff --git a/paimon-python/pypaimon/read/table_read.py
b/paimon-python/pypaimon/read/table_read.py
index 305826e976..e697b51d2f 100644
--- a/paimon-python/pypaimon/read/table_read.py
+++ b/paimon-python/pypaimon/read/table_read.py
@@ -219,6 +219,9 @@ class TableRead:
table reads do not guarantee row order. Python fallback reads remain
serial.
"""
+ # Cleanup ownership follows the uncapped concurrency decision. LIMIT
+ # may reduce the actual native worker count to one, but the returned
+ # PyArrow reader still does not close its suspended batch iterator.
effective = self._resolve_parallelism(parallelism, len(splits))
reader, batch_iterator = self._new_arrow_batch_reader(
splits, blob_parallelism, parallelism)
@@ -237,6 +240,8 @@ class TableRead:
``from_batches``. Retain it explicitly for a parallel read so native
split workers are stopped when a caller ends the read early.
"""
+ # Keep the explicit iterator cleanup chain even when LIMIT caps the
+ # native worker count inside _new_arrow_batch_reader to one.
effective = self._resolve_parallelism(parallelism, len(splits))
reader, batch_iterator = self._new_arrow_batch_reader(
splits, blob_parallelism, parallelism)
@@ -255,7 +260,7 @@ class TableRead:
schema = PyarrowFieldParser.from_paimon_schema(self.read_type)
if self.include_row_kind:
schema = self._add_row_kind_to_schema(schema)
- effective = self._resolve_parallelism(parallelism, len(splits))
+ effective = self._effective_parallelism(parallelism, len(splits))
effective_bp = self._resolve_blob_parallelism(blob_parallelism)
if self._should_run_parallel(splits, effective):
effective_bp = self._cap_blob_parallelism(
@@ -330,8 +335,10 @@ class TableRead:
back to the table option; when that is also unset the read
auto-scales to ``min(number of splits, CPU count)``. ``1``
keeps reads serial; ``>= 2`` caps the thread pool that reads
- splits concurrently and assembles the final table in input
- order. Must be ``>= 1``. Note that with ``>= 2`` (or auto)
+ byte-balanced contiguous split groups concurrently and
+ assembles the final table in input order. An unfiltered row
+ limit caps this further to avoid speculative work. Must be
+ ``>= 1``. Note that with ``>= 2`` (or auto)
and a ``limit`` set, the returned rows are an arbitrary
subset of the requested size, since which splits fill the row
quota first is non-deterministic. Data-evolution reads with
@@ -346,7 +353,7 @@ class TableRead:
shrunk to stay within it.
"""
effective_bp = self._resolve_blob_parallelism(blob_parallelism)
- effective = self._resolve_parallelism(parallelism, len(splits))
+ effective = self._effective_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)
@@ -423,6 +430,7 @@ class TableRead:
"Native read failed, falling back to the Python reader: %s", e)
return None
rust_splits = []
+ split_weights = []
for split in splits:
if isinstance(split, QueryAuthSplit):
return None
@@ -432,11 +440,12 @@ class TableRead:
if rust_split is None:
return None
rust_splits.append(rust_split)
+ split_weights.append(self._native_split_weight(split))
if (parallelism is not None
and self._should_run_parallel(splits, parallelism)):
if streaming:
groups = self._native_split_groups(
- rust_splits, parallelism)
+ rust_splits, parallelism, split_weights)
read_kwargs = {
'predicate': self.predicate,
'limit': self.limit,
@@ -469,7 +478,7 @@ class TableRead:
try:
return self._native_batches_parallel(
native_read, rust_splits, schema, parallelism,
- blob_parallelism)
+ blob_parallelism, split_weights)
except _NativeReadSetupError as e:
logger.warning(
"Native read failed, falling back to the Python reader:
%s", e)
@@ -490,22 +499,71 @@ class TableRead:
return self._convert_native_batches(batches, schema)
@staticmethod
- def _native_split_groups(rust_splits, effective):
+ def _native_split_weight(split):
+ """Estimate native split work from its physical data-file bytes."""
+ try:
+ sizes = [data_file.file_size for data_file in split.files]
+ except (AttributeError, TypeError):
+ return 1
+ if any(not isinstance(size, int) for size in sizes):
+ return 1
+ return max(1, sum(max(0, size) for size in sizes))
+
+ @staticmethod
+ def _native_split_groups(rust_splits, effective, weights=None):
workers = min(effective, len(rust_splits))
- base_size, larger_groups = divmod(len(rust_splits), workers)
+ if weights is None:
+ weights = [1] * len(rust_splits)
+ if len(weights) != len(rust_splits):
+ raise ValueError("split weights must match splits")
+ weights = [max(1, int(weight)) for weight in weights]
+
+ # Find the smallest possible maximum group weight for a contiguous
+ # partition. Keeping groups contiguous preserves split order when the
+ # non-streaming path flattens worker results.
+ low, high = max(weights), sum(weights)
+ while low < high:
+ capacity = (low + high) // 2
+ groups_needed = 1
+ current = 0
+ for weight in weights:
+ if current + weight > capacity:
+ groups_needed += 1
+ current = weight
+ else:
+ current += weight
+ if groups_needed <= workers:
+ high = capacity
+ else:
+ low = capacity + 1
+
+ capacity = low
groups = []
offset = 0
for index in range(workers):
- size = base_size + (1 if index < larger_groups else 0)
- groups.append(rust_splits[offset:offset + size])
- offset += size
+ remaining_groups = workers - index
+ if remaining_groups == 1:
+ groups.append(rust_splits[offset:])
+ break
+ end = offset
+ group_weight = 0
+ max_end = len(rust_splits) - (remaining_groups - 1)
+ while end < max_end:
+ weight = weights[end]
+ if end > offset and group_weight + weight > capacity:
+ break
+ group_weight += weight
+ end += 1
+ groups.append(rust_splits[offset:end])
+ offset = end
return groups
def _native_batches_parallel(
self, native_read, rust_splits, schema, effective,
- blob_parallelism):
+ blob_parallelism, split_weights=None):
"""Read contiguous split groups with independent Rust readers."""
- groups = self._native_split_groups(rust_splits, effective)
+ groups = self._native_split_groups(
+ rust_splits, effective, split_weights)
workers = len(groups)
remaining_state = _RemainingRows(self.limit)
@@ -598,6 +656,22 @@ class TableRead:
raise value
finally:
stop.set()
+ # A worker may be blocked inside its first native next(), before
+ # it has produced a batch or observed stop. Close every reader to
+ # interrupt that in-flight Rust operation before waiting for the
+ # executor. Native close is idempotent because workers also close
+ # their own readers in read_group's finally block.
+ for reader in readers:
+ close = getattr(reader, 'close', None)
+ if close is not None:
+ try:
+ close()
+ except Exception:
+ logger.debug(
+ "Failed to close a native reader while stopping "
+ "parallel read",
+ exc_info=True,
+ )
for future in futures:
future.cancel()
executor.shutdown(wait=True)
@@ -749,6 +823,14 @@ class TableRead:
raise ValueError(f"{source} must be >= 1, got {value}")
return value
+ def _effective_parallelism(
+ self, runtime: Optional[int], num_splits: int) -> int:
+ """Resolve parallelism and bound speculative unfiltered LIMIT reads."""
+ effective = self._resolve_parallelism(runtime, num_splits)
+ if self.limit is None or (self.limit > 0 and self.predicate is not
None):
+ return effective
+ return min(effective, max(1, self.limit))
+
@staticmethod
def _resolve_blob_parallelism(runtime: Optional[int]) -> int:
if runtime is None:
diff --git a/paimon-python/pypaimon/tests/native_read_test.py
b/paimon-python/pypaimon/tests/native_read_test.py
index e1ff4006a0..cae8c4deb5 100644
--- a/paimon-python/pypaimon/tests/native_read_test.py
+++ b/paimon-python/pypaimon/tests/native_read_test.py
@@ -26,8 +26,8 @@ from pypaimon.schema.data_types import AtomicType, DataField
class _Split:
- def __init__(self, file_name='data.parquet'):
- self.files = [Mock(file_name=file_name)]
+ def __init__(self, file_name='data.parquet', file_size=1):
+ self.files = [Mock(file_name=file_name, file_size=file_size)]
def _table_read(limit=None):
@@ -290,6 +290,152 @@ def
test_parallel_native_stream_close_does_not_start_another_read():
assert not reader.second_read_started.is_set()
+def test_parallel_native_stream_close_interrupts_other_in_flight_reader():
+ read = _table_read()
+ both_reading = threading.Barrier(2)
+
+ class FirstReader:
+ def __init__(self):
+ self.closed = threading.Event()
+
+ def __iter__(self):
+ return self
+
+ def __next__(self):
+ both_reading.wait(timeout=5)
+ return _id_batch([1])
+
+ def close(self):
+ self.closed.set()
+
+ class BlockingReader:
+ def __init__(self):
+ self.closed = threading.Event()
+
+ def __iter__(self):
+ return self
+
+ def __next__(self):
+ both_reading.wait(timeout=5)
+ self.closed.wait(timeout=5)
+ raise StopIteration
+
+ def close(self):
+ self.closed.set()
+
+ first = FirstReader()
+ blocked = BlockingReader()
+ batches = read._native_batches_parallel_streaming([first, blocked])
+ assert next(batches).column('id').to_pylist() == [1]
+
+ close_finished = threading.Event()
+ close_thread = threading.Thread(
+ target=lambda: (batches.close(), close_finished.set()), daemon=True)
+ close_thread.start()
+ closed_in_flight_reader = close_finished.wait(timeout=1)
+ try:
+ assert closed_in_flight_reader
+ finally:
+ blocked.close()
+ close_thread.join(timeout=5)
+
+ assert first.closed.is_set()
+ assert blocked.closed.is_set()
+
+
+def test_native_read_limit_caps_split_reader_fanout():
+ read = _table_read(limit=1)
+ splits = [_Split() for _ in range(16)]
+ for index, split in enumerate(splits):
+ split._native_split = index
+
+ with patch(
+ 'pypaimon.read.native_plan.native_read',
+ side_effect=lambda table, group, **kwargs: [
+ _id_batch(group)]) as native:
+ result = read.to_arrow_batch_reader(
+ splits, parallelism=16).read_all()
+
+ assert result.num_rows == 1
+ native.assert_called_once()
+ assert native.call_args.args[1] == list(range(16))
+
+
+def test_native_read_limit_close_reaches_capped_native_reader():
+ read = _table_read(limit=1)
+ splits = [_Split(), _Split()]
+ for index, split in enumerate(splits):
+ split._native_split = index
+
+ class CloseTrackingReader:
+ def __init__(self):
+ self._batch = _id_batch([0, 1])
+ self.closed = False
+
+ def __iter__(self):
+ return self
+
+ def __next__(self):
+ if self._batch is None:
+ raise StopIteration
+ batch, self._batch = self._batch, None
+ return batch
+
+ def close(self):
+ self.closed = True
+
+ native_reader = CloseTrackingReader()
+ with patch(
+ 'pypaimon.read.native_plan.native_read',
+ return_value=native_reader) as native:
+ batch_reader = read.to_arrow_batch_reader(
+ splits, parallelism=2)
+ assert batch_reader.read_next_batch().column('id').to_pylist() == [0]
+ assert not native_reader.closed
+ batch_reader.close()
+
+ native.assert_called_once()
+ assert native_reader.closed
+
+
+def test_filtered_native_read_limit_keeps_split_parallelism():
+ read = _table_read(limit=1)
+ read.predicate = Mock()
+
+ assert read._effective_parallelism(16, 16) == 16
+
+
+def test_native_split_groups_balance_contiguous_file_bytes():
+ groups = TableRead._native_split_groups(
+ list(range(4)), 2, weights=[8, 8, 1, 1])
+
+ assert groups == [[0], [1, 2, 3]]
+ assert [sum([8, 8, 1, 1][value] for value in group)
+ for group in groups] == [8, 10]
+
+
+def test_native_split_groups_keep_equal_splits_evenly_distributed():
+ assert TableRead._native_split_groups(
+ list(range(5)), 2, weights=[1] * 5) == [[0, 1, 2], [3, 4]]
+
+
+def test_native_read_groups_splits_by_file_bytes():
+ read = _table_read()
+ splits = [_Split(file_size=size) for size in [8, 8, 1, 1]]
+ for index, split in enumerate(splits):
+ split._native_split = index
+
+ with patch(
+ 'pypaimon.read.native_plan.native_read',
+ side_effect=lambda table, group, **kwargs: [
+ _id_batch(group)]) as native:
+ result = read.to_arrow(splits, parallelism=2)
+
+ assert result.num_rows == 4
+ assert sorted(call.args[1] for call in native.call_args_list) == [
+ [0], [1, 2, 3]]
+
+
def test_native_read_runtime_parallelism_overrides_table_option():
read = _table_read()
read._read_parallelism = 4
