[Python] Execute Ray bucket routing and postpone planning in Rust

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"`.

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

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)

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,21 +151,17 @@ 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:
             return batch.append_column(
                 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

paimon-python/pypaimon/table/row/generic_row.py

@@ -455,8 +455,11 @@ def to_bytes(cls, row: Union[GenericRow, BinaryRow]) -> bytes:
             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 @@ def to_bytes(cls, row: Union[GenericRow, BinaryRow]) -> bytes:
                 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 @@ def to_bytes(cls, row: Union[GenericRow, BinaryRow]) -> bytes:
         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