From 40383074dbf2d56780e9d71058ec5116c6ea05b9 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Fri, 23 Oct 2026 19:56:21 +0800 Subject: [PATCH] [python] Dispatch captured row ID and bucket execution --- paimon-python/dev/test_packaging.py | 2 +- .../execution/native/write/row_id_plan.py | 35 ++ .../read/scanner/bucket_select_converter.py | 4 +- .../execution/python/write/dynamic_bucket.py | 333 ++++++++++++++++- .../execution/python/write/input_routing.py | 330 ++++++++++++++++- .../execution/python/write/row_id_plan.py | 109 +++++- .../python/write/row_key_extractor.py | 14 +- .../pypaimon/index/dynamic_bucket.py | 334 +---------------- .../pypaimon/ray/data_evolution_merge_join.py | 138 +------ paimon-python/pypaimon/read/plan.py | 39 +- .../dynamic_bucket_execution_boundary_test.py | 61 +++- .../native_dynamic_input_routing_test.py | 10 +- .../tests/native_input_routing_test.py | 2 +- .../tests/native_ray_row_id_plan_test.py | 104 ++++++ .../tests/noncompact_null_row_test.py | 2 +- .../pypaimon/tests/pushdown_bucket_test.py | 2 +- .../write/distributed_input_execution_test.py | 11 +- .../tests/write/dynamic_bucket_test.py | 12 +- .../tests/write/dynamic_input_routing_test.py | 4 +- .../tests/write/native_input_routing_test.py | 2 +- .../tests/write/row_id_routing_test.py | 275 ++++++++++++++ .../pypaimon/write/backend_dispatch.py | 14 +- paimon-python/pypaimon/write/input_routing.py | 343 ------------------ paimon-python/pypaimon/write/row_id_plan.py | 47 ++- 24 files changed, 1327 insertions(+), 900 deletions(-) create mode 100644 paimon-python/pypaimon/tests/write/row_id_routing_test.py delete mode 100644 paimon-python/pypaimon/write/input_routing.py diff --git a/paimon-python/dev/test_packaging.py b/paimon-python/dev/test_packaging.py index 428fbac0eef6..02fe0efdca8e 100644 --- a/paimon-python/dev/test_packaging.py +++ b/paimon-python/dev/test_packaging.py @@ -211,7 +211,7 @@ def test_distribution_contents(self): "pypaimon/globalindex/vindex_reader.py", "pypaimon/write/bucket_planning.py", "pypaimon/write/commit_options.py", - "pypaimon/write/input_routing.py", + "pypaimon/execution/python/write/input_routing.py", ) for name in execution_files: self.assertIn(name, names) diff --git a/paimon-python/pypaimon/execution/native/write/row_id_plan.py b/paimon-python/pypaimon/execution/native/write/row_id_plan.py index d040b5b2059d..3e2eea650b16 100644 --- a/paimon-python/pypaimon/execution/native/write/row_id_plan.py +++ b/paimon-python/pypaimon/execution/native/write/row_id_plan.py @@ -64,3 +64,38 @@ def python_groups(table, plan): return [(start, end, deserialize_split_v1(split.serialize(), _partition_fields(table), table.trimmed_primary_keys_fields)) for start, end, split in plan.groups()] + + +class NativeRowIdPlan: + """Transport Arrow row IDs to Core without reading or replanning files.""" + + def __init__(self, table, plan, native=None, groups=None): + self.table = table + self.plan = plan + self._native = native + self.groups = groups if groups is not None else python_groups(table, self._ensure_native()) + + def _ensure_native(self): + if self._native is None: + from pypaimon_rust.datafusion import RowIdPlan + self._native = RowIdPlan(native_plan(self.table, self.plan)) + return self._native + + def route(self, row_ids): + return self._ensure_native().route(row_ids) + + def route_by_bucket(self, row_ids): + return self._ensure_native().route_by_bucket(row_ids) + + def select(self, row_ids): + from pypaimon.execution.native.read.plan import _from_native_plan + selected = _from_native_plan(self.table, self._ensure_native().select(row_ids)) + return self.plan._with_splits(selected.splits()) + + def for_plan(self, plan, groups): + return NativeRowIdPlan(self.table, plan, groups=groups) + + def __getstate__(self): + state = self.__dict__.copy() + state['_native'] = None + return state diff --git a/paimon-python/pypaimon/execution/python/read/scanner/bucket_select_converter.py b/paimon-python/pypaimon/execution/python/read/scanner/bucket_select_converter.py index 279adbc70999..eb7ad1d00e2b 100644 --- a/paimon-python/pypaimon/execution/python/read/scanner/bucket_select_converter.py +++ b/paimon-python/pypaimon/execution/python/read/scanner/bucket_select_converter.py @@ -30,7 +30,7 @@ happen — that would be silent data loss. The hashing routine reuses ``_hash_bytes_by_words`` / -``_bucket_from_hash`` from ``pypaimon.write.input_routing`` — the same +``_bucket_from_hash`` from ``pypaimon.execution.python.write.input_routing`` — the same code path the writer uses to assign rows to buckets. Reusing it (rather than copying) is what guarantees read/write hash agreement in the face of future routine changes. @@ -70,7 +70,7 @@ from pypaimon.schema.data_types import DataField from pypaimon.table.row.generic_row import GenericRow, GenericRowSerializer from pypaimon.table.row.internal_row import InternalRow, RowKind -from pypaimon.write.input_routing import ( +from pypaimon.execution.python.write.input_routing import ( _bucket_from_hash, _hash_bytes_by_words, ) diff --git a/paimon-python/pypaimon/execution/python/write/dynamic_bucket.py b/paimon-python/pypaimon/execution/python/write/dynamic_bucket.py index 759412e3d827..8a38ac72a2af 100644 --- a/paimon-python/pypaimon/execution/python/write/dynamic_bucket.py +++ b/paimon-python/pypaimon/execution/python/write/dynamic_bucket.py @@ -16,14 +16,17 @@ # limitations under the License. ################################################################################ -"""Python writer lifecycle for persistent dynamic-bucket HASH indexes.""" +"""Python dynamic-bucket allocation and persistent HASH index execution.""" +import random import struct import uuid from dataclasses import dataclass -from typing import Dict, List, Set, Tuple +from typing import Dict, Iterator, List, Set, Tuple -from pypaimon.index.dynamic_bucket import HASH_INDEX, _iter_hashes, to_signed_int32 +from pypaimon.index.dynamic_bucket import ( + HASH_INDEX, MAX_DYNAMIC_BUCKETS, to_signed_int32, validate_bucket_id, validate_max_buckets, +) from pypaimon.index.index_file_handler import IndexFileHandler from pypaimon.index.index_file_meta import IndexFileMeta from pypaimon.manifest.index_manifest_entry import IndexManifestEntry @@ -33,6 +36,330 @@ _SNAPSHOT_UNSET = object() +def _java_remainder(value: int, divisor: int) -> int: + remainder = abs(value) % divisor + return -remainder if value < 0 else remainder + + +def compute_assigner( + partition_hash: int, + key_hash: int, + num_channels: int, + num_assigners: int, +) -> int: + """Mirror Java ``BucketAssigner.computeAssigner``.""" + if num_channels <= 0 or num_assigners <= 0: + raise ValueError("num_channels and num_assigners must be positive") + start = abs(_java_remainder(to_signed_int32(partition_hash), num_channels)) + assigner = abs(_java_remainder(to_signed_int32(key_hash), num_assigners)) + return _java_remainder(to_signed_int32(start + assigner), num_channels) + + +def is_my_bucket(bucket: int, num_assigners: int, assign_id: int) -> bool: + return bucket % num_assigners == assign_id % num_assigners + + +def _iter_hashes(table, entry: IndexManifestEntry) -> Iterator[int]: + meta = entry.index_file + path = table.path_factory().bucket_index_path( + tuple(entry.partition.values), entry.bucket, meta + ) + with table.file_io.new_input_stream(path) as stream: + remainder = b"" + while True: + payload = stream.read(64 * 1024) + if not payload: + break + payload = remainder + payload + aligned_size = len(payload) - len(payload) % 4 + for value in struct.iter_unpack(">i", payload[:aligned_size]): + yield value[0] + remainder = payload[aligned_size:] + if remainder: + raise RuntimeError( + f"Corrupt dynamic-bucket HASH index {path}: " + "expected a multiple of 4 bytes" + ) + + +class _PartitionIndex: + def __init__( + self, + hash_to_bucket: Dict[int, int], + bucket_information: Dict[int, int], + target_bucket_row_number: int, + ) -> None: + self.hash_to_bucket = hash_to_bucket + self.non_full_bucket_information = bucket_information + self.total_bucket_set = set(bucket_information) + self.total_bucket_list = list(bucket_information) + self.target_bucket_row_number = target_bucket_row_number + self._bucket_upper_bound_reached = False + + def assign( + self, + key_hash: int, + bucket_filter, + max_buckets_num: int, + max_bucket_id: int, + ) -> int: + return self.assign_with_status( + key_hash, bucket_filter, max_buckets_num, max_bucket_id + )[0] + + def assign_with_status( + self, + key_hash: int, + bucket_filter, + max_buckets_num: int, + max_bucket_id: int, + ) -> Tuple[int, bool]: + validate_max_buckets(max_buckets_num) + if key_hash in self.hash_to_bucket: + return self.hash_to_bucket[key_hash], False + + for bucket, row_count in list(self.non_full_bucket_information.items()): + if row_count < self.target_bucket_row_number: + self.non_full_bucket_information[bucket] = row_count + 1 + self.hash_to_bucket[key_hash] = bucket + return bucket, True + del self.non_full_bucket_information[bucket] + + global_max_bucket_id = ( + MAX_DYNAMIC_BUCKETS if max_buckets_num == -1 else max_buckets_num + ) - 1 + # Java's occupancy and bucket limit are partition-local. A larger + # bucket restored from another partition cannot hide a free local ID. + if not self._bucket_upper_bound_reached: + for bucket in range(global_max_bucket_id + 1): + if bucket_filter(bucket) and bucket not in self.total_bucket_set: + self.non_full_bucket_information[bucket] = 1 + self.total_bucket_set.add(bucket) + self.total_bucket_list.append(bucket) + self.hash_to_bucket[key_hash] = bucket + return bucket, True + if max_buckets_num == -1: + raise RuntimeError( + "No dynamic bucket id remains below Java Short.MAX_VALUE. " + "Increase dynamic-bucket.target-row-num." + ) + # The assigner's ownership and limit are fixed for this partition, + # so Java remembers exhaustion and reuses an owned bucket directly. + self._bucket_upper_bound_reached = True + + if not self.total_bucket_list: + raise RuntimeError( + "No dynamic bucket is available for this assigner. Reduce the " + "assigner parallelism or increase dynamic-bucket.max-buckets." + ) + bucket = random.choice(self.total_bucket_list) + self.hash_to_bucket[key_hash] = bucket + return bucket, True + + +class DynamicBucketAssigner: + """Restore dynamic-bucket mappings from HASH indexes before assigning.""" + + def __init__( + self, + table, + num_channels: int, + num_assigners: int, + assign_id: int, + target_bucket_row_number: int, + max_buckets_num: int, + ignore_existing: bool = False, + snapshot=_SNAPSHOT_UNSET, + ) -> None: + validate_max_buckets(max_buckets_num) + if not 0 <= assign_id < num_channels: + raise ValueError( + f"assign_id must be in [0, {num_channels}), got {assign_id}" + ) + self.table = table + self.num_channels = num_channels + self.num_assigners = num_assigners + self.assign_id = assign_id + self.target_bucket_row_number = target_bucket_row_number + self.max_buckets_num = max_buckets_num + self.ignore_existing = ignore_existing + self.snapshot = ( + table.snapshot_manager().get_latest_snapshot() + if snapshot is _SNAPSHOT_UNSET + else snapshot + ) + self.max_bucket_id = 0 + self._partition_indexes: Dict[Tuple, _PartitionIndex] = {} + + def assign(self, partition: Tuple, partition_hash: int, key_hash: int) -> int: + return self.assign_with_status(partition, partition_hash, key_hash)[0] + + def assign_with_status( + self, partition: Tuple, partition_hash: int, key_hash: int + ) -> Tuple[int, bool]: + result = self.assign_batch( + [partition], [partition_hash], [key_hash] + ) + return result[0] + + def assign_batch( + self, + partitions: List[Tuple], + partition_hashes: List[int], + key_hashes: List[int], + ) -> List[Tuple[int, bool]]: + if not ( + len(partitions) == len(partition_hashes) == len(key_hashes) + ): + raise ValueError( + "Partition, partition-hash, and key-hash counts must match" + ) + requested_by_partition: Dict[Tuple, Set[int]] = {} + normalized_hashes = [to_signed_int32(value) for value in key_hashes] + for partition, partition_hash, key_hash in zip( + partitions, partition_hashes, normalized_hashes + ): + self._validate_assigner(partition_hash, key_hash) + requested_by_partition.setdefault(partition, set()).add(key_hash) + + for partition, requested in requested_by_partition.items(): + index = self._partition_indexes.get(partition) + if index is None: + index = self._load_partition( + partition, + requested, + ) + self._partition_indexes[partition] = index + else: + missing = requested.difference(index.hash_to_bucket) + if missing: + self._restore_requested_hashes( + partition, + missing, + index.hash_to_bucket, + ) + + results = [] + for partition, key_hash in zip(partitions, normalized_hashes): + index = self._partition_indexes[partition] + bucket, is_new = index.assign_with_status( + key_hash, + lambda value: is_my_bucket( + value, self.num_assigners, self.assign_id + ), + self.max_buckets_num, + self.max_bucket_id, + ) + self.max_bucket_id = max(self.max_bucket_id, bucket) + results.append((bucket, is_new)) + return results + + def _validate_assigner(self, partition_hash: int, key_hash: int) -> None: + expected_assigner = compute_assigner( + partition_hash, + key_hash, + self.num_channels, + self.num_assigners, + ) + if expected_assigner != self.assign_id: + raise ValueError( + f"Record assigner {expected_assigner} does not match writer " + f"assigner {self.assign_id}" + ) + + def _load_partition( + self, partition: Tuple, requested_hashes: Set[int] + ) -> _PartitionIndex: + if self.ignore_existing: + return _PartitionIndex({}, {}, self.target_bucket_row_number) + snapshot = self.snapshot + if snapshot is None: + return _PartitionIndex({}, {}, self.target_bucket_row_number) + entries = IndexFileHandler(self.table).scan( + snapshot, + lambda entry: entry.index_file.index_type == HASH_INDEX + and tuple(entry.partition.values) == partition, + ) + entries_by_bucket = {} + for entry in entries: + validate_bucket_id(entry.bucket) + previous = entries_by_bucket.get(entry.bucket) + if previous is not None: + raise RuntimeError( + "Found multiple dynamic-bucket HASH indexes for partition " + f"{partition}, bucket {entry.bucket}: " + f"{[previous.index_file.file_name, entry.index_file.file_name]}" + ) + entries_by_bucket[entry.bucket] = entry + + data_buckets = self._active_data_buckets(partition, snapshot) + missing_buckets = data_buckets.difference(entries_by_bucket) + if missing_buckets: + raise RuntimeError( + f"Dynamic-bucket partition {partition} has data files but no " + "complete HASH index for buckets " + f"{sorted(missing_buckets)}. Rewrite the partition before " + "performing incremental writes." + ) + + hash_to_bucket: Dict[int, int] = {} + bucket_information: Dict[int, int] = {} + for entry in entries: + self.max_bucket_id = max(self.max_bucket_id, entry.bucket) + if is_my_bucket( + entry.bucket, self.num_assigners, self.assign_id + ): + bucket_information[entry.bucket] = entry.index_file.row_count + remaining = set(requested_hashes) + for entry in entries: + if not remaining: + break + for key_hash in _iter_hashes(self.table, entry): + if key_hash in remaining: + hash_to_bucket[key_hash] = entry.bucket + remaining.remove(key_hash) + if not remaining: + break + return _PartitionIndex( + hash_to_bucket, + bucket_information, + self.target_bucket_row_number, + ) + + def _restore_requested_hashes( + self, + partition: Tuple, + requested_hashes: Set[int], + hash_to_bucket: Dict[int, int], + ) -> None: + if self.ignore_existing or not requested_hashes: + return + snapshot = self.snapshot + if snapshot is None: + return + entries = IndexFileHandler(self.table).scan( + snapshot, + lambda entry: entry.index_file.index_type == HASH_INDEX + and tuple(entry.partition.values) == partition, + ) + remaining = set(requested_hashes) + for entry in entries: + validate_bucket_id(entry.bucket) + for key_hash in _iter_hashes(self.table, entry): + if key_hash not in remaining: + continue + hash_to_bucket[key_hash] = entry.bucket + remaining.remove(key_hash) + if not remaining: + return + + def _active_data_buckets(self, partition: Tuple, snapshot) -> Set[int]: + if snapshot is None: + return set() + from pypaimon.write.backend_dispatch import dynamic_bucket_data_buckets + return dynamic_bucket_data_buckets(self.table, partition, snapshot) + + @dataclass class DynamicBucketIndexChanges: additions: List[IndexManifestEntry] diff --git a/paimon-python/pypaimon/execution/python/write/input_routing.py b/paimon-python/pypaimon/execution/python/write/input_routing.py index 133e3f2378a2..1f546986f0d0 100644 --- a/paimon-python/pypaimon/execution/python/write/input_routing.py +++ b/paimon-python/pypaimon/execution/python/write/input_routing.py @@ -15,14 +15,334 @@ # specific language governing permissions and limitations # under the License. -"""Arrow contracts for Python distributed dynamic-bucket execution.""" +"""Python hashing and routing for distributed bucket execution.""" + +import math +import struct +from typing import Any, Dict, List, Tuple import pyarrow as pa -from pypaimon.index.dynamic_bucket import ( - DynamicBucketAssigner, compute_assigner, resolve_bucket_snapshot, to_signed_int32, -) -from pypaimon.write.input_routing import DynamicBucketInputRouter +from pypaimon.common.options.core_options import CoreOptions +from pypaimon.execution.python.write.dynamic_bucket import DynamicBucketAssigner, compute_assigner +from pypaimon.index.dynamic_bucket import resolve_bucket_snapshot, to_signed_int32 +from pypaimon.schema.data_types import AtomicType +from pypaimon.schema.table_schema import TableSchema +from pypaimon.table.bucket_mode import BucketMode +from pypaimon.table.row.generic_row import GenericRow, GenericRowSerializer +from pypaimon.table.row.internal_row import RowKind + + +_MURMUR_C1 = 0xCC9E2D51 +_MURMUR_C2 = 0x1B873593 +_DEFAULT_SEED = 42 + + +def _mix_k1(k1: int) -> int: + k1 = (k1 * _MURMUR_C1) & 0xFFFFFFFF + k1 = ((k1 << 15) | (k1 >> 17)) & 0xFFFFFFFF + k1 = (k1 * _MURMUR_C2) & 0xFFFFFFFF + return k1 + + +def _mix_h1(h1: int, k1: int) -> int: + h1 = (h1 ^ k1) & 0xFFFFFFFF + h1 = ((h1 << 13) | (h1 >> 19)) & 0xFFFFFFFF + h1 = (h1 * 5 + 0xE6546B64) & 0xFFFFFFFF + return h1 + + +def _fmix(h1: int, length: int) -> int: + h1 = (h1 ^ length) & 0xFFFFFFFF + h1 ^= h1 >> 16 + h1 = (h1 * 0x85EBCA6B) & 0xFFFFFFFF + h1 ^= h1 >> 13 + h1 = (h1 * 0xC2B2AE35) & 0xFFFFFFFF + h1 ^= h1 >> 16 + return h1 + + +def _hash_bytes_by_words(data: bytes, seed: int = _DEFAULT_SEED) -> int: + n = len(data) + length_aligned = n - (n % 4) + h1 = seed + for i in range(0, length_aligned, 4): + k1 = struct.unpack_from(" int: + if hash_unsigned >= 0x80000000: + hash_signed = hash_unsigned - 0x100000000 + else: + hash_signed = hash_unsigned + rem = hash_signed - math.trunc(hash_signed / num_buckets) * num_buckets + return abs(rem) + + +class PartitionKeyExtractor: + """Resolve input routing fields and hash their Java binary layout.""" + + def __init__(self, table_schema: TableSchema): + self.table_schema = table_schema + self.partition_indices = self._get_field_indices(table_schema.partition_keys) + + def extract_partition_bucket_batch(self, data: pa.RecordBatch) -> Tuple[List[Tuple], List[int]]: + return self._extract_partitions_batch(data), self._extract_buckets_batch(data) + + def extract_partition_bucket_row( + self, values_by_name: Dict[str, Any]) -> Tuple[Tuple, int]: + partition = tuple( + values_by_name[self.table_schema.fields[i].name] + for i in self.partition_indices + ) + bucket = self._extract_bucket_row(values_by_name) + return partition, bucket + + def _get_field_indices(self, field_names: List[str]) -> List[int]: + if not field_names: + return [] + field_map = {field.name: i for i, field in enumerate(self.table_schema.fields)} + return [field_map[name] for name in field_names if name in field_map] + + def _extract_partitions_batch(self, data: pa.RecordBatch) -> List[Tuple]: + if not self.partition_indices: + return [() for _ in range(data.num_rows)] + + partition_columns = self._input_columns(data, self.partition_indices) + + partitions = [] + for row_idx in range(data.num_rows): + partition_values = tuple(col[row_idx].as_py() for col in partition_columns) + partitions.append(partition_values) + + return partitions + + def _input_columns(self, data: pa.RecordBatch, field_indices: List[int]): + """Resolve routing fields in the actual write type, as Java's converter does.""" + columns = [] + for field_index in field_indices: + name = self.table_schema.fields[field_index].name + index = data.schema.get_field_index(name) + if index < 0: + raise ValueError(f"Input must contain exactly one routing column '{name}'.") + columns.append(data.column(index)) + return columns + + @staticmethod + def _binary_row_hash_code(values: Tuple, fields: List) -> int: + return _hash_bytes_by_words( + GenericRowSerializer.to_bytes( + GenericRow(list(values), fields, RowKind.INSERT) + )[4:] + ) + + +class PostponeFixedBucketInputRouter(PartitionKeyExtractor): + """Group distributed inputs by the same partition and bucket as writers.""" + + def __init__(self, table, bucket_plan): + super().__init__(table.table_schema) + if table.options.bucket() != BucketMode.POSTPONE_BUCKET.value: + raise ValueError( + "Postpone fixed bucket writes require bucket = -2, got {}".format( + table.options.bucket() + ) + ) + bucket_function = str( + table.table_schema.options.get("bucket-function.type", "default") + ).strip().lower() + if bucket_function != "default": + raise ValueError( + "Postpone fixed bucket writes only support " + "bucket-function.type=default, got {}" + .format(bucket_function) + ) + self.bucket_keys = table.table_schema.bucket_keys + self.bucket_key_indices = self._get_field_indices(self.bucket_keys) + self._bucket_key_fields = table.table_schema.logical_bucket_key_fields + self._bucket_plan = bucket_plan + + def with_bucket_plan(self, bucket_plan) -> None: + self._bucket_plan = bucket_plan + + def num_buckets(self, partition: Tuple) -> int: + return self._bucket_plan.num_buckets(partition) + + def route(self, data) -> pa.Int32Array: + """Return bucket IDs in input order, as the Native router does.""" + return pa.array(self.extract_partition_bucket_batch(data)[1], type=pa.int32()) + + def extract_partition_bucket_batch( + self, data: pa.RecordBatch + ) -> Tuple[List[Tuple], List[int]]: + partitions = self._extract_partitions_batch(data) + columns = self._input_columns(data, self.bucket_key_indices) + buckets = [ + _bucket_from_hash( + self._binary_row_hash_code( + tuple(col[row_idx].as_py() for col in columns), + self._bucket_key_fields, + ), + self.num_buckets(partition), + ) + for row_idx, partition in enumerate(partitions) + ] + return partitions, buckets + + def _extract_buckets_batch(self, data: pa.RecordBatch) -> List[int]: + return self.extract_partition_bucket_batch(data)[1] + + def _extract_bucket_row(self, values_by_name: Dict[str, Any]) -> int: + partition = tuple( + values_by_name[self.table_schema.fields[i].name] + for i in self.partition_indices + ) + return _bucket_from_hash( + self._binary_row_hash_code( + tuple(values_by_name[name] for name in self.bucket_keys), + self._bucket_key_fields, + ), + self.num_buckets(partition), + ) + + +class FixedBucketInputRouter(PartitionKeyExtractor): + """Route distributed inputs without constructing a Python file writer.""" + + def __init__(self, table_schema: TableSchema): + super().__init__(table_schema) + options = CoreOptions.from_dict(table_schema.options) + self.num_buckets = options.bucket() + if self.num_buckets <= 0: + raise ValueError(f"Fixed bucket mode requires bucket > 0, got {self.num_buckets}") + + # Bucket-key resolution lives on TableSchema (mirrors Java + # ``TableSchema.bucketKeys()`` / ``logicalBucketKeyType()``); reuse + # it so any reader path that walks the same logic stays in sync. + self.bucket_keys = table_schema.bucket_keys + self.bucket_key_indices = self._get_field_indices(self.bucket_keys) + self._bucket_key_fields = table_schema.logical_bucket_key_fields + bucket_function = table_schema.options.get('bucket-function.type') + self.bucket_function = 'default' if bucket_function is None else str(bucket_function).lower() + if self.bucket_function == 'mod': + if len(self._bucket_key_fields) != 1: + raise ValueError('bucket key must have exactly one field in mod bucket function') + key_type = self._bucket_key_fields[0].type + if not isinstance(key_type, AtomicType) or key_type.type.upper() not in ('INT', 'INTEGER', 'BIGINT'): + raise ValueError( + 'bucket key type must be INT or BIGINT in mod bucket function, but got %s' % key_type) + elif self.bucket_function not in ('default', 'hive'): + raise ValueError('Unsupported bucket function type: %s' % self.bucket_function) + + def route(self, data) -> pa.Int32Array: + """Return bucket IDs in input order, as the Native router does.""" + return pa.array(self.extract_partition_bucket_batch(data)[1], type=pa.int32()) + + def _extract_buckets_batch(self, data: pa.RecordBatch) -> List[int]: + self._check_automatic_bucket_function() + columns = self._input_columns(data, self.bucket_key_indices) + if self.bucket_function == 'mod': + return [self._mod_bucket(value.as_py()) for value in columns[0]] + return [ + _bucket_from_hash( + self._binary_row_hash_code( + tuple(col[row_idx].as_py() for col in columns), + self._bucket_key_fields, + ), + self.num_buckets, + ) + for row_idx in range(data.num_rows) + ] + + def _extract_bucket_row(self, values_by_name: Dict[str, Any]) -> int: + self._check_automatic_bucket_function() + if self.bucket_function == 'mod': + return self._mod_bucket(values_by_name[self.bucket_keys[0]]) + return _bucket_from_hash( + self._binary_row_hash_code( + tuple(values_by_name[name] for name in self.bucket_keys), + self._bucket_key_fields, + ), + self.num_buckets, + ) + + def _check_automatic_bucket_function(self): + # Explicit buckets do not use this method. A single bucket also keeps + # the existing zero result without needing the Hive hash algorithm. + if self.bucket_function == 'hive' and self.num_buckets > 1: + raise NotImplementedError( + 'Hive bucket functions require the native writer; enable write.native.enabled') + + def _mod_bucket(self, value): + # Java reads the zeroed fixed-width slot for NULL; Python's modulo + # with a positive divisor otherwise matches Math.floorMod, even for + # Long.MIN_VALUE. Avoid floating-point division for BIGINT values. + return 0 if value is None else value % self.num_buckets + + +class DynamicBucketInputRouter(PartitionKeyExtractor): + """Hash dynamic-bucket input without assigning buckets or maintaining indexes.""" + + def __init__(self, table_schema: TableSchema): + super().__init__(table_schema) + num_buckets = int(table_schema.options.get(CoreOptions.BUCKET.key(), -1)) + if num_buckets != -1: + raise ValueError( + "Only 'bucket' = '-1' is allowed for " + f"'DynamicBucketRowKeyExtractor', but found: {num_buckets}" + ) + if table_schema.options.get(CoreOptions.BUCKET_KEY.key()) is not None: + raise ValueError("Cannot define 'bucket-key' in dynamic bucket mode") + self.bucket_keys = table_schema.bucket_keys + self.bucket_key_indices = self._get_field_indices(self.bucket_keys) + self._bucket_key_fields = table_schema.logical_bucket_key_fields + self._partition_fields = [ + table_schema.fields[index] for index in self.partition_indices + ] + + def extract_hashes_batch( + self, data: pa.RecordBatch + ) -> Tuple[List[Tuple], List[int], List[int]]: + """Return partitions, BinaryRow partition hashes, and key hashes.""" + partitions, partition_hashes = self.extract_partition_hashes_batch(data) + return partitions, partition_hashes, self._extract_key_hashes_batch(data) + + def extract_partition_hashes_batch(self, data: pa.RecordBatch) -> Tuple[List[Tuple], List[int]]: + """Hash partitions independently when key hashes arrived from upstream.""" + partitions = self._extract_partitions_batch(data) + partition_hash_cache = {} + partition_hashes = [] + for partition in partitions: + if partition not in partition_hash_cache: + partition_hash_cache[partition] = self._binary_row_hash_code( + partition, self._partition_fields + ) + partition_hashes.append(partition_hash_cache[partition]) + return partitions, partition_hashes + + def _extract_key_hashes_batch(self, data: pa.RecordBatch) -> List[int]: + key_columns = self._input_columns(data, self.bucket_key_indices) + return [ + self._binary_row_hash_code( + tuple(column[row_idx].as_py() for column in key_columns), + self._bucket_key_fields, + ) + for row_idx in range(data.num_rows) + ] + + def extract_assigners_batch( + self, data: pa.RecordBatch, num_channels: int, num_assigners: int + ) -> List[int]: + _, partition_hashes, key_hashes = self.extract_hashes_batch(data) + return [ + compute_assigner( + partition_hash, key_hash, num_channels, num_assigners + ) + for partition_hash, key_hash in zip(partition_hashes, key_hashes) + ] class PythonDynamicBucketAssignerRouter: diff --git a/paimon-python/pypaimon/execution/python/write/row_id_plan.py b/paimon-python/pypaimon/execution/python/write/row_id_plan.py index 72f486afafdc..c364187a940e 100644 --- a/paimon-python/pypaimon/execution/python/write/row_id_plan.py +++ b/paimon-python/pypaimon/execution/python/write/row_id_plan.py @@ -18,9 +18,13 @@ """Python grouping for distributed row-ID plans, selected before execution.""" -from pypaimon.execution.python.write.table_update import PythonTableUpdate +import pyarrow as pa + from pypaimon.manifest.schema.data_file_meta import DataFileMeta from pypaimon.read.query_auth_split import QueryAuthSplit +from pypaimon.utils.data_evolution_utils import retrieve_anchor_file +from pypaimon.utils.range import Range +from pypaimon.utils.range_helper import RangeHelper def _file_groups(split): @@ -28,15 +32,108 @@ def _file_groups(split): for group in _file_groups(split.split): yield QueryAuthSplit(group, split.auth_result) else: - yield from PythonTableUpdate._predicate_update_file_groups([split]) + normal, dedicated = [], [] + names = set() + for file in split.files: + file.non_null_row_id_range() + if file.file_name in names: + raise ValueError('A row-ID split contains duplicate file metadata') + names.add(file.file_name) + if DataFileMeta.is_blob_file(file.file_name) or DataFileMeta.is_vector_file(file.file_name): + dedicated.append(file) + else: + normal.append(file) + if not normal: + raise ValueError('A row-ID split requires a normal anchor file') + if split.data_deletion_files is not None and len(split.data_deletion_files) != len(split.files): + raise ValueError('Row-ID deletion metadata must match the captured data files') + matched = set() + for files in RangeHelper(lambda file: file.non_null_row_id_range()).merge_overlapping_ranges(normal): + anchor_range = retrieve_anchor_file(files).non_null_row_id_range() + if any(file.non_null_row_id_range() != anchor_range for file in files): + raise ValueError('Overlapping row-ID anchors must have identical ranges') + files += [file for file in dedicated if anchor_range.overlaps(file.non_null_row_id_range())] + selected_names = {file.file_name for file in files} + matched.update(selected_names) + # Keep the captured physical paths and aligned deletion metadata. + # Dedicated files may serve more than one normal anchor, but never + # merge those owners or make the gaps between them addressable. + group = split.filter_file(lambda file: file.file_name in selected_names) + group.raw_convertible = len(group.files) == 1 + yield group + if matched != names: + raise ValueError('A row-ID split contains a dedicated file without its normal anchor') def row_id_groups(plan): groups = [] for planned in plan.splits(): for split in _file_groups(planned): - ranges = [file.row_id_range() for file in split.files - if file.first_row_id is not None and not DataFileMeta.is_blob_file(file.file_name)] - if ranges: - groups.append((min(r.from_ for r in ranges), max(r.to for r in ranges), split)) + anchor_range = retrieve_anchor_file(split.files).non_null_row_id_range() + groups.append((anchor_range.from_, anchor_range.to, split)) return sorted(groups, key=lambda group: group[0]) + + +class PythonRowIdPlan: + """Route captured metadata with no file access or mutation state.""" + + def __init__(self, plan, groups=None): + import numpy as np + from pypaimon.table.row.generic_row import GenericRowSerializer + + self.plan = plan + self.groups = row_id_groups(plan) if groups is None else groups + self._starts = np.asarray([first for first, _, _ in self.groups], dtype=np.int64) + self._ends = np.asarray([last for _, last, _ in self.groups], dtype=np.int64) + if any(left[1] >= right[0] for left, right in zip(self.groups, self.groups[1:])): + raise ValueError('A row-ID plan contains overlapping file groups') + owners = {} + bucket_owners = [] + for first, _, split in self.groups: + identity = (GenericRowSerializer.to_bytes(split.partition), split.bucket) + bucket_owners.append(owners.setdefault(identity, first)) + self._bucket_owners = np.asarray(bucket_owners, dtype=np.int64) + + def _positions(self, row_ids): + import numpy as np + + values = row_ids.to_numpy(zero_copy_only=False) + positions = np.searchsorted(self._starts, values, side='right') - 1 + if len(values) == 0: + return positions + if not self.groups: + valid = np.zeros(len(values), dtype=bool) + else: + valid = (values >= 0) & (positions >= 0) & (values <= self._ends[np.maximum(positions, 0)]) + if not valid.all(): + bad = values[~valid][0] + ranges = ['[{}, {}]'.format(first, last) for first, last, _ in self.groups] + raise ValueError('_ROW_ID {} does not belong to any valid range {}; planner snapshot is stale ' + 'or matched rows come from a different table.'.format(bad, ranges)) + return positions + + def route(self, row_ids): + return pa.array(self._starts[self._positions(row_ids)], type=pa.int64()) + + def route_by_bucket(self, row_ids): + return pa.array(self._bucket_owners[self._positions(row_ids)], type=pa.int64()) + + def select(self, row_ids): + from pypaimon.globalindex.indexed_split import IndexedSplit + + positions = self._positions(row_ids) + requested = {} + for position, row_id in zip(positions, row_ids.to_pylist()): + requested.setdefault(int(position), set()).add(row_id) + splits = [] + for position in sorted(requested): + split = self.groups[position][2] + ranges = Range.to_ranges(list(requested[position])) + if isinstance(split, QueryAuthSplit): + splits.append(QueryAuthSplit(IndexedSplit(split.split, ranges), split.auth_result)) + else: + splits.append(IndexedSplit(split, ranges)) + return self.plan._with_splits(splits) + + def for_plan(self, plan, groups): + return PythonRowIdPlan(plan, groups) diff --git a/paimon-python/pypaimon/execution/python/write/row_key_extractor.py b/paimon-python/pypaimon/execution/python/write/row_key_extractor.py index 996f440c723d..81589e09262a 100644 --- a/paimon-python/pypaimon/execution/python/write/row_key_extractor.py +++ b/paimon-python/pypaimon/execution/python/write/row_key_extractor.py @@ -24,14 +24,11 @@ import pyarrow as pa from pypaimon.common.options.core_options import CoreOptions -from pypaimon.index.dynamic_bucket import ( - MAX_DYNAMIC_BUCKETS, - is_my_bucket, - validate_max_buckets, -) +from pypaimon.index.dynamic_bucket import MAX_DYNAMIC_BUCKETS, validate_max_buckets +from pypaimon.execution.python.write.dynamic_bucket import is_my_bucket from pypaimon.schema.table_schema import TableSchema from pypaimon.table.bucket_mode import BucketMode -from pypaimon.write.input_routing import ( +from pypaimon.execution.python.write.input_routing import ( DynamicBucketInputRouter, FixedBucketInputRouter, PartitionKeyExtractor, PostponeFixedBucketInputRouter, _bucket_from_hash as _shared_bucket_from_hash, _hash_bytes_by_words as _shared_hash_bytes_by_words, ) @@ -338,8 +335,9 @@ def __init__( ) self._index_maintainer = None else: - from pypaimon.execution.python.write.dynamic_bucket import DynamicBucketIndexMaintainer - from pypaimon.index.dynamic_bucket import DynamicBucketAssigner + from pypaimon.execution.python.write.dynamic_bucket import ( + DynamicBucketAssigner, DynamicBucketIndexMaintainer, + ) from pypaimon.index.dynamic_bucket import resolve_bucket_snapshot snapshot = resolve_bucket_snapshot(table, base_snapshot_id) diff --git a/paimon-python/pypaimon/index/dynamic_bucket.py b/paimon-python/pypaimon/index/dynamic_bucket.py index ad338d006504..0a7a999c7947 100644 --- a/paimon-python/pypaimon/index/dynamic_bucket.py +++ b/paimon-python/pypaimon/index/dynamic_bucket.py @@ -16,19 +16,11 @@ # limitations under the License. ################################################################################ -"""Shared dynamic-bucket input assignment and HASH mapping restoration.""" - -import random -import struct -from typing import Dict, Iterator, List, Set, Tuple - -from pypaimon.index.index_file_handler import IndexFileHandler -from pypaimon.manifest.index_manifest_entry import IndexManifestEntry +"""Shared dynamic-bucket limits, validation and read-view configuration.""" HASH_INDEX = "HASH" SHORT_MAX_VALUE = 32767 MAX_DYNAMIC_BUCKETS = SHORT_MAX_VALUE + 1 -_SNAPSHOT_UNSET = object() def resolve_bucket_snapshot(table, base_snapshot_id=None): @@ -67,29 +59,6 @@ def to_signed_int32(value: int) -> int: return value - 0x100000000 if value >= 0x80000000 else value -def _java_remainder(value: int, divisor: int) -> int: - remainder = abs(value) % divisor - return -remainder if value < 0 else remainder - - -def compute_assigner( - partition_hash: int, - key_hash: int, - num_channels: int, - num_assigners: int, -) -> int: - """Mirror Java ``BucketAssigner.computeAssigner``.""" - if num_channels <= 0 or num_assigners <= 0: - raise ValueError("num_channels and num_assigners must be positive") - start = abs(_java_remainder(to_signed_int32(partition_hash), num_channels)) - assigner = abs(_java_remainder(to_signed_int32(key_hash), num_assigners)) - return _java_remainder(to_signed_int32(start + assigner), num_channels) - - -def is_my_bucket(bucket: int, num_assigners: int, assign_id: int) -> bool: - return bucket % num_assigners == assign_id % num_assigners - - def validate_max_buckets(max_buckets_num: int) -> None: if max_buckets_num != -1 and not 1 <= max_buckets_num <= MAX_DYNAMIC_BUCKETS: raise ValueError( @@ -104,304 +73,3 @@ def validate_bucket_id(bucket: int) -> None: "Dynamic bucket id must be between 0 and " f"{MAX_DYNAMIC_BUCKETS - 1}, but was {bucket}." ) - - -def _iter_hashes(table, entry: IndexManifestEntry) -> Iterator[int]: - meta = entry.index_file - path = table.path_factory().bucket_index_path( - tuple(entry.partition.values), entry.bucket, meta - ) - with table.file_io.new_input_stream(path) as stream: - remainder = b"" - while True: - payload = stream.read(64 * 1024) - if not payload: - break - payload = remainder + payload - aligned_size = len(payload) - len(payload) % 4 - for value in struct.iter_unpack(">i", payload[:aligned_size]): - yield value[0] - remainder = payload[aligned_size:] - if remainder: - raise RuntimeError( - f"Corrupt dynamic-bucket HASH index {path}: " - "expected a multiple of 4 bytes" - ) - - -class _PartitionIndex: - def __init__( - self, - hash_to_bucket: Dict[int, int], - bucket_information: Dict[int, int], - target_bucket_row_number: int, - ) -> None: - self.hash_to_bucket = hash_to_bucket - self.non_full_bucket_information = bucket_information - self.total_bucket_set = set(bucket_information) - self.total_bucket_list = list(bucket_information) - self.target_bucket_row_number = target_bucket_row_number - self._bucket_upper_bound_reached = False - - def assign( - self, - key_hash: int, - bucket_filter, - max_buckets_num: int, - max_bucket_id: int, - ) -> int: - return self.assign_with_status( - key_hash, bucket_filter, max_buckets_num, max_bucket_id - )[0] - - def assign_with_status( - self, - key_hash: int, - bucket_filter, - max_buckets_num: int, - max_bucket_id: int, - ) -> Tuple[int, bool]: - validate_max_buckets(max_buckets_num) - if key_hash in self.hash_to_bucket: - return self.hash_to_bucket[key_hash], False - - for bucket, row_count in list(self.non_full_bucket_information.items()): - if row_count < self.target_bucket_row_number: - self.non_full_bucket_information[bucket] = row_count + 1 - self.hash_to_bucket[key_hash] = bucket - return bucket, True - del self.non_full_bucket_information[bucket] - - global_max_bucket_id = ( - MAX_DYNAMIC_BUCKETS if max_buckets_num == -1 else max_buckets_num - ) - 1 - # Java's occupancy and bucket limit are partition-local. A larger - # bucket restored from another partition cannot hide a free local ID. - if not self._bucket_upper_bound_reached: - for bucket in range(global_max_bucket_id + 1): - if bucket_filter(bucket) and bucket not in self.total_bucket_set: - self.non_full_bucket_information[bucket] = 1 - self.total_bucket_set.add(bucket) - self.total_bucket_list.append(bucket) - self.hash_to_bucket[key_hash] = bucket - return bucket, True - if max_buckets_num == -1: - raise RuntimeError( - "No dynamic bucket id remains below Java Short.MAX_VALUE. " - "Increase dynamic-bucket.target-row-num." - ) - # The assigner's ownership and limit are fixed for this partition, - # so Java remembers exhaustion and reuses an owned bucket directly. - self._bucket_upper_bound_reached = True - - if not self.total_bucket_list: - raise RuntimeError( - "No dynamic bucket is available for this assigner. Reduce the " - "assigner parallelism or increase dynamic-bucket.max-buckets." - ) - bucket = random.choice(self.total_bucket_list) - self.hash_to_bucket[key_hash] = bucket - return bucket, True - - -class DynamicBucketAssigner: - """Restore dynamic-bucket mappings from HASH indexes before assigning.""" - - def __init__( - self, - table, - num_channels: int, - num_assigners: int, - assign_id: int, - target_bucket_row_number: int, - max_buckets_num: int, - ignore_existing: bool = False, - snapshot=_SNAPSHOT_UNSET, - ) -> None: - validate_max_buckets(max_buckets_num) - if not 0 <= assign_id < num_channels: - raise ValueError( - f"assign_id must be in [0, {num_channels}), got {assign_id}" - ) - self.table = table - self.num_channels = num_channels - self.num_assigners = num_assigners - self.assign_id = assign_id - self.target_bucket_row_number = target_bucket_row_number - self.max_buckets_num = max_buckets_num - self.ignore_existing = ignore_existing - self.snapshot = ( - table.snapshot_manager().get_latest_snapshot() - if snapshot is _SNAPSHOT_UNSET - else snapshot - ) - self.max_bucket_id = 0 - self._partition_indexes: Dict[Tuple, _PartitionIndex] = {} - - def assign(self, partition: Tuple, partition_hash: int, key_hash: int) -> int: - return self.assign_with_status(partition, partition_hash, key_hash)[0] - - def assign_with_status( - self, partition: Tuple, partition_hash: int, key_hash: int - ) -> Tuple[int, bool]: - result = self.assign_batch( - [partition], [partition_hash], [key_hash] - ) - return result[0] - - def assign_batch( - self, - partitions: List[Tuple], - partition_hashes: List[int], - key_hashes: List[int], - ) -> List[Tuple[int, bool]]: - if not ( - len(partitions) == len(partition_hashes) == len(key_hashes) - ): - raise ValueError( - "Partition, partition-hash, and key-hash counts must match" - ) - requested_by_partition: Dict[Tuple, Set[int]] = {} - normalized_hashes = [to_signed_int32(value) for value in key_hashes] - for partition, partition_hash, key_hash in zip( - partitions, partition_hashes, normalized_hashes - ): - self._validate_assigner(partition_hash, key_hash) - requested_by_partition.setdefault(partition, set()).add(key_hash) - - for partition, requested in requested_by_partition.items(): - index = self._partition_indexes.get(partition) - if index is None: - index = self._load_partition( - partition, - requested, - ) - self._partition_indexes[partition] = index - else: - missing = requested.difference(index.hash_to_bucket) - if missing: - self._restore_requested_hashes( - partition, - missing, - index.hash_to_bucket, - ) - - results = [] - for partition, key_hash in zip(partitions, normalized_hashes): - index = self._partition_indexes[partition] - bucket, is_new = index.assign_with_status( - key_hash, - lambda value: is_my_bucket( - value, self.num_assigners, self.assign_id - ), - self.max_buckets_num, - self.max_bucket_id, - ) - self.max_bucket_id = max(self.max_bucket_id, bucket) - results.append((bucket, is_new)) - return results - - def _validate_assigner(self, partition_hash: int, key_hash: int) -> None: - expected_assigner = compute_assigner( - partition_hash, - key_hash, - self.num_channels, - self.num_assigners, - ) - if expected_assigner != self.assign_id: - raise ValueError( - f"Record assigner {expected_assigner} does not match writer " - f"assigner {self.assign_id}" - ) - - def _load_partition( - self, partition: Tuple, requested_hashes: Set[int] - ) -> _PartitionIndex: - if self.ignore_existing: - return _PartitionIndex({}, {}, self.target_bucket_row_number) - snapshot = self.snapshot - if snapshot is None: - return _PartitionIndex({}, {}, self.target_bucket_row_number) - entries = IndexFileHandler(self.table).scan( - snapshot, - lambda entry: entry.index_file.index_type == HASH_INDEX - and tuple(entry.partition.values) == partition, - ) - entries_by_bucket = {} - for entry in entries: - validate_bucket_id(entry.bucket) - previous = entries_by_bucket.get(entry.bucket) - if previous is not None: - raise RuntimeError( - "Found multiple dynamic-bucket HASH indexes for partition " - f"{partition}, bucket {entry.bucket}: " - f"{[previous.index_file.file_name, entry.index_file.file_name]}" - ) - entries_by_bucket[entry.bucket] = entry - - data_buckets = self._active_data_buckets(partition, snapshot) - missing_buckets = data_buckets.difference(entries_by_bucket) - if missing_buckets: - raise RuntimeError( - f"Dynamic-bucket partition {partition} has data files but no " - "complete HASH index for buckets " - f"{sorted(missing_buckets)}. Rewrite the partition before " - "performing incremental writes." - ) - - hash_to_bucket: Dict[int, int] = {} - bucket_information: Dict[int, int] = {} - for entry in entries: - self.max_bucket_id = max(self.max_bucket_id, entry.bucket) - if is_my_bucket( - entry.bucket, self.num_assigners, self.assign_id - ): - bucket_information[entry.bucket] = entry.index_file.row_count - remaining = set(requested_hashes) - for entry in entries: - if not remaining: - break - for key_hash in _iter_hashes(self.table, entry): - if key_hash in remaining: - hash_to_bucket[key_hash] = entry.bucket - remaining.remove(key_hash) - if not remaining: - break - return _PartitionIndex( - hash_to_bucket, - bucket_information, - self.target_bucket_row_number, - ) - - def _restore_requested_hashes( - self, - partition: Tuple, - requested_hashes: Set[int], - hash_to_bucket: Dict[int, int], - ) -> None: - if self.ignore_existing or not requested_hashes: - return - snapshot = self.snapshot - if snapshot is None: - return - entries = IndexFileHandler(self.table).scan( - snapshot, - lambda entry: entry.index_file.index_type == HASH_INDEX - and tuple(entry.partition.values) == partition, - ) - remaining = set(requested_hashes) - for entry in entries: - validate_bucket_id(entry.bucket) - for key_hash in _iter_hashes(self.table, entry): - if key_hash not in remaining: - continue - hash_to_bucket[key_hash] = entry.bucket - remaining.remove(key_hash) - if not remaining: - return - - def _active_data_buckets(self, partition: Tuple, snapshot) -> Set[int]: - if snapshot is None: - return set() - from pypaimon.write.backend_dispatch import dynamic_bucket_data_buckets - return dynamic_bucket_data_buckets(self.table, partition, snapshot) diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py b/paimon-python/pypaimon/ray/data_evolution_merge_join.py index edbb845d395b..5544353b967c 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py @@ -617,7 +617,6 @@ def distributed_update_apply( data_context=None, materialize_before_routing: bool = False, ) -> Tuple[list, int, list]: - import numpy as np import pickle import pyarrow.compute as pc @@ -667,42 +666,9 @@ def distributed_update_apply( frid_col = "_FIRST_ROW_ID" sentinel_row_id = -1 captured_sorted = sorted_first_row_ids - captured_sorted_arr = np.asarray(captured_sorted, dtype=np.int64) - valid_ranges = planner.valid_row_id_ranges - range_starts = np.asarray([r.from_ for r in valid_ranges], dtype=np.int64) - range_ends = np.asarray([r.to for r in valid_ranges], dtype=np.int64) def _assign_frid(batch: pa.Table) -> pa.Table: - if batch.num_rows == 0: - return batch.append_column( - frid_col, pa.array([], type=pa.int64()) - ) - rid_col = batch.column(row_id_name) - if rid_col.null_count: - raise ValueError( - "_ROW_ID is null; planner snapshot is stale " - "or matched rows come from a different table." - ) - rids = rid_col.to_numpy(zero_copy_only=False) - # Check each row_id belongs to a valid range (vectorized). - in_range = np.zeros(len(rids), dtype=bool) - for s, e in zip(range_starts, range_ends): - in_range |= (rids >= s) & (rids <= e) - if not in_range.all(): - bad = rids[~in_range][0] - raise ValueError( - f"_ROW_ID {bad} does not belong to any valid range " - f"{[f'[{r.from_}, {r.to}]' for r in valid_ranges]}; " - f"planner snapshot is stale or matched rows come " - f"from a different table." - ) - idx = np.searchsorted( - captured_sorted_arr, rids, side="right" - ) - 1 - frids = captured_sorted_arr[idx] - return batch.append_column( - frid_col, pa.array(frids, type=pa.int64()) - ) + return batch.append_column(frid_col, planner.route(batch.column(row_id_name))) map_kwargs = _map_kwargs(ray_remote_args) with_frid = update_ds.map_batches(_assign_frid, **map_kwargs) @@ -814,13 +780,9 @@ def distributed_read_by_row_id( ``ray.data.Dataset`` of ``(*projection, _ROW_ID)``, or ``None`` if the target is empty. Read-side mirror of ``distributed_update_apply``. """ - import numpy as np - import ray - from pypaimon.globalindex.indexed_split import IndexedSplit from pypaimon.table.special_fields import SpecialFields - from pypaimon.utils.range import Range row_id_name = SpecialFields.ROW_ID.name read_cols = list(projection) @@ -847,41 +809,9 @@ def distributed_read_by_row_id( precomputed_info_ref = ray.put(planner) frid_col = "_FIRST_ROW_ID" - sorted_arr = np.asarray(sorted_first_row_ids, dtype=np.int64) - valid_ranges = planner.valid_row_id_ranges - range_starts = np.asarray([r.from_ for r in valid_ranges], dtype=np.int64) - range_ends = np.asarray([r.to for r in valid_ranges], dtype=np.int64) def _assign_frid(batch: pa.Table) -> pa.Table: - if batch.num_rows == 0: - return batch.append_column(frid_col, pa.array([], type=pa.int64())) - rid_col = batch.column(row_id_name) - if rid_col.null_count: - raise ValueError( - "_ROW_ID is null; the planner snapshot is stale or the row ids " - "come from a different table." - ) - rids = rid_col.to_numpy(zero_copy_only=False) - # Foreign-id check: valid_ranges are sorted+merged, so one searchsorted finds - # the candidate range (O(rows log ranges), like distributed_delete_apply). - ridx = np.searchsorted(range_starts, rids, side="right") - 1 - safe = np.clip(ridx, 0, len(range_starts) - 1) - in_range = ( - (ridx >= 0) - & (rids >= range_starts[safe]) - & (rids <= range_ends[safe]) - ) - if not in_range.all(): - bad = rids[~in_range][0] - raise ValueError( - f"_ROW_ID {bad} does not belong to any valid range " - f"{[f'[{r.from_}, {r.to}]' for r in valid_ranges]}; the planner " - f"snapshot is stale or the row ids come from a different table." - ) - idx = np.searchsorted(sorted_arr, rids, side="right") - 1 - return batch.append_column( - frid_col, pa.array(sorted_arr[idx], type=pa.int64()) - ) + return batch.append_column(frid_col, planner.route(batch.column(row_id_name))) captured_table = planner.table # read the captured view used for routing captured_read_cols = read_cols @@ -890,18 +820,12 @@ def _assign_frid(batch: pa.Table) -> pa.Table: def _read_group(group: pa.Table) -> pa.Table: if group.num_rows == 0: return captured_empty - frid = int(group.column(frid_col)[0].as_py()) info = ray.get(precomputed_info_ref) - owning_split, target_files = info.first_row_id_index[frid] - target_names = {file.file_name for file in target_files} - origin_split = owning_split.filter_file(lambda file: file.file_name in target_names) - # Only matched rows (deduped, contiguous ids -> ranges); blob gets row-index pushdown. - wanted = set(group.column(row_id_name).to_pylist()) - indexed = IndexedSplit(origin_split, Range.to_ranges(list(wanted))) + selected = info.select(group.column(row_id_name)) read = captured_table.new_read_builder().with_projection( captured_read_cols ).new_read() - return read.to_arrow([indexed]) + return read.to_arrow(selected.splits()) map_kwargs = _map_kwargs(ray_remote_args) with_frid = row_ids_ds.map_batches(_assign_frid, **map_kwargs) @@ -921,8 +845,6 @@ def distributed_delete_apply( collect_row_ids: bool = False, materialize_before_routing: bool = False, ) -> Tuple[list, int, list]: - import base64 - import numpy as np import pickle import pyarrow.compute as pc @@ -949,58 +871,12 @@ def distributed_delete_apply( precomputed_info_ref = ray.put(anchor_info) - starts = np.asarray( - [a.row_range.from_ for a in anchor_info.anchors], dtype=np.int64 - ) - ends = np.asarray( - [a.row_range.to for a in anchor_info.anchors], dtype=np.int64 - ) - - def _group_key(anchor) -> str: - partition_blob = base64.b64encode( - pickle.dumps(tuple(anchor.partition.values)) - ).decode("ascii") - return f"{anchor.bucket}:{partition_blob}" - - group_keys = [_group_key(a) for a in anchor_info.anchors] - unique_group_count = len(set(group_keys)) + bucket_owners = anchor_info.route_by_bucket(pa.array(anchor_info.first_row_ids, type=pa.int64())) + unique_group_count = len(set(bucket_owners.to_pylist())) group_col = "_DELETE_GROUP_KEY" - valid_ranges = [ - f"[{a.row_range.from_}, {a.row_range.to}]" - for a in anchor_info.anchors - ] def _assign_group(batch: pa.Table) -> pa.Table: - if batch.num_rows == 0: - return batch.append_column( - group_col, pa.array([], type=pa.string()) - ) - rid_col = batch.column(row_id_name) - if rid_col.null_count: - raise ValueError( - "_ROW_ID is null; planner snapshot is stale " - "or matched rows come from a different table." - ) - rids = rid_col.to_numpy(zero_copy_only=False) - idx = np.searchsorted(starts, rids, side="right") - 1 - safe_idx = np.clip(idx, 0, len(starts) - 1) - in_range = ( - (idx >= 0) - & (idx < len(starts)) - & (rids >= starts[safe_idx]) - & (rids <= ends[safe_idx]) - ) - if not in_range.all(): - bad = rids[~in_range][0] - raise ValueError( - f"_ROW_ID {bad} does not belong to any valid range " - f"{valid_ranges}; planner snapshot is stale or matched " - f"rows come from a different table." - ) - return batch.append_column( - group_col, - pa.array([group_keys[i] for i in safe_idx], type=pa.string()), - ) + return batch.append_column(group_col, anchor_info.route_by_bucket(batch.column(row_id_name))) map_kwargs = _map_kwargs(ray_remote_args) with_group = delete_ds.map_batches(_assign_group, **map_kwargs) diff --git a/paimon-python/pypaimon/read/plan.py b/paimon-python/pypaimon/read/plan.py index 88746311dc07..966d2ce9f987 100644 --- a/paimon-python/pypaimon/read/plan.py +++ b/paimon-python/pypaimon/read/plan.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Any, List, Optional @@ -41,21 +41,34 @@ def stats(self) -> Optional[ScanStats]: native_stats = self._native_stats if native_stats is not None: self._scan_stats = ScanStats.from_native(native_stats) - # Python can filter or limit splits after Native planning. Keep the - # actual planning work, but summarize the final returned selection. - files, data_bytes = 0, 0 - for split in self._splits: - for file in split.files: - files += 1 - data_bytes += max(0, file.file_size) - counters = self._scan_stats.planning_counters - counters.update(final_splits=len(self._splits), splits_after_limit=len(self._splits), - final_files=files, planned_data_file_bytes=min(data_bytes, 2 ** 64 - 1)) - if self._stats_limit is not None: - counters['limit'] = self._stats_limit + self._update_split_counters() self._native_stats = None return self._scan_stats + def _update_split_counters(self): + # Keep the actual planning work while describing the returned splits. + if self._scan_stats is None or self._scan_stats.planning_counters is None: + return + files, data_bytes = 0, 0 + for split in self._splits: + for file in split.files: + files += 1 + data_bytes += max(0, file.file_size) + counters = self._scan_stats.planning_counters + counters.update(final_splits=len(self._splits), splits_after_limit=len(self._splits), + final_files=files, planned_data_file_bytes=min(data_bytes, 2 ** 64 - 1)) + if self._stats_limit is not None: + counters['limit'] = self._stats_limit + + def _with_splits(self, splits): + stats = self._scan_stats + if stats is not None: + counters = stats.planning_counters + stats = replace(stats, planning_counters=None if counters is None else dict(counters)) + selected = replace(self, _splits=list(splits), _scan_stats=stats) + selected._update_split_counters() + return selected + def __getstate__(self): # Native handles are process-local. Serialize counters as Python scalars. self.stats() diff --git a/paimon-python/pypaimon/tests/dynamic_bucket_execution_boundary_test.py b/paimon-python/pypaimon/tests/dynamic_bucket_execution_boundary_test.py index 76c47ffbf29f..c865150035cb 100644 --- a/paimon-python/pypaimon/tests/dynamic_bucket_execution_boundary_test.py +++ b/paimon-python/pypaimon/tests/dynamic_bucket_execution_boundary_test.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -"""Shared distributed bucket assignment does not load Python writer execution.""" +"""Shared bucket configuration and Native routing stay outside Python execution.""" from pathlib import Path import subprocess @@ -26,7 +26,7 @@ from pypaimon import CatalogFactory, Schema -def test_shared_hash_assignment_restores_nonempty_snapshot_without_python_execution(tmp_path): +def test_native_bucket_execution_does_not_load_python_routing_or_hash_restoration(tmp_path): warehouse = str(tmp_path / 'warehouse') catalog = CatalogFactory.create({'warehouse': warehouse}) catalog.create_database('default', True) @@ -52,37 +52,60 @@ def test_shared_hash_assignment_restores_nonempty_snapshot_without_python_execut import importlib.abc import sys from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import pyarrow as pa class BlockPythonExecution(importlib.abc.MetaPathFinder): def find_spec(self, fullname, path=None, target=None): if fullname.startswith('pypaimon.execution.python'): - raise AssertionError('Shared assignment loaded Python execution: ' + fullname) + raise AssertionError('Native assignment loaded Python execution: ' + fullname) sys.meta_path.insert(0, BlockPythonExecution()) from pypaimon import CatalogFactory -from pypaimon.index.dynamic_bucket import DynamicBucketAssigner, resolve_bucket_snapshot -from pypaimon.write.input_routing import DynamicBucketInputRouter +from pypaimon.index import dynamic_bucket as shared +from pypaimon.write.backend_dispatch import ( + create_dynamic_bucket_assigner, create_dynamic_bucket_assigner_router, +) + +# Only configuration, limits and validation remain in the shared module. +for name in ['DynamicBucketAssigner', '_PartitionIndex', '_iter_hashes', 'compute_assigner']: + assert not hasattr(shared, name), name +shared.validate_max_buckets(shared.MAX_DYNAMIC_BUCKETS) +shared.validate_bucket_id(shared.SHORT_MAX_VALUE) +assert shared.to_signed_int32(0xffffffff) == -1 catalog = CatalogFactory.create({'warehouse': sys.argv[1]}) table = catalog.get_table('default.records') -snapshot = resolve_bucket_snapshot(table) +snapshot = shared.resolve_bucket_snapshot(table) assert snapshot is not None and snapshot.index_manifest is not None +assert shared.resolve_bucket_snapshot(table, 0) is None +read_builder = shared.bucket_read_builder(table, ('p',), snapshot) +assert read_builder.table.options.scan_snapshot_id() == snapshot.id + +# Native adapters pass the captured baseline and Arrow inputs to the SDK; +# neither constructor nor execution may import the Python HASH implementation. rows = pa.record_batch([pa.array([2, 1], type=pa.int64()), pa.array(['other', 'new']), pa.array(['p', 'p'])], names=['id', 'value', 'p']) -router = DynamicBucketInputRouter(table.table_schema) -partitions, partition_hashes, key_hashes = router.extract_hashes_batch(rows) -# Keep the Rust metadata boundary deterministic. The assignment reads the -# actual nonempty HASH file produced above and restores the original key. -plan = SimpleNamespace(splits=lambda: [SimpleNamespace(bucket=0)]) -with patch('pypaimon.execution.native.read.plan.native_plan', return_value=plan) as native_plan: - assigner = DynamicBucketAssigner(table, 1, 1, 0, 1, -1, snapshot=snapshot) - assert assigner.assign_batch(partitions, partition_hashes, key_hashes) == [(1, True), (0, False)] - assert assigner.assign(partitions[1], partition_hashes[1], key_hashes[1]) == 0 -native_plan.assert_called_once() -assert native_plan.call_args[0][0].options.scan_snapshot_id() == snapshot.id -assert native_plan.call_args[1]['all_files'] is True +hashes = pa.array([42, -3], type=pa.int32()) +routing = pa.record_batch([pa.array([0, 0], type=pa.int32()), hashes], names=['assigner', 'key_hash']) +buckets = pa.array([1, 0], type=pa.int32()) +sdk = SimpleNamespace(DynamicBucketAssignerRouter=Mock(), DynamicBucketAssigner=Mock()) +sdk.DynamicBucketAssignerRouter.return_value.route.return_value = routing +sdk.DynamicBucketAssigner.return_value.restore_snapshot_id = snapshot.id +sdk.DynamicBucketAssigner.return_value.assign.return_value = buckets +native_table = Mock() +with patch.dict(sys.modules, {'pypaimon_rust.datafusion': sdk}), \ + patch('pypaimon.execution.native.write.input_routing.native_runtime_available', return_value=True), \ + patch('pypaimon.execution.native.write.input_routing.create_native_write_table', return_value=native_table): + router = create_dynamic_bucket_assigner_router(table, 1, 1) + assert router.route(rows).equals(routing) + assigner = create_dynamic_bucket_assigner(table, 1, 1, 0, base_snapshot_id=snapshot.id) + assert assigner.restore_snapshot_id == snapshot.id + assert assigner.assign(rows, hashes).equals(buckets) +sdk.DynamicBucketAssigner.assert_called_once_with( + native_table, 1, 1, 0, restore_snapshot_id=snapshot.id, overwrite=False) +sdk.DynamicBucketAssigner.return_value.assign.assert_called_once_with(rows, key_hashes=hashes) +native_table.new_batch_write_builder.assert_not_called() assert not any(name.startswith('pypaimon.execution.python') for name in sys.modules) ''' result = subprocess.run( diff --git a/paimon-python/pypaimon/tests/native_dynamic_input_routing_test.py b/paimon-python/pypaimon/tests/native_dynamic_input_routing_test.py index c23bb481448d..f57647446032 100644 --- a/paimon-python/pypaimon/tests/native_dynamic_input_routing_test.py +++ b/paimon-python/pypaimon/tests/native_dynamic_input_routing_test.py @@ -25,10 +25,10 @@ from pypaimon import Schema from pypaimon.execution.native.write.input_routing import NativeDynamicBucketAssigner, NativeDynamicBucketAssignerRouter -from pypaimon.index.dynamic_bucket import DynamicBucketAssigner +from pypaimon.execution.python.write.dynamic_bucket import DynamicBucketAssigner from pypaimon.tests.rest.rest_server import RESTCatalogServer from pypaimon.write.backend_dispatch import create_dynamic_bucket_assigner, create_dynamic_bucket_assigner_router -from pypaimon.write.input_routing import DynamicBucketInputRouter +from pypaimon.execution.python.write.input_routing import DynamicBucketInputRouter pytestmark = pytest.mark.native_plan @@ -62,9 +62,9 @@ def _commit(table, rows): def _forbid_python_input_execution(): stack = ExitStack() for target in [ - 'pypaimon.write.input_routing.FixedBucketInputRouter.__init__', - 'pypaimon.write.input_routing.DynamicBucketInputRouter.__init__', - 'pypaimon.index.dynamic_bucket.DynamicBucketAssigner.__init__', + 'pypaimon.execution.python.write.input_routing.FixedBucketInputRouter.__init__', + 'pypaimon.execution.python.write.input_routing.DynamicBucketInputRouter.__init__', + 'pypaimon.execution.python.write.dynamic_bucket.DynamicBucketAssigner.__init__', 'pypaimon.execution.python.write.table_write.PythonTableWrite.__init__', ]: stack.enter_context(patch(target, side_effect=AssertionError('Python input execution: ' + target))) diff --git a/paimon-python/pypaimon/tests/native_input_routing_test.py b/paimon-python/pypaimon/tests/native_input_routing_test.py index d0764b640917..a47360bf085b 100644 --- a/paimon-python/pypaimon/tests/native_input_routing_test.py +++ b/paimon-python/pypaimon/tests/native_input_routing_test.py @@ -30,7 +30,7 @@ from pypaimon.ray.shuffle import _make_bucket_udf from pypaimon.write.backend_dispatch import create_input_router, create_postpone_bucket_planner from pypaimon.write.bucket_planning import PostponeBucketPlanning -from pypaimon.write.input_routing import FixedBucketInputRouter, PostponeFixedBucketInputRouter +from pypaimon.execution.python.write.input_routing import FixedBucketInputRouter, PostponeFixedBucketInputRouter pytestmark = pytest.mark.native_plan diff --git a/paimon-python/pypaimon/tests/native_ray_row_id_plan_test.py b/paimon-python/pypaimon/tests/native_ray_row_id_plan_test.py index 1ff0a6c40022..4072c17a6bfb 100644 --- a/paimon-python/pypaimon/tests/native_ray_row_id_plan_test.py +++ b/paimon-python/pypaimon/tests/native_ray_row_id_plan_test.py @@ -18,6 +18,9 @@ """Exercise the actual SDK plan bridge without rebuilding the extension.""" from pathlib import Path +import pickle +import sys +from types import SimpleNamespace from unittest.mock import patch import pyarrow as pa @@ -25,6 +28,7 @@ from pypaimon import CatalogFactory, Schema from pypaimon.execution.native.runtime import native_runtime_available +from pypaimon.execution.python.write.row_id_plan import PythonRowIdPlan from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER from pypaimon.write.backend_dispatch import native_row_id_update_selected from pypaimon.write.row_id_plan import capture_row_id_plan @@ -110,3 +114,103 @@ def test_planned_update_keeps_resolved_write_schema_with_older_files(tmp_path, n read = table.new_read_builder() data = read.new_read().to_arrow(read.new_scan().plan().splits()).sort_by('id') assert data.to_pydict() == {'id': [1, 2], 'value': [0, 0], 'extra': [7, None]} + + +def _run_transported(function, data): + """Execute a serialized Ray callback in-process with the registered SDK.""" + cloudpickle = pytest.importorskip('cloudpickle') + restored = cloudpickle.loads(cloudpickle.dumps(function)) + with patch('pypaimon.read.table_scan.TableScan.plan_all', side_effect=AssertionError('worker replanned')): + return restored(data) + + +class _ArrowDataset: + """Exercise the real Ray callbacks without an installed worker extension.""" + + grouped_keys = [] + + def __init__(self, tables): + self.tables = tables + + def map_batches(self, function, **kwargs): + return _ArrowDataset([_run_transported(function, data) for data in self.tables]) + + def union(self, other): + return _ArrowDataset(self.tables + other.tables) + + def groupby(self, key, **kwargs): + return _ArrowGroups(self.tables, key) + + def iter_batches(self, **kwargs): + return iter(self.tables) + + +class _ArrowGroups: + + def __init__(self, tables, key): + self.data = pa.concat_tables(tables) + self.key = key + + def map_groups(self, function, **kwargs): + positions = {} + for offset, key in enumerate(self.data.column(self.key).to_pylist()): + positions.setdefault(key, []).append(offset) + _ArrowDataset.grouped_keys.append((self.key, self.data.schema.field(self.key).type, list(positions))) + return _ArrowDataset([ + _run_transported(function, self.data.take(pa.array(rows, type=pa.int64()))) + for rows in positions.values()]) + + +def _dataset(data): + return _ArrowDataset([data.slice(0, 0), data.slice(0, 1), data.slice(1)]) + + +@pytest.mark.parametrize('native', _ENGINES) +def test_ray_callbacks_route_read_update_and_delete_with_captured_core_plans(tmp_path, native, monkeypatch): + from contextlib import ExitStack + from pypaimon.ray.data_evolution_merge_join import ( + distributed_delete_apply, distributed_read_by_row_id, distributed_update_apply) + + table = _table(tmp_path) + _append(table, [1, 2]) + _append(table, [3, 4]) + table = table.copy({'write.native.enabled': str(native).lower(), 'read.native.enabled': str(native).lower()}) + snapshot_id = table.snapshot_manager().get_latest_snapshot().id + scheduler = SimpleNamespace( + put=lambda value: pickle.loads(pickle.dumps(value)), get=lambda value: value, + data=SimpleNamespace(from_arrow=lambda data: _ArrowDataset([data]))) + monkeypatch.setitem(sys.modules, 'ray', scheduler) + _ArrowDataset.grouped_keys = [] + + with ExitStack() as guards: + if native: + guards.enter_context(patch.object(PythonRowIdPlan, '_positions', + side_effect=AssertionError('Python row-ID routing'))) + source = _dataset(pa.table({'_ROW_ID': pa.array([3, 0, 3], type=pa.int64())})) + read = distributed_read_by_row_id(source, table, ['id', 'value'], num_partitions=2) + rows = pa.concat_tables(list(read.iter_batches())).sort_by('id') + assert rows.to_pydict() == {'id': [1, 4], 'value': [0, 0], '_ROW_ID': [0, 3]} + + update_schema = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False), ('value', pa.int32())]) + updates = _dataset(pa.table({'_ROW_ID': [3, 0], 'value': [9, 8]}, schema=update_schema)) + messages, count, row_ids = distributed_update_apply( + updates, table, ['value'], num_partitions=2, base_snapshot_id=snapshot_id, collect_row_ids=True) + assert count == 2 + assert sorted(row_ids) == [0, 3] + assert all(message.check_from_snapshot == snapshot_id for message in messages) + table.new_batch_write_builder().new_commit().commit(messages) + + deletes = _dataset(pa.table({'_ROW_ID': pa.array([1, 2], type=pa.int64())})) + messages, count, row_ids = distributed_delete_apply( + deletes, table, num_partitions=2, collect_row_ids=True) + assert count == 2 + assert sorted(row_ids) == [1, 2] + assert _ArrowDataset.grouped_keys[-1] == ('_DELETE_GROUP_KEY', pa.int64(), [0]) + table.new_batch_write_builder().new_commit().commit(messages) + + # Routing uses physical ranges; the selected read retains the anchor + # deletion vectors and removes the deleted IDs at the actual reader. + source = _dataset(pa.table({'_ROW_ID': pa.array([0, 1, 2, 3], type=pa.int64())})) + read = distributed_read_by_row_id(source, table, ['id', 'value'], num_partitions=2) + rows = pa.concat_tables(list(read.iter_batches())).sort_by('id') + assert rows.to_pydict() == {'id': [1, 4], 'value': [8, 9], '_ROW_ID': [0, 3]} diff --git a/paimon-python/pypaimon/tests/noncompact_null_row_test.py b/paimon-python/pypaimon/tests/noncompact_null_row_test.py index 24a0181dc725..ecc8102e3504 100644 --- a/paimon-python/pypaimon/tests/noncompact_null_row_test.py +++ b/paimon-python/pypaimon/tests/noncompact_null_row_test.py @@ -25,7 +25,7 @@ from pypaimon.schema.data_types import AtomicType, DataField, PyarrowFieldParser from pypaimon.schema.table_schema import TableSchema from pypaimon.table.row.generic_row import GenericRow, GenericRowDeserializer, GenericRowSerializer -from pypaimon.write.input_routing import FixedBucketInputRouter, _hash_bytes_by_words +from pypaimon.execution.python.write.input_routing import FixedBucketInputRouter, _hash_bytes_by_words # Current Java InternalRowSerializer for [NULL, BinaryString("abcdefghijk")]. diff --git a/paimon-python/pypaimon/tests/pushdown_bucket_test.py b/paimon-python/pypaimon/tests/pushdown_bucket_test.py index d2d72bb27cbf..3d3856c2363b 100644 --- a/paimon-python/pypaimon/tests/pushdown_bucket_test.py +++ b/paimon-python/pypaimon/tests/pushdown_bucket_test.py @@ -62,7 +62,7 @@ from pypaimon.execution.python.write.row_key_extractor import ( FixedBucketRowKeyExtractor, ) -from pypaimon.write.input_routing import ( +from pypaimon.execution.python.write.input_routing import ( _bucket_from_hash, _hash_bytes_by_words, ) diff --git a/paimon-python/pypaimon/tests/write/distributed_input_execution_test.py b/paimon-python/pypaimon/tests/write/distributed_input_execution_test.py index d45a511ac599..c141b7dad38e 100644 --- a/paimon-python/pypaimon/tests/write/distributed_input_execution_test.py +++ b/paimon-python/pypaimon/tests/write/distributed_input_execution_test.py @@ -24,11 +24,11 @@ import pytest from pypaimon import CatalogFactory, Schema -from pypaimon.index.dynamic_bucket import DynamicBucketAssigner +from pypaimon.execution.python.write.dynamic_bucket import DynamicBucketAssigner from pypaimon.read.plan import Plan from pypaimon.read.split import DataSplit from pypaimon.read.split_serializer import deserialize_split_v1, serialize_split_v1 -from pypaimon.write.input_routing import DynamicBucketInputRouter, FixedBucketInputRouter +from pypaimon.execution.python.write.input_routing import DynamicBucketInputRouter, FixedBucketInputRouter def _table(tmp_path, bucket): @@ -76,7 +76,8 @@ def test_ray_fixed_grouping_executes_input_router_without_python_extractor(tmp_p def test_daft_fixed_and_dynamic_routing_do_not_construct_python_extractors(tmp_path, forbid_python_execution): daft = pytest.importorskip('daft') from pypaimon.daft.daft_datasink import make_dynamic_routing_udf, make_fixed_bucket_udf - from pypaimon.index.dynamic_bucket import compute_assigner, to_signed_int32 + from pypaimon.index.dynamic_bucket import to_signed_int32 + from pypaimon.execution.python.write.dynamic_bucket import compute_assigner table = _table(tmp_path, '4') rows = {'id': [1, 2, 1], 'value': ['a', 'b', 'c'], 'p': ['p', 'p', 'q']} @@ -135,8 +136,8 @@ def test_daft_dynamic_assignment_dispatches_pinned_snapshot_and_arrow_hashes(tmp 'execution.python.write.row_key_extractor.SimpleDynamicBucketAssigner.__init__', 'execution.python.read.scanner.file_scanner.FileScanner.__init__', 'execution.python.write.dynamic_bucket.DynamicBucketIndexMaintainer.__init__', - 'write.input_routing.DynamicBucketInputRouter.__init__', - 'index.dynamic_bucket.DynamicBucketAssigner.__init__', + 'execution.python.write.input_routing.DynamicBucketInputRouter.__init__', + 'execution.python.write.dynamic_bucket.DynamicBucketAssigner.__init__', ]: stack.enter_context(patch('pypaimon.' + target, side_effect=AssertionError(target))) stack.enter_context(patch('pypaimon.daft.daft_datasink._load_worker_table', return_value=table)) diff --git a/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py b/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py index 8db7e598f236..100aabdf622e 100644 --- a/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py +++ b/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py @@ -26,13 +26,9 @@ import pytest from pypaimon import CatalogFactory, Schema -from pypaimon.index.dynamic_bucket import ( - DynamicBucketAssigner, - _PartitionIndex, - _iter_hashes, - compute_assigner, - to_signed_int32, - validate_bucket_id, +from pypaimon.index.dynamic_bucket import to_signed_int32, validate_bucket_id +from pypaimon.execution.python.write.dynamic_bucket import ( + DynamicBucketAssigner, _PartitionIndex, _iter_hashes, compute_assigner, ) from pypaimon.index.index_file_handler import IndexFileHandler from pypaimon.index.index_file_meta import IndexFileMeta @@ -261,7 +257,7 @@ def test_restore_rejects_bucket_id_above_java_short_range(self): entry.bucket = 32768 with patch( - 'pypaimon.index.dynamic_bucket.IndexFileHandler' + 'pypaimon.execution.python.write.dynamic_bucket.IndexFileHandler' ) as handler: handler.return_value.scan.return_value = [entry] with patch.object( diff --git a/paimon-python/pypaimon/tests/write/dynamic_input_routing_test.py b/paimon-python/pypaimon/tests/write/dynamic_input_routing_test.py index 24ec5c9e88a0..9210b2b65f90 100644 --- a/paimon-python/pypaimon/tests/write/dynamic_input_routing_test.py +++ b/paimon-python/pypaimon/tests/write/dynamic_input_routing_test.py @@ -161,7 +161,7 @@ def test_native_dynamic_router_pickle_rebuilds_without_python_hashing(tmp_path, with patch('pypaimon.execution.native.write.input_routing.native_runtime_available', return_value=True), \ patch('pypaimon.execution.native.write.input_routing.create_native_write_table', return_value=native_table), \ - patch('pypaimon.write.input_routing.DynamicBucketInputRouter.__init__', + patch('pypaimon.execution.python.write.input_routing.DynamicBucketInputRouter.__init__', side_effect=AssertionError('Python hashing')): router = create_dynamic_bucket_assigner_router(table, 2, 2) assert isinstance(router, NativeDynamicBucketAssignerRouter) @@ -264,7 +264,7 @@ def test_native_assignment_uses_dedicated_constructor_and_precomputed_arrow_hash patch('pypaimon.execution.native.write.input_routing.create_native_write_table', return_value=native_table), \ patch.object(table, 'new_batch_write_builder', side_effect=AssertionError('Python write builder')), \ - patch('pypaimon.index.dynamic_bucket.DynamicBucketAssigner.__init__', + patch('pypaimon.execution.python.write.dynamic_bucket.DynamicBucketAssigner.__init__', side_effect=AssertionError('Python index restoration')): assigner = create_dynamic_bucket_assigner(table, 2, 2, 1, overwrite, base_snapshot_id) assert isinstance(assigner, NativeDynamicBucketAssigner) diff --git a/paimon-python/pypaimon/tests/write/native_input_routing_test.py b/paimon-python/pypaimon/tests/write/native_input_routing_test.py index fdb5e70becf7..58efed762dfc 100644 --- a/paimon-python/pypaimon/tests/write/native_input_routing_test.py +++ b/paimon-python/pypaimon/tests/write/native_input_routing_test.py @@ -31,7 +31,7 @@ from pypaimon.write.backend_dispatch import create_input_router, create_postpone_bucket_planner from pypaimon.write.bucket_plan import PostponeBucketPlan from pypaimon.write.bucket_planning import PostponeBucketPlanning -from pypaimon.write.input_routing import FixedBucketInputRouter +from pypaimon.execution.python.write.input_routing import FixedBucketInputRouter def _table(tmp_path, bucket='4', partition='p', native=True): diff --git a/paimon-python/pypaimon/tests/write/row_id_routing_test.py b/paimon-python/pypaimon/tests/write/row_id_routing_test.py new file mode 100644 index 000000000000..171fdb1ab1a8 --- /dev/null +++ b/paimon-python/pypaimon/tests/write/row_id_routing_test.py @@ -0,0 +1,275 @@ +# 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. + +"""Captured row-ID execution preserves row order, ownership and read grants.""" + +from dataclasses import replace +import pickle +import sys +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pyarrow as pa +import pytest + +from pypaimon.catalog.table_query_auth import TableQueryAuthResult +from pypaimon.execution.native.write.row_id_plan import NativeRowIdPlan +from pypaimon.execution.python.write.row_id_plan import PythonRowIdPlan +from pypaimon.manifest.schema.data_file_meta import DataFileMeta +from pypaimon.manifest.schema.simple_stats import SimpleStats +from pypaimon.read.plan import Plan +from pypaimon.read.query_auth_split import QueryAuthSplit +from pypaimon.read.scan_stats import ScanStats, _NATIVE_COUNTER_FIELDS +from pypaimon.read.split import DataSplit +from pypaimon.schema.data_types import AtomicType, DataField +from pypaimon.table.row.generic_row import GenericRow +from pypaimon.table.source.deletion_file import DeletionFile +from pypaimon.utils.range import Range +from pypaimon.write.row_id_plan import CapturedRowIdPlan + +pytestmark = [pytest.mark.python_plan, pytest.mark.python_read, + pytest.mark.python_write, pytest.mark.python_commit] + + +def _file(name, first, count): + empty = GenericRow([], []) + return DataFileMeta.create( + name, 100, count, empty, empty, SimpleStats.empty_stats(), SimpleStats.empty_stats(), + 0, 0, 0, 0, [], first_row_id=first, external_path='memory:/external/' + name, + file_path='memory:/external/' + name) + + +def _split(first, count, partition, bucket=0, dedicated=False): + files = [_file('data-{}.parquet'.format(first), first, count)] + if dedicated: + files += [replace(files[0], file_name='delta.parquet', max_sequence_number=1, + write_cols=['value']), _file('payload.blob', first, count)] + field = DataField(0, 'p', AtomicType('STRING')) + return DataSplit(files, GenericRow([partition], [field]), bucket, + data_deletion_files=[DeletionFile('memory:/dv', 0, 8, 1)] * len(files), + snapshot_id=7, bucket_path='memory:/captured/bucket-{}'.format(bucket), + total_buckets=4) + + +@pytest.fixture +def captured(): + grant = TableQueryAuthResult(['value > 0'], {'value': 'value + 1'}) + splits = [_split(50, 2, 'west', 1), _split(20, 3, 'west'), + QueryAuthSplit(_split(10, 3, 'west', dedicated=True), grant), + _split(30, 2, 'east'), _split((1 << 63) - 1, 1, 'west')] + plan = Plan(splits, 7, _scan_stats=ScanStats(entries_total=17), _stats_limit=3) + execution = PythonRowIdPlan(plan) + return CapturedRowIdPlan(None, plan, execution.groups, execution) + + +def test_python_routing_keeps_unsorted_duplicates_and_inclusive_ends(captured): + row_ids = pa.array([22, 10, 31, 11, 20, 51, 22, (1 << 63) - 1], type=pa.int64()) + assert captured.route(row_ids).to_pylist() == [20, 10, 30, 10, 20, 50, 20, (1 << 63) - 1] + # Two independent anchors in one partition/bucket share the oldest owner. + # The same bucket in a different partition and a different bucket do not. + assert captured.route_by_bucket(row_ids).to_pylist() == [10, 10, 30, 10, 10, 50, 10, 10] + + +@pytest.mark.parametrize('operation', ['route', 'route_by_bucket', 'select']) +@pytest.mark.parametrize('row_id', [-1, 9, 13, 19, 23, 32, 49, 52]) +def test_python_routing_rejects_foreign_ids_and_gaps(captured, operation, row_id): + with pytest.raises(ValueError, match='valid range'): + getattr(captured, operation)(pa.array([10, row_id], type=pa.int64())) + + +@pytest.mark.parametrize('operation', ['route', 'route_by_bucket', 'select']) +@pytest.mark.parametrize('row_ids,error', [ + (pa.array([10.5]), TypeError), (pa.array([True]), TypeError), (pa.array(['10']), TypeError), + (pa.array([10, None], type=pa.int64()), ValueError), + (pa.array([1 << 63], type=pa.uint64()), ValueError), ([10], TypeError), +]) +def test_public_routing_validates_ids_before_backend_execution(captured, operation, row_ids, error): + execution = Mock() + captured = replace(captured, _execution=execution) + with pytest.raises(error): + getattr(captured, operation)(row_ids) + getattr(execution, operation).assert_not_called() + + +def test_chunked_integer_input_is_normalized_without_float_coercion(captured): + row_ids = pa.chunked_array([pa.array([10], type=pa.int32()), pa.array([22, 31], type=pa.int32())]) + assert captured.route(row_ids).to_pylist() == [10, 20, 30] + + +def test_selection_deduplicates_absolute_ids_and_preserves_complete_groups(captured): + original_splits = captured.plan.splits()[:] + selected = captured.select(pa.array([21, 12, 10, 11, 21, (1 << 63) - 1], type=pa.int64())) + assert selected.snapshot_id == 7 + assert selected._scan_stats is not captured.plan._scan_stats + assert selected._scan_stats == captured.plan._scan_stats + assert selected._stats_limit == 3 + assert [split.row_ranges() for split in selected.splits()] == [ + [Range(10, 12)], [Range(21, 21)], [Range((1 << 63) - 1, (1 << 63) - 1)]] + authorized = selected.splits()[0] + original = next(split for split in original_splits if isinstance(split, QueryAuthSplit)) + assert isinstance(authorized, QueryAuthSplit) + assert authorized.auth_result is original.auth_result + assert authorized.files == original.files + assert len(authorized.files) == 3 + assert authorized.data_deletion_files == original.data_deletion_files + assert authorized.split.data_split().bucket_path == original.bucket_path + assert authorized.split.data_split().total_buckets == original.total_buckets + assert authorized.snapshot_id == original.snapshot_id + assert captured.plan.splits() == original_splits + assert captured.groups[0][2].bucket_path == original.bucket_path + + +def test_empty_routing_and_selection_keep_snapshot(captured): + ids = pa.array([], type=pa.int64()) + for view in (captured, CapturedRowIdPlan(None, Plan([], None), [], PythonRowIdPlan(Plan([], None)))): + assert view.route(ids).to_pylist() == [] + assert view.route_by_bucket(ids).to_pylist() == [] + assert view.select(ids).splits() == [] + assert view.select(ids).snapshot_id == view.plan.snapshot_id + + +def test_python_captured_routing_pickle_and_complete_group_scope(captured): + restored = pickle.loads(pickle.dumps(captured)) + assert restored.snapshot_id == 7 + assert restored.route(pa.array([22, 10])).to_pylist() == [20, 10] + single = restored.for_split(restored.groups[1][2]) + assert single.route_by_bucket(pa.array([20, 22])).to_pylist() == [20, 20] + with pytest.raises(ValueError, match='valid range'): + single.route(pa.array([10])) + + +@pytest.mark.parametrize('operation', ['route', 'route_by_bucket', 'select']) +def test_native_routing_is_arrow_transport_and_never_replays_python(captured, operation): + core = Mock() + expected = (Plan([], 7) if operation == 'select' else pa.array([10, 20], type=pa.int64())) + getattr(core, operation).return_value = expected + execution = NativeRowIdPlan(None, captured.plan, native=core, groups=captured.groups) + captured = replace(captured, _execution=execution) + source = pa.chunked_array([pa.array([12], type=pa.int32()), pa.array([20], type=pa.int32())]) + with patch.object(PythonRowIdPlan, '__init__', side_effect=AssertionError('Python replay')), \ + patch('pypaimon.execution.native.read.plan._from_native_plan', return_value=expected) as decode: + actual = getattr(captured, operation)(source) + if operation == 'select': + assert actual == expected + assert actual._scan_stats is not captured.plan._scan_stats + assert actual._scan_stats == captured.plan._scan_stats + assert actual._stats_limit == 3 + else: + assert actual is expected + transported = getattr(core, operation).call_args.args[0] + assert transported.type == pa.int64() + assert transported.to_pylist() == [12, 20] + assert decode.call_count == int(operation == 'select') + failure = PermissionError('injected Core routing error') + getattr(core, operation).side_effect = failure + with patch.object(PythonRowIdPlan, '__init__', side_effect=AssertionError('Python replay')): + with pytest.raises(PermissionError) as raised: + getattr(captured, operation)(source) + assert raised.value is failure + + +def test_native_routing_pickle_rebuilds_only_the_captured_plan(captured, monkeypatch): + sdk = SimpleNamespace(RowIdPlan=Mock()) + monkeypatch.setitem(sys.modules, 'pypaimon_rust.datafusion', sdk) + sdk.RowIdPlan.return_value.route.return_value = pa.array([20, 10], type=pa.int64()) + execution = NativeRowIdPlan(None, captured.plan, native=lambda: None, groups=captured.groups) + restored = pickle.loads(pickle.dumps(replace(captured, _execution=execution))) + assert restored._execution._native is None + assert restored.snapshot_id == 7 + wire = object() + with patch('pypaimon.execution.native.write.row_id_plan.native_plan', return_value=wire) as transport, \ + patch('pypaimon.execution.native.write.row_id_plan.create_native_row_id_plan', + side_effect=AssertionError('backend reselected')), \ + patch('pypaimon.read.table_scan.TableScan.plan_all', side_effect=AssertionError('worker replanned')), \ + patch.object(PythonRowIdPlan, '__init__', side_effect=AssertionError('Python replay')): + assert restored.route(pa.array([22, 10])).to_pylist() == [20, 10] + assert restored.route(pa.array([22, 10])).to_pylist() == [20, 10] + sdk.RowIdPlan.assert_called_once_with(wire) + assert transport.call_args.args[1].snapshot_id == 7 + assert transport.call_args.args[1].splits() == restored.plan.splits() + scoped = restored.for_split(restored.groups[0][2]) + assert isinstance(scoped._execution, NativeRowIdPlan) + assert len(scoped.plan.splits()) == 1 + + +@pytest.mark.parametrize('dedicated_name', ['payload.blob', 'embedding.vector.parquet']) +def test_dedicated_files_do_not_merge_normal_owners_or_fill_row_id_gaps(dedicated_name): + original = _split(10, 3, 'west') + other = _file('data-20.parquet', 20, 3) + dedicated = _file(dedicated_name, 10, 13) + files = original.files + [other, dedicated] + deletion_files = [DeletionFile('memory:/dv-' + str(index), 0, 8, 1) for index in range(3)] + split = DataSplit(files, original.partition, 0, data_deletion_files=deletion_files, + snapshot_id=7, bucket_path=original.bucket_path, total_buckets=4) + grant = TableQueryAuthResult(['value > 0'], {}) + plan = Plan([QueryAuthSplit(split, grant)], 7) + execution = PythonRowIdPlan(plan) + captured = CapturedRowIdPlan(None, plan, execution.groups, execution) + assert [(first, last) for first, last, _ in captured.groups] == [(10, 12), (20, 22)] + assert captured.route(pa.array([10, 21])).to_pylist() == [10, 20] + assert captured.route_by_bucket(pa.array([10, 21])).to_pylist() == [10, 10] + for first, _, group in captured.groups: + assert isinstance(group, QueryAuthSplit) + assert group.auth_result is grant + assert group.bucket_path == split.bucket_path + assert group.total_buckets == 4 + assert {file.file_name for file in group.files} == {'data-{}.parquet'.format(first), dedicated_name} + assert group.data_deletion_files == [deletion_files[0 if first == 10 else 1], deletion_files[2]] + with pytest.raises(ValueError, match='valid range'): + captured.route(pa.array([15])) + selected = captured.select(pa.array([21, 11])) + assert [group.row_ranges() for group in selected.splits()] == [[Range(11, 11)], [Range(21, 21)]] + assert all(dedicated in group.files for group in selected.splits()) + + +def test_overlapping_normal_anchors_with_different_ranges_are_rejected(): + original = _split(10, 3, 'west') + overlap = _file('overlap.parquet', 11, 3) + split = DataSplit(original.files + [overlap], original.partition, 0, snapshot_id=7) + with pytest.raises(ValueError, match='identical ranges'): + PythonRowIdPlan(Plan([split], 7)) + + +@pytest.mark.parametrize('state', ['lazy', 'materialized', 'pickled']) +@pytest.mark.parametrize('selected_ids', [[], [10, 21]]) +def test_selected_plan_stats_keep_work_and_update_output_without_mutating_original(captured, state, selected_ids): + counters = {field: 0 for field in _NATIVE_COUNTER_FIELDS} + counters.update(manifest_entries_read=99, manifest_entries_after_merge=27) + captured.plan._scan_stats = None + captured.plan._native_stats = SimpleNamespace(**counters) + if state == 'materialized': + captured.plan.stats() + elif state == 'pickled': + captured = pickle.loads(pickle.dumps(captured)) + selected = captured.select(pa.array(selected_ids, type=pa.int64())) + source = dict(captured.plan.stats().planning_counters) + selected_stats = selected.stats().planning_counters + assert selected_stats['manifest_entries_read'] == 99 + assert selected_stats['manifest_entries_after_merge'] == 27 + assert selected_stats['limit'] == 3 + assert selected_stats['final_splits'] == selected_stats['splits_after_limit'] == len(selected_ids) + assert selected_stats['final_files'] == (4 if selected_ids else 0) + assert selected_stats['planned_data_file_bytes'] == (400 if selected_ids else 0) + assert captured.plan.stats().planning_counters == source + assert selected.stats() is not captured.plan.stats() + + +def test_selecting_python_stats_does_not_invent_native_counters(captured): + assert captured.plan.stats().planning_counters is None + selected = captured.select(pa.array([10])) + assert selected.stats().planning_counters is None + assert selected.stats() is not captured.plan.stats() diff --git a/paimon-python/pypaimon/write/backend_dispatch.py b/paimon-python/pypaimon/write/backend_dispatch.py index f6f20f269e26..0181174f5a60 100644 --- a/paimon-python/pypaimon/write/backend_dispatch.py +++ b/paimon-python/pypaimon/write/backend_dispatch.py @@ -214,9 +214,9 @@ def prepare(): if native is not None: return native if bucket_plan is not None: - from pypaimon.write.input_routing import PostponeFixedBucketInputRouter + from pypaimon.execution.python.write.input_routing import PostponeFixedBucketInputRouter return PostponeFixedBucketInputRouter(table, bucket_plan) - from pypaimon.write.input_routing import FixedBucketInputRouter + from pypaimon.execution.python.write.input_routing import FixedBucketInputRouter return FixedBucketInputRouter(table.table_schema) @@ -395,6 +395,16 @@ def create_shard_update_execution(owner, table, projection, columns, commit_user return create_python_shard_updator(owner, table, projection, columns, commit_user, shard, shard_count) +def create_row_id_plan_execution(table, plan): + """Select pure routing once; execution errors never replay another backend.""" + from pypaimon.execution.native.write.row_id_plan import create_native_row_id_plan, NativeRowIdPlan + native = create_native_row_id_plan(table, plan) + if native is not None: + return NativeRowIdPlan(table, plan, native=native) + from pypaimon.execution.python.write.row_id_plan import PythonRowIdPlan + return PythonRowIdPlan(plan) + + def create_row_id_update_execution(table, commit_user, commit_identifier, precomputed=None): """Select one incremental updater before capturing target files.""" from pypaimon.write.update_validation import _authorize_update diff --git a/paimon-python/pypaimon/write/input_routing.py b/paimon-python/pypaimon/write/input_routing.py deleted file mode 100644 index d2d049690252..000000000000 --- a/paimon-python/pypaimon/write/input_routing.py +++ /dev/null @@ -1,343 +0,0 @@ -# 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 math -import struct -from typing import Any, Dict, List, Tuple - -import pyarrow as pa - -from pypaimon.common.options.core_options import CoreOptions -from pypaimon.schema.data_types import AtomicType -from pypaimon.schema.table_schema import TableSchema -from pypaimon.table.bucket_mode import BucketMode -from pypaimon.table.row.generic_row import GenericRow, GenericRowSerializer -from pypaimon.table.row.internal_row import RowKind - - -_MURMUR_C1 = 0xCC9E2D51 -_MURMUR_C2 = 0x1B873593 -_DEFAULT_SEED = 42 - - -def _mix_k1(k1: int) -> int: - k1 = (k1 * _MURMUR_C1) & 0xFFFFFFFF - k1 = ((k1 << 15) | (k1 >> 17)) & 0xFFFFFFFF - k1 = (k1 * _MURMUR_C2) & 0xFFFFFFFF - return k1 - - -def _mix_h1(h1: int, k1: int) -> int: - h1 = (h1 ^ k1) & 0xFFFFFFFF - h1 = ((h1 << 13) | (h1 >> 19)) & 0xFFFFFFFF - h1 = (h1 * 5 + 0xE6546B64) & 0xFFFFFFFF - return h1 - - -def _fmix(h1: int, length: int) -> int: - h1 = (h1 ^ length) & 0xFFFFFFFF - h1 ^= h1 >> 16 - h1 = (h1 * 0x85EBCA6B) & 0xFFFFFFFF - h1 ^= h1 >> 13 - h1 = (h1 * 0xC2B2AE35) & 0xFFFFFFFF - h1 ^= h1 >> 16 - return h1 - - -def _hash_bytes_by_words(data: bytes, seed: int = _DEFAULT_SEED) -> int: - n = len(data) - length_aligned = n - (n % 4) - h1 = seed - for i in range(0, length_aligned, 4): - k1 = struct.unpack_from(" int: - if hash_unsigned >= 0x80000000: - hash_signed = hash_unsigned - 0x100000000 - else: - hash_signed = hash_unsigned - rem = hash_signed - math.trunc(hash_signed / num_buckets) * num_buckets - return abs(rem) - - -class PartitionKeyExtractor: - """Resolve input routing fields and hash their Java binary layout.""" - - def __init__(self, table_schema: TableSchema): - self.table_schema = table_schema - self.partition_indices = self._get_field_indices(table_schema.partition_keys) - - def extract_partition_bucket_batch(self, data: pa.RecordBatch) -> Tuple[List[Tuple], List[int]]: - return self._extract_partitions_batch(data), self._extract_buckets_batch(data) - - def extract_partition_bucket_row( - self, values_by_name: Dict[str, Any]) -> Tuple[Tuple, int]: - partition = tuple( - values_by_name[self.table_schema.fields[i].name] - for i in self.partition_indices - ) - bucket = self._extract_bucket_row(values_by_name) - return partition, bucket - - def _get_field_indices(self, field_names: List[str]) -> List[int]: - if not field_names: - return [] - field_map = {field.name: i for i, field in enumerate(self.table_schema.fields)} - return [field_map[name] for name in field_names if name in field_map] - - def _extract_partitions_batch(self, data: pa.RecordBatch) -> List[Tuple]: - if not self.partition_indices: - return [() for _ in range(data.num_rows)] - - partition_columns = self._input_columns(data, self.partition_indices) - - partitions = [] - for row_idx in range(data.num_rows): - partition_values = tuple(col[row_idx].as_py() for col in partition_columns) - partitions.append(partition_values) - - return partitions - - def _input_columns(self, data: pa.RecordBatch, field_indices: List[int]): - """Resolve routing fields in the actual write type, as Java's converter does.""" - columns = [] - for field_index in field_indices: - name = self.table_schema.fields[field_index].name - index = data.schema.get_field_index(name) - if index < 0: - raise ValueError(f"Input must contain exactly one routing column '{name}'.") - columns.append(data.column(index)) - return columns - - @staticmethod - def _binary_row_hash_code(values: Tuple, fields: List) -> int: - return _hash_bytes_by_words( - GenericRowSerializer.to_bytes( - GenericRow(list(values), fields, RowKind.INSERT) - )[4:] - ) - - -class PostponeFixedBucketInputRouter(PartitionKeyExtractor): - """Group distributed inputs by the same partition and bucket as writers.""" - - def __init__(self, table, bucket_plan): - super().__init__(table.table_schema) - if table.options.bucket() != BucketMode.POSTPONE_BUCKET.value: - raise ValueError( - "Postpone fixed bucket writes require bucket = -2, got {}".format( - table.options.bucket() - ) - ) - bucket_function = str( - table.table_schema.options.get("bucket-function.type", "default") - ).strip().lower() - if bucket_function != "default": - raise ValueError( - "Postpone fixed bucket writes only support " - "bucket-function.type=default, got {}" - .format(bucket_function) - ) - self.bucket_keys = table.table_schema.bucket_keys - self.bucket_key_indices = self._get_field_indices(self.bucket_keys) - self._bucket_key_fields = table.table_schema.logical_bucket_key_fields - self._bucket_plan = bucket_plan - - def with_bucket_plan(self, bucket_plan) -> None: - self._bucket_plan = bucket_plan - - def num_buckets(self, partition: Tuple) -> int: - return self._bucket_plan.num_buckets(partition) - - def route(self, data) -> pa.Int32Array: - """Return bucket IDs in input order, as the Native router does.""" - return pa.array(self.extract_partition_bucket_batch(data)[1], type=pa.int32()) - - def extract_partition_bucket_batch( - self, data: pa.RecordBatch - ) -> Tuple[List[Tuple], List[int]]: - partitions = self._extract_partitions_batch(data) - columns = self._input_columns(data, self.bucket_key_indices) - buckets = [ - _bucket_from_hash( - self._binary_row_hash_code( - tuple(col[row_idx].as_py() for col in columns), - self._bucket_key_fields, - ), - self.num_buckets(partition), - ) - for row_idx, partition in enumerate(partitions) - ] - return partitions, buckets - - def _extract_buckets_batch(self, data: pa.RecordBatch) -> List[int]: - return self.extract_partition_bucket_batch(data)[1] - - def _extract_bucket_row(self, values_by_name: Dict[str, Any]) -> int: - partition = tuple( - values_by_name[self.table_schema.fields[i].name] - for i in self.partition_indices - ) - return _bucket_from_hash( - self._binary_row_hash_code( - tuple(values_by_name[name] for name in self.bucket_keys), - self._bucket_key_fields, - ), - self.num_buckets(partition), - ) - - -class FixedBucketInputRouter(PartitionKeyExtractor): - """Route distributed inputs without constructing a Python file writer.""" - - def __init__(self, table_schema: TableSchema): - super().__init__(table_schema) - options = CoreOptions.from_dict(table_schema.options) - self.num_buckets = options.bucket() - if self.num_buckets <= 0: - raise ValueError(f"Fixed bucket mode requires bucket > 0, got {self.num_buckets}") - - # Bucket-key resolution lives on TableSchema (mirrors Java - # ``TableSchema.bucketKeys()`` / ``logicalBucketKeyType()``); reuse - # it so any reader path that walks the same logic stays in sync. - self.bucket_keys = table_schema.bucket_keys - self.bucket_key_indices = self._get_field_indices(self.bucket_keys) - self._bucket_key_fields = table_schema.logical_bucket_key_fields - bucket_function = table_schema.options.get('bucket-function.type') - self.bucket_function = 'default' if bucket_function is None else str(bucket_function).lower() - if self.bucket_function == 'mod': - if len(self._bucket_key_fields) != 1: - raise ValueError('bucket key must have exactly one field in mod bucket function') - key_type = self._bucket_key_fields[0].type - if not isinstance(key_type, AtomicType) or key_type.type.upper() not in ('INT', 'INTEGER', 'BIGINT'): - raise ValueError( - 'bucket key type must be INT or BIGINT in mod bucket function, but got %s' % key_type) - elif self.bucket_function not in ('default', 'hive'): - raise ValueError('Unsupported bucket function type: %s' % self.bucket_function) - - def route(self, data) -> pa.Int32Array: - """Return bucket IDs in input order, as the Native router does.""" - return pa.array(self.extract_partition_bucket_batch(data)[1], type=pa.int32()) - - def _extract_buckets_batch(self, data: pa.RecordBatch) -> List[int]: - self._check_automatic_bucket_function() - columns = self._input_columns(data, self.bucket_key_indices) - if self.bucket_function == 'mod': - return [self._mod_bucket(value.as_py()) for value in columns[0]] - return [ - _bucket_from_hash( - self._binary_row_hash_code( - tuple(col[row_idx].as_py() for col in columns), - self._bucket_key_fields, - ), - self.num_buckets, - ) - for row_idx in range(data.num_rows) - ] - - def _extract_bucket_row(self, values_by_name: Dict[str, Any]) -> int: - self._check_automatic_bucket_function() - if self.bucket_function == 'mod': - return self._mod_bucket(values_by_name[self.bucket_keys[0]]) - return _bucket_from_hash( - self._binary_row_hash_code( - tuple(values_by_name[name] for name in self.bucket_keys), - self._bucket_key_fields, - ), - self.num_buckets, - ) - - def _check_automatic_bucket_function(self): - # Explicit buckets do not use this method. A single bucket also keeps - # the existing zero result without needing the Hive hash algorithm. - if self.bucket_function == 'hive' and self.num_buckets > 1: - raise NotImplementedError( - 'Hive bucket functions require the native writer; enable write.native.enabled') - - def _mod_bucket(self, value): - # Java reads the zeroed fixed-width slot for NULL; Python's modulo - # with a positive divisor otherwise matches Math.floorMod, even for - # Long.MIN_VALUE. Avoid floating-point division for BIGINT values. - return 0 if value is None else value % self.num_buckets - - -class DynamicBucketInputRouter(PartitionKeyExtractor): - """Hash dynamic-bucket input without assigning buckets or maintaining indexes.""" - - def __init__(self, table_schema: TableSchema): - super().__init__(table_schema) - num_buckets = int(table_schema.options.get(CoreOptions.BUCKET.key(), -1)) - if num_buckets != -1: - raise ValueError( - "Only 'bucket' = '-1' is allowed for " - f"'DynamicBucketRowKeyExtractor', but found: {num_buckets}" - ) - if table_schema.options.get(CoreOptions.BUCKET_KEY.key()) is not None: - raise ValueError("Cannot define 'bucket-key' in dynamic bucket mode") - self.bucket_keys = table_schema.bucket_keys - self.bucket_key_indices = self._get_field_indices(self.bucket_keys) - self._bucket_key_fields = table_schema.logical_bucket_key_fields - self._partition_fields = [ - table_schema.fields[index] for index in self.partition_indices - ] - - def extract_hashes_batch( - self, data: pa.RecordBatch - ) -> Tuple[List[Tuple], List[int], List[int]]: - """Return partitions, BinaryRow partition hashes, and key hashes.""" - partitions, partition_hashes = self.extract_partition_hashes_batch(data) - return partitions, partition_hashes, self._extract_key_hashes_batch(data) - - def extract_partition_hashes_batch(self, data: pa.RecordBatch) -> Tuple[List[Tuple], List[int]]: - """Hash partitions independently when key hashes arrived from upstream.""" - partitions = self._extract_partitions_batch(data) - partition_hash_cache = {} - partition_hashes = [] - for partition in partitions: - if partition not in partition_hash_cache: - partition_hash_cache[partition] = self._binary_row_hash_code( - partition, self._partition_fields - ) - partition_hashes.append(partition_hash_cache[partition]) - return partitions, partition_hashes - - def _extract_key_hashes_batch(self, data: pa.RecordBatch) -> List[int]: - key_columns = self._input_columns(data, self.bucket_key_indices) - return [ - self._binary_row_hash_code( - tuple(column[row_idx].as_py() for column in key_columns), - self._bucket_key_fields, - ) - for row_idx in range(data.num_rows) - ] - - def extract_assigners_batch( - self, data: pa.RecordBatch, num_channels: int, num_assigners: int - ) -> List[int]: - from pypaimon.index.dynamic_bucket import compute_assigner - - _, partition_hashes, key_hashes = self.extract_hashes_batch(data) - return [ - compute_assigner( - partition_hash, key_hash, num_channels, num_assigners - ) - for partition_hash, key_hash in zip(partition_hashes, key_hashes) - ] diff --git a/paimon-python/pypaimon/write/row_id_plan.py b/paimon-python/pypaimon/write/row_id_plan.py index 3b859b54fef4..63de3f02d87d 100644 --- a/paimon-python/pypaimon/write/row_id_plan.py +++ b/paimon-python/pypaimon/write/row_id_plan.py @@ -22,7 +22,9 @@ Workers reuse complete logical groups and the original snapshot. """ -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace + +import pyarrow as pa from pypaimon.read.plan import Plan from pypaimon.utils.data_evolution_utils import retrieve_anchor_file @@ -42,6 +44,7 @@ class CapturedRowIdPlan: table: object plan: Plan groups: list + _execution: object = field(default=None, repr=False, compare=False) @property def snapshot_id(self): @@ -66,6 +69,24 @@ def anchors(self): for _, _, split in self.groups for anchor in [retrieve_anchor_file(split.files)]] + def route(self, row_ids): + """Return each input row's captured logical file-group owner.""" + return self._executor().route(_row_id_array(row_ids)) + + def route_by_bucket(self, row_ids): + """Keep each captured partition/bucket under one deletion writer.""" + return self._executor().route_by_bucket(_row_id_array(row_ids)) + + def select(self, row_ids): + """Select absolute row IDs while retaining complete captured groups.""" + return self._executor().select(_row_id_array(row_ids)) + + def _executor(self): + if self._execution is None: + from pypaimon.write.backend_dispatch import create_row_id_plan_execution + object.__setattr__(self, '_execution', create_row_id_plan_execution(self.table, self.plan)) + return self._execution + def for_split(self, split): """Keep one complete logical group and its captured snapshot metadata.""" names = {file.file_name for file in split.files} @@ -74,7 +95,18 @@ def for_split(self, split): and {file.file_name for file in group[2].files} == names] if len(groups) != 1: raise ValueError('Split is not a complete captured row-ID group') - return CapturedRowIdPlan(self.table, Plan([groups[0][2]], self.plan.snapshot_id), groups) + plan = self.plan._with_splits([groups[0][2]]) + return CapturedRowIdPlan(self.table, plan, groups, self._executor().for_plan(plan, groups)) + + +def _row_id_array(row_ids): + """Accept integer Arrow inputs without silently truncating floating IDs.""" + if not isinstance(row_ids, (pa.Array, pa.ChunkedArray)) or not pa.types.is_integer(row_ids.type): + raise TypeError('_ROW_ID input must be an integer Arrow array') + if row_ids.null_count: + raise ValueError('_ROW_ID is null; planner snapshot is stale or row IDs come from a different table.') + row_ids = row_ids.cast(pa.int64(), safe=True) + return row_ids.combine_chunks() if isinstance(row_ids, pa.ChunkedArray) else row_ids def _mutation_routing_split(split): @@ -102,11 +134,6 @@ def capture_row_id_plan(table, plan=None): if plan is None: plan = table.new_read_builder().new_scan().plan_all() plan = replace(plan, _splits=[_mutation_routing_split(split) for split in plan.splits()]) - from pypaimon.execution.native.write.row_id_plan import create_native_row_id_plan, python_groups - native = create_native_row_id_plan(table, plan) - if native is not None: - groups = python_groups(table, native) - else: - from pypaimon.execution.python.write.row_id_plan import row_id_groups - groups = row_id_groups(plan) - return CapturedRowIdPlan(table, plan, groups) + from pypaimon.write.backend_dispatch import create_row_id_plan_execution + execution = create_row_id_plan_execution(table, plan) + return CapturedRowIdPlan(table, plan, execution.groups, execution)