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 eeb765d405 [Python] Delegate Native split execution to Rust Core
(#10691)
eeb765d405 is described below
commit eeb765d405963b619cc0194ecf79e2378df3b77e
Author: Jingsong Lee <[email protected]>
AuthorDate: Tue Oct 20 17:18:49 2026 +0800
[Python] Delegate Native split execution to Rust Core (#10691)
---
paimon-python/README.md | 18 +-
paimon-python/pypaimon/read/native_plan.py | 10 +-
paimon-python/pypaimon/read/table_read.py | 342 +-----------
.../pypaimon/tests/native_iterator_test.py | 98 ++--
.../tests/native_parquet_vector_update_test.py | 20 +-
paimon-python/pypaimon/tests/native_read_test.py | 604 ++++++---------------
6 files changed, 292 insertions(+), 800 deletions(-)
diff --git a/paimon-python/README.md b/paimon-python/README.md
index 5f0e4c81d5..d72c46d2c0 100644
--- a/paimon-python/README.md
+++ b/paimon-python/README.md
@@ -419,11 +419,19 @@ Python-planned or serialized authorization rules 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. An uninstalled extension,
+is executed by Rust Core. Core owns byte-balanced split groups, bounded
+prefetch, the shared row limit and cancellation. The Rust reader's `ordered`
+argument defaults to `False`. PyPaimon requests `ordered=True` for materialized
+reads, emitting groups in supplied split order; streaming reads use
+`ordered=False` and emit batches as groups become ready. Both modes retain each
+split reader's row order.
+Ordered parallel reads disable background Parquet row-group prefetch to prevent
+later groups from holding scheduling permits needed by earlier groups; split
+workers still read concurrently.
+An unfiltered row limit also caps reader fan-out to avoid speculative work.
+The split scheduler buffers at most one batch per worker, and closing the batch
+reader interrupts native reads that are still in flight. Limited BLOB reads
visit
+splits lazily with the remaining quota, before resolving output payloads. An
uninstalled extension,
unsupported route, or native-reader construction failure falls back to Python;
I/O and data errors raised after streaming starts surface to the caller.
diff --git a/paimon-python/pypaimon/read/native_plan.py
b/paimon-python/pypaimon/read/native_plan.py
index fc8a638f76..336c807dfb 100644
--- a/paimon-python/pypaimon/read/native_plan.py
+++ b/paimon-python/pypaimon/read/native_plan.py
@@ -500,7 +500,7 @@ def _prepare_native_read(table, predicate:
Optional[Predicate] = None,
blob_parallelism: Optional[int] = None,
read_type: Optional[List[DataField]] = None,
include_row_kind: bool = False):
- """Create one Rust reader reusable across split groups."""
+ """Create the configured Core reader entrypoint."""
if not native_runtime_available():
raise RuntimeError(
"read.native.enabled needs the pypaimon-rust extension")
@@ -527,8 +527,10 @@ def native_read(table, splits, predicate:
Optional[Predicate] = None,
projection: Optional[List[str]] = None,
blob_parallelism: Optional[int] = None,
read_type: Optional[List[DataField]] = None,
- include_row_kind: bool = False):
- """Read Rust ``Split`` objects into PyArrow ``RecordBatch`` objects."""
+ include_row_kind: bool = False,
+ parallelism: int = 1,
+ ordered: bool = False):
+ """Run one Core read; ``ordered`` retains the supplied split order."""
read_splits = _prepare_native_read(
table,
predicate=predicate,
@@ -537,7 +539,7 @@ def native_read(table, splits, predicate:
Optional[Predicate] = None,
blob_parallelism=blob_parallelism,
read_type=read_type,
include_row_kind=include_row_kind)
- return read_splits(splits)
+ return read_splits(splits, parallelism=parallelism, ordered=ordered)
def native_plan(
diff --git a/paimon-python/pypaimon/read/table_read.py
b/paimon-python/pypaimon/read/table_read.py
index 6e6b9326d7..6159f17ced 100644
--- a/paimon-python/pypaimon/read/table_read.py
+++ b/paimon-python/pypaimon/read/table_read.py
@@ -17,7 +17,6 @@
import logging
import os
-import queue
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Callable, Dict, Iterator, List, Optional
@@ -138,10 +137,6 @@ class _RemainingRows:
return self._remaining <= 0
-class _NativeReadSetupError(Exception):
- """Native reader construction failed before a batch was exposed."""
-
-
class TableRead:
"""Implementation of TableRead for native Python reading."""
@@ -244,8 +239,7 @@ class TableRead:
read._row_kind_column = 'rowkind'
schema =
read._add_row_kind_to_schema(read._output_arrow_schema(splits),
read._row_kind_column)
fields = output_fields(self.read_type, self.output_projection)
- batches = read._try_native_batches(splits, schema, parallelism=1,
streaming=True,
- preserve_split_order=True)
+ batches = read._try_native_batches(splits, schema, parallelism=1,
ordered=True)
if batches is not None:
try:
yield from native_rows(batches, self.table, fields)
@@ -316,8 +310,8 @@ class TableRead:
"""Return a closeable batch reader supporting context management.
PyArrow readers do not close the Python iterator passed to
- ``from_batches``. Retain it explicitly for a parallel read so native
- split workers are stopped when a caller ends the read early.
+ ``from_batches``. Retain it explicitly so native reads release their
+ resources when a caller ends the read early, including serial reads.
"""
# Keep the explicit iterator cleanup chain even when LIMIT caps the
# native worker count inside _new_arrow_batch_reader to one.
@@ -339,22 +333,23 @@ class TableRead:
schema = self._output_arrow_schema(splits)
if self.include_row_kind:
schema = self._add_row_kind_to_schema(schema)
- effective = self._effective_parallelism(parallelism, len(splits))
+ effective = self._resolve_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(
- min(effective, len(splits)), effective_bp)
native_batches = self._try_native_batches(
splits,
schema,
parallelism=effective,
blob_parallelism=effective_bp,
- streaming=True,
+ ordered=False,
)
if native_batches is not None:
batch_iterator = iter(native_batches)
else:
self._check_python_merge_supported()
+ effective = self._effective_parallelism(parallelism, len(splits))
+ if self._should_run_parallel(splits, effective):
+ effective_bp = self._cap_blob_parallelism(
+ min(effective, len(splits)), effective_bp)
batch_iterator = self._arrow_batch_generator(
splits, schema, effective_bp)
reader_type = pyarrow.ipc.RecordBatchReader
@@ -436,26 +431,23 @@ class TableRead:
shrunk to stay within it.
"""
effective_bp = self._resolve_blob_parallelism(blob_parallelism)
- effective = self._effective_parallelism(parallelism, len(splits))
+ effective = self._resolve_parallelism(parallelism, len(splits))
schema = self._output_arrow_schema(splits)
if self.include_row_kind:
schema = self._add_row_kind_to_schema(schema)
- native_parallel = self._should_run_parallel(splits, effective)
- native_bp = effective_bp
- if native_parallel:
- native_bp = self._cap_blob_parallelism(
- min(effective, len(splits)), effective_bp)
native_batches = self._try_native_batches(
splits,
schema,
parallelism=effective,
- blob_parallelism=native_bp,
+ blob_parallelism=effective_bp,
+ ordered=True,
)
if native_batches is not None:
return self._batches_to_arrow(native_batches, schema)
self._check_python_merge_supported()
+ effective = self._effective_parallelism(parallelism, len(splits))
if self._should_run_parallel(splits, effective):
return self._to_arrow_parallel(splits, schema, effective,
effective_bp)
@@ -597,8 +589,7 @@ class TableRead:
schema: pyarrow.Schema,
parallelism: Optional[int] = None,
blob_parallelism: Optional[int] = None,
- streaming: bool = False,
- preserve_split_order: bool = False):
+ ordered: bool = False):
"""Return Rust-read batches, or ``None`` when this read must fall
back."""
if not self.table.options.native_read_enabled():
return self._native_fallback("read.native.enabled is false")
@@ -612,12 +603,10 @@ class TableRead:
"the table uses an unsupported BLOB view")
try:
from pypaimon.read.native_plan import (
- _prepare_native_read, _raise_if_native_fork_safety_error,
- native_read, native_split_from_python)
+ _raise_if_native_fork_safety_error, native_read,
native_split_from_python)
except Exception as e:
return self._native_fallback(str(e), e, log=True)
rust_splits = []
- split_weights = []
for split in splits:
if not self._native_split_files_supported(split):
return self._native_fallback(
@@ -633,86 +622,20 @@ class TableRead:
e,
log=True)
rust_splits.append(rust_split)
- split_weights.append(self._native_split_weight(split))
- if preserve_split_order:
- read_kwargs = self._native_read_kwargs(blob_parallelism)
- try:
- read_splits = _prepare_native_read(self.table, **read_kwargs)
- except Exception as e:
- _raise_if_native_fork_safety_error(e)
- return self._native_fallback(str(e), e, log=True)
- return self._convert_native_batches(
- self._native_batches_ordered(read_splits, rust_splits,
read_kwargs), schema)
- 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:
- _raise_if_native_fork_safety_error(e)
- return self._native_fallback(str(e), e, log=True)
- if streaming:
- groups = self._native_split_groups(
- rust_splits, parallelism, split_weights)
- readers = []
- try:
- for group in groups:
- readers.append(read_splits(group))
- except Exception as e:
- 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 after "
- "another reader failed to start",
- exc_info=True,
- )
- _raise_if_native_fork_safety_error(e)
- return self._native_fallback(str(e), e, log=True)
- batches = self._native_batches_parallel_streaming(readers)
- return self._convert_native_batches(batches, schema)
- try:
- return self._native_batches_parallel(
- read_splits, rust_splits, schema, parallelism,
- split_weights)
- except _NativeReadSetupError as e:
- _raise_if_native_fork_safety_error(e)
- return self._native_fallback(str(e), e, log=True)
+ # Python adapts the public configuration; Core owns split grouping,
+ # bounded concurrency, the row quota and cancellation of every stream.
+ effective = 1 if parallelism is None else parallelism
try:
- read_kwargs = self._native_read_kwargs(blob_parallelism)
- batches = native_read(self.table, rust_splits, **read_kwargs)
- except Exception as e:
- _raise_if_native_fork_safety_error(e)
- return self._native_fallback(str(e), e, log=True)
+ batches = native_read(
+ self.table, rust_splits,
+ parallelism=effective,
+ ordered=ordered,
+ **self._native_read_kwargs(blob_parallelism))
+ except Exception as error:
+ _raise_if_native_fork_safety_error(error)
+ return self._native_fallback(str(error), error, log=True)
return self._convert_native_batches(batches, schema)
- def _native_batches_ordered(self, read_splits, splits, read_kwargs):
- from pypaimon.read.native_plan import _prepare_native_read
-
- remaining = self.limit
- for index, split in enumerate(splits):
- if remaining is not None and remaining <= 0:
- break
- if index and remaining is not None:
- # BLOB readers must see the remaining quota before opening any
- # payload, including when earlier splits already spent it.
- read_splits = _prepare_native_read(self.table,
**dict(read_kwargs, limit=remaining))
- batches = read_splits([split])
- try:
- for batch in batches:
- yield batch
- if remaining is not None:
- remaining -= batch.num_rows
- if remaining <= 0:
- break
- finally:
- close = getattr(batches, 'close', None)
- if close is not None:
- close()
-
def _native_read_kwargs(self, blob_parallelism=None):
kwargs = {
'predicate': self.predicate,
@@ -725,207 +648,6 @@ class TableRead:
kwargs['include_row_kind'] = True
return kwargs
- @staticmethod
- 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))
- 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):
- 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, 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)
-
- remaining_state = _RemainingRows(self.limit)
- results = [None] * len(groups)
- with ThreadPoolExecutor(
- max_workers=workers,
- thread_name_prefix="pypaimon-native-read") as executor:
- futures = {
- executor.submit(
- self._read_native_split_group,
- read_splits,
- group,
- schema,
- remaining_state,
- ): index
- for index, group in enumerate(groups)
- }
- for future in as_completed(futures):
- results[futures[future]] = future.result()
-
- return [batch for group_batches in results for batch in group_batches]
-
- def _native_batches_parallel_streaming(self, readers):
- """Stream ready batches with one buffered batch per Rust reader."""
- stop = threading.Event()
- results = queue.Queue(maxsize=len(readers))
- capacities = [threading.Semaphore(1) for _ in readers]
- executor = ThreadPoolExecutor(
- max_workers=len(readers),
- thread_name_prefix="pypaimon-native-read",
- )
-
- def put(index, item):
- while not stop.is_set():
- try:
- results.put((index, item), timeout=0.1)
- return True
- except queue.Full:
- pass
- return False
-
- def read_group(index, batches):
- try:
- iterator = iter(batches)
- while not stop.is_set():
- if not capacities[index].acquire(timeout=0.1):
- continue
- if stop.is_set():
- capacities[index].release()
- return
- try:
- batch = next(iterator)
- except StopIteration:
- put(index, ('done', None))
- return
- if not put(index, ('batch', batch)):
- return
- except Exception as error:
- put(index, ('error', error))
- finally:
- close = getattr(batches, 'close', None)
- if close is not None:
- close()
-
- futures = [
- executor.submit(read_group, index, reader)
- for index, reader in enumerate(readers)
- ]
- try:
- completed = 0
- while completed < len(readers):
- index, (kind, value) = results.get()
- if kind == 'batch':
- try:
- yield value
- except BaseException:
- # Stop workers before releasing this reader's permit.
- # Otherwise an early close can let it enter another
- # blocking ``next`` before the outer finally runs.
- stop.set()
- raise
- finally:
- capacities[index].release()
- elif kind == 'done':
- capacities[index].release()
- completed += 1
- else:
- capacities[index].release()
- 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)
-
- def _read_native_split_group(
- self, read_splits, rust_splits, schema, remaining_state):
- if remaining_state.exhausted():
- return []
- try:
- batches = read_splits(rust_splits)
- except Exception as e:
- raise _NativeReadSetupError(str(e)) from e
- result = []
- for batch in batches:
- if batch.num_rows == 0:
- continue
- allowed = remaining_state.try_consume(batch.num_rows)
- if allowed == 0:
- break
- if allowed < batch.num_rows:
- batch = batch.slice(0, allowed)
- batch = self._prepare_native_batch(batch, schema)
- batch = self._project_batch_to_output(batch)
- result.append(self._try_to_pad_batch_by_schema(batch, schema))
- if remaining_state.exhausted():
- break
- return result
-
@staticmethod
def _native_split_files_supported(split):
for data_file in split.files:
@@ -936,20 +658,13 @@ class TableRead:
return True
def _convert_native_batches(self, batches, schema):
- remaining = self.limit
try:
for batch in batches:
if batch.num_rows == 0:
continue
batch = self._prepare_native_batch(batch, schema)
- if remaining is not None and batch.num_rows > remaining:
- batch = batch.slice(0, remaining)
batch = self._project_batch_to_output(batch)
yield self._try_to_pad_batch_by_schema(batch, schema)
- if remaining is not None:
- remaining -= batch.num_rows
- if remaining <= 0:
- break
finally:
close = getattr(batches, 'close', None)
if close is not None:
@@ -1124,7 +839,10 @@ class TableRead:
and not self._limit_covers_all_splits(splits))
def _requires_batch_iterator_cleanup(self, splits: List[Split], effective:
int) -> bool:
- return (self._should_run_parallel(splits, effective)
+ # A serial Core stream owns resources too. PyArrow's reader does not
+ # close the Python iterator supplied to from_batches on an early close.
+ return (self.table.options.native_read_enabled()
+ or self._should_run_parallel(splits, effective)
or self._deferred_blob_limit_may_prune(splits))
def _native_blob_view_supported(self) -> bool:
diff --git a/paimon-python/pypaimon/tests/native_iterator_test.py
b/paimon-python/pypaimon/tests/native_iterator_test.py
index 4f20936b75..76eb987950 100644
--- a/paimon-python/pypaimon/tests/native_iterator_test.py
+++ b/paimon-python/pypaimon/tests/native_iterator_test.py
@@ -107,6 +107,46 @@ def test_iterator_uses_native_read(tmp_path, mode,
projection, limit):
assert all(kind == RowKind.INSERT for _, kind in actual)
[email protected]('method_name', [
+ 'to_arrow', 'to_arrow_batch_reader', '_to_managed_arrow_batch_reader',
+])
+def
test_native_arrow_modes_share_limit_and_materialized_keeps_split_prefix(tmp_path,
method_name):
+ table, _ = _table(tmp_path)
+ predicate =
table.new_read_builder().new_predicate_builder().greater_than('v', 10)
+ builder =
table.new_read_builder().with_projection(['id']).with_filter(predicate)
+ splits = list(reversed(builder.new_scan().plan().splits()))
+ assert len(splits) >= 2
+ python = table.copy({'read.native.enabled': 'false',
'scan.native-plan.enabled': 'false'})
+ python_read =
python.new_read_builder().with_projection(['id']).with_filter(predicate).new_read()
+ expected = []
+ for split in splits:
+ expected.extend(python_read.to_arrow([split],
parallelism=1).column('id').to_pylist())
+ assert len(expected) == 2
+ for limit in [None, 1]:
+ for parallelism in [1, 4]:
+ limited =
table.new_read_builder().with_projection(['id']).with_filter(predicate)
+ if limit is not None:
+ limited.with_limit(limit)
+ read = limited.new_read()
+ with patch.object(read, '_create_split_read',
side_effect=AssertionError('Python read fallback')):
+ result = getattr(read, method_name)(splits,
parallelism=parallelism)
+ if method_name != 'to_arrow':
+ try:
+ result_table = result.read_all()
+ finally:
+ result.close()
+ else:
+ result_table = result
+ values = result_table.column('id').to_pylist()
+ count = len(expected) if limit is None else limit
+ assert len(values) == count
+ if method_name == 'to_arrow' or parallelism == 1:
+ assert values == expected[:count]
+ else:
+ assert len(set(values)) == len(values)
+ assert set(values).issubset(expected)
+
+
def test_append_iterator_row_kind_is_insert(tmp_path):
table, _ = _table(tmp_path)
table = table.copy({'read.native.enabled': 'false',
'scan.native-plan.enabled': 'false'})
@@ -143,14 +183,19 @@ def _planned_read(tmp_path, limit=None):
builder = table.new_read_builder().with_projection(['id'])
if limit is not None:
builder.with_limit(limit)
- return builder.new_read(), builder.new_scan().plan().splits()
+ splits = builder.new_scan().plan().splits()
+ # Transport tests mock native_read itself; opaque retained split tokens
+ # keep their lifetime assertions independent of Native split construction.
+ for split in splits:
+ split._native_split = object()
+ return builder.new_read(), splits
@pytest.mark.parametrize('consume', [False, True])
def test_iterator_is_lazy_and_close_stops_native_input(tmp_path, consume):
read, splits = _planned_read(tmp_path)
batches = _Batches([_batch([1, 2]), RuntimeError('must not fetch next
batch')])
- with patch('pypaimon.read.native_plan._prepare_native_read',
return_value=lambda _: batches) as prepare, \
+ with patch('pypaimon.read.native_plan.native_read', return_value=batches)
as prepare, \
patch.object(read, '_create_split_read',
side_effect=AssertionError('Python fallback')):
iterator = read.to_iterator(splits)
prepare.assert_not_called()
@@ -171,7 +216,7 @@ def
test_iterator_closes_native_input_on_failure_without_replaying_python(tmp_pa
pa.record_batch([pa.array(['+I']), pa.array(['wrong type'])],
names=['rowkind', 'id'])
if failure == 'schema' else _batch([2], ['invalid']))
batches = _Batches([_batch([1]), bad])
- with patch('pypaimon.read.native_plan._prepare_native_read',
return_value=lambda _: batches), \
+ with patch('pypaimon.read.native_plan.native_read', return_value=batches),
\
patch.object(read, '_create_split_read',
side_effect=AssertionError('Python replay')):
iterator = read.to_iterator(splits)
assert next(iterator).get_field(0) == 1
@@ -180,32 +225,23 @@ def
test_iterator_closes_native_input_on_failure_without_replaying_python(tmp_pa
assert batches.closed
-def test_iterator_spends_limit_before_opening_next_split(tmp_path):
+def test_iterator_delegates_ordered_global_limit_in_one_native_read(tmp_path):
read, splits = _planned_read(tmp_path, limit=3)
assert len(splits) >= 2
- prepared = []
- opened = []
- readers = []
-
- def prepare(table, **kwargs):
- prepared.append(kwargs['limit'])
-
- def open_split(group):
- assert len(group) == 1
- assert not readers or readers[-1].closed
- opened.append(group[0])
- reader = _Batches([_batch(list(range(min(2, kwargs['limit']))))])
- readers.append(reader)
- return reader
- return open_split
-
- with patch('pypaimon.read.native_plan._prepare_native_read',
side_effect=prepare), \
+ batches = _Batches([_batch([1, 2]), _batch([3])])
+ with patch('pypaimon.read.native_plan.native_read', return_value=batches)
as native, \
+ patch('pypaimon.read.native_plan.native_split_from_python',
side_effect=lambda split, **kwargs: split), \
patch.object(read, '_create_split_read',
side_effect=AssertionError('Python fallback')):
result = _rows(read, splits + splits)
- assert len(result) == 3
- assert prepared == [3, 1]
- assert opened == [splits[0]._native_split, splits[1]._native_split]
- assert all(reader.closed for reader in readers)
+ assert result == [((value,), RowKind.INSERT) for value in [1, 2, 3]]
+ native.assert_called_once()
+ transported = [getattr(split, '_native_split', split) for split in splits
+ splits]
+ assert native.call_args.args == (read.table, transported)
+ assert native.call_args.kwargs['limit'] == 3
+ assert native.call_args.kwargs['parallelism'] == 1
+ assert native.call_args.kwargs['ordered'] is True
+ assert native.call_args.kwargs['include_row_kind'] is True
+ assert batches.closed
assert read.limit == 3
@@ -213,14 +249,14 @@ def
test_iterator_setup_failure_falls_back_before_native_rows(tmp_path):
read, splits = _planned_read(tmp_path)
expected = _rows(read.table.copy({'read.native.enabled':
'false'}).new_read_builder()
.with_projection(['id']).new_read(), splits)
- with patch('pypaimon.read.native_plan._prepare_native_read',
side_effect=RuntimeError('setup failed')):
+ with patch('pypaimon.read.native_plan.native_read',
side_effect=RuntimeError('setup failed')):
assert _rows(read, splits) == expected
def test_iterator_qualifies_every_split_before_starting_native(tmp_path):
read, splits = _planned_read(tmp_path)
with patch.object(read, '_native_split_files_supported',
side_effect=[True, False]), \
- patch('pypaimon.read.native_plan._prepare_native_read') as prepare:
+ patch('pypaimon.read.native_plan.native_read') as prepare:
assert len(_rows(read, splits)) == 3
prepare.assert_not_called()
@@ -229,10 +265,10 @@ def
test_iterator_qualifies_every_split_before_starting_native(tmp_path):
def test_iterator_preserves_all_row_kinds_and_hides_transport_column(tmp_path,
limit):
read, splits = _planned_read(tmp_path, limit)
kinds = ['+I', '-U', '+U', '-D']
- batches = _Batches([_batch([1, 2, 3, 4], kinds)])
- with patch('pypaimon.read.native_plan._prepare_native_read',
return_value=lambda _: batches) as prepare:
- rows = _rows(read, splits[:1])
count = 4 if limit is None else limit
+ batches = _Batches([_batch([1, 2, 3, 4][:count], kinds[:count])])
+ with patch('pypaimon.read.native_plan.native_read', return_value=batches)
as prepare:
+ rows = _rows(read, splits[:1])
assert rows == [((i + 1,), RowKind.from_string(kind)) for i, kind in
enumerate(kinds[:count])]
assert prepare.call_count == int(limit != 0)
assert batches.closed == (limit != 0)
@@ -439,7 +475,7 @@ def
test_iterator_keeps_plan_authorization_across_splits_and_fallback(rest_catal
if mode == 'wire-fallback':
splits = [deserialize_split_v1(serialize_split_v1(split), []) for
split in splits]
failure = AssertionError('Serialized rules must not mint a native
grant')
- context = patch('pypaimon.read.native_plan._prepare_native_read',
side_effect=failure)
+ context = patch('pypaimon.read.native_plan.native_read',
side_effect=failure)
with context:
assert sorted(_rows(read, splits), key=lambda row: row[0]) == [
((2,), RowKind.INSERT), ((3,), RowKind.INSERT)]
diff --git a/paimon-python/pypaimon/tests/native_parquet_vector_update_test.py
b/paimon-python/pypaimon/tests/native_parquet_vector_update_test.py
index 90a6c1b4f8..57a0688bb9 100644
--- a/paimon-python/pypaimon/tests/native_parquet_vector_update_test.py
+++ b/paimon-python/pypaimon/tests/native_parquet_vector_update_test.py
@@ -25,7 +25,7 @@ import pyarrow.compute as pc
import pytest
from pypaimon import CatalogFactory, Schema
-from pypaimon.read.table_read import TableRead
+from pypaimon.read.native_plan import native_read as read_native_splits
from pypaimon.schema.data_types import PyarrowFieldParser
from pypaimon.table.data_evolution_merge_into import WhenMatched,
WhenNotMatched, source_col
from pypaimon.write.table_update import TableUpdate
@@ -172,17 +172,9 @@ def
test_shard_normal_updates_preserve_parquet_vector_and_blob_sidecars(tmp_path
update = builder.new_update().with_update_type(['score'])
update.with_read_projection(['id', 'score'])
- def reversed_groups(self, readers):
- # A valid completion order for parallel readers must not change which
- # row receives an update. Read actual files, but finish the last group
first.
- for reader in reversed(readers):
- try:
- yield from reader
- finally:
- reader.close()
-
with ExitStack() as stack:
- stack.enter_context(patch.object(TableRead,
'_native_batches_parallel_streaming', reversed_groups))
+ native_calls = stack.enter_context(patch(
+ 'pypaimon.read.native_plan.native_read', wraps=read_native_splits))
if native:
stack.enter_context(patch('pypaimon.write.table_update.ShardTableUpdator.__init__',
side_effect=AssertionError('Python shard
fallback')))
@@ -195,6 +187,12 @@ def
test_shard_normal_updates_preserve_parquet_vector_and_blob_sidecars(tmp_path
finally:
reader.close()
messages = shard.prepare_commit()
+ if native_read and not native:
+ # Pure shard writers consume their ranges in split order. Their
+ # Native read delegates every split to one serial Core stream.
+ native_calls.assert_called_once()
+ assert native_calls.call_args.kwargs['parallelism'] == 1
+ assert native_calls.call_args.kwargs['ordered'] is False
_commit(builder, messages, stream)
for rows in _rows(table):
assert rows == [dict(row, score=row['score'] + 10) for row in before]
diff --git a/paimon-python/pypaimon/tests/native_read_test.py
b/paimon-python/pypaimon/tests/native_read_test.py
index e82e996a81..20966ef9fc 100644
--- a/paimon-python/pypaimon/tests/native_read_test.py
+++ b/paimon-python/pypaimon/tests/native_read_test.py
@@ -200,12 +200,12 @@ def test_floating_sequence_uses_native_read(type_):
native.assert_called_once()
-def test_native_read_consumes_retained_rust_splits_and_enforces_limit():
+def test_native_read_consumes_retained_rust_splits_and_transports_limit():
read = _table_read(limit=2)
first, second = _Split(), _Split()
first._native_split = object()
second._native_split = object()
- batches = [_id_batch([1, 2, 3])]
+ batches = [_id_batch([1, 2])]
with patch(
'pypaimon.read.native_plan.native_read',
@@ -220,6 +220,8 @@ def
test_native_read_consumes_retained_rust_splits_and_enforces_limit():
limit=2,
read_type=read.read_type,
blob_parallelism=1,
+ parallelism=1,
+ ordered=True,
)
@@ -339,6 +341,8 @@ def test_native_read_bridges_python_planned_split():
predicate=None,
limit=None,
read_type=read.read_type,
+ parallelism=1,
+ ordered=False,
)
@@ -366,95 +370,68 @@ def
test_native_query_auth_transports_rules_for_core_authorization():
assert native.call_args.args[1] == [converted]
-def test_native_read_uses_effective_parallelism_from_table_option():
[email protected]('method_name, ordered', [
+ ('to_arrow', True), ('to_arrow_batch_reader', False),
+ ('_to_managed_arrow_batch_reader', False),
+])
[email protected]('option, runtime, expected', [
+ (2, None, 2), (4, 2, 2), (1, 4, 4), (4, 1, 1), (4, 16, 16),
+])
+def test_native_read_delegates_all_splits_and_parallelism_once(method_name,
ordered, option, runtime, expected):
read = _table_read()
- read._read_parallelism = 2
- splits = [_Split() for _ in range(4)]
+ read._read_parallelism = option
+ splits = [_Split(file_size=size) for size in [8, 8, 1, 1]]
for index, split in enumerate(splits):
split._native_split = index
-
- barrier = threading.Barrier(2)
- worker_names = set()
- worker_names_lock = threading.Lock()
-
- groups = []
-
- def read_group(rust_splits):
- groups.append(rust_splits)
-
- def batches():
- with worker_names_lock:
- worker_names.add(threading.current_thread().name)
- barrier.wait(timeout=5)
- yield _id_batch(rust_splits)
- return batches()
-
- with patch('pypaimon.read.native_plan._prepare_native_read',
- return_value=read_group) as prepare:
- result = read.to_arrow(splits)
-
+ with patch('pypaimon.read.native_plan.native_read',
return_value=[_id_batch(range(4))]) as native, \
+ patch('pypaimon.read.table_read.ThreadPoolExecutor',
side_effect=AssertionError('Python executor')):
+ result = getattr(read, method_name)(splits, parallelism=runtime)
+ if method_name != 'to_arrow':
+ result = result.read_all()
assert result.to_pydict() == {'id': [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]]
+ native.assert_called_once()
+ assert native.call_args.args == (read.table, list(range(4)))
+ assert native.call_args.kwargs['parallelism'] == expected
+ assert native.call_args.kwargs['ordered'] == ordered
-def test_native_batch_reader_uses_effective_parallelism():
[email protected]('ordered', [False, True])
+def test_native_read_transports_explicit_order_without_opening_groups(ordered):
read = _table_read()
- read._read_parallelism = 2
- splits = [_Split() for _ in range(4)]
+ splits = [_Split(), _Split()]
for index, split in enumerate(splits):
split._native_split = index
+ with patch('pypaimon.read.native_plan.native_read',
return_value=[_id_batch([0, 1])]) as native:
+ batches = list(read._try_native_batches(
+ splits, pa.schema([('id', pa.int32())]), parallelism=2,
ordered=ordered))
+ assert batches[0].column('id').to_pylist() == [0, 1]
+ native.assert_called_once()
+ assert native.call_args.args[1] == [0, 1]
+ assert native.call_args.kwargs['ordered'] == ordered
+ assert native.call_args.kwargs['parallelism'] == 2
- barrier = threading.Barrier(2)
- worker_names = set()
- worker_names_lock = threading.Lock()
-
- groups = []
- def read_group(rust_splits):
- groups.append(rust_splits)
[email protected]('ordered', [None, False, True])
+def test_native_read_bridge_transports_current_ordered_api(ordered):
+ from pypaimon.read.native_plan import native_read
- def batches():
- with worker_names_lock:
- worker_names.add(threading.current_thread().name)
- barrier.wait(timeout=5)
- yield _id_batch(rust_splits)
- return batches()
+ table = object()
+ splits = [object(), object()]
+ batches = [_id_batch([2]), _id_batch([1])]
+ calls = []
- with patch('pypaimon.read.native_plan._prepare_native_read',
- return_value=read_group) as prepare:
- result = read.to_arrow_batch_reader(splits).read_all()
+ def read_splits(values, *, parallelism, ordered=False):
+ calls.append((values, parallelism, ordered))
+ return batches
- assert sorted(result.column('id').to_pylist()) == [0, 1, 2, 3]
- prepare.assert_called_once()
- assert sorted(groups) == [[0, 1], [2, 3]]
- assert len(worker_names) == 2
+ kwargs = {} if ordered is None else {'ordered': ordered}
+ with patch('pypaimon.read.native_plan._prepare_native_read',
return_value=read_splits) as prepare:
+ result = native_read(table, splits, parallelism=2, limit=2, **kwargs)
+ assert result is batches
+ assert calls == [(splits, 2, False if ordered is None else ordered)]
+ prepare.assert_called_once_with(
+ table, predicate=None, limit=2, projection=None, blob_parallelism=None,
+ read_type=None, include_row_kind=False)
def test_native_batch_reader_close_closes_batch_iterator():
@@ -482,433 +459,187 @@ def
test_native_batch_reader_close_closes_batch_iterator():
self.closed = True
reader = Reader()
- with patch.object(
- read,
- '_new_arrow_batch_reader',
- return_value=(reader, batch_iterator)):
- batch_reader = read.to_arrow_batch_reader(
- [_Split(), _Split()])
+ with patch.object(read, '_new_arrow_batch_reader', return_value=(reader,
batch_iterator)):
+ batch_reader = read.to_arrow_batch_reader([_Split(), _Split()])
assert batch_reader.read_next_batch().column('id').to_pylist() == [1]
batch_reader.close()
-
assert reader.closed
assert iterator_closed.wait(timeout=1)
-def test_native_batch_reader_caps_blob_parallelism_across_rust_readers():
- read = _table_read()
- splits = [_Split() for _ in range(16)]
- 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._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
- 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():
- read = _table_read()
- all_started = threading.Barrier(3)
- produced = [0, 0, 0]
- produced_lock = threading.Lock()
-
- def reader_batches(index):
- for value in range(3):
- with produced_lock:
- produced[index] += 1
- if value == 0:
- all_started.wait(timeout=2)
- yield _id_batch([index * 10 + value])
-
- batches = read._native_batches_parallel_streaming([
- reader_batches(0),
- reader_batches(1),
- reader_batches(2),
- ])
- try:
- next(batches)
- with produced_lock:
- assert sum(produced) <= 4
- assert all(value <= 2 for value in produced)
- finally:
- batches.close()
-
-
-def test_parallel_native_stream_close_does_not_start_another_read():
- read = _table_read()
-
- class BlockingReader:
- def __init__(self):
- self._first = True
- self._release = threading.Event()
- self.second_read_started = threading.Event()
- self.closed = threading.Event()
-
- def __iter__(self):
- return self
-
- def __next__(self):
- if self._first:
- self._first = False
- return _id_batch([1])
- self.second_read_started.set()
- self._release.wait(timeout=5)
- raise StopIteration
-
- def close(self):
- self.closed.set()
- self._release.set()
-
- reader = BlockingReader()
- batches = read._native_batches_parallel_streaming([reader])
- assert next(batches).column('id').to_pylist() == [1]
- reader.second_read_started.wait(timeout=1)
-
- close_finished = threading.Event()
-
- def close_batches():
- batches.close()
- close_finished.set()
-
- close_thread = threading.Thread(target=close_batches, daemon=True)
- close_thread.start()
- closed_without_unblocking = close_finished.wait(timeout=1)
- try:
- assert closed_without_unblocking
- finally:
- reader.close()
- close_thread.join(timeout=5)
-
- assert reader.closed.is_set()
- 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)
[email protected]('method_name', ['to_arrow_batch_reader',
'_to_managed_arrow_batch_reader'])
[email protected]('blob_limit', [False, True])
+def test_native_batch_reader_close_reaches_core_input(method_name, blob_limit):
+ read = _blob_table_read(limit=2) if blob_limit else _table_read()
+ read.table.is_primary_key_table = blob_limit
splits = [_Split(), _Split()]
for index, split in enumerate(splits):
split._native_split = index
+ split.merged_row_count = Mock(return_value=None)
+ expected = b'payload' if blob_limit else 1
+ batch = (pa.record_batch([pa.array([expected], type=pa.large_binary())],
names=['payload'])
+ if blob_limit else _id_batch([expected]))
class CloseTrackingReader:
- def __init__(self):
- self._batch = _id_batch([0, 1])
- self.closed = False
+ 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
+ core_reader = CloseTrackingReader()
+ with patch('pypaimon.read.native_plan.native_read',
return_value=core_reader):
+ reader = getattr(read, method_name)(splits, parallelism=4)
+ assert reader.read_next_batch().column(0).to_pylist() == [expected]
+ assert not core_reader.closed
+ reader.close()
+ reader.close()
+ assert core_reader.closed
[email protected]('managed_reader', [False, True])
-def
test_pruning_pk_blob_limit_close_reaches_serial_native_reader(managed_reader):
- read = _blob_table_read(limit=2)
- read.table.is_primary_key_table = True
- splits = [_Split(), _Split()]
[email protected]('method_name', ['to_arrow_batch_reader',
'_to_managed_arrow_batch_reader'])
[email protected]('split_count, parallelism', [(1, 1), (3, 1)])
+def
test_serial_native_batch_reader_close_releases_suspended_core_stream(method_name,
split_count, parallelism):
+ read = _table_read()
+ splits = [_Split() for _ in range(split_count)]
for index, split in enumerate(splits):
split._native_split = index
- split.merged_row_count = Mock(return_value=None)
- class CloseTrackingReader:
- closed = False
+ class CoreReader:
+ def __init__(self):
+ self.reads = 0
+ self.closes = 0
def __iter__(self):
return self
def __next__(self):
- return pa.record_batch([pa.array([b'payload'],
type=pa.large_binary())], names=['payload'])
+ self.reads += 1
+ return _id_batch([self.reads])
def close(self):
- self.closed = True
+ self.closes += 1
- native_reader = CloseTrackingReader()
- method = read._to_managed_arrow_batch_reader if managed_reader else
read.to_arrow_batch_reader
- with patch('pypaimon.read.native_plan.native_read',
return_value=native_reader):
- reader = method(splits, parallelism=4)
- assert reader.read_next_batch().column('payload').to_pylist() ==
[b'payload']
- assert not native_reader.closed
+ core_reader = CoreReader()
+ with patch('pypaimon.read.native_plan.native_read',
return_value=core_reader):
+ reader = getattr(read, method_name)(splits, parallelism=parallelism)
+ assert reader.read_next_batch().column('id').to_pylist() == [1]
+ assert core_reader.closes == 0
+ reader.close()
reader.close()
- assert native_reader.closed
+ assert core_reader.closes == 1
+ assert core_reader.reads == 1
-def test_filtered_native_read_limit_keeps_split_parallelism():
- read = _table_read(limit=1)
[email protected]('method_name', ['to_arrow', 'to_arrow_batch_reader'])
[email protected]('limit', [0, 1, 3, None])
+def test_native_read_forwards_one_global_limit(method_name, limit):
+ read = _table_read(limit=limit)
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
-
- groups = []
-
- def read_group(group):
- groups.append(group)
- return [_id_batch(group)]
-
- with patch(
- '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(groups) == [[0], [1, 2, 3]]
-
-
-def test_native_read_runtime_parallelism_overrides_table_option():
- read = _table_read()
- read._read_parallelism = 4
+ read._read_parallelism = 2
splits = [_Split() for _ in range(4)]
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._prepare_native_read',
- return_value=read_group) as prepare:
- result = read.to_arrow(splits, parallelism=2)
-
- assert result.num_rows == 4
- prepare.assert_called_once()
- assert len(groups) == 2
+ # Core owns quota consumption and returns the already limited batches.
+ values = [0, 1, 2] if limit is None else [0, 1, 2][:limit]
+ with patch('pypaimon.read.native_plan.native_read',
+ return_value=[_id_batch([]), _id_batch(values)]) as native:
+ result = getattr(read, method_name)(splits)
+ if method_name != 'to_arrow':
+ result = result.read_all()
+ assert result.to_pydict() == {'id': values}
+ native.assert_called_once()
+ assert native.call_args.args[1] == list(range(4))
+ assert native.call_args.kwargs['limit'] == limit
+ assert native.call_args.kwargs['predicate'] is read.predicate
+ assert native.call_args.kwargs['parallelism'] == 2
-def test_native_read_caps_blob_parallelism_across_split_readers():
[email protected]('method_name', ['to_arrow', 'to_arrow_batch_reader'])
+def
test_native_read_transports_blob_parallelism_for_one_core_read(method_name):
read = _table_read()
splits = [_Split() for _ in range(16)]
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._prepare_native_read',
- return_value=read_group) as prepare:
- result = read.to_arrow(
- splits, parallelism=16, blob_parallelism=16)
-
+ with patch('pypaimon.read.native_plan.native_read',
return_value=[_id_batch(range(16))]) as native:
+ result = getattr(read, method_name)(splits, parallelism=16,
blob_parallelism=16)
+ if method_name != 'to_arrow':
+ result = result.read_all()
assert result.num_rows == 16
- 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():
- read = _table_read(limit=3)
- read._read_parallelism = 2
- splits = [_Split() for _ in range(4)]
- for index, split in enumerate(splits):
- split._native_split = index
-
- with patch(
- 'pypaimon.read.native_plan._prepare_native_read',
- return_value=lambda group: [_id_batch(group)]):
- result = read.to_arrow(splits)
-
- assert result.num_rows == 3
+ native.assert_called_once()
+ assert native.call_args.args[1] == list(range(16))
+ assert native.call_args.kwargs['parallelism'] == 16
+ assert native.call_args.kwargs['blob_parallelism'] == 16
-def test_parallel_native_batch_reader_shares_limit_across_readers():
- read = _table_read(limit=3)
- read._read_parallelism = 2
- splits = [_Split() for _ in range(4)]
+def test_native_batch_reader_setup_fallback_preserves_python_blob_worker_cap():
+ read = _table_read()
+ splits = [_Split() for _ in range(16)]
for index, split in enumerate(splits):
split._native_split = index
-
- with patch(
- '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
+ schema = pa.schema([('id', pa.int32())])
+ with patch('pypaimon.read.native_plan.native_read',
side_effect=RuntimeError('setup failed')) as native, \
+ patch.object(read, '_arrow_batch_generator',
return_value=iter([_id_batch([1])])) as python_read:
+ result = read.to_arrow_batch_reader(splits, parallelism=16,
blob_parallelism=16).read_all()
+ assert result.to_pydict() == {'id': [1]}
+ assert native.call_args.kwargs['parallelism'] == 16
+ assert native.call_args.kwargs['blob_parallelism'] == 16
+ python_read.assert_called_once_with(splits, schema, 4)
[email protected]('streaming', [False, True])
-def test_parallel_native_reader_setup_failure_falls_back(streaming):
[email protected]('ordered', [False, True])
+def test_native_reader_setup_failure_falls_back_before_batches(ordered):
read = _table_read()
- splits = [_Split() for _ in range(2)]
+ splits = [_Split(), _Split()]
for split in splits:
split._native_split = object()
-
- with patch('pypaimon.read.native_plan._prepare_native_read',
- side_effect=RuntimeError('setup failed')):
+ with patch('pypaimon.read.native_plan.native_read',
side_effect=RuntimeError('setup failed')) as native:
assert read._try_native_batches(
- splits,
- pa.schema([('id', pa.int32())]),
- parallelism=2,
- streaming=streaming,
- ) is None
+ splits, pa.schema([('id', pa.int32())]), parallelism=2,
ordered=ordered) is None
+ native.assert_called_once()
-def test_parallel_native_stream_setup_failure_closes_started_readers():
+def test_native_reader_qualifies_all_splits_before_starting_core():
read = _table_read()
- splits = [_Split() for _ in range(2)]
+ splits = [_Split(), _Split('unsupported.vortex')]
for split in splits:
split._native_split = object()
- started = Mock()
-
- read_splits = Mock(side_effect=[started, RuntimeError('setup failed')])
- with patch(
- 'pypaimon.read.native_plan._prepare_native_read',
- return_value=read_splits):
+ with patch('pypaimon.read.native_plan.native_read') as native:
assert read._try_native_batches(
- splits,
- pa.schema([('id', pa.int32())]),
- parallelism=2,
- streaming=True,
- ) is None
-
- started.close.assert_called_once_with()
+ splits, pa.schema([('id', pa.int32())]), parallelism=2) is None
+ native.assert_not_called()
-def test_parallel_native_stream_error_propagates():
[email protected]('method_name', ['to_arrow', 'to_arrow_batch_reader'])
+def
test_native_stream_error_closes_input_without_replaying_python(method_name):
read = _table_read()
- splits = [_Split() for _ in range(2)]
+ read._read_parallelism = 2
+ splits = [_Split(), _Split()]
for split in splits:
split._native_split = object()
+ closed = []
- def broken_stream(group):
- def batches():
+ def batches():
+ try:
+ yield _id_batch([1])
raise RuntimeError('stream failed')
- yield
- return batches()
-
- 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)
+ finally:
+ closed.append(True)
+
+ with patch('pypaimon.read.native_plan.native_read',
return_value=batches()) as native, \
+ patch.object(read, '_create_split_read',
side_effect=AssertionError('Python replay')):
+ with pytest.raises(RuntimeError, match='stream failed'):
+ result = getattr(read, method_name)(splits)
+ if method_name != 'to_arrow':
+ assert result.read_next_batch().column('id').to_pylist() == [1]
+ result.read_next_batch()
+ assert closed == [True]
+ native.assert_called_once()
def test_native_read_failure_falls_back():
@@ -1014,9 +745,9 @@ def
test_native_read_preserves_parallel_large_binary_blob_schema():
split._native_split = index
with patch(
- 'pypaimon.read.native_plan._prepare_native_read',
- return_value=lambda group: [pa.record_batch(
- [pa.array([bytes(group)], type=pa.large_binary())],
+ 'pypaimon.read.native_plan.native_read',
+ return_value=[pa.record_batch(
+ [pa.array([b'payload'], type=pa.large_binary())],
names=['payload'])]):
result = read.to_arrow(splits, parallelism=2)
@@ -1040,7 +771,6 @@ def
test_native_primary_key_blob_limit_passes_quota_to_rust(descriptor):
blob_parallelism=1)) == [batch]
native.assert_called_once()
assert native.call_args.kwargs['limit'] == 1
- assert not read._should_run_parallel([split, split], 2)
@pytest.mark.parametrize('descriptor', [False, True])
