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 9301269c15 [python] Avoid repeated catalog setup in parallel native
reads (#10118)
9301269c15 is described below
commit 9301269c15b59600881c78fd643d8209dbc97f59
Author: XiaoHongbo <[email protected]>
AuthorDate: Tue Sep 22 22:27:41 2026 +0800
[python] Avoid repeated catalog setup in parallel native reads (#10118)
---
paimon-python/pypaimon/read/native_plan.py | 30 +++--
paimon-python/pypaimon/read/table_read.py | 32 ++---
.../pypaimon/tests/native_plan_integration_test.py | 24 ++--
paimon-python/pypaimon/tests/native_read_test.py | 144 ++++++++++++++-------
4 files changed, 147 insertions(+), 83 deletions(-)
diff --git a/paimon-python/pypaimon/read/native_plan.py
b/paimon-python/pypaimon/read/native_plan.py
index e457bccb28..6eac710aaa 100644
--- a/paimon-python/pypaimon/read/native_plan.py
+++ b/paimon-python/pypaimon/read/native_plan.py
@@ -340,13 +340,13 @@ def _configure_native_read_builder(builder, predicate,
limit, projection,
return builder
-def native_read(table, splits, predicate: Optional[Predicate] = None,
- limit: Optional[int] = None,
- projection: Optional[List[str]] = None,
- blob_parallelism: Optional[int] = None,
- nested_projection: Optional[List[List[str]]] = None,
- include_row_kind: bool = False):
- """Read Rust ``Split`` objects into PyArrow ``RecordBatch`` objects."""
+def _prepare_native_read(table, predicate: Optional[Predicate] = None,
+ limit: Optional[int] = None,
+ projection: Optional[List[str]] = None,
+ blob_parallelism: Optional[int] = None,
+ nested_projection: Optional[List[List[str]]] = None,
+ include_row_kind: bool = False):
+ """Create one Rust reader reusable across split groups."""
if not native_reader_available():
raise RuntimeError(
"read.native.enabled needs the pypaimon-rust native reader API")
@@ -358,8 +358,20 @@ def native_read(table, splits, predicate:
Optional[Predicate] = None,
builder = builder.with_blob_parallelism(blob_parallelism)
reader = builder.new_read()
read_arrow = getattr(reader, 'read_arrow', None)
- return (read_arrow(splits) if callable(read_arrow)
- else reader.read(splits))
+ return read_arrow if callable(read_arrow) else reader.read
+
+
+def native_read(table, splits, predicate: Optional[Predicate] = None,
+ limit: Optional[int] = None,
+ projection: Optional[List[str]] = None,
+ blob_parallelism: Optional[int] = None,
+ nested_projection: Optional[List[List[str]]] = None,
+ include_row_kind: bool = False):
+ """Read Rust ``Split`` objects into PyArrow ``RecordBatch`` objects."""
+ read_splits = _prepare_native_read(
+ table, predicate, limit, projection, blob_parallelism,
+ nested_projection, include_row_kind)
+ return read_splits(splits)
def native_plan(
diff --git a/paimon-python/pypaimon/read/table_read.py
b/paimon-python/pypaimon/read/table_read.py
index fb2926a467..eb5e64752d 100644
--- a/paimon-python/pypaimon/read/table_read.py
+++ b/paimon-python/pypaimon/read/table_read.py
@@ -441,7 +441,7 @@ class TableRead:
return None
try:
from pypaimon.read.native_plan import (
- native_read, native_split_from_python)
+ _prepare_native_read, native_read, native_split_from_python)
except Exception as e:
logger.warning(
"Native read failed, falling back to the Python reader: %s", e)
@@ -464,15 +464,20 @@ class TableRead:
split_weights.append(self._native_split_weight(split))
if (parallelism is not None
and self._should_run_parallel(splits, parallelism)):
+ read_kwargs = self._native_read_kwargs(blob_parallelism)
+ try:
+ read_splits = _prepare_native_read(self.table, **read_kwargs)
+ except Exception as e:
+ logger.warning(
+ "Native read failed, falling back to the Python reader:
%s", e)
+ return None
if streaming:
groups = self._native_split_groups(
rust_splits, parallelism, split_weights)
- read_kwargs = self._native_read_kwargs(blob_parallelism)
readers = []
try:
for group in groups:
- readers.append(
- native_read(self.table, group, **read_kwargs))
+ readers.append(read_splits(group))
except Exception as e:
for reader in readers:
close = getattr(reader, 'close', None)
@@ -492,8 +497,8 @@ class TableRead:
return self._convert_native_batches(batches, schema)
try:
return self._native_batches_parallel(
- native_read, rust_splits, schema, parallelism,
- blob_parallelism, split_weights)
+ read_splits, rust_splits, schema, parallelism,
+ split_weights)
except _NativeReadSetupError as e:
logger.warning(
"Native read failed, falling back to the Python reader:
%s", e)
@@ -582,9 +587,9 @@ class TableRead:
return groups
def _native_batches_parallel(
- self, native_read, rust_splits, schema, effective,
- blob_parallelism, split_weights=None):
- """Read contiguous split groups with independent Rust readers."""
+ self, read_splits, rust_splits, schema, effective,
+ split_weights=None):
+ """Read contiguous split groups with independent Rust streams."""
groups = self._native_split_groups(
rust_splits, effective, split_weights)
workers = len(groups)
@@ -597,11 +602,10 @@ class TableRead:
futures = {
executor.submit(
self._read_native_split_group,
- native_read,
+ read_splits,
group,
schema,
remaining_state,
- blob_parallelism,
): index
for index, group in enumerate(groups)
}
@@ -700,13 +704,11 @@ class TableRead:
executor.shutdown(wait=True)
def _read_native_split_group(
- self, native_read, rust_splits, schema, remaining_state,
- blob_parallelism):
+ self, read_splits, rust_splits, schema, remaining_state):
if remaining_state.exhausted():
return []
try:
- read_kwargs = self._native_read_kwargs(blob_parallelism)
- batches = native_read(self.table, rust_splits, **read_kwargs)
+ batches = read_splits(rust_splits)
except Exception as e:
raise _NativeReadSetupError(str(e)) from e
result = []
diff --git a/paimon-python/pypaimon/tests/native_plan_integration_test.py
b/paimon-python/pypaimon/tests/native_plan_integration_test.py
index 5fb207ea4d..646ea5eb68 100644
--- a/paimon-python/pypaimon/tests/native_plan_integration_test.py
+++ b/paimon-python/pypaimon/tests/native_plan_integration_test.py
@@ -32,9 +32,9 @@ from pypaimon import CatalogFactory, Schema
from pypaimon.catalog.table_query_auth import TableQueryAuthResult
from pypaimon.globalindex.global_index_result import GlobalIndexResult
from pypaimon.read.native_plan import (
- native_family_search_modes_available, native_method_available, native_read,
- native_reader_available, native_split_bridge_available,
- native_split_from_python,
+ _prepare_native_read, native_family_search_modes_available,
+ native_method_available, native_read, native_reader_available,
+ native_split_bridge_available, native_split_from_python,
)
from pypaimon.schema.schema_change import SchemaChange
from pypaimon.table.row.blob import BlobDescriptor, BlobViewStruct
@@ -373,8 +373,8 @@ class NativePlanIntegrationTest(unittest.TestCase):
plan = builder.new_scan().plan()
self.assertEqual(len(plan.splits()), 4)
- with patch('pypaimon.read.native_plan.native_read',
- wraps=native_read) as rust_reads, \
+ with patch('pypaimon.read.native_plan._prepare_native_read',
+ wraps=_prepare_native_read) as prepare, \
patch(
'pypaimon.read.table_read.TableRead._create_split_read',
side_effect=AssertionError('Python reader was used')):
@@ -386,10 +386,10 @@ class NativePlanIntegrationTest(unittest.TestCase):
{'k': 3, 'v': 'c', 'dt': 'p3'},
{'k': 4, 'v': 'd', 'dt': 'p4'},
])
- self.assertEqual(rust_reads.call_count, 2)
+ prepare.assert_called_once()
- with patch('pypaimon.read.native_plan.native_read',
- wraps=native_read) as rust_reads, \
+ with patch('pypaimon.read.native_plan._prepare_native_read',
+ wraps=_prepare_native_read) as prepare, \
patch(
'pypaimon.read.table_read.TableRead._create_split_read',
side_effect=AssertionError('Python reader was used')):
@@ -402,7 +402,7 @@ class NativePlanIntegrationTest(unittest.TestCase):
{'k': 3, 'v': 'c', 'dt': 'p3'},
{'k': 4, 'v': 'd', 'dt': 'p4'},
])
- self.assertEqual(rust_reads.call_count, 2)
+ prepare.assert_called_once()
@unittest.skipUnless(native_reader_available(),
"pypaimon-rust native reader API not installed")
@@ -1527,10 +1527,10 @@ class NativePlanIntegrationTest(unittest.TestCase):
})
builder = native_table.new_read_builder()
plan = builder.new_scan().plan()
- with patch('pypaimon.read.native_plan.native_read',
- wraps=native_read) as read:
+ with patch('pypaimon.read.native_plan._prepare_native_read',
+ wraps=_prepare_native_read) as prepare:
rows = builder.new_read().to_arrow(plan.splits()).to_pylist()
- self.assertEqual(read.call_count, len(plan.splits()))
+ prepare.assert_called_once()
self.assertEqual(sorted(rows, key=lambda row: row['k']), [
{'k': 1, 'p': 'a/b'},
{'k': 2, 'p': 'a/b'},
diff --git a/paimon-python/pypaimon/tests/native_read_test.py
b/paimon-python/pypaimon/tests/native_read_test.py
index 63e49ed50b..8aa85566f4 100644
--- a/paimon-python/pypaimon/tests/native_read_test.py
+++ b/paimon-python/pypaimon/tests/native_read_test.py
@@ -223,7 +223,11 @@ def
test_native_read_uses_effective_parallelism_from_table_option():
worker_names = set()
worker_names_lock = threading.Lock()
- def read_group(table, rust_splits, **kwargs):
+ groups = []
+
+ def read_group(rust_splits):
+ groups.append(rust_splits)
+
def batches():
with worker_names_lock:
worker_names.add(threading.current_thread().name)
@@ -231,17 +235,41 @@ def
test_native_read_uses_effective_parallelism_from_table_option():
yield _id_batch(rust_splits)
return batches()
- with patch('pypaimon.read.native_plan.native_read',
- side_effect=read_group) as native:
+ with patch('pypaimon.read.native_plan._prepare_native_read',
+ return_value=read_group) as prepare:
result = read.to_arrow(splits)
assert result.to_pydict() == {'id': [0, 1, 2, 3]}
- assert native.call_count == 2
- groups = sorted(call.args[1] for call in native.call_args_list)
- assert groups == [[0, 1], [2, 3]]
+ prepare.assert_called_once()
+ assert sorted(groups) == [[0, 1], [2, 3]]
assert len(worker_names) == 2
+def test_parallel_native_read_prepares_rust_reader_once():
+ read = _table_read()
+ splits = [_Split() for _ in range(4)]
+ for index, split in enumerate(splits):
+ split._native_split = index
+
+ groups = []
+
+ def read_group(rust_splits):
+ groups.append(rust_splits)
+ return [_id_batch(rust_splits)]
+
+ with patch(
+ 'pypaimon.read.native_plan._prepare_native_read',
+ create=True,
+ return_value=read_group) as prepare, patch(
+ 'pypaimon.read.native_plan.native_read',
+ side_effect=AssertionError('rebuilt native reader')):
+ result = read.to_arrow(splits, parallelism=2)
+
+ assert result.to_pydict() == {'id': [0, 1, 2, 3]}
+ prepare.assert_called_once()
+ assert sorted(groups) == [[0, 1], [2, 3]]
+
+
def test_native_batch_reader_uses_effective_parallelism():
read = _table_read()
read._read_parallelism = 2
@@ -253,7 +281,11 @@ def test_native_batch_reader_uses_effective_parallelism():
worker_names = set()
worker_names_lock = threading.Lock()
- def read_group(table, rust_splits, **kwargs):
+ groups = []
+
+ def read_group(rust_splits):
+ groups.append(rust_splits)
+
def batches():
with worker_names_lock:
worker_names.add(threading.current_thread().name)
@@ -261,14 +293,13 @@ def test_native_batch_reader_uses_effective_parallelism():
yield _id_batch(rust_splits)
return batches()
- with patch('pypaimon.read.native_plan.native_read',
- side_effect=read_group) as native:
+ with patch('pypaimon.read.native_plan._prepare_native_read',
+ return_value=read_group) as prepare:
result = read.to_arrow_batch_reader(splits).read_all()
assert sorted(result.column('id').to_pylist()) == [0, 1, 2, 3]
- assert native.call_count == 2
- assert sorted(call.args[1] for call in native.call_args_list) == [
- [0, 1], [2, 3]]
+ prepare.assert_called_once()
+ assert sorted(groups) == [[0, 1], [2, 3]]
assert len(worker_names) == 2
@@ -316,17 +347,22 @@ def
test_native_batch_reader_caps_blob_parallelism_across_rust_readers():
for index, split in enumerate(splits):
split._native_split = index
+ groups = []
+
+ def read_group(group):
+ groups.append(group)
+ return [_id_batch(group)]
+
with patch(
- 'pypaimon.read.native_plan.native_read',
- side_effect=lambda table, group, **kwargs: [
- _id_batch(group)]) as native:
+ 'pypaimon.read.native_plan._prepare_native_read',
+ return_value=read_group) as prepare:
result = read.to_arrow_batch_reader(
splits, blob_parallelism=16, parallelism=16).read_all()
assert result.num_rows == 16
- assert native.call_count == 16
- assert {call.kwargs['blob_parallelism']
- for call in native.call_args_list} == {4}
+ prepare.assert_called_once()
+ assert prepare.call_args.kwargs['blob_parallelism'] == 4
+ assert len(groups) == 16
def test_parallel_native_stream_bounds_prefetch_per_reader():
@@ -541,15 +577,19 @@ def test_native_read_groups_splits_by_file_bytes():
for index, split in enumerate(splits):
split._native_split = index
+ groups = []
+
+ def read_group(group):
+ groups.append(group)
+ return [_id_batch(group)]
+
with patch(
- 'pypaimon.read.native_plan.native_read',
- side_effect=lambda table, group, **kwargs: [
- _id_batch(group)]) as native:
+ 'pypaimon.read.native_plan._prepare_native_read',
+ return_value=read_group):
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]]
+ assert sorted(groups) == [[0], [1, 2, 3]]
def test_native_read_runtime_parallelism_overrides_table_option():
@@ -559,14 +599,20 @@ def
test_native_read_runtime_parallelism_overrides_table_option():
for index, split in enumerate(splits):
split._native_split = index
+ groups = []
+
+ def read_group(group):
+ groups.append(group)
+ return [_id_batch(group)]
+
with patch(
- 'pypaimon.read.native_plan.native_read',
- side_effect=lambda table, group, **kwargs: [
- _id_batch(group)]) as native:
+ 'pypaimon.read.native_plan._prepare_native_read',
+ return_value=read_group) as prepare:
result = read.to_arrow(splits, parallelism=2)
assert result.num_rows == 4
- assert native.call_count == 2
+ prepare.assert_called_once()
+ assert len(groups) == 2
def test_native_read_caps_blob_parallelism_across_split_readers():
@@ -575,17 +621,22 @@ def
test_native_read_caps_blob_parallelism_across_split_readers():
for index, split in enumerate(splits):
split._native_split = index
+ groups = []
+
+ def read_group(group):
+ groups.append(group)
+ return [_id_batch(group)]
+
with patch(
- 'pypaimon.read.native_plan.native_read',
- side_effect=lambda table, group, **kwargs: [
- _id_batch(group)]) as native:
+ 'pypaimon.read.native_plan._prepare_native_read',
+ return_value=read_group) as prepare:
result = read.to_arrow(
splits, parallelism=16, blob_parallelism=16)
assert result.num_rows == 16
- assert native.call_count == 16
- assert {call.kwargs['blob_parallelism']
- for call in native.call_args_list} == {4}
+ prepare.assert_called_once()
+ assert prepare.call_args.kwargs['blob_parallelism'] == 4
+ assert len(groups) == 16
def test_parallel_native_read_shares_limit_across_readers():
@@ -596,9 +647,8 @@ def test_parallel_native_read_shares_limit_across_readers():
split._native_split = index
with patch(
- 'pypaimon.read.native_plan.native_read',
- side_effect=lambda table, group, **kwargs: [
- _id_batch(group)]):
+ 'pypaimon.read.native_plan._prepare_native_read',
+ return_value=lambda group: [_id_batch(group)]):
result = read.to_arrow(splits)
assert result.num_rows == 3
@@ -612,9 +662,8 @@ def
test_parallel_native_batch_reader_shares_limit_across_readers():
split._native_split = index
with patch(
- 'pypaimon.read.native_plan.native_read',
- side_effect=lambda table, group, **kwargs: [
- _id_batch(group)]):
+ 'pypaimon.read.native_plan._prepare_native_read',
+ return_value=lambda group: [_id_batch(group)]):
result = read.to_arrow_batch_reader(splits).read_all()
assert result.num_rows == 3
@@ -627,7 +676,7 @@ def
test_parallel_native_reader_setup_failure_falls_back(streaming):
for split in splits:
split._native_split = object()
- with patch('pypaimon.read.native_plan.native_read',
+ with patch('pypaimon.read.native_plan._prepare_native_read',
side_effect=RuntimeError('setup failed')):
assert read._try_native_batches(
splits,
@@ -644,9 +693,10 @@ def
test_parallel_native_stream_setup_failure_closes_started_readers():
split._native_split = object()
started = Mock()
+ read_splits = Mock(side_effect=[started, RuntimeError('setup failed')])
with patch(
- 'pypaimon.read.native_plan.native_read',
- side_effect=[started, RuntimeError('setup failed')]):
+ 'pypaimon.read.native_plan._prepare_native_read',
+ return_value=read_splits):
assert read._try_native_batches(
splits,
pa.schema([('id', pa.int32())]),
@@ -663,14 +713,14 @@ def test_parallel_native_stream_error_propagates():
for split in splits:
split._native_split = object()
- def broken_stream(table, group, **kwargs):
+ def broken_stream(group):
def batches():
raise RuntimeError('stream failed')
yield
return batches()
- with patch('pypaimon.read.native_plan.native_read',
- side_effect=broken_stream), \
+ with patch('pypaimon.read.native_plan._prepare_native_read',
+ return_value=broken_stream), \
pytest.raises(RuntimeError, match='stream failed'):
read._try_native_batches(
splits, pa.schema([('id', pa.int32())]), parallelism=2)
@@ -777,8 +827,8 @@ def
test_native_read_preserves_parallel_large_binary_blob_schema():
split._native_split = index
with patch(
- 'pypaimon.read.native_plan.native_read',
- side_effect=lambda table, group, **kwargs: [pa.record_batch(
+ 'pypaimon.read.native_plan._prepare_native_read',
+ return_value=lambda group: [pa.record_batch(
[pa.array([bytes(group)], type=pa.large_binary())],
names=['payload'])]):
result = read.to_arrow(splits, parallelism=2)
