This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit a74d86960bf6b949670db99edea392997bf259ed Author: Jingsong Lee <[email protected]> AuthorDate: Fri Sep 18 21:45:01 2026 +0800 [python] Add opt-in Rust native data reads (#9954) --- paimon-python/README.md | 21 ++ .../pypaimon/common/options/core_options.py | 12 + paimon-python/pypaimon/read/native_plan.py | 81 ++++- paimon-python/pypaimon/read/table_read.py | 321 ++++++++++++++++-- paimon-python/pypaimon/read/table_scan.py | 13 +- .../pypaimon/tests/data_evolution_formats_test.py | 4 +- .../pypaimon/tests/data_evolution_test.py | 4 +- .../pypaimon/tests/native_plan_integration_test.py | 373 ++++++++++++++++++++- paimon-python/pypaimon/tests/native_plan_test.py | 30 +- paimon-python/pypaimon/tests/native_read_test.py | 349 +++++++++++++++++++ paimon-python/pypaimon/tests/table_update_test.py | 3 +- 11 files changed, 1151 insertions(+), 60 deletions(-) diff --git a/paimon-python/README.md b/paimon-python/README.md index b3541b1bdb..c8ee07f1d3 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -124,6 +124,27 @@ Python planner for unsupported scans. New bindings preserve `plan.snapshot_id` even when pruning removes every split. Native explain output includes snapshot and split metadata; native pruning counters are not exposed. +To run both split planning and data-file reading in Rust, enable the independent +native-read option: + +```python +native_table = table.copy({"read.native.enabled": "true"}) +builder = native_table.new_read_builder().with_projection(["id", "name"]) +plan = builder.new_scan().plan() +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. + With Rust main's `Table.from_resolved_schema()` binding, filesystem and JDBC catalog tables preserve the Python table's resolved schema and complete effective options. Stale table objects, historical schemas, and `copy()` overrides or diff --git a/paimon-python/pypaimon/common/options/core_options.py b/paimon-python/pypaimon/common/options/core_options.py index 7df8ee0ef2..50e48ab3ca 100644 --- a/paimon-python/pypaimon/common/options/core_options.py +++ b/paimon-python/pypaimon/common/options/core_options.py @@ -686,6 +686,15 @@ class CoreOptions: "reader still reads the files.") ) + READ_NATIVE_ENABLED: ConfigOption[bool] = ( + ConfigOptions.key("read.native.enabled") + .boolean_type() + .default_value(False) + .with_description("Read data via pypaimon_rust and return PyArrow batches. " + "This also enables native split planning; unsupported " + "routes fall back to pypaimon.") + ) + CHANGELOG_PRODUCER: ConfigOption[ChangelogProducer] = ( ConfigOptions.key("changelog-producer") .enum_type(ChangelogProducer) @@ -1582,6 +1591,9 @@ class CoreOptions: def native_plan_enabled(self, default=None): return self.options.get(CoreOptions.SCAN_NATIVE_PLAN_ENABLED, default) + def native_read_enabled(self, default=None): + return self.options.get(CoreOptions.READ_NATIVE_ENABLED, default) + def changelog_producer(self, default=None): return self.options.get(CoreOptions.CHANGELOG_PRODUCER, default) diff --git a/paimon-python/pypaimon/read/native_plan.py b/paimon-python/pypaimon/read/native_plan.py index ed12a98948..6a15b41939 100644 --- a/paimon-python/pypaimon/read/native_plan.py +++ b/paimon-python/pypaimon/read/native_plan.py @@ -49,6 +49,13 @@ def native_runtime_available() -> bool: return hasattr(PaimonCatalog, 'get_table') and hasattr(Split, 'serialize') +def native_reader_available() -> bool: + """Whether pypaimon-rust exposes the batch-read API used by PyPaimon.""" + return (native_runtime_available() + and native_method_available('ReadBuilder', 'new_read') + and native_method_available('TableRead', 'read')) + + def native_family_search_modes_available() -> bool: """Whether Rust supports family-specific global-index search modes.""" return native_version_at_least(0, 4) @@ -260,25 +267,13 @@ def _restore_python_partition_paths(table, splits: List[Split]) -> None: for data_file, python_path in candidates: if data_file.file_name in bucket_files[bucket_path]: data_file.file_path = python_path + # The retained Rust split still points at its canonical path. + # Invalidate it so native reading cannot bypass this repair. + split._native_split = None -def native_plan( - table, - predicate: Optional[Predicate] = None, - limit: Optional[int] = None, - projection: Optional[List[str]] = None, - row_ranges: Optional[List[Tuple[int, int]]] = None, - incremental_range: Optional[Tuple[int, int]] = None, - row_position_slice: Optional[Tuple[int, int]] = None, - row_position_shard: Optional[Tuple[int, int]] = None) -> Plan: - """Plan with pypaimon_rust, preserving snapshot metadata. - - Native conversion or planning failures are handled by TableScan, which - falls back to the Python planner. - """ - if not native_runtime_available(): - raise RuntimeError( - "scan.native-plan.enabled needs pypaimon-rust>=0.3.0 (split planning API)") +def _native_read_builder(table): + """Reconstruct the Rust table and return a builder for the same schema.""" file_io_options = _resolved_schema_file_io_options(table) if file_io_options is not None: from pypaimon_rust.datafusion import Table @@ -307,12 +302,56 @@ def native_plan( branch = getattr(rt, 'branch', None) if not callable(branch) or branch() != table.current_branch(): raise RuntimeError("Native table did not resolve the requested branch") + return builder + + +def _configure_native_read_builder(builder, predicate, limit, projection): if projection is not None: builder = builder.with_projection(projection) if predicate is not None: builder = builder.with_filter(_predicate_to_native(predicate)) if limit is not None: builder = builder.with_limit(limit) + 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): + """Read Rust ``Split`` objects into PyArrow ``RecordBatch`` objects.""" + if not native_reader_available(): + raise RuntimeError( + "read.native.enabled needs the pypaimon-rust native reader API") + builder = _configure_native_read_builder( + _native_read_builder(table), predicate, limit, projection) + if blob_parallelism is not 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)) + + +def native_plan( + table, + predicate: Optional[Predicate] = None, + limit: Optional[int] = None, + projection: Optional[List[str]] = None, + row_ranges: Optional[List[Tuple[int, int]]] = None, + incremental_range: Optional[Tuple[int, int]] = None, + row_position_slice: Optional[Tuple[int, int]] = None, + row_position_shard: Optional[Tuple[int, int]] = None) -> Plan: + """Plan with pypaimon_rust, preserving snapshot metadata. + + Native conversion or planning failures are handled by TableScan, which + falls back to the Python planner. + """ + if not native_runtime_available(): + raise RuntimeError( + "scan.native-plan.enabled needs pypaimon-rust>=0.3.0 (split planning API)") + builder = _configure_native_read_builder( + _native_read_builder(table), predicate, limit, projection) if row_ranges is not None: builder = builder.with_row_ranges(row_ranges) scan = (builder.new_scan() if incremental_range is None @@ -327,6 +366,14 @@ def native_plan( # Trimmed primary keys decode per-file min/max keys (PK merge-on-read). kfields = table.trimmed_primary_keys_fields splits = [deserialize_split_v1(split.serialize(), pfields, kfields) for split in rust_splits] + if table.options.native_read_enabled(): + # Retain the opaque Rust split next to the Python metadata view. The + # normal planner/reader contract remains a Python Split list, while + # native reads can consume the exact Rust split without a second lossy + # conversion. Any Python split transformation creates a fresh object + # without this marker and thus safely falls back to the Python reader. + for split, rust_split in zip(splits, rust_splits): + split._native_split = rust_split _restore_python_partition_paths(table, splits) snapshot_id = getattr(rust_plan, 'snapshot_id', None) if callable(snapshot_id): diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index 5b8407874e..166b8d4aac 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import logging import os import threading from concurrent.futures import ThreadPoolExecutor, as_completed @@ -23,6 +24,7 @@ from typing import Any, Callable, Dict, Iterator, List, Optional import pandas import pyarrow +from pypaimon.common.options.core_options import CoreOptions from pypaimon.common.predicate import Predicate from pypaimon.common.predicate_json_parser import extract_referenced_fields from pypaimon.data.map_shared_shredding import map_selected_keys_field @@ -42,8 +44,19 @@ from pypaimon.schema.data_types import ( from pypaimon.table.row.offset_row import OffsetRow ROW_KIND_COLUMN = "_row_kind" +logger = logging.getLogger(__name__) _RECORD_BATCH_READER_FROM_STREAM = getattr( pyarrow.ipc.RecordBatchReader, "from_stream", None) +_NATIVE_READ_FILE_FORMATS = frozenset({ + CoreOptions.FILE_FORMAT_PARQUET, + CoreOptions.FILE_FORMAT_ORC, + CoreOptions.FILE_FORMAT_AVRO, + CoreOptions.FILE_FORMAT_ROW, + CoreOptions.FILE_FORMAT_MOSAIC, +}) +_NATIVE_READ_FILE_SUFFIXES = tuple( + '.%s' % file_format for file_format in _NATIVE_READ_FILE_FORMATS) +_NATIVE_BLOB_FILE_SUFFIX = '.blob' class _ClosableArrowBatchReader: @@ -113,12 +126,16 @@ 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.""" - # Cap on peak concurrent blob reads across the whole parallel read: split - # workers (P) each spin up blob_parallelism (B) blob threads, so peak - # connections ~= P*B. Shrink per-split B to keep the product bounded. + # Cap peak concurrent BLOB range reads across the whole parallel read: + # split workers (P) each allow blob_parallelism (B) in-flight reads, so + # peak connections ~= P*B. Shrink per-reader B to keep the product bounded. _MAX_TOTAL_BLOB_WORKERS = 64 def __init__( @@ -213,11 +230,17 @@ class TableRead: self, splits: List[Split], blob_parallelism: Optional[int] = None): - effective_bp = self._resolve_blob_parallelism(blob_parallelism) 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) + effective_bp = self._resolve_blob_parallelism(blob_parallelism) + native_batches = self._try_native_batches( + splits, schema, blob_parallelism=effective_bp) + if native_batches is not None: + batch_iterator = iter(native_batches) + else: + batch_iterator = self._arrow_batch_generator( + splits, schema, effective_bp) reader_type = pyarrow.ipc.RecordBatchReader reader = reader_type.from_batches(schema, batch_iterator) return reader, batch_iterator @@ -229,8 +252,18 @@ class TableRead: return pyarrow.schema([row_kind_field] + list(schema)) @staticmethod - def _try_to_pad_batch_by_schema(batch: pyarrow.RecordBatch, target_schema): - if batch.schema.names == target_schema.names: + def _try_to_pad_batch_by_schema( + batch: pyarrow.RecordBatch, + target_schema, + allow_type_cast: bool = False): + """Align a batch with the requested schema. + + Native batches stay strict so an outdated Rust BLOB representation + cannot be silently converted. Python query-auth masking is allowed + to cast internal dependency columns because masking intentionally may + change their logical type. + """ + if batch.schema.equals(target_schema, check_metadata=True): return batch columns = [] @@ -239,6 +272,12 @@ class TableRead: for field in target_schema: if field.name in batch.schema.names: col = batch.column(field.name) + if col.type != field.type: + if not allow_type_cast: + raise TypeError( + "Batch field '%s' has type %s, expected %s" % ( + field.name, col.type, field.type)) + col = col.cast(field.type) else: col = pyarrow.nulls(num_rows, type=field.type) columns.append(col) @@ -265,12 +304,11 @@ class TableRead: 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 - deferred BLOB resolution run serially when a limit may discard - rows, so payloads are not materialized from discarded splits. - blob_parallelism: number of threads for concurrent blob reads - within each batch. ``None`` or ``1`` (default) reads blobs - serially; ``>= 2`` uses a thread pool with ``pread`` for - concurrent ranged reads. GIL is released during I/O. On the + deferred BLOB resolution falls back to the Python reader when + a limit may discard rows, so discarded payloads are not read. + blob_parallelism: maximum concurrent blob range reads within each + split reader. ``None`` or ``1`` (default) reads blobs serially; + ``>= 2`` enables concurrent ranged reads. On the parallel path, peak blob threads (``parallelism`` * ``blob_parallelism``) are capped at ``_MAX_TOTAL_BLOB_WORKERS``; per-split ``blob_parallelism`` is @@ -282,21 +320,202 @@ class TableRead: 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, + ) + if native_batches is not None: + return self._batches_to_arrow(native_batches, schema) + 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_iterator = self._arrow_batch_generator( + splits, schema, effective_bp) + batch_reader = pyarrow.ipc.RecordBatchReader.from_batches( + schema, batch_iterator) table_list = [] for batch in iter(batch_reader.read_next_batch, None): if batch.num_rows == 0: continue - table_list.append(self._try_to_pad_batch_by_schema(batch, schema)) + table_list.append(self._try_to_pad_batch_by_schema( + batch, schema, allow_type_cast=True)) - if not table_list: - return pyarrow.Table.from_arrays([pyarrow.array([], type=field.type) for field in schema], schema=schema) - else: - return pyarrow.Table.from_batches(table_list) + return self._batches_to_arrow( + table_list, schema, allow_type_cast=True) + + @staticmethod + def _batches_to_arrow(batches, schema, allow_type_cast: bool = False): + batches = [ + TableRead._try_to_pad_batch_by_schema( + batch, schema, allow_type_cast=allow_type_cast) + for batch in batches + if batch.num_rows > 0 + ] + if not batches: + return pyarrow.Table.from_arrays( + [pyarrow.array([], type=field.type) for field in schema], + schema=schema) + return pyarrow.Table.from_batches(batches) + + def _try_native_batches( + self, + splits: List[Split], + schema: pyarrow.Schema, + parallelism: Optional[int] = None, + blob_parallelism: Optional[int] = None): + """Return Rust-read batches, or ``None`` when this read must fall back.""" + if not self.table.options.native_read_enabled(): + return None + # These Python-only output controls do not yet have native equivalents. + if self.include_row_kind or self.nested_name_paths: + return None + if self.table.options.file_format() not in _NATIVE_READ_FILE_FORMATS: + return None + if not splits: + return [] + if self._deferred_blob_limit_may_prune(splits): + return None + if self._native_schema_needs_python_fallback(schema): + return None + try: + from pypaimon.read.native_plan import native_read + except Exception as e: + logger.warning( + "Native read failed, falling back to the Python reader: %s", e) + return None + rust_splits = [] + for split in splits: + if isinstance(split, QueryAuthSplit): + return None + if not self._native_split_files_supported(split): + return None + rust_split = getattr(split, '_native_split', None) + if rust_split is None: + return None + rust_splits.append(rust_split) + if (parallelism is not None + and self._should_run_parallel(splits, parallelism)): + try: + return self._native_batches_parallel( + native_read, rust_splits, schema, parallelism, + blob_parallelism) + except _NativeReadSetupError as e: + logger.warning( + "Native read failed, falling back to the Python reader: %s", e) + return None + try: + read_kwargs = { + 'predicate': self.predicate, + 'limit': self.limit, + 'projection': [field.name for field in self.read_type], + } + if blob_parallelism is not None: + read_kwargs['blob_parallelism'] = blob_parallelism + batches = native_read(self.table, rust_splits, **read_kwargs) + except Exception as e: + logger.warning( + "Native read failed, falling back to the Python reader: %s", e) + return None + return self._convert_native_batches(batches, schema) + + def _native_batches_parallel( + self, native_read, rust_splits, schema, effective, + blob_parallelism): + """Read contiguous split groups with independent Rust readers.""" + workers = min(effective, len(rust_splits)) + base_size, larger_groups = divmod(len(rust_splits), workers) + 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_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, + native_read, + group, + schema, + remaining_state, + blob_parallelism, + ): 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 _read_native_split_group( + self, native_read, rust_splits, schema, remaining_state, + blob_parallelism): + if remaining_state.exhausted(): + return [] + try: + read_kwargs = { + 'predicate': self.predicate, + 'limit': self.limit, + 'projection': [field.name for field in self.read_type], + } + if blob_parallelism is not None: + read_kwargs['blob_parallelism'] = blob_parallelism + batches = native_read(self.table, rust_splits, **read_kwargs) + 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._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: + file_name = data_file.file_name.lower() + if ('.vector.' not in file_name + and not file_name.endswith(_NATIVE_READ_FILE_SUFFIXES) + and not file_name.endswith(_NATIVE_BLOB_FILE_SUFFIX)): + return False + return True + + def _convert_native_batches(self, batches, schema): + """Apply PyPaimon's exact output limit lazily to native batches.""" + remaining = self.limit + for batch in batches: + if batch.num_rows == 0: + continue + 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 def _arrow_batch_generator(self, splits: List[Split], schema: pyarrow.Schema, blob_parallelism: int = 1) -> Iterator[pyarrow.RecordBatch]: @@ -400,9 +619,9 @@ class TableRead: return max(1, cls._MAX_TOTAL_BLOB_WORKERS // workers) def _should_run_parallel( - self, - splits: List[Split], - effective: int, + self, + splits: List[Split], + effective: int, ) -> bool: """Decide whether to take the parallel read path. @@ -410,14 +629,59 @@ class TableRead: overhead, no behavior change). A single split is never parallelized since there is nothing to fan out across. """ - deferred_limit_may_prune = ( - self.limit is not None - and self._deferred_blob_fields - and not self._limit_covers_all_splits(splits) - ) + deferred_limit_may_prune = self._deferred_blob_limit_may_prune(splits) return (effective >= 2 and len(splits) >= 2 and not deferred_limit_may_prune) + def _deferred_blob_limit_may_prune(self, splits: List[Split]) -> bool: + return (self.limit is not None + and (self._deferred_blob_fields + or self._native_inline_blob_fields()) + and not self._limit_covers_all_splits(splits)) + + def _native_inline_blob_fields(self) -> set: + """Return configured BLOB fields that native reads resolve eagerly.""" + options = self.table.options + if options.blob_as_descriptor(): + return set() + inline_fields = ( + options.blob_descriptor_fields() + | options.blob_view_fields() + ) + read_names = { + field.name + for field in getattr(self, '_scan_read_type', self.read_type) + } + return inline_fields & read_names + + @classmethod + def _native_schema_needs_python_fallback(cls, schema: pyarrow.Schema) -> bool: + """Rust currently reads precision-zero timestamps as milliseconds.""" + return any( + cls._native_type_needs_python_fallback(field.type) + for field in schema + ) + + @classmethod + def _native_type_needs_python_fallback(cls, data_type) -> bool: + if pyarrow.types.is_timestamp(data_type): + return data_type.unit == 's' + if pyarrow.types.is_struct(data_type): + return any( + cls._native_type_needs_python_fallback(field.type) + for field in data_type + ) + if (pyarrow.types.is_list(data_type) + or pyarrow.types.is_large_list(data_type) + or pyarrow.types.is_fixed_size_list(data_type)): + return cls._native_type_needs_python_fallback(data_type.value_type) + if pyarrow.types.is_map(data_type): + return ( + cls._native_type_needs_python_fallback(data_type.key_type) + or cls._native_type_needs_python_fallback(data_type.item_type) + ) + return False + def _limit_covers_all_splits(self, splits: List[Split]) -> bool: """Return whether split metadata proves that LIMIT cannot drop rows.""" total_rows = 0 @@ -473,7 +737,8 @@ class TableRead: for batch in split_batches: if batch.num_rows == 0: continue - table_list.append(self._try_to_pad_batch_by_schema(batch, schema)) + table_list.append(self._try_to_pad_batch_by_schema( + batch, schema, allow_type_cast=True)) if not table_list: return pyarrow.Table.from_arrays( diff --git a/paimon-python/pypaimon/read/table_scan.py b/paimon-python/pypaimon/read/table_scan.py index 2a21527242..6d269957e6 100755 --- a/paimon-python/pypaimon/read/table_scan.py +++ b/paimon-python/pypaimon/read/table_scan.py @@ -57,6 +57,7 @@ _NATIVE_FORWARDED_OPTIONS = frozenset({ }) | _NATIVE_SEARCH_MODE_OPTIONS _NATIVE_PLAN_INDEPENDENT_OPTIONS = frozenset({ CoreOptions.BLOB_AS_DESCRIPTOR.key(), + CoreOptions.READ_NATIVE_ENABLED.key(), CoreOptions.READ_BATCH_SIZE.key(), CoreOptions.READ_PARALLELISM.key(), }) @@ -95,7 +96,7 @@ class TableScan: auth_result = self.__auth_query() # The native planner bypasses the auth-aware file scanner. Resolve auth # before selecting a planning backend. - if (auth_result is None and self.table.options.native_plan_enabled() + if (auth_result is None and self._native_requested() and self._native_plan_supported()): native = self._try_native_plan() if native is not None: @@ -105,6 +106,14 @@ class TableScan: plan = self.file_scanner.scan() return wrap_plan_with_auth(auth_result, plan) + def _native_requested(self) -> bool: + if self.table.options.native_plan_enabled(): + return True + if not self.table.options.native_read_enabled(): + return False + from pypaimon.read.native_plan import native_reader_available + return native_reader_available() + def _native_plan_supported(self) -> bool: # Any probe failure (e.g. a remote schema/metadata read) must fall back, not fail the scan. try: @@ -393,7 +402,7 @@ class TableScan: explain reports their split metadata without pruning counters. """ auth_result = self.__auth_query() - if (auth_result is None and self.table.options.native_plan_enabled() + if (auth_result is None and self._native_requested() and self._native_plan_supported()): native = self._try_native_plan() if native is not None: diff --git a/paimon-python/pypaimon/tests/data_evolution_formats_test.py b/paimon-python/pypaimon/tests/data_evolution_formats_test.py index c970387995..73057665ee 100644 --- a/paimon-python/pypaimon/tests/data_evolution_formats_test.py +++ b/paimon-python/pypaimon/tests/data_evolution_formats_test.py @@ -494,7 +494,9 @@ class DataEvolutionFormatsTest(unittest.TestCase): {'a': [1, 2, 3, 4], 'b': ['x', 'y', 'z', 'w'], 'c': [0.1, 0.2, 0.3, 0.4]}, schema=pa_schema) - self.assertEqual(actual, expect) + # Table scans are unordered; native planning may emit merged groups + # before raw groups. + self.assertEqual(actual.sort_by('a'), expect.sort_by('a')) # ------------------------------------------------------------------ # Blob-format data evolution diff --git a/paimon-python/pypaimon/tests/data_evolution_test.py b/paimon-python/pypaimon/tests/data_evolution_test.py index 8460f6ef2f..99af2e4933 100644 --- a/paimon-python/pypaimon/tests/data_evolution_test.py +++ b/paimon-python/pypaimon/tests/data_evolution_test.py @@ -497,7 +497,9 @@ class DataEvolutionTest(unittest.TestCase): 'f1': ['a'] * 100 + ['x'] + ['c'], 'f2': ['b'] * 100 + ['y'] + ['d'], }, schema=simple_pa_schema) - self.assertEqual(actual, expect) + # A table scan does not guarantee row-id order. Native planning may + # pack merge-required and raw groups in a different split order. + self.assertEqual(actual.sort_by('f0'), expect.sort_by('f0')) self.assertEqual(len(actual.schema), len(expect.schema), 'Merge read output column count must match schema') self.assertEqual(actual.schema.names, expect.schema.names, 'Merge read output column names must match schema') diff --git a/paimon-python/pypaimon/tests/native_plan_integration_test.py b/paimon-python/pypaimon/tests/native_plan_integration_test.py index 10cc5d3471..bbface6217 100644 --- a/paimon-python/pypaimon/tests/native_plan_integration_test.py +++ b/paimon-python/pypaimon/tests/native_plan_integration_test.py @@ -15,6 +15,9 @@ # specific language governing permissions and limitations # under the License. +import datetime +import os +import pickle import tempfile import unittest from unittest.mock import patch @@ -25,9 +28,10 @@ import pytest from pypaimon import CatalogFactory, Schema from pypaimon.globalindex.global_index_result import GlobalIndexResult from pypaimon.read.native_plan import ( - native_family_search_modes_available, native_method_available, + native_family_search_modes_available, native_method_available, native_read, + native_reader_available, ) -from pypaimon.table.row.blob import BlobDescriptor +from pypaimon.table.row.blob import BlobDescriptor, BlobRef from pypaimon.utils.range import Range @@ -92,6 +96,13 @@ class NativePlanIntegrationTest(unittest.TestCase): self.assertEqual( native_table.new_read_builder().explain().native_planned, expect_native) + @staticmethod + def _native_rows(builder, plan): + with patch( + 'pypaimon.read.table_read.TableRead._create_split_read', + side_effect=AssertionError('Python reader was used')): + return builder.new_read().to_arrow(plan.splits()).to_pylist() + def test_primary_key_matches_normal_plan(self): self.cat.create_table('default.pk_t', Schema.from_pyarrow_schema( self.schema, primary_keys=['k'], options={'bucket': '1'}), False) @@ -145,6 +156,102 @@ class NativePlanIntegrationTest(unittest.TestCase): self._write('ap_t', [{'k': 3, 'v': 'c'}]) self._assert_matches('ap_t') + @unittest.skipUnless(native_reader_available(), + "pypaimon-rust native reader API not installed") + def test_native_read_bypasses_python_split_reader(self): + self.cat.create_table( + 'default.native_read_t', Schema.from_pyarrow_schema(self.schema), False) + self._write('native_read_t', [ + {'k': 1, 'v': 'a'}, {'k': 2, 'v': 'b'}, {'k': 3, 'v': 'c'}]) + + table = self.cat.get_table('default.native_read_t').copy( + {'read.native.enabled': 'true'}) + builder = table.new_read_builder().with_projection(['k']).with_limit(2) + builder.with_filter(builder.new_predicate_builder().greater_or_equal('k', 2)) + plan = builder.new_scan().plan() + splits = [pickle.loads(pickle.dumps(split)) for split in plan.splits()] + + with patch( + 'pypaimon.read.table_read.TableRead._create_split_read', + side_effect=AssertionError('Python reader was used')): + read = builder.new_read() + rows = read.to_arrow(splits).to_pylist() + streamed = pa.Table.from_batches( + list(read.to_arrow_batch_reader(splits))).to_pylist() + + self.assertEqual(rows, [{'k': 2}, {'k': 3}]) + self.assertEqual(streamed, rows) + self.assertTrue(builder.explain().native_planned) + + @unittest.skipUnless(native_reader_available(), + "pypaimon-rust native reader API not installed") + def test_native_read_uses_effective_split_parallelism(self): + schema = pa.schema([ + ('k', pa.int64()), + ('v', pa.string()), + ('dt', pa.string()), + ]) + self.cat.create_table('default.native_parallel_t', + Schema.from_pyarrow_schema( + schema, + partition_keys=['dt']), False) + table = self.cat.get_table('default.native_parallel_t') + write_builder = table.new_batch_write_builder() + write = write_builder.new_write() + write.write_arrow(pa.Table.from_pylist([ + {'k': 1, 'v': 'a', 'dt': 'p1'}, + {'k': 2, 'v': 'b', 'dt': 'p2'}, + {'k': 3, 'v': 'c', 'dt': 'p3'}, + {'k': 4, 'v': 'd', 'dt': 'p4'}, + ], schema=schema)) + write_builder.new_commit().commit(write.prepare_commit()) + write.close() + + native_table = table.copy({ + 'read.native.enabled': 'true', + 'read.parallelism': '2', + }) + builder = native_table.new_read_builder() + 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, \ + patch( + 'pypaimon.read.table_read.TableRead._create_split_read', + side_effect=AssertionError('Python reader was used')): + rows = builder.new_read().to_arrow(plan.splits()).to_pylist() + + self.assertEqual(sorted(rows, key=lambda row: row['k']), [ + {'k': 1, 'v': 'a', 'dt': 'p1'}, + {'k': 2, 'v': 'b', 'dt': 'p2'}, + {'k': 3, 'v': 'c', 'dt': 'p3'}, + {'k': 4, 'v': 'd', 'dt': 'p4'}, + ]) + self.assertEqual(rust_reads.call_count, 2) + + @unittest.skipUnless(native_reader_available(), + "pypaimon-rust native reader API not installed") + def test_native_read_primary_key_matches_python(self): + self.cat.create_table('default.native_read_pk', Schema.from_pyarrow_schema( + self.schema, primary_keys=['k'], options={'bucket': '1'}), False) + self._write('native_read_pk', [{'k': 1, 'v': 'a'}, {'k': 2, 'v': 'old'}]) + self._write('native_read_pk', [{'k': 2, 'v': 'new'}, {'k': 3, 'v': 'c'}]) + + normal = self._plan_and_read('native_read_pk', native=False)[1] + table = self.cat.get_table('default.native_read_pk').copy( + {'read.native.enabled': 'true'}) + builder = table.new_read_builder() + plan = builder.new_scan().plan() + with patch( + 'pypaimon.read.table_read.TableRead._create_split_read', + side_effect=AssertionError('Python reader was used')): + native = sorted( + builder.new_read().to_arrow(plan.splits()).to_pylist(), + key=lambda row: row['k']) + + self.assertEqual(native, normal) + def test_append_distribution_matches_interleaved_partition_buckets(self): self.schema = pa.schema([('k', pa.int64()), ('v', pa.string()), ('p', pa.string())]) self.cat.create_table('default.interleaved_t', Schema.from_pyarrow_schema( @@ -238,7 +345,7 @@ class NativePlanIntegrationTest(unittest.TestCase): self._assert_matches('de_t') native_table = self.cat.get_table('default.de_t').copy( - {'scan.native-plan.enabled': 'true'}) + {'read.native.enabled': 'true'}) predicate = native_table.new_read_builder().new_predicate_builder().equal( 'v', 'b2') builder = (native_table.new_read_builder() @@ -246,7 +353,7 @@ class NativePlanIntegrationTest(unittest.TestCase): .with_filter(predicate) .with_limit(1)) plan = builder.new_scan().plan() - rows = builder.new_read().to_arrow(plan.splits()).to_pylist() + rows = self._native_rows(builder, plan) self.assertEqual(rows, [{'k': 2}]) self.assertTrue(builder.explain().native_planned) @@ -255,6 +362,9 @@ class NativePlanIntegrationTest(unittest.TestCase): .with_projection(['media.camera']) .with_limit(1)) blob_plan = blob_builder.new_scan().plan() + # A dotted top-level field currently uses PyPaimon's nested-projection + # machinery, so this part deliberately exercises the documented + # Python fallback while the scalar projection above is native. blob_rows = blob_builder.new_read().to_arrow( blob_plan.splits()).to_pylist() self.assertEqual(blob_rows, [{'media.camera': b'a'}]) @@ -279,6 +389,241 @@ class NativePlanIntegrationTest(unittest.TestCase): ) self.assertTrue(descriptor_builder.explain().native_planned) + @unittest.skipUnless(native_reader_available(), + "pypaimon-rust native reader API not installed") + def test_native_read_data_evolution_partial_appends(self): + schema = pa.schema([ + ('f0', pa.int32()), + ('f1', pa.string()), + ('f2', pa.string()), + ]) + self.cat.create_table('default.de_order_t', Schema.from_pyarrow_schema( + schema, options={ + 'row-tracking.enabled': 'true', + 'data-evolution.enabled': 'true', + }), False) + table = self.cat.get_table('default.de_order_t') + write_builder = table.new_batch_write_builder() + + base = write_builder.new_write() + base.write_arrow(pa.Table.from_pydict({ + 'f0': [0, 1], + 'f1': ['z', 'a'], + 'f2': ['q', 'b'], + }, schema=schema)) + write_builder.new_commit().commit(base.prepare_commit()) + base.close() + + left = write_builder.new_write().with_write_type(['f0', 'f1']) + right = write_builder.new_write().with_write_type(['f2']) + left.write_arrow(pa.Table.from_pydict( + {'f0': [2], 'f1': ['x']}, + schema=pa.schema([('f0', pa.int32()), ('f1', pa.string())]))) + right.write_arrow(pa.Table.from_pydict( + {'f2': ['y']}, schema=pa.schema([('f2', pa.string())]))) + messages = left.prepare_commit() + right.prepare_commit() + for message in messages: + for data_file in message.new_files: + data_file.first_row_id = 2 + write_builder.new_commit().commit(messages) + left.close() + right.close() + + native_table = table.copy({'read.native.enabled': 'true'}) + builder = native_table.new_read_builder() + plan = builder.new_scan().plan() + rows = self._native_rows(builder, plan) + + self.assertEqual(sorted(rows, key=lambda row: row['f0']), [ + {'f0': 0, 'f1': 'z', 'f2': 'q'}, + {'f0': 1, 'f1': 'a', 'f2': 'b'}, + {'f0': 2, 'f1': 'x', 'f2': 'y'}, + ]) + self.assertTrue(builder.explain().native_planned) + + @unittest.skipUnless(native_reader_available(), + "pypaimon-rust native reader API not installed") + def test_native_read_data_evolution_blob_parallelism(self): + schema = pa.schema([ + ('id', pa.int32()), + ('img', pa.large_binary()), + ]) + self.cat.create_table('default.native_blob_t', Schema.from_pyarrow_schema( + schema, options={ + 'row-tracking.enabled': 'true', + 'data-evolution.enabled': 'true', + }), False) + table = self.cat.get_table('default.native_blob_t') + payloads = [b'a', b'bb', b'ccc'] + write = table.new_batch_write_builder().new_write() + write.write_arrow(pa.Table.from_pydict({ + 'id': [1, 2, 3], + 'img': payloads, + }, schema=schema)) + table.new_batch_write_builder().new_commit().commit(write.prepare_commit()) + write.close() + + native_table = table.copy({'read.native.enabled': 'true'}) + builder = native_table.new_read_builder().with_projection(['id', 'img']) + plan = builder.new_scan().plan() + self.assertTrue(any( + data_file.file_name.endswith('.blob') + for split in plan.splits() + for data_file in split.files + )) + + with patch( + 'pypaimon.read.table_read.TableRead._create_split_read', + side_effect=AssertionError('Python reader was used')): + table_result = builder.new_read().to_arrow( + plan.splits(), parallelism=1, blob_parallelism=3) + batch_table = builder.new_read().to_arrow_batch_reader( + plan.splits(), blob_parallelism=2).read_all() + + expected = {'id': [1, 2, 3], 'img': payloads} + self.assertEqual(table_result.schema, schema) + self.assertEqual(batch_table.schema, schema) + self.assertEqual(table_result.to_pydict(), expected) + self.assertEqual(batch_table.to_pydict(), expected) + self.assertTrue(builder.explain().native_planned) + + @unittest.skipUnless(native_reader_available(), + "pypaimon-rust native reader API not installed") + def test_native_read_pruning_limit_defers_blob_payload_io(self): + schema = pa.schema([ + ('id', pa.int32()), + ('payload', pa.large_binary()), + ]) + self.cat.create_table('default.native_blob_limit_t', + Schema.from_pyarrow_schema(schema, options={ + 'row-tracking.enabled': 'true', + 'data-evolution.enabled': 'true', + }), False) + table = self.cat.get_table('default.native_blob_limit_t') + write = table.new_batch_write_builder().new_write() + write.write_arrow(pa.Table.from_pydict({ + 'id': [1, 2, 3], + 'payload': [b'a', b'bb', b'ccc'], + }, schema=schema)) + table.new_batch_write_builder().new_commit().commit(write.prepare_commit()) + write.close() + + native_table = table.copy({'read.native.enabled': 'true'}) + builder = native_table.new_read_builder().with_limit(1) + plan = builder.new_scan().plan() + fetched = [] + original_to_data = BlobRef.to_data + + def tracked_to_data(blob): + fetched.append(blob) + return original_to_data(blob) + + with patch.object(BlobRef, 'to_data', tracked_to_data), patch( + 'pypaimon.read.native_plan.native_read', + return_value=[]) as native: + result = builder.new_read().to_arrow( + plan.splits(), parallelism=1) + + native.assert_not_called() + self.assertEqual(result.to_pydict(), {'id': [1], 'payload': [b'a']}) + self.assertEqual(len(fetched), 1) + + @unittest.skipUnless(native_reader_available(), + "pypaimon-rust native reader API not installed") + def test_native_read_pruning_limit_defers_descriptor_blob_payload_io(self): + schema = pa.schema([ + ('id', pa.int32()), + ('payload', pa.large_binary()), + ]) + with tempfile.TemporaryDirectory() as payload_dir: + first_path = os.path.join(payload_dir, 'first') + with open(first_path, 'wb') as output: + output.write(b'first') + missing_path = os.path.join(payload_dir, 'missing') + descriptors = [ + BlobDescriptor( + 'file://' + first_path, 0, 5).serialize(), + BlobDescriptor( + 'file://' + missing_path, 0, 7).serialize(), + ] + + self.cat.create_table( + 'default.native_descriptor_limit_t', + Schema.from_pyarrow_schema(schema, options={ + 'row-tracking.enabled': 'true', + 'data-evolution.enabled': 'true', + 'blob-descriptor-field': 'payload', + }), False) + table = self.cat.get_table('default.native_descriptor_limit_t') + write = table.new_batch_write_builder().new_write() + write.write_arrow(pa.Table.from_pydict({ + 'id': [1, 2], + 'payload': descriptors, + }, schema=schema)) + table.new_batch_write_builder().new_commit().commit( + write.prepare_commit()) + write.close() + + native_table = table.copy({'read.native.enabled': 'true'}) + builder = native_table.new_read_builder().with_limit(1) + plan = builder.new_scan().plan() + fetched = [] + original_to_data = BlobRef.to_data + + def tracked_to_data(blob): + fetched.append(blob) + return original_to_data(blob) + + with patch.object(BlobRef, 'to_data', tracked_to_data), patch( + 'pypaimon.read.native_plan.native_read', + return_value=[]) as native: + result = builder.new_read().to_arrow(plan.splits()) + + native.assert_not_called() + self.assertEqual( + result.to_pydict(), {'id': [1], 'payload': [b'first']}) + self.assertEqual(len(fetched), 1) + + @unittest.skipUnless(native_reader_available(), + "pypaimon-rust native reader API not installed") + def test_native_read_falls_back_for_precision_zero_timestamps(self): + cases = [ + ('timestamp', pa.timestamp('s'), datetime.datetime(1970, 1, 1)), + ('timestamp_ltz', pa.timestamp('s', tz='UTC'), + datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc)), + ] + for name, timestamp_type, first_value in cases: + with self.subTest(name=name): + schema = pa.schema([ + ('id', pa.int32()), + ('ts', timestamp_type), + ]) + table_name = 'native_%s_zero_t' % name + self.cat.create_table( + 'default.%s' % table_name, + Schema.from_pyarrow_schema(schema), False) + table = self.cat.get_table('default.%s' % table_name) + write = table.new_batch_write_builder().new_write() + write.write_arrow(pa.Table.from_pydict({ + 'id': [1, 2], + 'ts': [first_value, first_value], + }, schema=schema)) + table.new_batch_write_builder().new_commit().commit( + write.prepare_commit()) + write.close() + + native_table = table.copy({'read.native.enabled': 'true'}) + builder = native_table.new_read_builder() + plan = builder.new_scan().plan() + with patch( + 'pypaimon.read.native_plan.native_read', + return_value=[]) as native: + result = builder.new_read().to_arrow(plan.splits()) + + native.assert_not_called() + self.assertEqual(result.schema, schema) + self.assertEqual(result.num_rows, 2) + @unittest.skipUnless(_has_native_row_ranges(), "pypaimon_rust row-range API not installed") def test_data_evolution_global_index_row_ranges(self): @@ -293,7 +638,7 @@ class NativePlanIntegrationTest(unittest.TestCase): {'k': 3, 'v': 'c'}, ]) table = self.cat.get_table('default.de_range_t').copy( - {'scan.native-plan.enabled': 'true'}) + {'read.native.enabled': 'true'}) builder = table.new_read_builder() scan = builder.new_scan().with_global_index_result( GlobalIndexResult.from_range(Range(1, 1))) @@ -302,7 +647,7 @@ class NativePlanIntegrationTest(unittest.TestCase): with patch.object( scan.file_scanner, 'scan', side_effect=AssertionError("fallback")): plan = scan.plan() - rows = builder.new_read().to_arrow(plan.splits()).to_pylist() + rows = self._native_rows(builder, plan) self.assertEqual(rows, [{'k': 2, 'v': 'b'}]) self.assertEqual( @@ -471,6 +816,22 @@ class NativePlanIntegrationTest(unittest.TestCase): self._assert_matches('pt_t') + native_table = t.copy({ + 'scan.native-plan.enabled': 'true', + 'read.native.enabled': 'true', + }) + builder = native_table.new_read_builder() + plan = builder.new_scan().plan() + with patch('pypaimon.read.native_plan.native_read', + wraps=native_read) as read: + rows = builder.new_read().to_arrow(plan.splits()).to_pylist() + read.assert_not_called() + self.assertEqual(sorted(rows, key=lambda row: row['k']), [ + {'k': 1, 'p': 'a/b'}, + {'k': 2, 'p': 'a/b'}, + {'k': 3, 'p': 'c'}, + ]) + def test_explain_reflects_native_plan(self): self.cat.create_table( 'default.ex_t', Schema.from_pyarrow_schema(self.schema), False) diff --git a/paimon-python/pypaimon/tests/native_plan_test.py b/paimon-python/pypaimon/tests/native_plan_test.py index b7892c1185..b62177d8bd 100644 --- a/paimon-python/pypaimon/tests/native_plan_test.py +++ b/paimon-python/pypaimon/tests/native_plan_test.py @@ -50,6 +50,7 @@ def _scan(native_enabled, file_scanner): scan = TableScan.__new__(TableScan) scan.table = Mock() scan.table.options.native_plan_enabled.return_value = native_enabled + scan.table.options.native_read_enabled.return_value = False scan.table.options.options.contains_key.return_value = False # no time-travel scan.table.options.options.contains.return_value = False # no incremental scan.table.options.merge_engine.return_value = None # not first-row @@ -145,9 +146,13 @@ class NativePlanTest(unittest.TestCase): self.assertIsNone(_resolved_schema_file_io_options(table)) def test_switch_defaults_off(self): - self.assertFalse(CoreOptions(Options({})).native_plan_enabled()) + defaults = CoreOptions(Options({})) + self.assertFalse(defaults.native_plan_enabled()) + self.assertFalse(defaults.native_read_enabled()) self.assertTrue( CoreOptions(Options({"scan.native-plan.enabled": "true"})).native_plan_enabled()) + self.assertTrue( + CoreOptions(Options({"read.native.enabled": "true"})).native_read_enabled()) def test_catalogless_standard_file_io_options_are_preserved(self): from pypaimon.catalog.catalog_environment import CatalogEnvironment @@ -173,6 +178,19 @@ class NativePlanTest(unittest.TestCase): self.assertIs(scan.plan(), sentinel) fs.scan.assert_called_once_with() + def test_native_read_switch_also_requests_native_plan(self): + fs = Mock(partition_key_predicate=None) + scan = _scan(native_enabled=False, file_scanner=fs) + scan.table.options.native_read_enabled.return_value = True + expected = Plan([], 1) + with patch('pypaimon.read.native_plan.native_reader_available', + return_value=True), \ + patch('pypaimon.read.native_plan.native_plan', + return_value=expected) as np: + self.assertEqual(scan.plan(), expected) + np.assert_called_once() + fs.scan.assert_not_called() + def test_plan_routes_to_native_and_prunes_partitions(self): # Native planner returns every partition; the predicate keeps only [2026, 7]. keep = Mock(partition=Mock(values=[2026, 7])) @@ -664,6 +682,7 @@ class NativePlanTest(unittest.TestCase): ) split = Mock( partition=Mock(values=['a/b']), bucket=0, files=[data_file]) + split._native_split = object() _restore_python_partition_paths(table, [split]) @@ -671,6 +690,7 @@ class NativePlanTest(unittest.TestCase): data_file.file_path, '/warehouse/t/p=a/b/bucket-0/data.parquet', ) + self.assertIsNone(split._native_split) def test_partition_path_keeps_existing_rust_path(self): table = Mock(partition_keys=['p']) @@ -762,11 +782,13 @@ class NativePlanTest(unittest.TestCase): with patch.dict(sys.modules, {'pypaimon_rust': fake_mod, 'pypaimon_rust.datafusion': fake_df}), \ patch('pypaimon.read.native_plan._catalog_options', return_value={}), \ - patch('pypaimon.read.native_plan.deserialize_split_v1', - return_value='decoded') as des: + patch('pypaimon.read.native_plan.deserialize_split_v1') as des: + decoded = Mock() + des.return_value = decoded result = native_plan(table, row_ranges=[(1, 2)]) - self.assertEqual(result.splits(), ['decoded']) + self.assertEqual(result.splits(), [decoded]) + self.assertIs(decoded._native_split, split) self.assertEqual(result.snapshot_id, 3) rt.new_read_builder.assert_called_once_with({ CoreOptions.SOURCE_SPLIT_TARGET_SIZE.key(): '1024', diff --git a/paimon-python/pypaimon/tests/native_read_test.py b/paimon-python/pypaimon/tests/native_read_test.py new file mode 100644 index 0000000000..11282480b0 --- /dev/null +++ b/paimon-python/pypaimon/tests/native_read_test.py @@ -0,0 +1,349 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import threading +from unittest.mock import Mock, patch + +import pyarrow as pa +import pytest + +from pypaimon.read.table_read import TableRead +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 _table_read(limit=None): + read = TableRead.__new__(TableRead) + read.table = Mock() + read.table.options.native_read_enabled.return_value = True + read.table.options.file_format.return_value = 'parquet' + read.table.options.blob_as_descriptor.return_value = False + read.table.options.blob_descriptor_fields.return_value = set() + read.table.options.blob_view_fields.return_value = set() + read.predicate = None + read.read_type = [DataField(0, 'id', AtomicType('INT'))] + read.include_row_kind = False + read.nested_name_paths = None + read.limit = limit + read._read_parallelism = 1 + read._deferred_blob_fields = set() + read._predicate_extra_fields = [] + read._output_column_names = ['id'] + return read + + +def _blob_table_read(limit=None): + read = _table_read(limit) + read.read_type = [DataField(0, 'payload', AtomicType('BLOB'))] + read._output_column_names = ['payload'] + return read + + +def _id_batch(values): + return pa.record_batch( + [pa.array(values, type=pa.int32())], names=['id']) + + +def test_native_read_consumes_retained_rust_splits_and_enforces_limit(): + read = _table_read(limit=2) + first, second = _Split(), _Split() + first._native_split = object() + second._native_split = object() + batches = [_id_batch([1, 2, 3])] + + with patch( + 'pypaimon.read.native_plan.native_read', + return_value=batches) as native: + result = read.to_arrow([first, second]) + + assert result.to_pydict() == {'id': [1, 2]} + native.assert_called_once_with( + read.table, + [first._native_split, second._native_split], + predicate=None, + limit=2, + projection=['id'], + blob_parallelism=1, + ) + + +def test_native_read_falls_back_for_transformed_python_split(): + read = _table_read() + schema = pa.schema([('id', pa.int32())]) + split = _Split() + + with patch('pypaimon.read.native_plan.native_read') as native: + assert read._try_native_batches([split], schema) is None + + native.assert_not_called() + + +def test_native_read_uses_effective_parallelism_from_table_option(): + read = _table_read() + read._read_parallelism = 2 + splits = [_Split() for _ in range(4)] + for index, split in enumerate(splits): + split._native_split = index + + barrier = threading.Barrier(2) + worker_names = set() + worker_names_lock = threading.Lock() + + def read_group(table, rust_splits, **kwargs): + with worker_names_lock: + worker_names.add(threading.current_thread().name) + barrier.wait(timeout=5) + return [_id_batch(rust_splits)] + + with patch('pypaimon.read.native_plan.native_read', + side_effect=read_group) as native: + 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]] + assert len(worker_names) == 2 + + +def test_native_read_runtime_parallelism_overrides_table_option(): + read = _table_read() + read._read_parallelism = 4 + splits = [_Split() for _ in range(4)] + 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 native.call_count == 2 + + +def test_native_read_caps_blob_parallelism_across_split_readers(): + 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.native_read', + side_effect=lambda table, group, **kwargs: [ + _id_batch(group)]) as native: + 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} + + +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.native_read', + side_effect=lambda table, group, **kwargs: [ + _id_batch(group)]): + result = read.to_arrow(splits) + + assert result.num_rows == 3 + + +def test_parallel_native_reader_setup_failure_falls_back(): + read = _table_read() + splits = [_Split() for _ in range(2)] + for split in splits: + split._native_split = object() + + with patch('pypaimon.read.native_plan.native_read', + side_effect=RuntimeError('setup failed')): + assert read._try_native_batches( + splits, pa.schema([('id', pa.int32())]), parallelism=2) is None + + +def test_parallel_native_stream_error_propagates(): + read = _table_read() + splits = [_Split() for _ in range(2)] + for split in splits: + split._native_split = object() + + def broken_stream(table, group, **kwargs): + def batches(): + raise RuntimeError('stream failed') + yield + return batches() + + with patch('pypaimon.read.native_plan.native_read', + side_effect=broken_stream), \ + pytest.raises(RuntimeError, match='stream failed'): + read._try_native_batches( + splits, pa.schema([('id', pa.int32())]), parallelism=2) + + +def test_native_read_failure_falls_back(): + read = _table_read() + schema = pa.schema([('id', pa.int32())]) + split = _Split() + split._native_split = object() + + with patch('pypaimon.read.native_plan.native_read', + side_effect=RuntimeError('unsupported')): + assert read._try_native_batches([split], schema) is None + + +def test_native_read_falls_back_for_unsupported_file_format(): + read = _table_read() + read.table.options.file_format.return_value = 'vortex' + schema = pa.schema([('id', pa.int32())]) + split = _Split() + split._native_split = object() + + with patch('pypaimon.read.native_plan.native_read') as native: + assert read._try_native_batches([split], schema) is None + + native.assert_not_called() + + +def test_native_read_falls_back_for_unsupported_dedicated_file(): + read = _table_read() + schema = pa.schema([('id', pa.int32())]) + split = _Split('camera.video') + split._native_split = object() + + with patch('pypaimon.read.native_plan.native_read') as native: + assert read._try_native_batches([split], schema) is None + + native.assert_not_called() + + +def test_native_read_supports_blob_file_and_forwards_parallelism(): + read = _table_read() + schema = pa.schema([('id', pa.int32())]) + split = _Split('picture.blob') + split._native_split = object() + + with patch( + 'pypaimon.read.native_plan.native_read', + return_value=[_id_batch([1])]) as native: + batches = list(read._try_native_batches( + [split], schema, blob_parallelism=3)) + + assert batches[0].column('id').to_pylist() == [1] + assert native.call_args.kwargs['blob_parallelism'] == 3 + + +def test_native_read_preserves_large_binary_blob_schema_for_batch_reader(): + read = _blob_table_read() + split = _Split('payload.blob') + split._native_split = object() + native_batch = pa.record_batch( + [pa.array([b'a'], type=pa.large_binary())], names=['payload']) + + with patch( + 'pypaimon.read.native_plan.native_read', + return_value=[native_batch]): + result = read.to_arrow_batch_reader([split]).read_all() + + assert result.schema == pa.schema([('payload', pa.large_binary())]) + assert result.to_pydict() == {'payload': [b'a']} + + +def test_native_read_does_not_cast_legacy_binary_blob_schema(): + batch = pa.record_batch( + [pa.array([b'a'], type=pa.binary())], names=['payload']) + + with pytest.raises( + TypeError, + match="Batch field 'payload' has type binary, expected large_binary"): + TableRead._try_to_pad_batch_by_schema( + batch, pa.schema([('payload', pa.large_binary())])) + + +def test_native_read_preserves_parallel_large_binary_blob_schema(): + read = _blob_table_read() + splits = [_Split('payload.blob') for _ in range(2)] + for index, split in enumerate(splits): + split._native_split = index + + with patch( + 'pypaimon.read.native_plan.native_read', + side_effect=lambda table, group, **kwargs: [pa.record_batch( + [pa.array([bytes(group)], type=pa.large_binary())], + names=['payload'])]): + result = read.to_arrow(splits, parallelism=2) + + assert result.schema == pa.schema([('payload', pa.large_binary())]) + + +def test_native_read_defers_to_python_for_pruning_blob_limit(): + read = _blob_table_read(limit=1) + read._deferred_blob_fields = {'payload'} + split = _Split('payload.blob') + split._native_split = object() + split.merged_row_count = Mock(return_value=2) + + with patch( + 'pypaimon.read.native_plan.native_read') as native: + assert read._try_native_batches( + [split], pa.schema([('payload', pa.large_binary())]), + blob_parallelism=1) is None + + native.assert_not_called() + + +def test_native_read_defers_to_python_for_pruning_descriptor_blob_limit(): + read = _blob_table_read(limit=1) + read.table.options.blob_descriptor_fields.return_value = {'payload'} + split = _Split('payload.parquet') + split._native_split = object() + split.merged_row_count = Mock(return_value=2) + + with patch('pypaimon.read.native_plan.native_read') as native: + assert read._try_native_batches( + [split], pa.schema([('payload', pa.large_binary())]), + blob_parallelism=1) is None + + native.assert_not_called() + + [email protected]('data_type', [ + pa.timestamp('s'), + pa.timestamp('s', tz='UTC'), +]) +def test_native_read_falls_back_for_precision_zero_timestamps(data_type): + read = _table_read() + split = _Split() + split._native_split = object() + + with patch('pypaimon.read.native_plan.native_read') as native: + assert read._try_native_batches( + [split], pa.schema([('ts', data_type)])) is None + + native.assert_not_called() diff --git a/paimon-python/pypaimon/tests/table_update_test.py b/paimon-python/pypaimon/tests/table_update_test.py index 130dd53686..092990ab98 100644 --- a/paimon-python/pypaimon/tests/table_update_test.py +++ b/paimon-python/pypaimon/tests/table_update_test.py @@ -288,7 +288,8 @@ class _TableUpdateTestBase(DataEvolutionTestBase): self._apply_commit(tc, msgs, cid) tc.close() - result = self._read_all(table) + # Table scans are unordered; compare in stable business-key order. + result = self._read_all(table).sort_by('id') self.assertEqual( [25, 30, 99, 99, 99], result['age'].to_pylist(),
