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 c0a939a96c [Python] Execute Ray bucket routing and postpone planning
in Rust (#10770)
c0a939a96c is described below
commit c0a939a96c252ac22abb62220879c7aa3815327b
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Oct 23 17:44:49 2026 +0800
[Python] Execute Ray bucket routing and postpone planning in Rust (#10770)
---
docs/docs/pypaimon/ray-data.md | 6 +
.../execution/native/write/input_routing.py | 73 ++++++++
.../execution/native/write/postpone_bucket.py | 102 +++++++++--
paimon-python/pypaimon/ray/shuffle.py | 20 +--
paimon-python/pypaimon/table/row/generic_row.py | 31 +++-
.../pypaimon/tests/native_input_routing_test.py | 192 +++++++++++++++++++++
.../pypaimon/tests/noncompact_null_row_test.py | 81 +++++++++
.../pypaimon/tests/test_ray_shuffle_helper.py | 38 ++--
.../tests/write/native_input_routing_test.py | 180 +++++++++++++++++++
.../write/postpone_execution_boundary_test.py | 44 ++++-
paimon-python/pypaimon/write/backend_dispatch.py | 16 ++
paimon-python/pypaimon/write/input_routing.py | 8 +
paimon-python/pypaimon/write/ray_datasink.py | 9 +-
13 files changed, 741 insertions(+), 59 deletions(-)
diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md
index 5566ee9e36..135f83f654 100644
--- a/docs/docs/pypaimon/ray-data.md
+++ b/docs/docs/pypaimon/ray-data.md
@@ -278,6 +278,12 @@ including the default dynamic-bucket primary-key table
(`bucket = -1`).
Ray write tasks create independent Paimon writers, which can assign
overlapping buckets or sequence numbers for those modes.
+With `write.native.enabled=true`, fixed and postpone input bucket routing runs
+in Rust, including POSTPONE per-partition input statistics and bucket-count
+inference. Ray still materializes, groups, and schedules the input; workers
+reconstruct native routing state from the shared plan. The grouping and
+single-writer ownership rules above remain the same.
+
**Parameters:**
- `dataset`: the Ray Dataset to write.
- `table_identifier`: full table name, e.g. `"db_name.table_name"`.
diff --git a/paimon-python/pypaimon/execution/native/write/input_routing.py
b/paimon-python/pypaimon/execution/native/write/input_routing.py
new file mode 100644
index 0000000000..461cba52f1
--- /dev/null
+++ b/paimon-python/pypaimon/execution/native/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 transport for Rust's distributed input bucket routing."""
+
+import pyarrow as pa
+
+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):
+ if isinstance(data, pa.RecordBatch):
+ return data
+ if isinstance(data, pa.Table):
+ return 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')
+
+
+def create_native_input_router(table, bucket_plan=None):
+ if not native_runtime_available():
+ return None
+ native_table = create_native_write_table(table)
+ if native_table is None:
+ return None
+ router = NativeInputRouter(table, bucket_plan)
+ router._prepare(native_table)
+ return router
+
+
+class NativeInputRouter:
+ """Keep serializable configuration while Rust owns every input-row hash."""
+
+ def __init__(self, table, bucket_plan=None):
+ self.table = table
+ self._bucket_plan = bucket_plan
+ self._router = None
+
+ def _prepare(self, native_table=None):
+ from pypaimon_rust.datafusion import BucketRouter
+ if native_table is None:
+ native_table = create_native_write_table(self.table)
+ if native_table is None:
+ raise RuntimeError('Native input bucket routing is unavailable')
+ router = BucketRouter(native_table.schema())
+ if self._bucket_plan is not None:
+ router =
router.with_bucket_plan(self._bucket_plan.to_arrow(self.table))
+ self._router = router
+
+ def route(self, data) -> pa.Int32Array:
+ if self._router is None:
+ self._prepare()
+ return self._router.route(as_record_batch(data))
+
+ def __getstate__(self):
+ state = dict(self.__dict__)
+ state['_router'] = None
+ return state
diff --git a/paimon-python/pypaimon/execution/native/write/postpone_bucket.py
b/paimon-python/pypaimon/execution/native/write/postpone_bucket.py
index 2bb38944b5..acb3823aa4 100644
--- a/paimon-python/pypaimon/execution/native/write/postpone_bucket.py
+++ b/paimon-python/pypaimon/execution/native/write/postpone_bucket.py
@@ -15,32 +15,114 @@
# specific language governing permissions and limitations
# under the License.
-"""Read existing bucket layouts through Rust before distributed input work."""
+"""Arrow transport for Rust's postpone bucket statistics and planning."""
+
+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.native_commit import
create_native_write_table
-from pypaimon.write.bucket_planning import PostponeBucketPlanning
+from pypaimon.schema.data_types import PyarrowFieldParser
+from pypaimon.write.bucket_plan import PostponeBucketPlan
def create_native_postpone_bucket_planner(table, known_num_buckets=None):
if not native_runtime_available():
return None
- return NativePostponeBucketPlanner(table, known_num_buckets)
+ native_table = create_native_write_table(table)
+ if native_table is None:
+ return None
+ return NativePostponeBucketPlanner(table, known_num_buckets, native_table)
+
+
+def _plan_dict(plan):
+ # Both Core and Python use positional columns. Partition keys may be
+ # called total_buckets, row_count or data_size.
+ columns = [column.to_pylist() for column in plan.columns]
+ return {tuple(column[row] for column in columns[:-1]): columns[-1][row]
+ for row in range(plan.num_rows)}
+
+
+class NativePostponeBucketPlanner:
+ """Keep Python plans serializable while Rust evaluates all input rows."""
+
+ def __init__(self, table, known_num_buckets=None, native_table=None):
+ self.table = table
+ self._planner = None
+ # Preserve the attributes exposed by the existing public planner.
+ options = table.options
+ self.default_bucket_num = options.postpone_default_bucket_num()
+ self.target_row_num_per_bucket = None
+ self.target_size_per_bucket = None
+ if self.default_bucket_num is None:
+ self.max_num_buckets =
options.postpone_batch_write_fixed_bucket_max_parallelism()
+ self.target_row_num_per_bucket =
options.postpone_target_row_num_per_bucket()
+ if self.target_row_num_per_bucket is None:
+ self.target_size_per_bucket =
options.postpone_target_size_per_bucket()
+ if native_table is None:
+ native_table = create_native_write_table(table)
+ if native_table is None:
+ raise RuntimeError('Native postpone bucket planning is
unavailable')
+ if known_num_buckets is None:
+ known_num_buckets = self._load_bucket_metadata(table, native_table)
+ self._known_num_buckets = dict(known_num_buckets)
+ self._prepare(native_table)
+
+ def _prepare(self, native_table=None):
+ from pypaimon_rust.datafusion import PostponeBucketPlanner
+ if native_table is None:
+ native_table = create_native_write_table(self.table)
+ if native_table is None:
+ raise RuntimeError('Native postpone bucket planning is
unavailable')
+ self._planner = PostponeBucketPlanner(
+ native_table.schema(),
PostponeBucketPlan(self._known_num_buckets).to_arrow(self.table))
+ def _ensure_planner(self):
+ if self._planner is None:
+ self._prepare()
+ return self._planner
-class NativePostponeBucketPlanner(PostponeBucketPlanning):
- """Use Rust planning metadata and shared input-only bucket arithmetic."""
+ def current_plan(self):
+ self._known_num_buckets =
_plan_dict(self._ensure_planner().current_plan())
+ 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)}
+
+ def plan(self, partition_stats):
+ fields = [self.table.field_dict[name] for name in
self.table.partition_keys]
+ schema = PyarrowFieldParser.from_paimon_schema(fields)
+ schema = schema.append(pa.field('row_count', pa.int64(),
nullable=False))
+ schema = schema.append(pa.field('data_size', pa.int64(),
nullable=False))
+ partitions = list(partition_stats)
+ if any(len(partition) != len(fields) for partition in partitions):
+ raise ValueError('Input statistics partition arity does not match
the table schema')
+ columns = [pa.array([partition[i] for partition in partitions],
type=field.type)
+ for i, field in enumerate(list(schema)[:-2])]
+ columns.extend([
+ pa.array([partition_stats[partition][i] for partition in
partitions], type=pa.int64())
+ for i in range(2)])
+ plan = self._ensure_planner().plan(pa.RecordBatch.from_arrays(columns,
schema=schema))
+ self._known_num_buckets = _plan_dict(plan)
+ return PostponeBucketPlan(self._known_num_buckets)
+
+ def __getstate__(self):
+ state = dict(self.__dict__)
+ state['_planner'] = None
+ return state
@staticmethod
- def _load_bucket_metadata(table):
+ def _load_bucket_metadata(table, native_table=None):
# Writer restoration uses the catalog's write identity. Query plans
# would add SELECT rules and can combine different physical layouts.
- native_table = create_native_write_table(table)
+ if native_table is None:
+ native_table = create_native_write_table(table)
if native_table is None:
raise RuntimeError('Native bucket metadata planning is
unavailable')
plan =
native_table.new_postpone_fixed_bucket_write_builder().resolve_bucket_plan()
# A partition key may itself be called total_buckets. Keep the
# positional schema used by the Core plan and with_bucket_plan input.
- columns = [column.to_pylist() for column in plan.columns]
- return {tuple(column[row] for column in columns[:-1]): columns[-1][row]
- for row in range(plan.num_rows)}
+ return _plan_dict(plan)
diff --git a/paimon-python/pypaimon/ray/shuffle.py
b/paimon-python/pypaimon/ray/shuffle.py
index 4dfb63a379..71ea575adb 100644
--- a/paimon-python/pypaimon/ray/shuffle.py
+++ b/paimon-python/pypaimon/ray/shuffle.py
@@ -130,15 +130,15 @@ def maybe_apply_repartition(
def _group_by_partition_bucket(
dataset: "ray.data.Dataset",
table: "Table",
- extractor=None,
+ router=None,
):
partition_keys = list(table.table_schema.partition_keys or [])
- if extractor is None:
- from pypaimon.write.input_routing import FixedBucketInputRouter
- extractor = FixedBucketInputRouter(table.table_schema)
+ if router is None:
+ from pypaimon.write.backend_dispatch import create_input_router
+ router = create_input_router(table)
col_names = set(f.name for f in table.table_schema.fields)
bucket_col = _pick_bucket_col_name(col_names)
- bucket_udf = _make_bucket_udf(extractor, bucket_col)
+ bucket_udf = _make_bucket_udf(router, bucket_col)
ds_with_bucket = dataset.map_batches(
bucket_udf, batch_format="pyarrow", zero_copy_batch=True,
@@ -151,11 +151,10 @@ def _identity_batch(batch: pa.Table) -> pa.Table:
return batch
-def _make_bucket_udf(extractor, bucket_col):
+def _make_bucket_udf(router, bucket_col):
"""Build a map_batches UDF that appends a transient bucket column.
- The bucket value comes from ``extract_partition_bucket_batch`` so it
- matches the writer's bucket assignment for the same row exactly.
+ The prepared router uses the writer's bucket assignment for each row.
"""
def _udf(batch: pa.Table) -> pa.Table:
if batch.num_rows == 0:
@@ -163,9 +162,6 @@ def _make_bucket_udf(extractor, bucket_col):
bucket_col, pa.array([], type=pa.int32())
)
record_batch = batch.combine_chunks().to_batches()[0]
- _, buckets = extractor.extract_partition_bucket_batch(record_batch)
- return batch.append_column(
- bucket_col, pa.array(buckets, type=pa.int32())
- )
+ return batch.append_column(bucket_col, router.route(record_batch))
return _udf
diff --git a/paimon-python/pypaimon/table/row/generic_row.py
b/paimon-python/pypaimon/table/row/generic_row.py
index 07659ba503..31a7701edb 100644
--- a/paimon-python/pypaimon/table/row/generic_row.py
+++ b/paimon-python/pypaimon/table/row/generic_row.py
@@ -455,8 +455,11 @@ class GenericRowSerializer:
field_fixed_offset = null_bits_size_in_bytes + i * 8
if value is None:
- cls._set_null_bit(fixed_part, 0, i)
- struct.pack_into('<q', fixed_part, field_fixed_offset, 0)
+ null_data = cls._serialize_null_field(
+ fixed_part, field_fixed_offset, i, field.type,
+ fixed_part_size + current_variable_offset)
+ variable_part_data.append(null_data)
+ current_variable_offset += len(null_data)
continue
if not isinstance(field.type, AtomicType):
@@ -486,8 +489,11 @@ class GenericRowSerializer:
d = value if isinstance(value, Decimal) else
Decimal(str(value))
unscaled_value, overflow = _decimal_to_unscaled_with_check(d,
precision, scale)
if overflow:
- cls._set_null_bit(fixed_part, 0, i)
- struct.pack_into('<q', fixed_part, field_fixed_offset, 0)
+ null_data = cls._serialize_null_field(
+ fixed_part, field_fixed_offset, i, field.type,
+ fixed_part_size + current_variable_offset)
+ variable_part_data.append(null_data)
+ current_variable_offset += len(null_data)
continue
if is_non_compact_timestamp:
@@ -566,6 +572,23 @@ class GenericRowSerializer:
arity_prefix = struct.pack('>i', arity)
return arity_prefix + row_data
+ @classmethod
+ def _serialize_null_field(cls, fixed_part, field_offset, index, data_type,
variable_offset):
+ cls._set_null_bit(fixed_part, 0, index)
+ reserved_size = 0
+ if isinstance(data_type, AtomicType):
+ type_name = data_type.type.upper()
+ if type_name.startswith(('DECIMAL', 'NUMERIC')):
+ if _parse_type_precision_scale(data_type)[0] > 18:
+ reserved_size = 16
+ elif type_name.startswith('TIMESTAMP'):
+ if _parse_type_precision_scale(data_type)[0] > 3:
+ reserved_size = 8
+ # Java reserves updateable space even for NULL noncompact values and
+ # keeps its cursor in the fixed slot, with a zero length or nanos word.
+ struct.pack_into('<q', fixed_part, field_offset, variable_offset << 32
if reserved_size else 0)
+ return b'\x00' * reserved_size
+
@classmethod
def _calculate_bit_set_width_in_bytes(cls, arity: int) -> int:
return ((arity + 63 + cls.HEADER_SIZE_IN_BITS) // 64) * 8
diff --git a/paimon-python/pypaimon/tests/native_input_routing_test.py
b/paimon-python/pypaimon/tests/native_input_routing_test.py
new file mode 100644
index 0000000000..d0764b6409
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_input_routing_test.py
@@ -0,0 +1,192 @@
+# 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 input hashes and postpone statistics execute in the real SDK."""
+
+from decimal import Decimal
+from datetime import datetime, timezone
+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 NativeInputRouter
+from pypaimon.execution.native.write.postpone_bucket import
NativePostponeBucketPlanner
+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
+
+pytestmark = pytest.mark.native_plan
+
+
+def _table(catalog, schema, options, partition='p'):
+ table_options = {
+ 'file.format': 'parquet', 'write.native.enabled': 'true',
+ 'postpone.target-row-num-per-bucket': '2',
+ 'postpone.batch-write-fixed-bucket.max-parallelism': '8',
+ }
+ table_options.update(options)
+ table_options = {key: value for key, value in table_options.items() if
value is not None}
+ catalog.create_table('default.input_routing', Schema.from_pyarrow_schema(
+ schema, primary_keys=['id', partition], partition_keys=[partition],
options=table_options), False)
+ return catalog.get_table('default.input_routing')
+
+
[email protected]('bucket_function', ['default', 'mod'])
+def
test_native_ray_bucket_hashes_match_writer_for_reordered_subset(native_rest_catalog,
bucket_function):
+ schema = pa.schema([('id', pa.int64()), ('p', pa.string()), ('value',
pa.string())])
+ table = _table(native_rest_catalog, schema, {
+ 'bucket': '7', 'bucket-function.type': bucket_function})
+ batch = pa.record_batch([['west', 'east', 'west', 'east'], [-9, 0, 11, -(1
<< 63)]], names=['p', 'id'])
+ expected =
FixedBucketInputRouter(table.table_schema).extract_partition_bucket_batch(batch)[1]
+ with patch.object(FixedBucketInputRouter, '_extract_partitions_batch',
+ side_effect=AssertionError('Python partitions')), \
+ patch.object(FixedBucketInputRouter, '_binary_row_hash_code',
side_effect=AssertionError('Python hashes')):
+ router = create_input_router(table)
+ assert isinstance(router, NativeInputRouter)
+ restored = pickle.loads(pickle.dumps(router))
+ chunked = pa.concat_tables([pa.Table.from_batches([batch.slice(0, 2)]),
+ pa.Table.from_batches([batch.slice(2,
2)])])
+ result = _make_bucket_udf(restored, '__bucket__')(chunked)
+ assert result.column('__bucket__').to_pylist() == expected
+ assert result.column('__bucket__').type == pa.int32()
+ assert result.drop_columns(['__bucket__']).equals(chunked)
+ with pytest.raises(Exception, match='id|routing|column|field'):
+ restored.route(pa.record_batch([['west']], names=['p']))
+
+
+def
test_native_postpone_statistics_inference_and_routing_skip_python(native_rest_catalog):
+ schema = pa.schema([
+ ('id', pa.int64()), ('row_count', pa.string()), ('value', pa.string()),
+ ('items', pa.list_(pa.int64())), ('attrs', pa.map_(pa.string(),
pa.int64())),
+ ('nested', pa.struct([('label', pa.string()), ('count', pa.int64())])),
+ ('amount', pa.decimal128(24, 3)),
+ ])
+ table = _table(native_rest_catalog, schema, {
+ 'bucket': '-2', 'postpone.target-size-per-bucket': '64',
+ 'postpone.target-row-num-per-bucket': None}, partition='row_count')
+ rows = [dict(id=i, row_count='known' if i == 0 else 'west' if i < 4 else
'east',
+ value='abcdefghi' * i, items=[1, None, i], attrs=[('key', i)],
+ nested={'label': 'nested-value', 'count': i},
amount=Decimal('12345678901234567890.123'))
+ for i in range(6)]
+ batch = pa.RecordBatch.from_pylist(rows, schema=schema)
+ reference = PostponeBucketPlanning(table, {('known',): 3})
+ expected_stats = reference.input_partition_stats(batch)
+ expected = reference.plan(expected_stats)
+ with patch.object(PostponeBucketPlanning, 'input_partition_stats',
+ side_effect=AssertionError('Python row sizes')), \
+ patch.object(PostponeBucketPlanning, 'plan',
side_effect=AssertionError('Python bucket inference')):
+ planner = create_postpone_bucket_planner(table, {('known',): 3})
+ assert isinstance(planner, NativePostponeBucketPlanner)
+ restored = pickle.loads(pickle.dumps(planner))
+ stats =
restored.input_partition_stats(pa.Table.from_batches([batch.slice(0, 2),
batch.slice(2)]))
+ assert stats == expected_stats
+ plan = restored.plan(stats)
+ assert plan.as_dict() == expected.as_dict()
+ assert restored.current_plan().as_dict() == expected.as_dict()
+ assert restored.input_partition_stats(batch) == {}
+ reference_router = PostponeFixedBucketInputRouter(table, expected)
+ expected_buckets =
reference_router.extract_partition_bucket_batch(batch)[1]
+ with patch.object(PostponeFixedBucketInputRouter,
'extract_partition_bucket_batch',
+ side_effect=AssertionError('Python postpone hashes')):
+ router = create_input_router(table, plan)
+ assert isinstance(router, NativeInputRouter)
+ assert router.route(batch).to_pylist() == expected_buckets
+ # A worker reconstructs from the immutable plan without SDK pickling.
+ assert pickle.loads(pickle.dumps(router)).route(batch).to_pylist() ==
expected_buckets
+
+
+def
test_native_postpone_known_empty_skips_metadata_and_stats_failure_is_terminal(native_rest_catalog):
+ table = _table(native_rest_catalog, pa.schema([('id', pa.int64()), ('p',
pa.string())]), {'bucket': '-2'})
+ with patch.object(NativePostponeBucketPlanner, '_load_bucket_metadata',
side_effect=AssertionError('metadata')):
+ planner = create_postpone_bucket_planner(table, {})
+ assert isinstance(planner, NativePostponeBucketPlanner)
+ assert planner.current_plan().as_dict() == {}
+ assert planner.plan({('west',): (9, 0)}).as_dict() == {('west',): 8}
+ with patch.object(PostponeBucketPlanning, 'input_partition_stats',
side_effect=AssertionError('Python replay')):
+ with pytest.raises(Exception, match='p|partition|column|field'):
+ planner.input_partition_stats(pa.record_batch([[1]], names=['id']))
+
+
[email protected]('key_type,value,null_bucket', [
+ (pa.decimal128(19, 2), Decimal('123.40'), 1),
+ (pa.timestamp('us'), datetime(2024, 1, 2, 3, 4, 5, 123456), 28),
+ (pa.timestamp('us', tz='UTC'), datetime(2024, 1, 2, 3, 4, 5, 123456,
tzinfo=timezone.utc), 28),
+])
+def
test_native_nullable_noncompact_bucket_keys_match_java(native_rest_catalog,
key_type, value, null_bucket):
+ schema = pa.schema([('id', key_type), ('label', pa.string()), ('p',
pa.string())])
+ native_rest_catalog.create_table('default.null_bucket',
Schema.from_pyarrow_schema(
+ schema, primary_keys=['id', 'label', 'p'], partition_keys=['p'],
options={
+ 'bucket': '100', 'primary-key.nullable': 'true',
'write.native.enabled': 'true',
+ }), False)
+ table = native_rest_catalog.get_table('default.null_bucket')
+ batch = pa.RecordBatch.from_arrays([
+ pa.array([None, value, None], type=key_type),
+ pa.array(['abcdefghijk'] * 3), pa.array(['west'] * 3),
+ ], schema=schema)
+ expected =
FixedBucketInputRouter(table.table_schema).extract_partition_bucket_batch(batch)[1]
+ # Current Java hashes [NULL, "abcdefghijk"] as 1558048601 for DECIMAL
+ # and -684602128 for noncompact TIMESTAMP/LTZ; Java abs(hash % 100)
follows.
+ assert expected[0] == expected[2] == null_bucket
+ with patch.object(FixedBucketInputRouter, '_binary_row_hash_code',
side_effect=AssertionError('Python hashes')):
+ router = create_input_router(table)
+ assert isinstance(router, NativeInputRouter)
+ assert router.route(batch).to_pylist() == expected
+
+
+def
test_native_postpone_fixed_binary_partition_slices_match_python(native_rest_catalog):
+ schema = pa.schema([('id', pa.int64()), ('p', pa.binary(4)), ('value',
pa.string())])
+ table = _table(native_rest_catalog, schema, {
+ 'bucket': '-2', 'partition.legacy-name': 'false',
'primary-key.nullable': 'true',
+ 'postpone.target-row-num-per-bucket': None,
'postpone.target-size-per-bucket': '128',
+ })
+ assert table.field_dict['p'].type.type == 'BINARY(4)'
+ payloads = ['payload-' * 4] * 9
+ payloads[3] = None
+ source = pa.RecordBatch.from_arrays([
+ pa.array(range(9), type=pa.int64()),
+ pa.array([b'pad!', b'abcd', None, b'abcd', b'abcd', None, b'efgh',
b'efgh', b'end!'], type=pa.binary(4)),
+ pa.array(payloads),
+ ], schema=schema)
+ batch = source.slice(1, 7)
+ assert batch.column(1).offset == 1
+ chunked = pa.Table.from_batches([batch.slice(0, 2), batch.slice(2)])
+ reference = PostponeBucketPlanning(table, {})
+ expected_stats = reference.input_partition_stats(batch)
+ assert expected_stats == {(b'abcd',): (3, 160), (None,): (2, 128),
(b'efgh',): (2, 128)}
+ expected_plan = reference.plan(expected_stats)
+ assert expected_plan.as_dict() == {(b'abcd',): 2, (None,): 1, (b'efgh',):
1}
+ expected_buckets = PostponeFixedBucketInputRouter(table,
expected_plan).extract_partition_bucket_batch(batch)[1]
+
+ with patch.object(PostponeBucketPlanning, 'input_partition_stats',
+ side_effect=AssertionError('Python statistics')), \
+ patch.object(PostponeBucketPlanning, 'plan',
side_effect=AssertionError('Python bucket inference')):
+ planner = create_postpone_bucket_planner(table, {})
+ assert isinstance(planner, NativePostponeBucketPlanner)
+ assert planner.input_partition_stats(batch) == expected_stats
+ assert planner.input_partition_stats(chunked) == expected_stats
+ plan = planner.plan(expected_stats)
+ assert plan.as_dict() == expected_plan.as_dict()
+ assert planner.input_partition_stats(chunked) == {}
+ with patch.object(PostponeFixedBucketInputRouter,
'extract_partition_bucket_batch',
+ side_effect=AssertionError('Python bucket hashing')):
+ router = create_input_router(table, plan)
+ assert isinstance(router, NativeInputRouter)
+ assert router.route(batch).to_pylist() == expected_buckets
+ assert router.route(chunked).to_pylist() == expected_buckets
diff --git a/paimon-python/pypaimon/tests/noncompact_null_row_test.py
b/paimon-python/pypaimon/tests/noncompact_null_row_test.py
new file mode 100644
index 0000000000..24a0181dc7
--- /dev/null
+++ b/paimon-python/pypaimon/tests/noncompact_null_row_test.py
@@ -0,0 +1,81 @@
+# 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.
+
+"""NULL noncompact fields retain Java's updateable BinaryRow layout."""
+
+from decimal import Decimal
+import struct
+
+import pyarrow as pa
+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
+
+
+# Current Java InternalRowSerializer for [NULL, BinaryString("abcdefghijk")].
+# Java BinaryRow.hashCode() hashes these raw bytes without Python's arity
prefix.
+_COMPACT =
'000100000000000000000000000000000b000000180000006162636465666768696a6b0000000000'
+_DECIMAL = ('000100000000000000000000180000000b00000028000000'
+ '000000000000000000000000000000006162636465666768696a6b0000000000')
+_TIMESTAMP = ('000100000000000000000000180000000b00000020000000'
+ '00000000000000006162636465666768696a6b0000000000')
+
+
[email protected]('type_name,body,java_hash', [
+ ('DECIMAL(18, 2)', _COMPACT, 1425286886),
+ ('DECIMAL(19, 2)', _DECIMAL, 1558048601),
+ ('DECIMAL(20, 2)', _DECIMAL, 1558048601),
+ ('TIMESTAMP(3)', _COMPACT, 1425286886),
+ ('TIMESTAMP(4)', _TIMESTAMP, -684602128),
+ ('TIMESTAMP(9)', _TIMESTAMP, -684602128),
+ ('TIMESTAMP(3) WITH LOCAL TIME ZONE', _COMPACT, 1425286886),
+ ('TIMESTAMP(4) WITH LOCAL TIME ZONE', _TIMESTAMP, -684602128),
+ ('TIMESTAMP(9) WITH LOCAL TIME ZONE', _TIMESTAMP, -684602128),
+ ('TIMESTAMP_LTZ(6)', _TIMESTAMP, -684602128),
+])
+def
test_null_variable_field_bytes_hash_and_following_string_match_java(type_name,
body, java_hash):
+ fields = [DataField(0, 'key', AtomicType(type_name)), DataField(1,
'label', AtomicType('STRING'))]
+ values = [None, 'abcdefghijk']
+ encoded = GenericRowSerializer.to_bytes(GenericRow(values, fields))
+ assert encoded == struct.pack('>i', 2) + bytes.fromhex(body)
+ assert _hash_bytes_by_words(encoded[4:]) == java_hash & 0xFFFFFFFF
+ assert GenericRowDeserializer.from_bytes(encoded, fields).values == values
+
+ # Real Python input routing must retain both the NULL slot and the next
+ # variable field's offset when hashing a nullable composite bucket key.
+ partition = DataField(2, 'p', AtomicType('STRING'))
+ schema = TableSchema(fields=fields + [partition], primary_keys=['key',
'label', 'p'], partition_keys=['p'],
+ options={'bucket': '100', 'primary-key.nullable':
'true'})
+ arrow_schema = PyarrowFieldParser.from_paimon_schema(schema.fields)
+ batch = pa.RecordBatch.from_arrays([
+ pa.array([None], type=arrow_schema.field('key').type),
+ pa.array(['abcdefghijk']), pa.array(['west']),
+ ], schema=arrow_schema)
+ partitions, buckets =
FixedBucketInputRouter(schema).extract_partition_bucket_batch(batch)
+ assert partitions == [('west',)]
+ assert buckets == [abs(java_hash) % 100]
+
+
[email protected]('type_name', ['DECIMAL(19, 2)', 'NUMERIC(19, 2)'])
+def test_decimal_overflow_uses_the_same_noncompact_null_layout(type_name):
+ fields = [DataField(0, 'key', AtomicType(type_name)), DataField(1,
'label', AtomicType('STRING'))]
+ encoded =
GenericRowSerializer.to_bytes(GenericRow([Decimal('99999999999999999.999'),
'abcdefghijk'], fields))
+ assert encoded == struct.pack('>i', 2) + bytes.fromhex(_DECIMAL)
+ assert _hash_bytes_by_words(encoded[4:]) == 1558048601
+ assert GenericRowDeserializer.from_bytes(encoded, fields).values == [None,
'abcdefghijk']
diff --git a/paimon-python/pypaimon/tests/test_ray_shuffle_helper.py
b/paimon-python/pypaimon/tests/test_ray_shuffle_helper.py
index 7401eb1e3a..b6535b15f2 100644
--- a/paimon-python/pypaimon/tests/test_ray_shuffle_helper.py
+++ b/paimon-python/pypaimon/tests/test_ray_shuffle_helper.py
@@ -19,7 +19,7 @@
"""Unit tests for the Ray pre-shuffle helper in pypaimon/ray/shuffle.py.
These tests exercise the helper in isolation: the bucket-key UDF (with a
-stub extractor), the collision-safe column name picker, the
+stub router), the collision-safe column name picker, the
large-type coercion, and the bucket-mode dispatch in
``maybe_apply_repartition``. Ray-based end-to-end behaviour is covered
in ``pypaimon/tests/ray_repartition_test.py``.
@@ -53,17 +53,14 @@ maybe_apply_repartition = _SHUFFLE.maybe_apply_repartition
class BucketUdfTest(unittest.TestCase):
"""The bucket-key UDF appends a deterministic int32 column."""
- def _make_extractor(self, buckets_per_row):
- extractor = MagicMock()
- extractor.extract_partition_bucket_batch.return_value = (
- [() for _ in buckets_per_row],
- list(buckets_per_row),
- )
- return extractor
+ def _make_router(self, buckets_per_row):
+ router = MagicMock(spec=["route"])
+ router.route.return_value = pa.array(buckets_per_row, type=pa.int32())
+ return router
def test_appends_int32_bucket_column(self):
- extractor = self._make_extractor([0, 1, 0])
- udf = _make_bucket_udf(extractor, BUCKET_KEY_COL)
+ router = self._make_router([0, 1, 0])
+ udf = _make_bucket_udf(router, BUCKET_KEY_COL)
batch = pa.table({"id": [10, 11, 12]})
out = udf(batch)
@@ -73,24 +70,24 @@ class BucketUdfTest(unittest.TestCase):
self.assertEqual(out.column(BUCKET_KEY_COL).to_pylist(), [0, 1, 0])
def test_empty_batch_appends_empty_column(self):
- extractor = self._make_extractor([])
- udf = _make_bucket_udf(extractor, BUCKET_KEY_COL)
+ router = self._make_router([])
+ udf = _make_bucket_udf(router, BUCKET_KEY_COL)
batch = pa.table({"id": pa.array([], type=pa.int32())})
out = udf(batch)
self.assertEqual(out.num_rows, 0)
self.assertEqual(out.column_names, ["id", BUCKET_KEY_COL])
- # The extractor is short-circuited on empty input — we don't pay
+ # The router is short-circuited on empty input — we don't pay
# the cost of combining empty chunks just to call into it.
- extractor.extract_partition_bucket_batch.assert_not_called()
+ router.route.assert_not_called()
- def test_multichunk_batch_combines_before_extracting(self):
+ def test_multichunk_batch_combines_before_routing(self):
# Two record batches in the same table — the UDF must combine
- # before calling the extractor, otherwise the extractor sees
+ # before calling the router, otherwise the router sees
# half the rows.
- extractor = self._make_extractor([0, 1, 2, 3])
- udf = _make_bucket_udf(extractor, BUCKET_KEY_COL)
+ router = self._make_router([0, 1, 2, 3])
+ udf = _make_bucket_udf(router, BUCKET_KEY_COL)
rb1 = pa.record_batch({"id": [1, 2]})
rb2 = pa.record_batch({"id": [3, 4]})
batch = pa.Table.from_batches([rb1, rb2])
@@ -99,8 +96,9 @@ class BucketUdfTest(unittest.TestCase):
self.assertEqual(out.num_rows, 4)
self.assertEqual(out.column(BUCKET_KEY_COL).to_pylist(), [0, 1, 2, 3])
- # Extractor is called exactly once with all four rows.
- call = extractor.extract_partition_bucket_batch.call_args
+ # Router is called exactly once with all four rows.
+ router.route.assert_called_once()
+ call = router.route.call_args
passed_batch = call.args[0]
self.assertEqual(passed_batch.num_rows, 4)
diff --git a/paimon-python/pypaimon/tests/write/native_input_routing_test.py
b/paimon-python/pypaimon/tests/write/native_input_routing_test.py
new file mode 100644
index 0000000000..da3fc398ea
--- /dev/null
+++ b/paimon-python/pypaimon/tests/write/native_input_routing_test.py
@@ -0,0 +1,180 @@
+# 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.
+
+"""Native input execution is chosen before data and remains native on
workers."""
+
+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.execution.native.write.input_routing import NativeInputRouter
+from pypaimon.execution.native.write.postpone_bucket import
NativePostponeBucketPlanner
+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
+
+
+def _table(tmp_path, bucket='4', partition='p', 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()), (partition, pa.string()), ('value',
pa.string())]),
+ primary_keys=['id', partition], partition_keys=[partition], options={
+ 'bucket': bucket, 'write.native.enabled': str(native).lower(),
+ 'postpone.target-row-num-per-bucket': '2',
+ 'postpone.batch-write-fixed-bucket.max-parallelism': '8',
+ }), False)
+ return catalog.get_table('default.t')
+
+
[email protected]
+def sdk(monkeypatch):
+ module = SimpleNamespace(BucketRouter=Mock(), PostponeBucketPlanner=Mock())
+ monkeypatch.setitem(sys.modules, 'pypaimon_rust.datafusion', module)
+ return module
+
+
[email protected]('native', [False, True])
+def test_bucket_router_selection_finishes_before_input(tmp_path, native):
+ table = _table(tmp_path, native=native)
+ with
patch('pypaimon.execution.native.write.input_routing.create_native_input_router',
+ side_effect=RuntimeError('unsupported table environment')) as
prepare:
+ router = create_input_router(table)
+ assert isinstance(router, FixedBucketInputRouter)
+ assert prepare.call_count == int(native)
+
+
[email protected]('bucket', ['4', '-2'])
[email protected]('chunked', [False, True])
+def
test_python_router_returns_arrow_buckets_for_the_shared_route_interface(tmp_path,
bucket, chunked):
+ table = _table(tmp_path, bucket=bucket, native=False)
+ plan = PostponeBucketPlan({('west',): 3, ('east',): 2}) if bucket == '-2'
else None
+ router = create_input_router(table, plan)
+ batch = pa.record_batch([['west', 'east', 'west'], [1, 2, 3]], names=['p',
'id'])
+ expected = router.extract_partition_bucket_batch(batch)[1]
+ data = pa.Table.from_batches([batch.slice(0, 1), batch.slice(1)]) if
chunked else batch
+ buckets = router.route(data)
+ assert isinstance(buckets, pa.Int32Array)
+ assert buckets.to_pylist() == expected
+ assert buckets.null_count == 0
+ empty = router.route(data.slice(0, 0))
+ assert isinstance(empty, pa.Int32Array)
+ assert len(empty) == 0
+
+
+def test_native_router_passes_arrow_without_python_row_hashing(tmp_path, sdk):
+ table = _table(tmp_path)
+ batch = pa.record_batch([['west', 'east'], [1, 2]], names=['p', 'id'])
+ buckets = pa.array([2, 1], type=pa.int32())
+ sdk.BucketRouter.return_value.route.return_value = buckets
+ native_table = Mock()
+ 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(FixedBucketInputRouter, '_binary_row_hash_code',
+ side_effect=AssertionError('Python input hashing')):
+ router = create_input_router(table)
+ assert isinstance(router, NativeInputRouter)
+ assert router.route(batch) is buckets
+ sdk.BucketRouter.assert_called_once_with(native_table.schema.return_value)
+ sdk.BucketRouter.return_value.route.assert_called_once_with(batch)
+
+
+def test_native_router_failure_propagates_without_selecting_python(tmp_path,
sdk):
+ table = _table(tmp_path)
+ router = NativeInputRouter(table)
+ failure = ValueError('Input must contain a routing column')
+ router._router = Mock()
+ router._router.route.side_effect = failure
+ with patch.object(FixedBucketInputRouter, '__init__',
side_effect=AssertionError('Python replay')):
+ with pytest.raises(ValueError) as raised:
+ router.route(pa.record_batch([[1]], names=['value']))
+ assert raised.value is failure
+
+
+def test_native_router_pickle_discards_sdk_and_restores_bucket_plan(tmp_path,
sdk):
+ table = _table(tmp_path, bucket='-2')
+ plan = PostponeBucketPlan({('west',): 3})
+ router = NativeInputRouter(table, plan)
+ # An SDK object is deliberately not serializable.
+ router._router = lambda: None
+ restored = pickle.loads(pickle.dumps(router))
+ assert restored._router is None
+ native = sdk.BucketRouter.return_value
+ native.with_bucket_plan.return_value = native
+ native.route.return_value = pa.array([1], type=pa.int32())
+ with
patch('pypaimon.execution.native.write.input_routing.create_native_write_table',
return_value=Mock()):
+ assert restored.route(pa.record_batch([[7], ['west']], names=['id',
'p'])).to_pylist() == [1]
+ assert native.with_bucket_plan.call_args[0][0].equals(plan.to_arrow(table))
+
+
+def test_native_planner_transports_group_statistics_positionally(tmp_path,
sdk):
+ table = _table(tmp_path, bucket='-2', partition='row_count')
+ native_table = Mock()
+ core = sdk.PostponeBucketPlanner.return_value
+ core.input_statistics.return_value = pa.record_batch(
+ [['west'], pa.array([3], type=pa.int64()), pa.array([80],
type=pa.int64())],
+ names=['row_count', 'row_count', 'data_size'])
+ expected = PostponeBucketPlan({('west',): 2})
+ core.plan.return_value = expected.to_arrow(table)
+ with
patch('pypaimon.execution.native.write.postpone_bucket.native_runtime_available',
return_value=True), \
+
patch('pypaimon.execution.native.write.postpone_bucket.create_native_write_table',
+ return_value=native_table), \
+ patch.object(PostponeBucketPlanning, 'input_partition_stats',
+ side_effect=AssertionError('Python input
statistics')), \
+ patch.object(PostponeBucketPlanning, 'plan',
side_effect=AssertionError('Python bucket inference')):
+ planner = create_postpone_bucket_planner(table, known_num_buckets={})
+ assert isinstance(planner, NativePostponeBucketPlanner)
+ assert not isinstance(planner, PostponeBucketPlanning)
+ batch = pa.table({'id': [1, 2, 3], 'row_count': ['west'] * 3, 'value':
['data'] * 3})
+ stats = planner.input_partition_stats(batch)
+ assert stats == {('west',): (3, 80)}
+ assert planner.plan(stats).as_dict() == expected.as_dict()
+ native_table.new_postpone_fixed_bucket_write_builder.assert_not_called()
+ transported = core.input_statistics.call_args[0][0]
+ assert isinstance(transported, pa.RecordBatch) and transported.num_rows ==
3
+ grouped = core.plan.call_args[0][0]
+ assert grouped.schema.names == ['row_count', 'row_count', 'data_size']
+ assert [column.to_pylist() for column in grouped.columns] == [['west'],
[3], [80]]
+ assert grouped.column(1).type == pa.int64() and grouped.column(2).type ==
pa.int64()
+
+
+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
+ expected = PostponeBucketPlan({('west',): 7})
+ core.plan.return_value = expected.to_arrow(table)
+ planner = NativePostponeBucketPlanner(table, {}, Mock())
+ planner.plan({('west',): (13, 0)})
+ planner._planner = lambda: None
+ restored = pickle.loads(pickle.dumps(planner))
+ assert restored._planner is None
+ sdk.PostponeBucketPlanner.reset_mock()
+ failure = RuntimeError('worker statistics failed')
+ core.input_statistics.side_effect = failure
+ with
patch('pypaimon.execution.native.write.postpone_bucket.create_native_write_table',
return_value=Mock()), \
+ patch.object(PostponeBucketPlanning, 'input_partition_stats',
side_effect=AssertionError('Python replay')):
+ with pytest.raises(RuntimeError) as raised:
+ restored.input_partition_stats(pa.table({'id': [1], 'p': ['west'],
'value': ['x']}))
+ assert raised.value is failure
+ assert
sdk.PostponeBucketPlanner.call_args[0][1].equals(expected.to_arrow(table))
diff --git
a/paimon-python/pypaimon/tests/write/postpone_execution_boundary_test.py
b/paimon-python/pypaimon/tests/write/postpone_execution_boundary_test.py
index 8453fecbd5..447658a0df 100644
--- a/paimon-python/pypaimon/tests/write/postpone_execution_boundary_test.py
+++ b/paimon-python/pypaimon/tests/write/postpone_execution_boundary_test.py
@@ -20,6 +20,7 @@
from pathlib import Path
import subprocess
import sys
+from types import SimpleNamespace
from unittest.mock import Mock, patch
import pyarrow as pa
@@ -102,6 +103,26 @@ def _native_table(table):
return native
[email protected](autouse=True)
+def native_input_sdk(monkeypatch):
+ """These boundary tests supply SDK results without a native
installation."""
+ def planner(_schema, known):
+ core = Mock()
+ core.current_plan.return_value = known
+ core.plan.return_value = known
+ return core
+
+ def router(_schema):
+ core = Mock()
+ core.with_bucket_plan.return_value = core
+ core.route.side_effect = lambda batch: pa.array([0] * batch.num_rows,
type=pa.int32())
+ return core
+
+ sdk = SimpleNamespace(PostponeBucketPlanner=Mock(side_effect=planner),
BucketRouter=Mock(side_effect=router))
+ monkeypatch.setitem(sys.modules, 'pypaimon_rust.datafusion', sdk)
+ return sdk
+
+
@pytest.fixture
def forbid_python_execution():
with
patch('pypaimon.execution.python.write.postpone_bucket.PythonPostponeBucketPlanner.__init__',
@@ -116,32 +137,37 @@ def forbid_python_execution():
@pytest.mark.parametrize('overwrite', [False, True])
-def test_native_ray_planning_uses_native_metadata_and_shared_input_routing(
+def test_native_ray_planning_uses_native_metadata_and_input_routing(
tmp_path, forbid_python_execution, overwrite):
pytest.importorskip('ray')
table = _table(tmp_path)
core = Mock()
writer = NativePostponeFixedBucketTableWrite(table, 'user', None, None,
core)
dataset = Mock()
+ expected = {('existing',): 1 if overwrite else 3, ('new',): 8,
('pending',): 1}
+ native_table = _native_table(table)
with
patch('pypaimon.execution.native.write.postpone_bucket.native_runtime_available',
return_value=True), \
patch('pypaimon.execution.native.write.postpone_bucket.create_native_write_table',
- return_value=_native_table(table)) as plan, \
+ return_value=native_table), \
+
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.native.write.postpone_bucket.NativePostponeBucketPlanner.plan',
+ return_value=PostponeBucketPlan(expected)), \
patch('pypaimon.write.ray_datasink._collect_partition_stats',
return_value=(
dataset, {('existing',): (2, 0), ('new',): (9, 0),
('pending',): (1, 0)})), \
patch('pypaimon.write.ray_datasink._write_primary_key_groups') as
write:
writer.write_ray(dataset, overwrite=overwrite)
args = write.call_args[1]
- expected = {('existing',): 1 if overwrite else 3, ('new',): 8,
('pending',): 1}
assert args['postpone_bucket_plan'].as_dict() == expected
batch = pa.record_batch([[1, 2], ['existing', 'new'], [10, 20]],
names=['id', 'p', 'value'])
- partitions, buckets =
args['bucket_extractor'].extract_partition_bucket_batch(batch)
- assert partitions == [('existing',), ('new',)]
+ buckets = args['bucket_router'].route(batch).to_pylist()
assert 0 <= buckets[0] < expected[('existing',)]
assert 0 <= buckets[1] < 8
if overwrite:
- plan.assert_not_called()
+
native_table.new_postpone_fixed_bucket_write_builder.assert_not_called()
else:
- plan.assert_called_once_with(table)
+
native_table.new_postpone_fixed_bucket_write_builder.return_value.resolve_bucket_plan.assert_called_once()
assert writer._python_writer is None
core.write_arrow.assert_not_called()
writer.close()
@@ -153,13 +179,15 @@ def
test_native_ray_explicit_plan_does_not_load_metadata(tmp_path, forbid_python
explicit = PostponeBucketPlan({('existing',): 7})
writer = NativePostponeFixedBucketTableWrite(table, 'user', None,
explicit, Mock())
dataset = Mock()
+ native_table = _native_table(table)
with
patch('pypaimon.execution.native.write.postpone_bucket.native_runtime_available',
return_value=True), \
patch('pypaimon.execution.native.write.postpone_bucket.create_native_write_table',
- side_effect=AssertionError('explicit bucket plan must not
read metadata')), \
+ return_value=native_table), \
patch('pypaimon.write.ray_datasink._collect_partition_stats',
return_value=(dataset, {})), \
patch('pypaimon.write.ray_datasink._write_primary_key_groups') as
write:
writer.write_ray(dataset)
assert write.call_args[1]['postpone_bucket_plan'].as_dict() ==
explicit.as_dict()
+
native_table.new_postpone_fixed_bucket_write_builder.assert_not_called()
writer.close()
diff --git a/paimon-python/pypaimon/write/backend_dispatch.py
b/paimon-python/pypaimon/write/backend_dispatch.py
index 5961d946a2..f62e553fc2 100644
--- a/paimon-python/pypaimon/write/backend_dispatch.py
+++ b/paimon-python/pypaimon/write/backend_dispatch.py
@@ -204,6 +204,22 @@ def create_postpone_bucket_planner(table,
known_num_buckets=None):
return PythonPostponeBucketPlanner(table,
known_num_buckets=known_num_buckets)
+def create_input_router(table, bucket_plan=None):
+ """Select bucket routing before a distributed task receives input rows."""
+ if table.options.native_write_enabled():
+ def prepare():
+ from pypaimon.execution.native.write.input_routing import
create_native_input_router
+ return create_native_input_router(table, bucket_plan)
+ native = try_prepare_native(prepare)
+ if native is not None:
+ return native
+ if bucket_plan is not None:
+ from pypaimon.write.input_routing import PostponeFixedBucketInputRouter
+ return PostponeFixedBucketInputRouter(table, bucket_plan)
+ from pypaimon.write.input_routing import FixedBucketInputRouter
+ return FixedBucketInputRouter(table.table_schema)
+
+
def load_postpone_bucket_metadata(table):
"""Select metadata planning before public bucket arithmetic consumes
input."""
if table.options.native_write_enabled():
diff --git a/paimon-python/pypaimon/write/input_routing.py
b/paimon-python/pypaimon/write/input_routing.py
index b8e700e840..d2d0496902 100644
--- a/paimon-python/pypaimon/write/input_routing.py
+++ b/paimon-python/pypaimon/write/input_routing.py
@@ -167,6 +167,10 @@ class
PostponeFixedBucketInputRouter(PartitionKeyExtractor):
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]]:
@@ -229,6 +233,10 @@ class FixedBucketInputRouter(PartitionKeyExtractor):
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)
diff --git a/paimon-python/pypaimon/write/ray_datasink.py
b/paimon-python/pypaimon/write/ray_datasink.py
index 4005798649..0da5eeb383 100644
--- a/paimon-python/pypaimon/write/ray_datasink.py
+++ b/paimon-python/pypaimon/write/ray_datasink.py
@@ -326,8 +326,7 @@ def write_paimon_dataset(
)
)
):
- from pypaimon.write.backend_dispatch import
create_postpone_bucket_planner
- from pypaimon.write.input_routing import PostponeFixedBucketInputRouter
+ from pypaimon.write.backend_dispatch import create_input_router,
create_postpone_bucket_planner
planner = (
postpone_bucket_planner
@@ -351,7 +350,7 @@ def write_paimon_dataset(
static_partition=static_partition,
concurrency=concurrency,
ray_remote_args=ray_remote_args,
- bucket_extractor=PostponeFixedBucketInputRouter(table, plan),
+ bucket_router=create_input_router(table, plan),
postpone_bucket_plan=plan,
)
@@ -489,7 +488,7 @@ def _write_primary_key_groups(
static_partition: Optional[Dict[str, Any]],
concurrency: Optional[int],
ray_remote_args: Optional[Dict[str, Any]],
- bucket_extractor=None,
+ bucket_router=None,
postpone_bucket_plan=None,
) -> Optional[PaimonWriteResult]:
import pickle
@@ -499,7 +498,7 @@ def _write_primary_key_groups(
)
grouped, bucket_col = _group_by_partition_bucket(
- dataset, table, extractor=bucket_extractor
+ dataset, table, router=bucket_router
)
message_col = "__paimon_commit_messages__"
error_col = "__paimon_write_error__"
