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 b8af4bf2b1 [python] Use native Daft bucket routing and assignment
(#10775)
b8af4bf2b1 is described below
commit b8af4bf2b1a1bd6d2e1808174f669825eb6d88b1
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Oct 23 19:34:17 2026 +0800
[python] Use native Daft bucket routing and assignment (#10775)
---
docs/docs/pypaimon/daft.md | 7 +
paimon-python/pypaimon/daft/daft_datasink.py | 81 ++----
paimon-python/pypaimon/daft/daft_paimon.py | 2 +-
.../execution/native/write/input_routing.py | 124 +++++++-
.../execution/native/write/postpone_bucket.py | 16 +-
.../read/reader/outer_projection_record_reader.py | 24 +-
.../execution/python/write/input_routing.py | 73 +++++
.../execution/python/write/row_key_extractor.py | 8 +-
paimon-python/pypaimon/index/dynamic_bucket.py | 12 +-
.../pypaimon/read/reader/field_indices.py | 17 +-
.../tests/daft/daft_pk_distributed_write_test.py | 2 +-
.../dynamic_bucket_execution_boundary_test.py | 4 +-
.../tests/native_dynamic_input_routing_test.py | 151 ++++++++++
.../write/distributed_input_execution_test.py | 43 +--
.../pypaimon/tests/write/dynamic_bucket_test.py | 28 +-
.../tests/write/dynamic_input_routing_test.py | 317 +++++++++++++++++++++
.../tests/write/native_input_routing_test.py | 51 ++++
...t.py => simple_dynamic_bucket_assigner_test.py} | 18 +-
paimon-python/pypaimon/write/backend_dispatch.py | 31 ++
19 files changed, 880 insertions(+), 129 deletions(-)
diff --git a/docs/docs/pypaimon/daft.md b/docs/docs/pypaimon/daft.md
index d219912b87..000c2c23a8 100644
--- a/docs/docs/pypaimon/daft.md
+++ b/docs/docs/pypaimon/daft.md
@@ -174,6 +174,13 @@ write_paimon(
`write_paimon` opens its own catalog, resolves the table, and commits the
write through Daft's DataSink API.
+With `write.native.enabled=true`, primary-key writes use Rust for fixed-bucket
+routing and for both stages of dynamic-bucket assignment: input hashing and
+restoring HASH indexes from the captured writer snapshot. Daft still schedules
+and groups the input so one writer owns each partition/bucket. Overwrite starts
+with an empty bucket-assignment state; append restores the snapshot captured
+before dispatch. These write-side operations do not request SELECT permission.
+
**Overwrite mode:**
```python
diff --git a/paimon-python/pypaimon/daft/daft_datasink.py
b/paimon-python/pypaimon/daft/daft_datasink.py
index 675276efff..59dd539fd9 100644
--- a/paimon-python/pypaimon/daft/daft_datasink.py
+++ b/paimon-python/pypaimon/daft/daft_datasink.py
@@ -438,19 +438,16 @@ def make_fixed_bucket_udf(table: FileStoreTable):
from daft import DataType, Series
from pypaimon.schema.data_types import PyarrowFieldParser
- from pypaimon.write.input_routing import FixedBucketInputRouter
+ from pypaimon.write.backend_dispatch import create_input_router
schema = PyarrowFieldParser.from_paimon_schema(table.fields)
- extractor = FixedBucketInputRouter(table.table_schema)
+ router = create_input_router(table)
@daft.func.batch(return_dtype=DataType.int32())
def _fixed_bucket(*columns) -> Series:
arrow_table = _series_to_arrow_table(columns, schema)
- buckets = []
- for batch in arrow_table.to_batches():
- _, batch_buckets = extractor.extract_partition_bucket_batch(batch)
- buckets.extend(batch_buckets)
- return Series.from_arrow(pa.array(buckets, type=pa.int32()))
+ buckets = [router.route(batch) for batch in arrow_table.to_batches()]
+ return Series.from_arrow(pa.concat_arrays(buckets) if buckets else
pa.array([], type=pa.int32()))
return _fixed_bucket
@@ -466,12 +463,11 @@ def make_dynamic_routing_udf(
import daft
from daft import DataType, Series
- from pypaimon.index.dynamic_bucket import compute_assigner, to_signed_int32
from pypaimon.schema.data_types import PyarrowFieldParser
- from pypaimon.write.input_routing import DynamicBucketInputRouter
+ from pypaimon.write.backend_dispatch import
create_dynamic_bucket_assigner_router
schema = PyarrowFieldParser.from_paimon_schema(table.fields)
- extractor = DynamicBucketInputRouter(table.table_schema)
+ router = create_dynamic_bucket_assigner_router(table, num_channels,
num_assigners)
return_fields = {
assigner_column: DataType.int32(),
key_hash_column: DataType.int32(),
@@ -486,25 +482,14 @@ def make_dynamic_routing_udf(
assigners = []
key_hashes = []
for batch in arrow_table.to_batches():
- _, partition_hashes, batch_key_hashes = (
- extractor.extract_hashes_batch(batch)
- )
- assigners.extend([
- compute_assigner(
- partition_hash, key_hash, num_channels, num_assigners
- )
- for partition_hash, key_hash in zip(
- partition_hashes, batch_key_hashes
- )
- ])
- key_hashes.extend(
- to_signed_int32(value) for value in batch_key_hashes
- )
+ routed = router.route(batch)
+ assigners.append(routed.column(0))
+ key_hashes.append(routed.column(1))
return Series.from_arrow(
pa.StructArray.from_arrays(
[
- pa.array(assigners, type=pa.int32()),
- pa.array(key_hashes, type=pa.int32()),
+ pa.concat_arrays(assigners) if assigners else pa.array([],
type=pa.int32()),
+ pa.concat_arrays(key_hashes) if key_hashes else
pa.array([], type=pa.int32()),
],
names=[assigner_column, key_hash_column],
)
@@ -528,7 +513,7 @@ def make_dynamic_bucket_assignment_udf(
from daft import DataType, Series
from pypaimon.schema.data_types import PyarrowFieldParser
- from pypaimon.write.input_routing import DynamicBucketInputRouter
+ from pypaimon.write.backend_dispatch import create_dynamic_bucket_assigner
catalog_options = _extract_catalog_options(table)
table_identifier = _extract_identifier(table)
@@ -556,9 +541,10 @@ def make_dynamic_bucket_assignment_udf(
unnest=True,
)
def _assign_dynamic_bucket(*columns) -> Series:
- assigner_values = columns[data_column_count].to_pylist()
- key_hash_values = columns[data_column_count + 1].to_pylist()
- if not assigner_values:
+ import pyarrow.compute as pc
+
+ assigner_values = columns[data_column_count].to_arrow()
+ if len(assigner_values) == 0:
return Series.from_arrow(
pa.array([], type=pa.struct([
pa.field(field.name, field.type) for field in output_fields
@@ -567,8 +553,8 @@ def make_dynamic_bucket_assignment_udf(
pa.field(base_snapshot_column, pa.int64()),
]))
)
- assign_id = assigner_values[0]
- if any(value != assign_id for value in assigner_values):
+ assign_id = assigner_values[0].as_py()
+ if assigner_values.null_count or
pc.count_distinct(assigner_values).as_py() != 1:
raise RuntimeError(
"A dynamic-bucket group contained multiple assigners"
)
@@ -582,36 +568,17 @@ def make_dynamic_bucket_assignment_udf(
table_path,
schema_token,
)
- from pypaimon.index.dynamic_bucket import HashBucketAssigner,
resolve_bucket_snapshot
-
- router = DynamicBucketInputRouter(worker_table.table_schema)
- snapshot = resolve_bucket_snapshot(worker_table, base_snapshot_id)
- options = worker_table.options
- assigner = HashBucketAssigner(
- worker_table,
- num_channels=num_channels,
- num_assigners=num_assigners,
- assign_id=assign_id,
- ignore_existing=ignore_existing,
- target_bucket_row_number=options.dynamic_bucket_target_row_num(),
- max_buckets_num=options.dynamic_bucket_max_buckets(),
- snapshot=snapshot,
- )
- buckets = []
- offset = 0
- for batch in arrow_table.to_batches():
- batch_hashes = key_hash_values[offset:offset + batch.num_rows]
- partitions, partition_hashes =
router.extract_partition_hashes_batch(batch)
- buckets.extend(bucket for bucket, _ in assigner.assign_batch(
- partitions, partition_hashes, batch_hashes))
- offset += batch.num_rows
+ assigner = create_dynamic_bucket_assigner(
+ worker_table, num_channels, num_assigners, assign_id,
+ ignore_existing=ignore_existing, base_snapshot_id=base_snapshot_id)
+ buckets = assigner.assign(arrow_table,
key_hashes=columns[data_column_count + 1].to_arrow())
arrays = [
arrow_table.column(field.name).combine_chunks()
for field in output_fields
]
- arrays.append(pa.array(buckets, type=pa.int32()))
+ arrays.append(buckets)
arrays.append(pa.array(
- [snapshot.id if snapshot is not None else 0] *
arrow_table.num_rows,
+ [assigner.restore_snapshot_id] * arrow_table.num_rows,
type=pa.int64(),
))
return Series.from_arrow(
diff --git a/paimon-python/pypaimon/daft/daft_paimon.py
b/paimon-python/pypaimon/daft/daft_paimon.py
index 765c009a33..ec69834946 100644
--- a/paimon-python/pypaimon/daft/daft_paimon.py
+++ b/paimon-python/pypaimon/daft/daft_paimon.py
@@ -296,7 +296,7 @@ def _write_table(
return prepared.write_sink(commit_sink)
if bucket_mode == BucketMode.HASH_DYNAMIC:
- latest_snapshot = table.snapshot_manager().get_latest_snapshot()
+ latest_snapshot = None if mode == 'overwrite' else
table.snapshot_manager().get_latest_snapshot()
base_snapshot_id = (
latest_snapshot.id if latest_snapshot is not None else 0
)
diff --git a/paimon-python/pypaimon/execution/native/write/input_routing.py
b/paimon-python/pypaimon/execution/native/write/input_routing.py
index 461cba52f1..b5c9ff45bb 100644
--- a/paimon-python/pypaimon/execution/native/write/input_routing.py
+++ b/paimon-python/pypaimon/execution/native/write/input_routing.py
@@ -23,12 +23,15 @@ from pypaimon.execution.native.runtime import
native_runtime_available
from pypaimon.execution.native.write.native_commit import
create_native_write_table
-def as_record_batch(data):
+def input_batches(data):
if isinstance(data, pa.RecordBatch):
- return data
+ return [data]
if isinstance(data, pa.Table):
- return pa.RecordBatch.from_arrays(
- [column.combine_chunks() for column in data.columns],
schema=data.schema)
+ # Binary columns can retain multiple chunks even after combine_chunks.
+ # Preserve every batch and its row count, including zero-column input.
+ batches = data.to_batches()
+ return batches if batches else [pa.RecordBatch.from_arrays(
+ [column.combine_chunks() for column in data.columns],
schema=data.schema)]
raise TypeError('Input routing requires a pyarrow Table or RecordBatch')
@@ -65,9 +68,120 @@ class NativeInputRouter:
def route(self, data) -> pa.Int32Array:
if self._router is None:
self._prepare()
- return self._router.route(as_record_batch(data))
+ buckets = [self._router.route(batch) for batch in input_batches(data)]
+ return buckets[0] if len(buckets) == 1 else pa.concat_arrays(buckets)
def __getstate__(self):
state = dict(self.__dict__)
state['_router'] = None
return state
+
+
+def create_native_dynamic_bucket_assigner_router(table, num_channels,
num_assigners):
+ if not native_runtime_available():
+ return None
+ native_table = create_native_write_table(table)
+ if native_table is None:
+ return None
+ router = NativeDynamicBucketAssignerRouter(table, num_channels,
num_assigners)
+ router._prepare(native_table)
+ return router
+
+
+class NativeDynamicBucketAssignerRouter:
+ """Route input rows to dynamic bucket assigners using Rust's row hashes."""
+
+ def __init__(self, table, num_channels, num_assigners):
+ self.table = table
+ self._num_channels = num_channels
+ self._num_assigners = num_assigners
+ self._router = None
+
+ def _prepare(self, native_table=None):
+ from pypaimon_rust.datafusion import DynamicBucketAssignerRouter
+ if native_table is None:
+ native_table = create_native_write_table(self.table)
+ if native_table is None:
+ raise RuntimeError('Native dynamic bucket assigner routing is
unavailable')
+ self._router = DynamicBucketAssignerRouter(native_table.schema(),
self._num_channels, self._num_assigners)
+
+ def route(self, data) -> pa.RecordBatch:
+ if self._router is None:
+ self._prepare()
+ routed = [self._router.route(batch) for batch in input_batches(data)]
+ if len(routed) == 1:
+ return routed[0]
+ return pa.RecordBatch.from_arrays([
+ pa.concat_arrays([batch.column(index) for batch in routed]) for
index in range(2)
+ ], schema=routed[0].schema)
+
+ def __getstate__(self):
+ state = dict(self.__dict__)
+ state['_router'] = None
+ return state
+
+
+def create_native_dynamic_bucket_assigner(table, num_channels, num_assigners,
assign_id,
+ ignore_existing=False,
base_snapshot_id=None):
+ if not native_runtime_available():
+ return None
+ native_table = create_native_write_table(table)
+ if native_table is None:
+ return None
+ assigner = NativeDynamicBucketAssigner(
+ table, num_channels, num_assigners, assign_id, ignore_existing,
base_snapshot_id)
+ assigner._prepare(native_table)
+ return assigner
+
+
+class NativeDynamicBucketAssigner:
+
+ def __init__(self, table, num_channels, num_assigners, assign_id,
+ ignore_existing=False, base_snapshot_id=None):
+ self.table = table
+ self._num_channels = num_channels
+ self._num_assigners = num_assigners
+ self._assign_id = assign_id
+ self._ignore_existing = ignore_existing
+ self._base_snapshot_id = base_snapshot_id
+ self._assigner = None
+
+ def _prepare(self, native_table=None):
+ from pypaimon_rust.datafusion import DynamicBucketAssigner
+ if native_table is None:
+ native_table = create_native_write_table(self.table)
+ if native_table is None:
+ raise RuntimeError('Native dynamic bucket assignment is
unavailable')
+ self._assigner = DynamicBucketAssigner(
+ native_table, self._num_channels, self._num_assigners,
self._assign_id,
+ restore_snapshot_id=self._base_snapshot_id,
overwrite=self._ignore_existing)
+ # Core captures the restore point once. Final worker writers must
+ # restore that same snapshot even if another commit finishes later.
+ self._base_snapshot_id = self._assigner.restore_snapshot_id
+
+ @property
+ def restore_snapshot_id(self):
+ if self._assigner is None:
+ self._prepare()
+ return self._base_snapshot_id
+
+ def assign(self, data, key_hashes=None) -> pa.Int32Array:
+ batches = input_batches(data)
+ if isinstance(key_hashes, pa.ChunkedArray):
+ key_hashes = key_hashes.combine_chunks()
+ if key_hashes is not None:
+ if not isinstance(key_hashes, pa.Int32Array):
+ raise TypeError('Dynamic bucket key hashes must be an
Int32Array')
+ # Validate the whole input before the stateful assigner sees any
+ # batch, including a null or missing hash in a later chunk.
+ if len(key_hashes) != data.num_rows or key_hashes.null_count:
+ raise ValueError('Dynamic bucket key hashes must be non-null
and match the input row count')
+ if self._assigner is None:
+ self._prepare()
+ buckets = []
+ offset = 0
+ for batch in batches:
+ hashes = None if key_hashes is None else key_hashes.slice(offset,
batch.num_rows)
+ buckets.append(self._assigner.assign(batch, key_hashes=hashes))
+ offset += batch.num_rows
+ return buckets[0] if len(buckets) == 1 else pa.concat_arrays(buckets)
diff --git a/paimon-python/pypaimon/execution/native/write/postpone_bucket.py
b/paimon-python/pypaimon/execution/native/write/postpone_bucket.py
index acb3823aa4..98837940bc 100644
--- a/paimon-python/pypaimon/execution/native/write/postpone_bucket.py
+++ b/paimon-python/pypaimon/execution/native/write/postpone_bucket.py
@@ -20,7 +20,7 @@
import pyarrow as pa
from pypaimon.execution.native.runtime import native_runtime_available
-from pypaimon.execution.native.write.input_routing import as_record_batch
+from pypaimon.execution.native.write.input_routing import input_batches
from pypaimon.execution.native.write.native_commit import
create_native_write_table
from pypaimon.schema.data_types import PyarrowFieldParser
from pypaimon.write.bucket_plan import PostponeBucketPlan
@@ -87,10 +87,16 @@ class NativePostponeBucketPlanner:
return PostponeBucketPlan(self._known_num_buckets)
def input_partition_stats(self, data):
- stats = self._ensure_planner().input_statistics(as_record_batch(data))
- columns = [column.to_pylist() for column in stats.columns]
- return {tuple(column[row] for column in columns[:-2]):
- (columns[-2][row], columns[-1][row]) for row in
range(stats.num_rows)}
+ planner = self._ensure_planner()
+ partition_stats = {}
+ for batch in input_batches(data):
+ stats = planner.input_statistics(batch)
+ columns = [column.to_pylist() for column in stats.columns]
+ for row in range(stats.num_rows):
+ partition = tuple(column[row] for column in columns[:-2])
+ count, size = partition_stats.get(partition, (0, 0))
+ partition_stats[partition] = (count + columns[-2][row], size +
columns[-1][row])
+ return partition_stats
def plan(self, partition_stats):
fields = [self.table.field_dict[name] for name in
self.table.partition_keys]
diff --git
a/paimon-python/pypaimon/execution/python/read/reader/outer_projection_record_reader.py
b/paimon-python/pypaimon/execution/python/read/reader/outer_projection_record_reader.py
index e2e6d0f5c0..fa4d710917 100644
---
a/paimon-python/pypaimon/execution/python/read/reader/outer_projection_record_reader.py
+++
b/paimon-python/pypaimon/execution/python/read/reader/outer_projection_record_reader.py
@@ -25,9 +25,8 @@ partial-update / aggregation see the original sub-structure),
and this
wrapper extracts the user-visible flat columns afterwards.
"""
-from typing import Any, List, Optional
+from typing import Any, Iterable, List, Optional, Sequence, Set
-from pypaimon.read.reader.field_indices import project_top_level_field_indices
from pypaimon.read.reader.iface.record_iterator import RecordIterator
from pypaimon.read.reader.iface.record_reader import RecordReader
from pypaimon.table.row.internal_row import InternalRow
@@ -62,11 +61,11 @@ class
OuterProjectionRecordReader(RecordReader[InternalRow]):
self._inner = inner
self._flat_arity = len(name_paths)
self._file_io = file_io
- self._blob_field_indices = project_top_level_field_indices(
+ self._blob_field_indices = _project_top_level_field_indices(
blob_field_indices, self._specs)
- self._descriptor_field_indices = project_top_level_field_indices(
+ self._descriptor_field_indices = _project_top_level_field_indices(
descriptor_field_indices, self._specs)
- self._vector_field_indices = project_top_level_field_indices(
+ self._vector_field_indices = _project_top_level_field_indices(
vector_field_indices, self._specs)
self.file_io = self._file_io
self.blob_field_indices = self._blob_field_indices
@@ -135,6 +134,21 @@ class _PathSpec:
self.sub_names = sub_names
+def _project_top_level_field_indices(
+ source_indices: Optional[Iterable[int]],
+ path_specs: Sequence[_PathSpec],
+) -> Optional[Set[int]]:
+ """Carry top-level field markers through this reader's resolved paths."""
+ if source_indices is None:
+ return None
+ source = set(source_indices)
+ return {
+ proj_pos
+ for proj_pos, spec in enumerate(path_specs)
+ if not spec.sub_names and spec.top_idx in source
+ }
+
+
def _extract(row: InternalRow, spec: _PathSpec) -> Any:
cur = row.get_field(spec.top_idx)
for name in spec.sub_names:
diff --git a/paimon-python/pypaimon/execution/python/write/input_routing.py
b/paimon-python/pypaimon/execution/python/write/input_routing.py
new file mode 100644
index 0000000000..133e3f2378
--- /dev/null
+++ b/paimon-python/pypaimon/execution/python/write/input_routing.py
@@ -0,0 +1,73 @@
+# 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.
+
+"""Arrow contracts for Python distributed dynamic-bucket execution."""
+
+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
+
+
+class PythonDynamicBucketAssignerRouter:
+ """Route input rows to dynamic bucket assigners and carry their key
hashes."""
+
+ def __init__(self, table, num_channels, num_assigners):
+ if num_channels <= 0 or num_assigners <= 0:
+ raise ValueError('num_channels and num_assigners must be positive')
+ self._router = DynamicBucketInputRouter(table.table_schema)
+ self._num_channels = num_channels
+ self._num_assigners = num_assigners
+
+ def route(self, data) -> pa.RecordBatch:
+ _, partition_hashes, key_hashes =
self._router.extract_hashes_batch(data)
+ return pa.RecordBatch.from_arrays([
+ pa.array([compute_assigner(partition_hash, key_hash,
self._num_channels, self._num_assigners)
+ for partition_hash, key_hash in zip(partition_hashes,
key_hashes)], type=pa.int32()),
+ pa.array([to_signed_int32(key_hash) for key_hash in key_hashes],
type=pa.int32()),
+ ], names=['assigner', 'key_hash'])
+
+
+class PythonDynamicBucketAssigner:
+
+ def __init__(self, table, num_channels, num_assigners, assign_id,
+ ignore_existing=False, base_snapshot_id=None):
+ self._router = DynamicBucketInputRouter(table.table_schema)
+ snapshot = None if ignore_existing else resolve_bucket_snapshot(table,
base_snapshot_id)
+ self.restore_snapshot_id = snapshot.id if snapshot is not None else 0
+ self._assigner = DynamicBucketAssigner(
+ table, num_channels=num_channels, num_assigners=num_assigners,
assign_id=assign_id,
+
target_bucket_row_number=table.options.dynamic_bucket_target_row_num(),
+ max_buckets_num=table.options.dynamic_bucket_max_buckets(),
+ ignore_existing=ignore_existing, snapshot=snapshot)
+
+ def assign(self, data, key_hashes=None) -> pa.Int32Array:
+ if key_hashes is None:
+ partitions, partition_hashes, hashes =
self._router.extract_hashes_batch(data)
+ else:
+ if isinstance(key_hashes, pa.ChunkedArray):
+ key_hashes = key_hashes.combine_chunks()
+ if not isinstance(key_hashes, pa.Int32Array):
+ raise TypeError('Dynamic bucket key hashes must be an
Int32Array')
+ if len(key_hashes) != data.num_rows or key_hashes.null_count:
+ raise ValueError('Dynamic bucket key hashes must be non-null
and match the input row count')
+ partitions, partition_hashes =
self._router.extract_partition_hashes_batch(data)
+ hashes = key_hashes.to_pylist()
+ return pa.array([bucket for bucket, _ in self._assigner.assign_batch(
+ partitions, partition_hashes, hashes)], type=pa.int32())
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 76a65d0555..996f440c72 100644
--- a/paimon-python/pypaimon/execution/python/write/row_key_extractor.py
+++ b/paimon-python/pypaimon/execution/python/write/row_key_extractor.py
@@ -281,7 +281,7 @@ class _SimplePartitionIndex:
)
-class SimpleHashBucketAssigner:
+class SimpleDynamicBucketAssigner:
def __init__(self, num_assigners, assign_id, target_bucket_row_number,
max_buckets_num):
validate_max_buckets(max_buckets_num)
self.num_assigners = num_assigners
@@ -330,7 +330,7 @@ class
DynamicBucketRowKeyExtractor(DynamicBucketInputRouter, RowKeyExtractor):
max_buckets_num = opts.dynamic_bucket_max_buckets()
if table is None:
- self._assigner = SimpleHashBucketAssigner(
+ self._assigner = SimpleDynamicBucketAssigner(
num_assigners=num_assigners,
assign_id=assign_id,
target_bucket_row_number=target_bucket_row_number,
@@ -339,12 +339,12 @@ class
DynamicBucketRowKeyExtractor(DynamicBucketInputRouter, RowKeyExtractor):
self._index_maintainer = None
else:
from pypaimon.execution.python.write.dynamic_bucket import
DynamicBucketIndexMaintainer
- from pypaimon.index.dynamic_bucket import HashBucketAssigner
+ from pypaimon.index.dynamic_bucket import DynamicBucketAssigner
from pypaimon.index.dynamic_bucket import resolve_bucket_snapshot
snapshot = resolve_bucket_snapshot(table, base_snapshot_id)
self.base_snapshot_id = snapshot.id if snapshot is not None else 0
- self._assigner = HashBucketAssigner(
+ self._assigner = DynamicBucketAssigner(
table=table,
num_channels=num_channels,
num_assigners=num_assigners,
diff --git a/paimon-python/pypaimon/index/dynamic_bucket.py
b/paimon-python/pypaimon/index/dynamic_bucket.py
index a72d600b79..ad338d0065 100644
--- a/paimon-python/pypaimon/index/dynamic_bucket.py
+++ b/paimon-python/pypaimon/index/dynamic_bucket.py
@@ -83,7 +83,7 @@ def compute_assigner(
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 (start + assigner) % num_channels
+ return _java_remainder(to_signed_int32(start + assigner), num_channels)
def is_my_bucket(bucket: int, num_assigners: int, assign_id: int) -> bool:
@@ -141,6 +141,7 @@ class _PartitionIndex:
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,
@@ -174,7 +175,9 @@ class _PartitionIndex:
global_max_bucket_id = (
MAX_DYNAMIC_BUCKETS if max_buckets_num == -1 else max_buckets_num
) - 1
- if not self.total_bucket_set or max_bucket_id < global_max_bucket_id:
+ # 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
@@ -187,6 +190,9 @@ class _PartitionIndex:
"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(
@@ -198,7 +204,7 @@ class _PartitionIndex:
return bucket, True
-class HashBucketAssigner:
+class DynamicBucketAssigner:
"""Restore dynamic-bucket mappings from HASH indexes before assigning."""
def __init__(
diff --git a/paimon-python/pypaimon/read/reader/field_indices.py
b/paimon-python/pypaimon/read/reader/field_indices.py
index b5a6690a6b..7392aeafa4 100644
--- a/paimon-python/pypaimon/read/reader/field_indices.py
+++ b/paimon-python/pypaimon/read/reader/field_indices.py
@@ -17,7 +17,7 @@
"""Helpers for carrying typed field metadata through projection wrappers."""
-from typing import Any, Iterable, List, Optional, Sequence, Set
+from typing import Iterable, List, Set
from pypaimon.schema.data_types import DataField, VectorType
@@ -56,18 +56,3 @@ def descriptor_field_indices_for_table(table, fields:
List[DataField]) -> Set[in
def vector_field_indices(fields: List[DataField]) -> Set[int]:
return {i for i, f in enumerate(fields) if isinstance(f.type, VectorType)}
-
-
-def project_top_level_field_indices(
- source_indices: Optional[Iterable[int]],
- path_specs: Sequence[Any],
-) -> Optional[Set[int]]:
- """Map top-level field indices through path specs with
top_idx/sub_names."""
- if source_indices is None:
- return None
- source = set(source_indices)
- return {
- proj_pos
- for proj_pos, spec in enumerate(path_specs)
- if not spec.sub_names and spec.top_idx in source
- }
diff --git
a/paimon-python/pypaimon/tests/daft/daft_pk_distributed_write_test.py
b/paimon-python/pypaimon/tests/daft/daft_pk_distributed_write_test.py
index a9a0fb58c3..e5472cac25 100644
--- a/paimon-python/pypaimon/tests/daft/daft_pk_distributed_write_test.py
+++ b/paimon-python/pypaimon/tests/daft/daft_pk_distributed_write_test.py
@@ -289,7 +289,7 @@ def
test_hash_dynamic_pk_restores_bucket_mapping_across_commits(catalog):
_write_table(
daft.from_pydict(
{
- # A fresh task-local SimpleHashBucketAssigner puts id=2 in
+ # A fresh task-local SimpleDynamicBucketAssigner puts id=2 in
# bucket 0 and incorrectly moves the existing id=1 to bucket 1.
"id": [2, 1],
"value": ["other", "new"],
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 49c169dda5..76c47ffbf2 100644
--- a/paimon-python/pypaimon/tests/dynamic_bucket_execution_boundary_test.py
+++ b/paimon-python/pypaimon/tests/dynamic_bucket_execution_boundary_test.py
@@ -62,7 +62,7 @@ class BlockPythonExecution(importlib.abc.MetaPathFinder):
sys.meta_path.insert(0, BlockPythonExecution())
from pypaimon import CatalogFactory
-from pypaimon.index.dynamic_bucket import HashBucketAssigner,
resolve_bucket_snapshot
+from pypaimon.index.dynamic_bucket import DynamicBucketAssigner,
resolve_bucket_snapshot
from pypaimon.write.input_routing import DynamicBucketInputRouter
catalog = CatalogFactory.create({'warehouse': sys.argv[1]})
@@ -77,7 +77,7 @@ partitions, partition_hashes, key_hashes =
router.extract_hashes_batch(rows)
# 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 = HashBucketAssigner(table, 1, 1, 0, 1, -1, snapshot=snapshot)
+ 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()
diff --git a/paimon-python/pypaimon/tests/native_dynamic_input_routing_test.py
b/paimon-python/pypaimon/tests/native_dynamic_input_routing_test.py
new file mode 100644
index 0000000000..c23bb48144
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_dynamic_input_routing_test.py
@@ -0,0 +1,151 @@
+# 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.
+
+"""Daft distributed input execution uses the real SDK for routing and
restoration."""
+
+from contextlib import ExitStack
+import pickle
+from unittest.mock import patch
+
+import pyarrow as pa
+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.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
+
+pytestmark = pytest.mark.native_plan
+
+
+def _table(catalog, bucket='-1', partitioned=True):
+ fields = [('id', pa.int64()), ('value', pa.string())]
+ if partitioned:
+ fields.append(('p', pa.string()))
+ catalog.create_table('default.dynamic_input', Schema.from_pyarrow_schema(
+ pa.schema(fields), primary_keys=['id', 'p'] if partitioned else ['id'],
+ partition_keys=['p'] if partitioned else [], options={
+ 'bucket': bucket, 'file.format': 'parquet',
'dynamic-bucket.target-row-num': '1',
+ 'write.native.enabled': 'true', 'commit.native.enabled': 'true',
+ 'read.native.enabled': 'true', 'scan.native-plan.enabled': 'true',
+ 'query-auth.enabled': 'true',
+ }), False)
+ return catalog.get_table('default.dynamic_input')
+
+
+def _commit(table, rows):
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ writer.write_arrow(pa.table(rows))
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+
+
+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.table_write.PythonTableWrite.__init__',
+ ]:
+ stack.enter_context(patch(target, side_effect=AssertionError('Python
input execution: ' + target)))
+ return stack
+
+
+def
test_native_dynamic_arrow_route_and_assignment_keep_one_restore_snapshot(native_rest_catalog):
+ table = _table(native_rest_catalog)
+ _commit(table, {'id': [1], 'value': ['old'], 'p': ['west']})
+ snapshot_id = table.snapshot_manager().get_latest_snapshot().id
+ batch = pa.record_batch([['west', 'west', 'east'], [2, 1, 2]], names=['p',
'id'])
+ chunked_input = pa.Table.from_batches([batch.slice(0, 1), batch.slice(1)])
+ reference = DynamicBucketInputRouter(table.table_schema)
+ _, partition_hashes, hashes = reference.extract_hashes_batch(batch)
+ reference_assigner = DynamicBucketAssigner(
+ table, 1, 1, 0, 1, -1,
snapshot=table.snapshot_manager().get_snapshot_by_id(snapshot_id))
+ expected = [bucket for bucket, _ in reference_assigner.assign_batch(
+ reference._extract_partitions_batch(batch), partition_hashes, hashes)]
+ with _forbid_python_input_execution(), \
+ patch.object(RESTCatalogServer, '_table_query_auth_handle',
+ side_effect=AssertionError('Bucket restoration must
not request SELECT')) as query:
+ router = create_dynamic_bucket_assigner_router(table, 1, 1)
+ assert isinstance(router, NativeDynamicBucketAssignerRouter)
+ routed = pickle.loads(pickle.dumps(router)).route(chunked_input)
+ assert routed.schema.names == ['assigner', 'key_hash']
+ assert routed.column(0).to_pylist() == [0] * batch.num_rows
+ assert routed.column(1).to_pylist() == [value - (1 << 32) if value >=
1 << 31 else value for value in hashes]
+ assigner = create_dynamic_bucket_assigner(table, 1, 1, 0,
base_snapshot_id=snapshot_id)
+ assert isinstance(assigner, NativeDynamicBucketAssigner)
+ assert assigner.restore_snapshot_id == snapshot_id
+ assert assigner.assign(chunked_input, routed.column(1)).to_pylist() ==
expected
+ assert assigner.assign(chunked_input).to_pylist() == expected
+ # Overwrite ignores existing mappings and exposes an empty restore
point.
+ overwritten = create_dynamic_bucket_assigner(table, 1, 1, 0, True,
snapshot_id)
+ assert overwritten.restore_snapshot_id == 0
+ assert overwritten.assign(chunked_input, routed.column(1)).to_pylist()
== [0, 1, 0]
+ query.assert_not_called()
+ assert table.snapshot_manager().get_latest_snapshot().id == snapshot_id
+
+
[email protected]('partitioned', [False, True])
+def
test_native_dynamic_precomputed_hashes_need_only_partition_input(native_rest_catalog,
partitioned):
+ table = _table(native_rest_catalog, partitioned=partitioned)
+ source = {'id': [1, 2, 1]}
+ if partitioned:
+ source['p'] = ['west'] * 3
+ source = pa.table(source)
+ with _forbid_python_input_execution():
+ routed = create_dynamic_bucket_assigner_router(table, 1,
1).route(source)
+ assigner = create_dynamic_bucket_assigner(table, 1, 1, 0,
base_snapshot_id=0)
+ partitions = source.select(['p'] if partitioned else [])
+ assert partitions.num_rows == 3
+ hashes = pa.chunked_array([routed.column(1).slice(0, 1),
routed.column(1).slice(1)])
+ assert assigner.assign(partitions, hashes).to_pylist() == [0, 1, 0]
+ assert assigner.restore_snapshot_id == 0
+ assert table.snapshot_manager().get_latest_snapshot() is None
+
+
[email protected]('bucket', ['4', '-1'])
[email protected]('partitioned', [False, True])
+def
test_native_daft_fixed_and_dynamic_grouped_writes_skip_python_input_execution(
+ native_rest_catalog, bucket, partitioned):
+ daft = pytest.importorskip('daft')
+ from pypaimon.daft.daft_paimon import _write_table
+ from pypaimon.write.write_dispatch import NativeTableWrite
+
+ table = _table(native_rest_catalog, bucket, partitioned)
+ first = {'id': [1, 2], 'value': ['old', 'keep']}
+ second = {'id': [1, 3], 'value': ['new', 'added']}
+ overwritten = {'id': [1, 4], 'value': ['replaced', 'last']}
+ if partitioned:
+ first['p'], second['p'], overwritten['p'] = ['a', 'b'], ['a', 'c'],
['a', 'd']
+ with _forbid_python_input_execution(), \
+ patch.object(RESTCatalogServer, '_table_query_auth_handle',
+ side_effect=AssertionError('Distributed writes must
not request SELECT')) as query, \
+ patch.object(NativeTableWrite, '_switch_to_python',
side_effect=AssertionError('Writer replay')):
+ for rows, mode in [(first, 'append'), (second, 'append'),
(overwritten, 'overwrite')]:
+ assert _write_table(daft.from_pydict(rows), table,
mode).to_pydict()['rows']
+ query.assert_not_called()
+ builder = table.new_read_builder()
+ result =
builder.new_read().to_arrow(builder.new_scan().plan().splits()).sort_by([('id',
'ascending')])
+ expected = overwritten if not partitioned else {
+ 'id': [1, 2, 3, 4], 'value': ['replaced', 'keep', 'added', 'last'],
'p': ['a', 'b', 'c', 'd']}
+ assert result.to_pydict() == expected
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 d1ed56bdce..d45a511ac5 100644
--- a/paimon-python/pypaimon/tests/write/distributed_input_execution_test.py
+++ b/paimon-python/pypaimon/tests/write/distributed_input_execution_test.py
@@ -24,7 +24,7 @@ import pyarrow as pa
import pytest
from pypaimon import CatalogFactory, Schema
-from pypaimon.index.dynamic_bucket import HashBucketAssigner
+from pypaimon.index.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
@@ -49,7 +49,7 @@ def forbid_python_execution():
for target in [
'execution.python.write.row_key_extractor.FixedBucketRowKeyExtractor.__init__',
'execution.python.write.row_key_extractor.DynamicBucketRowKeyExtractor.__init__',
-
'execution.python.write.row_key_extractor.SimpleHashBucketAssigner.__init__',
+
'execution.python.write.row_key_extractor.SimpleDynamicBucketAssigner.__init__',
'execution.python.read.scanner.file_scanner.FileScanner.__init__',
'execution.python.write.table_write.PythonTableWrite.__init__',
'execution.python.write.dynamic_bucket.DynamicBucketIndexMaintainer.__init__',
@@ -85,7 +85,7 @@ def
test_daft_fixed_and_dynamic_routing_do_not_construct_python_extractors(tmp_p
batch = pa.table(rows).to_batches()[0]
assert fixed['bucket'] ==
FixedBucketInputRouter(table.table_schema).extract_partition_bucket_batch(batch)[1]
# Only input hashing belongs to this stage; it must not construct even a
- # transient SimpleHashBucketAssigner or persistent HASH-index maintainer.
+ # transient SimpleDynamicBucketAssigner or persistent HASH-index
maintainer.
dynamic = _table(tmp_path / 'dynamic', '-1')
router = DynamicBucketInputRouter(dynamic.table_schema)
_, partition_hashes, key_hashes = router.extract_hashes_batch(batch)
@@ -116,41 +116,48 @@ def _seed_snapshot(table):
@pytest.mark.parametrize('ignore_existing', [False, True])
-def
test_daft_dynamic_assignment_restores_pinned_native_metadata_without_python_executor(tmp_path,
ignore_existing):
+def
test_daft_dynamic_assignment_dispatches_pinned_snapshot_and_arrow_hashes(tmp_path,
ignore_existing):
daft = pytest.importorskip('daft')
from pypaimon.daft.daft_datasink import make_dynamic_bucket_assignment_udf
table = _table(tmp_path, '-1')
- snapshot, plan = _seed_snapshot(table)
+ snapshot, _ = _seed_snapshot(table)
rows = {'id': [2, 1], 'value': ['other', 'new'], 'p': ['p', 'p']}
router = DynamicBucketInputRouter(table.table_schema)
_, _, hashes = router.extract_hashes_batch(pa.table(rows).to_batches()[0])
rows.update(assigner=[0, 0], hash=hashes)
+ execution = Mock()
+ execution.restore_snapshot_id = 0 if ignore_existing else snapshot.id
+ execution.assign.return_value = pa.array([0, 1] if ignore_existing else
[1, 0], type=pa.int32())
with ExitStack() as stack:
for target in [
'execution.python.write.row_key_extractor.DynamicBucketRowKeyExtractor.__init__',
-
'execution.python.write.row_key_extractor.SimpleHashBucketAssigner.__init__',
+
'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__',
]:
stack.enter_context(patch('pypaimon.' + target,
side_effect=AssertionError(target)))
stack.enter_context(patch('pypaimon.daft.daft_datasink._load_worker_table',
return_value=table))
+ factory = stack.enter_context(patch(
+ 'pypaimon.write.backend_dispatch.create_dynamic_bucket_assigner',
return_value=execution))
native_plan = stack.enter_context(patch(
- 'pypaimon.execution.native.read.plan.native_plan',
return_value=plan))
+ 'pypaimon.execution.native.read.plan.native_plan',
side_effect=AssertionError('Read authorization')))
udf = make_dynamic_bucket_assignment_udf(
table, 1, 1, 'bucket', 'hash', 'base_snapshot',
ignore_existing=ignore_existing, base_snapshot_id=snapshot.id)
- result = daft.from_pydict(rows).select(udf(*[daft.col(name) for name
in rows])).to_pydict()
+ # Match the Int32 transport emitted by the first routing stage.
+ frame = pa.table(rows).set_column(4, 'hash', pa.array(hashes,
type=pa.int32()))
+ result = daft.from_arrow(frame).select(udf(*[daft.col(name) for name
in rows])).to_pydict()
assert result['bucket'] == ([0, 1] if ignore_existing else [1, 0])
- assert result['base_snapshot'] == [snapshot.id, snapshot.id]
- if ignore_existing:
- native_plan.assert_not_called()
- else:
- 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
- predicate = native_plan.call_args[1]['predicate']
- assert predicate.field == 'p' and predicate.literals == ['p']
+ assert result['base_snapshot'] == [0 if ignore_existing else snapshot.id]
* 2
+ native_plan.assert_not_called()
+ factory.assert_called_once_with(
+ table, 1, 1, 0, ignore_existing=ignore_existing,
base_snapshot_id=snapshot.id)
+ transported = execution.assign.call_args[0][0]
+ assert transported.to_pydict() == {key: rows[key] for key in ('id',
'value', 'p')}
+ assert execution.assign.call_args[1]['key_hashes'].equals(pa.array(hashes,
type=pa.int32()))
def
test_native_dynamic_metadata_preserves_missing_hash_index_validation(tmp_path):
@@ -161,7 +168,7 @@ def
test_native_dynamic_metadata_preserves_missing_hash_index_validation(tmp_pat
with patch('pypaimon.execution.native.read.plan.native_plan',
return_value=plan), \
patch('pypaimon.execution.python.read.scanner.file_scanner.FileScanner.__init__',
side_effect=AssertionError('Pure scanner')):
- assigner = HashBucketAssigner(table, 1, 1, 0, 1, -1, snapshot=snapshot)
+ assigner = DynamicBucketAssigner(table, 1, 1, 0, 1, -1,
snapshot=snapshot)
with pytest.raises(RuntimeError, match='no complete HASH index.*7'):
assigner.assign(('p',), 0, 123)
diff --git a/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
b/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
index fc5530b6e7..8db7e598f2 100644
--- a/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
+++ b/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py
@@ -27,7 +27,7 @@ import pytest
from pypaimon import CatalogFactory, Schema
from pypaimon.index.dynamic_bucket import (
- HashBucketAssigner,
+ DynamicBucketAssigner,
_PartitionIndex,
_iter_hashes,
compute_assigner,
@@ -122,6 +122,28 @@ class DynamicBucketTest(unittest.TestCase):
min_int = -(2 ** 31)
self.assertEqual(compute_assigner(min_int, 0, 5, 5), 3)
self.assertEqual(compute_assigner(2, min_int, 5, 5), 0)
+ # Java int addition wraps before the final remainder.
+ self.assertEqual(compute_assigner(max_int - 1, max_int - 1, max_int,
max_int), -4)
+
+ def test_global_maximum_does_not_hide_an_unused_partition_bucket(self):
+ index = _PartitionIndex({10: 0}, {0: 1}, 1)
+ # A different partition already restored the global maximum bucket.
+ # Java PartitionIndex still allocates the unused local bucket 1.
+ self.assertEqual(index.assign(20, lambda _: True, 2, 1), 1)
+ self.assertEqual(index.assign(20, lambda _: True, 2, 1), 1)
+
+ def
test_reached_partition_bucket_limit_is_cached_and_hash_assignments_stay_stable(self):
+ index = _PartitionIndex({10: 1}, {1: 1}, 1)
+ owns_odd = Mock(side_effect=lambda bucket: bucket % 2 == 1)
+ self.assertEqual(index.assign(20, owns_odd, 4, 3), 3)
+ first_overflow = index.assign(30, owns_odd, 4, 3)
+ self.assertIn(first_overflow, (1, 3))
+ owns_odd.reset_mock()
+ for key in range(31, 50):
+ bucket = index.assign(key, owns_odd, 4, 3)
+ self.assertIn(bucket, (1, 3))
+ self.assertEqual(index.assign(key, owns_odd, 4, 3), bucket)
+ owns_odd.assert_not_called()
def test_binary_row_hash_matches_java_for_bucket_key_types(self):
# Generated with Java
InternalRowSerializer.toBinaryRow(...).hashCode().
@@ -226,7 +248,7 @@ class DynamicBucketTest(unittest.TestCase):
validate_bucket_id(32768)
def test_restore_rejects_bucket_id_above_java_short_range(self):
- assigner = HashBucketAssigner(
+ assigner = DynamicBucketAssigner(
table=Mock(),
num_channels=1,
num_assigners=1,
@@ -260,7 +282,7 @@ class DynamicBucketTest(unittest.TestCase):
def test_assigner_rejects_record_owned_by_another_writer(self):
with tempfile.TemporaryDirectory() as root:
table = self._create_table(root, 'wrong_assigner')
- assigner = HashBucketAssigner(
+ assigner = DynamicBucketAssigner(
table=table,
num_channels=2,
num_assigners=2,
diff --git a/paimon-python/pypaimon/tests/write/dynamic_input_routing_test.py
b/paimon-python/pypaimon/tests/write/dynamic_input_routing_test.py
new file mode 100644
index 0000000000..24ec5c9e88
--- /dev/null
+++ b/paimon-python/pypaimon/tests/write/dynamic_input_routing_test.py
@@ -0,0 +1,317 @@
+# 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.
+
+"""Distributed dynamic routing keeps Native errors and Arrow batches intact."""
+
+import pickle
+import sys
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import pyarrow as pa
+import pytest
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.catalog.catalog_exception import TableNoPermissionException
+from pypaimon.execution.native.write.input_routing import (
+ NativeDynamicBucketAssigner, NativeDynamicBucketAssignerRouter,
NativeInputRouter, input_batches,
+)
+from pypaimon.execution.python.write.input_routing import (
+ PythonDynamicBucketAssigner, PythonDynamicBucketAssignerRouter,
+)
+from pypaimon.write.backend_dispatch import create_dynamic_bucket_assigner,
create_dynamic_bucket_assigner_router
+
+
[email protected]('row_count', [0, 3])
+def
test_arrow_transport_preserves_rows_when_all_columns_are_projected_out(row_count):
+ data = pa.table({'id': list(range(row_count))}).select([])
+ batches = input_batches(data)
+ assert sum(batch.num_rows for batch in batches) == row_count
+ assert all(batch.num_columns == 0 for batch in batches)
+
+
+def test_empty_arrow_transport_preserves_routing_schema_and_metadata():
+ schema = pa.schema([pa.field('id', pa.int64(), nullable=False)],
metadata={b'input': b'empty'})
+ data = pa.Table.from_batches([], schema=schema)
+ batches = input_batches(data)
+ assert len(batches) == 1
+ assert batches[0].num_rows == 0
+ assert batches[0].schema.equals(schema, check_metadata=True)
+
+
[email protected]('mode', ['dense', 'sparse'])
+def test_empty_arrow_transport_preserves_union_payload_columns(mode):
+ type_ids = pa.array([], type=pa.int8())
+ children = [pa.array([], type=pa.int64()), pa.array([], type=pa.string())]
+ payload = (pa.UnionArray.from_dense(type_ids, pa.array([],
type=pa.int32()), children)
+ if mode == 'dense' else pa.UnionArray.from_sparse(type_ids,
children))
+ # An ignored payload still travels with the batch; constructing it with
+ # pa.array([], type=union_type) is unsupported by Arrow.
+ data = pa.table({'id': pa.array([], type=pa.int64()), 'payload': payload})
+ batches = input_batches(data)
+ assert len(batches) == 1
+ assert batches[0].num_rows == 0
+ assert batches[0].schema.equals(data.schema, check_metadata=True)
+
+
[email protected]('operation', ['fixed', 'dynamic', 'assign'])
+def test_empty_input_is_still_validated_by_the_sdk(tmp_path, operation):
+ table = _table(tmp_path)
+ core = Mock()
+ if operation == 'fixed':
+ execution = NativeInputRouter(table)
+ execution._router = core
+ method = 'route'
+ elif operation == 'dynamic':
+ execution = NativeDynamicBucketAssignerRouter(table, 1, 1)
+ execution._router = core
+ method = 'route'
+ else:
+ execution = NativeDynamicBucketAssigner(table, 1, 1, 0)
+ execution._assigner = core
+ method = 'assign'
+ failure = ValueError('Missing routing column p')
+ getattr(core, method).side_effect = failure
+ schema = pa.schema([('id', pa.int64())], metadata={b'input': b'empty'})
+ data = pa.Table.from_batches([], schema=schema)
+ with pytest.raises(ValueError) as raised:
+ getattr(execution, method)(data)
+ assert raised.value is failure
+ transported = getattr(core, method).call_args.args[0]
+ assert transported.num_rows == 0
+ assert transported.schema.equals(schema, check_metadata=True)
+
+
+def _uncombined_table(batch):
+ chunked = pa.Table.from_batches([batch.slice(0, 1), batch.slice(1)])
+ data = Mock(spec=pa.Table)
+ data.schema = chunked.schema
+ data.num_rows = chunked.num_rows
+ data.to_batches.return_value = chunked.to_batches()
+ data.combine_chunks.return_value = chunked
+ return data
+
+
+def _table(tmp_path, native=True):
+ catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('default', True)
+ catalog.create_table('default.t', Schema.from_pyarrow_schema(
+ pa.schema([('id', pa.int64()), ('p', pa.string()), ('value',
pa.string())]),
+ primary_keys=['id', 'p'], partition_keys=['p'], options={
+ 'bucket': '-1', 'write.native.enabled': str(native).lower(),
+ 'dynamic-bucket.target-row-num': '1',
+ }), False)
+ return catalog.get_table('default.t')
+
+
[email protected]
+def sdk(monkeypatch):
+ module = SimpleNamespace(DynamicBucketAssignerRouter=Mock(),
DynamicBucketAssigner=Mock())
+ monkeypatch.setitem(sys.modules, 'pypaimon_rust.datafusion', module)
+ return module
+
+
+def test_python_dynamic_router_and_assigner_share_arrow_contract(tmp_path):
+ table = _table(tmp_path, False)
+ batch = pa.table({'p': ['west', 'west', 'east'], 'id': [1, 2, 1]})
+ router = create_dynamic_bucket_assigner_router(table, 1, 1)
+ assert isinstance(router, PythonDynamicBucketAssignerRouter)
+ routed = router.route(batch)
+ assert routed.schema == pa.schema([('assigner', pa.int32()), ('key_hash',
pa.int32())])
+ assert routed.column(0).to_pylist() == [0, 0, 0]
+ assigner = create_dynamic_bucket_assigner(table, 1, 1, 0,
base_snapshot_id=0)
+ assert isinstance(assigner, PythonDynamicBucketAssigner)
+ assert assigner.restore_snapshot_id == 0
+ assert assigner.assign(batch, routed.column(1)).to_pylist() == [0, 1, 0]
+ assert assigner.assign(batch).to_pylist() == [0, 1, 0]
+
+
[email protected]('hashes,error', [
+ (pa.array([1], type=pa.int32()), ValueError),
+ (pa.array([1, None], type=pa.int32()), ValueError),
+ (pa.array([1, 2], type=pa.int64()), TypeError),
+])
+def test_python_dynamic_hash_transport_rejects_invalid_arrays(tmp_path,
hashes, error):
+ assigner = create_dynamic_bucket_assigner(_table(tmp_path, False), 1, 1,
0, base_snapshot_id=0)
+ with pytest.raises(error):
+ assigner.assign(pa.table({'id': [1, 2], 'p': ['west', 'west']}),
hashes)
+
+
+def
test_native_dynamic_router_pickle_rebuilds_without_python_hashing(tmp_path,
sdk):
+ table = _table(tmp_path)
+ native_table = Mock()
+ batch = pa.record_batch([[1, 2], ['west', 'west']], names=['id', 'p'])
+ expected = pa.record_batch([pa.array([0, 1], type=pa.int32()),
pa.array([42, -3], type=pa.int32())],
+ names=['assigner', 'key_hash'])
+ sdk.DynamicBucketAssignerRouter.return_value.route.side_effect =
[expected.slice(0, 1), expected.slice(1)]
+ 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__',
+ side_effect=AssertionError('Python hashing')):
+ router = create_dynamic_bucket_assigner_router(table, 2, 2)
+ assert isinstance(router, NativeDynamicBucketAssignerRouter)
+ router._router = lambda: None
+ restored = pickle.loads(pickle.dumps(router))
+ assert restored._router is None
+ assert restored.route(pa.Table.from_batches([batch.slice(0, 1),
batch.slice(1)])).equals(expected)
+ assert sdk.DynamicBucketAssignerRouter.call_count == 2
+
sdk.DynamicBucketAssignerRouter.assert_called_with(native_table.schema.return_value,
2, 2)
+ assert [call.args[0].to_pydict() for call in
sdk.DynamicBucketAssignerRouter.return_value.route.call_args_list] == [
+ batch.slice(0, 1).to_pydict(), batch.slice(1).to_pydict()]
+
+
+def
test_native_dynamic_routing_preserves_all_batches_and_sdk_output_schema(tmp_path,
sdk):
+ router = NativeDynamicBucketAssignerRouter(_table(tmp_path), 2, 2)
+ router._router = sdk.DynamicBucketAssignerRouter.return_value
+ schema = pa.schema([
+ pa.field('assigner', pa.int32(), nullable=False),
+ pa.field('key_hash', pa.int32(), nullable=False),
+ ], metadata={b'output': b'routing'})
+ expected = pa.record_batch([[0, 1, 0], [42, -3, 42]], schema=schema)
+ router._router.route.side_effect = [expected.slice(0, 1),
expected.slice(1)]
+ data = _uncombined_table(pa.record_batch([[1, 2, 1]], names=['id']))
+ actual = router.route(data)
+ assert actual.equals(expected, check_metadata=True)
+ assert [call.args[0].num_rows for call in
router._router.route.call_args_list] == [1, 2]
+ data.combine_chunks.assert_not_called()
+
+
[email protected]('precomputed', [False, True])
+def
test_native_assignment_keeps_state_and_slices_hashes_across_input_batches(tmp_path,
precomputed):
+ assigner = NativeDynamicBucketAssigner(_table(tmp_path), 1, 1, 0)
+ core = assigner._assigner = Mock()
+ core.assign.side_effect = [pa.array([0], type=pa.int32()), pa.array([1,
0], type=pa.int32())]
+ batch = pa.record_batch([['west', 'west', 'west'], [1, 2, 1]], names=['p',
'id'])
+ data = _uncombined_table(batch)
+ hashes = pa.chunked_array([pa.array([42, -3], type=pa.int32()),
pa.array([42], type=pa.int32())])
+ actual = assigner.assign(data, hashes if precomputed else None)
+ assert actual.to_pylist() == [0, 1, 0]
+ assert [call.args[0].num_rows for call in core.assign.call_args_list] ==
[1, 2]
+ actual_hashes = [call.kwargs['key_hashes'] for call in
core.assign.call_args_list]
+ if precomputed:
+ assert [values.to_pylist() for values in actual_hashes] == [[42], [-3,
42]]
+ else:
+ assert actual_hashes == [None, None]
+ data.combine_chunks.assert_not_called()
+
+
[email protected]('hashes,error', [
+ (pa.chunked_array([pa.array([1], type=pa.int32()), pa.array([2, None],
type=pa.int32())]), ValueError),
+ (pa.array([1, 2], type=pa.int32()), ValueError),
+ (pa.array([1, 2, 3, 4], type=pa.int32()), ValueError),
+ (pa.array([1, 2, 3], type=pa.int64()), TypeError),
+ ([1, 2, 3], TypeError),
+])
+def
test_native_assignment_validates_all_hashes_before_accepting_any_input(tmp_path,
hashes, error):
+ assigner = NativeDynamicBucketAssigner(_table(tmp_path), 1, 1, 0)
+ core = assigner._assigner = Mock()
+ data = _uncombined_table(pa.record_batch([[1, 2, 3]], names=['id']))
+ with pytest.raises(error):
+ assigner.assign(data, hashes)
+ core.assign.assert_not_called()
+
+
[email protected]('operation', ['route', 'assign'])
+def test_native_dynamic_later_batch_errors_are_terminal(tmp_path, sdk,
operation):
+ if operation == 'route':
+ execution = NativeDynamicBucketAssignerRouter(_table(tmp_path), 1, 1)
+ core = execution._router = sdk.DynamicBucketAssignerRouter.return_value
+ first = pa.record_batch([pa.array([0], type=pa.int32()),
pa.array([42], type=pa.int32())],
+ names=['assigner', 'key_hash'])
+ else:
+ execution = NativeDynamicBucketAssigner(_table(tmp_path), 1, 1, 0)
+ core = execution._assigner = Mock()
+ first = pa.array([0], type=pa.int32())
+ failure = ValueError('later input batch failed')
+ getattr(core, operation).side_effect = [first, failure]
+ data = _uncombined_table(pa.record_batch([[1, 2, 3]], names=['id']))
+ with
patch('pypaimon.execution.python.write.input_routing.PythonDynamicBucketAssignerRouter.__init__',
+ side_effect=AssertionError('Python replay')), \
+
patch('pypaimon.execution.python.write.input_routing.PythonDynamicBucketAssigner.__init__',
+ side_effect=AssertionError('Python replay')):
+ with pytest.raises(ValueError) as raised:
+ getattr(execution, operation)(data)
+ assert raised.value is failure
+
+
[email protected]('overwrite', [False, True])
[email protected]('base_snapshot_id', [None, 0, 7])
+def
test_native_assignment_uses_dedicated_constructor_and_precomputed_arrow_hashes(
+ tmp_path, sdk, overwrite, base_snapshot_id):
+ table = _table(tmp_path)
+ native_table = Mock()
+ core = sdk.DynamicBucketAssigner.return_value
+ core.restore_snapshot_id = 0 if overwrite or base_snapshot_id == 0 else 7
+ core.assign.return_value = pa.array([1, 0], type=pa.int32())
+ batch = pa.record_batch([[2, 1], ['west', 'west']], names=['id', 'p'])
+ hashes = pa.chunked_array([pa.array([42], type=pa.int32()), pa.array([-3],
type=pa.int32())])
+ 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.object(table, 'new_batch_write_builder',
side_effect=AssertionError('Python write builder')), \
+
patch('pypaimon.index.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)
+ assert assigner.restore_snapshot_id == core.restore_snapshot_id
+ assert assigner.assign(batch, hashes) is core.assign.return_value
+ native_table.new_batch_write_builder.assert_not_called()
+ sdk.DynamicBucketAssigner.assert_called_once_with(
+ native_table, 2, 2, 1, restore_snapshot_id=base_snapshot_id,
overwrite=overwrite)
+ assert core.assign.call_args[0][0] is batch
+ assert
core.assign.call_args[1]['key_hashes'].equals(hashes.combine_chunks())
+
+
[email protected]('operation', ['route', 'assign'])
[email protected]('phase', ['prepare', 'execute'])
+def test_native_dynamic_errors_never_replay_python(tmp_path, sdk, operation,
phase):
+ table = _table(tmp_path)
+ failure = TableNoPermissionException(table.identifier)
+ native_table = Mock()
+ core = (sdk.DynamicBucketAssignerRouter.return_value if operation ==
'route' else
+ sdk.DynamicBucketAssigner.return_value)
+ core.restore_snapshot_id = 0
+ if phase == 'prepare':
+ if operation == 'route':
+ sdk.DynamicBucketAssignerRouter.side_effect = failure
+ else:
+ sdk.DynamicBucketAssigner.side_effect = failure
+ else:
+ getattr(core, operation).side_effect = failure
+ 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.execution.python.write.input_routing.PythonDynamicBucketAssignerRouter.__init__',
+ side_effect=AssertionError('Python replay')), \
+
patch('pypaimon.execution.python.write.input_routing.PythonDynamicBucketAssigner.__init__',
+ side_effect=AssertionError('Python replay')):
+ with pytest.raises(TableNoPermissionException) as raised:
+ execution = (create_dynamic_bucket_assigner_router(table, 1, 1) if
operation == 'route' else
+ create_dynamic_bucket_assigner(table, 1, 1, 0))
+ getattr(execution, operation)(pa.record_batch([[1], ['west']],
names=['id', 'p']))
+ assert raised.value is failure
+
+
[email protected]('operation', ['route', 'assign'])
+def test_missing_native_environment_selects_python_before_input(tmp_path,
operation):
+ table = _table(tmp_path)
+ with
patch('pypaimon.execution.native.write.input_routing.native_runtime_available',
return_value=False):
+ execution = (create_dynamic_bucket_assigner_router(table, 1, 1) if
operation == 'route' else
+ create_dynamic_bucket_assigner(table, 1, 1, 0,
base_snapshot_id=0))
+ expected_type = PythonDynamicBucketAssignerRouter if operation == 'route'
else PythonDynamicBucketAssigner
+ assert isinstance(execution, expected_type)
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 da3fc398ea..fdb5e70bec 100644
--- a/paimon-python/pypaimon/tests/write/native_input_routing_test.py
+++ b/paimon-python/pypaimon/tests/write/native_input_routing_test.py
@@ -100,6 +100,41 @@ def
test_native_router_passes_arrow_without_python_row_hashing(tmp_path, sdk):
sdk.BucketRouter.return_value.route.assert_called_once_with(batch)
+def _uncombined_table(batch):
+ # Arrow may keep multiple binary chunks to avoid a buffer overflow. Model
+ # that contract with small batches instead of allocating several GiB.
+ chunked = pa.Table.from_batches([batch.slice(0, 1), batch.slice(1)])
+ data = Mock(spec=pa.Table)
+ data.schema = chunked.schema
+ data.num_rows = chunked.num_rows
+ data.to_batches.return_value = chunked.to_batches()
+ data.combine_chunks.return_value = chunked
+ return data
+
+
+def test_native_router_preserves_every_batch_of_binary_input(tmp_path, sdk):
+ batch = pa.record_batch([['west', 'east', 'west'], [1, 2, 3]], names=['p',
'id'])
+ router = NativeInputRouter(_table(tmp_path))
+ router._router = sdk.BucketRouter.return_value
+ router._router.route.side_effect = [pa.array([2], type=pa.int32()),
pa.array([1, 3], type=pa.int32())]
+ data = _uncombined_table(batch)
+ assert router.route(data).to_pylist() == [2, 1, 3]
+ assert [call.args[0].num_rows for call in
router._router.route.call_args_list] == [1, 2]
+ data.combine_chunks.assert_not_called()
+
+
+def test_native_router_later_batch_failure_is_terminal(tmp_path, sdk):
+ router = NativeInputRouter(_table(tmp_path))
+ router._router = sdk.BucketRouter.return_value
+ failure = ValueError('later routing batch failed')
+ router._router.route.side_effect = [pa.array([0], type=pa.int32()),
failure]
+ data = _uncombined_table(pa.record_batch([[1, 2, 3]], names=['id']))
+ with patch.object(FixedBucketInputRouter, '__init__',
side_effect=AssertionError('Python replay')):
+ with pytest.raises(ValueError) as raised:
+ router.route(data)
+ assert raised.value is failure
+
+
def test_native_router_failure_propagates_without_selecting_python(tmp_path,
sdk):
table = _table(tmp_path)
router = NativeInputRouter(table)
@@ -159,6 +194,22 @@ def
test_native_planner_transports_group_statistics_positionally(tmp_path, sdk):
assert grouped.column(1).type == pa.int64() and grouped.column(2).type ==
pa.int64()
+def
test_native_planner_combines_partition_statistics_from_all_input_batches(tmp_path,
sdk):
+ table = _table(tmp_path, bucket='-2')
+ planner = NativePostponeBucketPlanner(table, {}, Mock())
+ schema = pa.schema([('p', pa.string()), ('row_count', pa.int64()),
('data_size', pa.int64())])
+ sdk.PostponeBucketPlanner.return_value.input_statistics.side_effect = [
+ pa.record_batch([['west'], [1], [10]], schema=schema),
+ pa.record_batch([['west', 'east'], [1, 1], [20, 30]], schema=schema),
+ ]
+ batch = pa.record_batch([['west', 'west', 'east'], [1, 2, 3]], names=['p',
'id'])
+ data = _uncombined_table(batch)
+ assert planner.input_partition_stats(data) == {('west',): (2, 30),
('east',): (1, 30)}
+ assert [call.args[0].num_rows
+ for call in
sdk.PostponeBucketPlanner.return_value.input_statistics.call_args_list] == [1,
2]
+ data.combine_chunks.assert_not_called()
+
+
def
test_native_planner_pickle_keeps_resolved_counts_and_worker_failures(tmp_path,
sdk):
table = _table(tmp_path, bucket='-2')
core = sdk.PostponeBucketPlanner.return_value
diff --git
a/paimon-python/pypaimon/tests/write/simple_hash_bucket_assigner_test.py
b/paimon-python/pypaimon/tests/write/simple_dynamic_bucket_assigner_test.py
similarity index 88%
rename from
paimon-python/pypaimon/tests/write/simple_hash_bucket_assigner_test.py
rename to
paimon-python/pypaimon/tests/write/simple_dynamic_bucket_assigner_test.py
index ee9405658c..d68c925f53 100644
--- a/paimon-python/pypaimon/tests/write/simple_hash_bucket_assigner_test.py
+++ b/paimon-python/pypaimon/tests/write/simple_dynamic_bucket_assigner_test.py
@@ -18,13 +18,13 @@
import unittest
from unittest.mock import patch
-from pypaimon.execution.python.write.row_key_extractor import
SimpleHashBucketAssigner
+from pypaimon.execution.python.write.row_key_extractor import
SimpleDynamicBucketAssigner
-class SimpleHashBucketAssignerTest(unittest.TestCase):
+class SimpleDynamicBucketAssignerTest(unittest.TestCase):
def test_assign(self):
- assigner = SimpleHashBucketAssigner(2, 0, 100, -1)
+ assigner = SimpleDynamicBucketAssigner(2, 0, 100, -1)
partition = ()
buckets = [assigner.assign(partition, h) for h in range(301)]
for b in buckets[:100]:
@@ -34,7 +34,7 @@ class SimpleHashBucketAssignerTest(unittest.TestCase):
self.assertEqual(buckets[200], 4)
def test_assign_with_upper_bound(self):
- assigner = SimpleHashBucketAssigner(1, 0, 100, 3)
+ assigner = SimpleDynamicBucketAssigner(1, 0, 100, 3)
partition = ()
buckets = [assigner.assign(partition, h) for h in range(400)]
for b in buckets[:100]:
@@ -49,7 +49,7 @@ class SimpleHashBucketAssignerTest(unittest.TestCase):
def test_assign_with_same_hash(self):
for max_buckets in [-1, 1, 2]:
with self.subTest(max_buckets=max_buckets):
- assigner = SimpleHashBucketAssigner(1, 0, 100, max_buckets)
+ assigner = SimpleDynamicBucketAssigner(1, 0, 100, max_buckets)
partition = ()
hashes = list(range(100)) + list(range(100))
buckets = [assigner.assign(partition, h) for h in hashes]
@@ -59,7 +59,7 @@ class SimpleHashBucketAssignerTest(unittest.TestCase):
def test_maximum_bucket_id(self):
for max_buckets in [-1, 32768]:
with self.subTest(max_buckets=max_buckets):
- assigner = SimpleHashBucketAssigner(
+ assigner = SimpleDynamicBucketAssigner(
32768, 32767, 1, max_buckets)
self.assertEqual(assigner.assign((), 1), 32767)
self.assertEqual(assigner.assign((), 1), 32767)
@@ -69,7 +69,7 @@ class SimpleHashBucketAssignerTest(unittest.TestCase):
ValueError,
"'dynamic-bucket.max-buckets' must be -1 or between 1 and 32768",
):
- SimpleHashBucketAssigner(1, 0, 100, 40000)
+ SimpleDynamicBucketAssigner(1, 0, 100, 40000)
def test_register_each_bucket_once(self):
for num_assigners, assign_id, max_buckets, expected in [
@@ -80,14 +80,14 @@ class SimpleHashBucketAssignerTest(unittest.TestCase):
]:
with self.subTest(num_assigners=num_assigners, assign_id=assign_id,
max_buckets=max_buckets):
- assigner = SimpleHashBucketAssigner(num_assigners, assign_id,
2, max_buckets)
+ assigner = SimpleDynamicBucketAssigner(num_assigners,
assign_id, 2, max_buckets)
for h, expected_bucket in enumerate(expected):
self.assertEqual(assigner.assign((), h), expected_bucket)
index = assigner._partition_index[()]
self.assertCountEqual(index.bucket_list,
index.bucket_information)
def test_overflow_uses_all_registered_buckets(self):
- assigner = SimpleHashBucketAssigner(1, 0, 2, 3)
+ assigner = SimpleDynamicBucketAssigner(1, 0, 2, 3)
initial = [assigner.assign((), h) for h in range(6)]
self.assertEqual(initial, [0, 0, 1, 1, 2, 2])
index = assigner._partition_index[()]
diff --git a/paimon-python/pypaimon/write/backend_dispatch.py
b/paimon-python/pypaimon/write/backend_dispatch.py
index f62e553fc2..f6f20f269e 100644
--- a/paimon-python/pypaimon/write/backend_dispatch.py
+++ b/paimon-python/pypaimon/write/backend_dispatch.py
@@ -220,6 +220,37 @@ def create_input_router(table, bucket_plan=None):
return FixedBucketInputRouter(table.table_schema)
+def create_dynamic_bucket_assigner_router(table, num_channels, num_assigners):
+ """Select routing to dynamic bucket assigners before a task receives
rows."""
+ if table.options.native_write_enabled():
+ from pypaimon.execution.native.write.input_routing import
create_native_dynamic_bucket_assigner_router
+ try:
+ native = create_native_dynamic_bucket_assigner_router(table,
num_channels, num_assigners)
+ except NotImplementedError:
+ native = None
+ if native is not None:
+ return native
+ from pypaimon.execution.python.write.input_routing import
PythonDynamicBucketAssignerRouter
+ return PythonDynamicBucketAssignerRouter(table, num_channels,
num_assigners)
+
+
+def create_dynamic_bucket_assigner(table, num_channels, num_assigners,
assign_id,
+ ignore_existing=False,
base_snapshot_id=None):
+ """Select restoration before assignment; authorization failures are
terminal."""
+ if table.options.native_write_enabled():
+ from pypaimon.execution.native.write.input_routing import
create_native_dynamic_bucket_assigner
+ # Resolving the writer snapshot can authorize and read HASH metadata.
+ # Only an explicitly unsupported environment returns None; never turn
+ # an authorization/restoration exception into another backend attempt.
+ native = create_native_dynamic_bucket_assigner(
+ table, num_channels, num_assigners, assign_id, ignore_existing,
base_snapshot_id)
+ if native is not None:
+ return native
+ from pypaimon.execution.python.write.input_routing import
PythonDynamicBucketAssigner
+ return PythonDynamicBucketAssigner(
+ table, num_channels, num_assigners, assign_id, ignore_existing,
base_snapshot_id)
+
+
def load_postpone_bucket_metadata(table):
"""Select metadata planning before public bucket arithmetic consumes
input."""
if table.options.native_write_enabled():
