diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md index 5566ee9e3672..135f83f65409 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 000000000000..461cba52f1d0 --- /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 2bb38944b5f5..acb3823aa452 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 4dfb63a3796b..71ea575adb1d 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 _udf(batch: pa.Table) -> pa.Table: 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 07659ba503f2..31a7701edb9b 100644 --- a/paimon-python/pypaimon/table/row/generic_row.py +++ b/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(' 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(' 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(' 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 000000000000..d0764b640917 --- /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') + + +@pytest.mark.parametrize('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'])) + + +@pytest.mark.parametrize('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 000000000000..24a0181dc725 --- /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') + + +@pytest.mark.parametrize('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] + + +@pytest.mark.parametrize('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 7401eb1e3aad..b6535b15f247 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 @@ 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 @@ def test_appends_int32_bucket_column(self): 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 @@ def test_multichunk_batch_combines_before_extracting(self): 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 000000000000..da3fc398eaf5 --- /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') + + +@pytest.fixture +def sdk(monkeypatch): + module = SimpleNamespace(BucketRouter=Mock(), PostponeBucketPlanner=Mock()) + monkeypatch.setitem(sys.modules, 'pypaimon_rust.datafusion', module) + return module + + +@pytest.mark.parametrize('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) + + +@pytest.mark.parametrize('bucket', ['4', '-2']) +@pytest.mark.parametrize('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 8453fecbd5d4..447658a0dfde 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 +@pytest.fixture(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 5961d946a2fd..f62e553fc2b7 100644 --- a/paimon-python/pypaimon/write/backend_dispatch.py +++ b/paimon-python/pypaimon/write/backend_dispatch.py @@ -204,6 +204,22 @@ def prepare(): 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 b8e700e84072..d2d049690252 100644 --- a/paimon-python/pypaimon/write/input_routing.py +++ b/paimon-python/pypaimon/write/input_routing.py @@ -167,6 +167,10 @@ def with_bucket_plan(self, bucket_plan) -> None: 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 @@ def __init__(self, table_schema: TableSchema): 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 400579864944..0da5eeb3830d 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__"