This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new d706289806 [python] Dispatch captured row ID and bucket execution
(#10778)
d706289806 is described below
commit d706289806e264daf6fb47759c1f35a9b979f248
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Oct 23 20:35:31 2026 +0800
[python] Dispatch captured row ID and bucket execution (#10778)
---
paimon-python/dev/test_packaging.py | 2 +-
.../pypaimon/execution/native/write/row_id_plan.py | 35 +++
.../python/read/scanner/bucket_select_converter.py | 4 +-
.../execution/python/write/dynamic_bucket.py | 333 +++++++++++++++++++-
.../execution/python/write/input_routing.py | 330 +++++++++++++++++++-
.../pypaimon/execution/python/write/row_id_plan.py | 109 ++++++-
.../execution/python/write/row_key_extractor.py | 14 +-
paimon-python/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 ++--
.../tests/native_dynamic_input_routing_test.py | 10 +-
.../pypaimon/tests/native_input_routing_test.py | 2 +-
.../pypaimon/tests/native_ray_row_id_plan_test.py | 104 +++++++
.../pypaimon/tests/noncompact_null_row_test.py | 2 +-
.../pypaimon/tests/pushdown_bucket_test.py | 2 +-
.../write/distributed_input_execution_test.py | 11 +-
.../pypaimon/tests/write/dynamic_bucket_test.py | 12 +-
.../tests/write/dynamic_input_routing_test.py | 4 +-
.../tests/write/native_input_routing_test.py | 2 +-
.../pypaimon/tests/write/row_id_routing_test.py | 275 +++++++++++++++++
paimon-python/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(-)
diff --git a/paimon-python/dev/test_packaging.py
b/paimon-python/dev/test_packaging.py
index 428fbac0ee..02fe0efdca 100644
--- a/paimon-python/dev/test_packaging.py
+++ b/paimon-python/dev/test_packaging.py
@@ -211,7 +211,7 @@ class PackagingTest(unittest.TestCase):
"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 d040b5b205..3e2eea650b 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 279adbc709..eb7ad1d00e 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 @@ allowed; false-negative (drop a bucket that has matching
rows) MUST never
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.common.predicate import Predicate
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 759412e3d8..8a38ac72a2 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 @@ _ADD = 0
_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 133e3f2378..1f546986f0 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("<I", data, i)[0]
+ k1 = _mix_k1(k1)
+ h1 = _mix_h1(h1, k1)
+ return _fmix(h1, n)
+
+
+def _bucket_from_hash(hash_unsigned: int, num_buckets: int) -> 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 72f486afaf..c364187a94 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 996f440c72..81589e0926 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 numpy as np
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 @@ class
DynamicBucketRowKeyExtractor(DynamicBucketInputRouter, RowKeyExtractor):
)
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 ad338d0065..0a7a999c79 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 edbb845d39..5544353b96 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 distributed_read_by_row_id(
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 88746311dc..966d2ce9f9 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 @@ class Plan:
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 76c47ffbf2..c865150035 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 @@ import pyarrow as pa
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 c23bb48144..f576474460 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 @@ import pytest
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 d0764b6409..a47360bf08 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.execution.native.write.postpone_bucket import
NativePostponeBucket
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 1ff0a6c400..4072c17a6b 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 @@ import pytest
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)])
+
+
[email protected]('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 24a0181dc7..ecc8102e35 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 @@ import pytest
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 d2d72bb27c..3d3856c236 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.schema.data_types import AtomicType, DataField
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 d45a511ac5..c141b7dad3 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 pyarrow as pa
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 8db7e598f2..100aabdf62 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 pyarrow as pa
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 @@ class DynamicBucketTest(unittest.TestCase):
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 24ec5c9e88..9210b2b65f 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 fdb5e70bec..58efed762d 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.execution.native.write.postpone_bucket import
NativePostponeBucket
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 0000000000..171fdb1ab1
--- /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)
+
+
[email protected]
+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]
+
+
[email protected]('operation', ['route', 'route_by_bucket', 'select'])
[email protected]('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()))
+
+
[email protected]('operation', ['route', 'route_by_bucket', 'select'])
[email protected]('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]))
+
+
[email protected]('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
+
+
[email protected]('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))
+
+
[email protected]('state', ['lazy', 'materialized', 'pickled'])
[email protected]('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 f6f20f269e..0181174f5a 100644
--- a/paimon-python/pypaimon/write/backend_dispatch.py
+++ b/paimon-python/pypaimon/write/backend_dispatch.py
@@ -214,9 +214,9 @@ def create_input_router(table, bucket_plan=None):
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 d2d0496902..0000000000
--- 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("<I", data, i)[0]
- k1 = _mix_k1(k1)
- h1 = _mix_h1(h1, k1)
- return _fmix(h1, n)
-
-
-def _bucket_from_hash(hash_unsigned: int, num_buckets: int) -> 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 3b859b54fe..63de3f02d8 100644
--- a/paimon-python/pypaimon/write/row_id_plan.py
+++ b/paimon-python/pypaimon/write/row_id_plan.py
@@ -22,7 +22,9 @@ Ray transports Python metadata; file grouping belongs to the
selected backend.
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 @@ class CapturedRowIdPlan:
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 @@ class CapturedRowIdPlan:
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)
