| # 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 input_batches(data): | |
| if isinstance(data, pa.RecordBatch): | |
| return [data] | |
| if isinstance(data, pa.Table): | |
| # Binary columns can retain multiple chunks even after combine_chunks. | |
| # Preserve every batch and its row count, including zero-column input. | |
| batches = data.to_batches() | |
| return batches if batches else [pa.RecordBatch.from_arrays( | |
| [column.combine_chunks() for column in data.columns], schema=data.schema)] | |
| raise TypeError('Input routing requires a pyarrow Table or RecordBatch') | |
| 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() | |
| buckets = [self._router.route(batch) for batch in input_batches(data)] | |
| return buckets[0] if len(buckets) == 1 else pa.concat_arrays(buckets) | |
| def __getstate__(self): | |
| state = dict(self.__dict__) | |
| state['_router'] = None | |
| return state | |
| def create_native_dynamic_bucket_assigner_router(table, num_channels, num_assigners): | |
| if not native_runtime_available(): | |
| return None | |
| native_table = create_native_write_table(table) | |
| if native_table is None: | |
| return None | |
| router = NativeDynamicBucketAssignerRouter(table, num_channels, num_assigners) | |
| router._prepare(native_table) | |
| return router | |
| class NativeDynamicBucketAssignerRouter: | |
| """Route input rows to dynamic bucket assigners using Rust's row hashes.""" | |
| def __init__(self, table, num_channels, num_assigners): | |
| self.table = table | |
| self._num_channels = num_channels | |
| self._num_assigners = num_assigners | |
| self._router = None | |
| def _prepare(self, native_table=None): | |
| from pypaimon_rust.datafusion import DynamicBucketAssignerRouter | |
| if native_table is None: | |
| native_table = create_native_write_table(self.table) | |
| if native_table is None: | |
| raise RuntimeError('Native dynamic bucket assigner routing is unavailable') | |
| self._router = DynamicBucketAssignerRouter(native_table.schema(), self._num_channels, self._num_assigners) | |
| def route(self, data) -> pa.RecordBatch: | |
| if self._router is None: | |
| self._prepare() | |
| routed = [self._router.route(batch) for batch in input_batches(data)] | |
| if len(routed) == 1: | |
| return routed[0] | |
| return pa.RecordBatch.from_arrays([ | |
| pa.concat_arrays([batch.column(index) for batch in routed]) for index in range(2) | |
| ], schema=routed[0].schema) | |
| def __getstate__(self): | |
| state = dict(self.__dict__) | |
| state['_router'] = None | |
| return state | |
| def create_native_dynamic_bucket_assigner(table, num_channels, num_assigners, assign_id, | |
| ignore_existing=False, base_snapshot_id=None): | |
| if not native_runtime_available(): | |
| return None | |
| native_table = create_native_write_table(table) | |
| if native_table is None: | |
| return None | |
| assigner = NativeDynamicBucketAssigner( | |
| table, num_channels, num_assigners, assign_id, ignore_existing, base_snapshot_id) | |
| assigner._prepare(native_table) | |
| return assigner | |
| class NativeDynamicBucketAssigner: | |
| def __init__(self, table, num_channels, num_assigners, assign_id, | |
| ignore_existing=False, base_snapshot_id=None): | |
| self.table = table | |
| self._num_channels = num_channels | |
| self._num_assigners = num_assigners | |
| self._assign_id = assign_id | |
| self._ignore_existing = ignore_existing | |
| self._base_snapshot_id = base_snapshot_id | |
| self._assigner = None | |
| def _prepare(self, native_table=None): | |
| from pypaimon_rust.datafusion import DynamicBucketAssigner | |
| if native_table is None: | |
| native_table = create_native_write_table(self.table) | |
| if native_table is None: | |
| raise RuntimeError('Native dynamic bucket assignment is unavailable') | |
| self._assigner = DynamicBucketAssigner( | |
| native_table, self._num_channels, self._num_assigners, self._assign_id, | |
| restore_snapshot_id=self._base_snapshot_id, overwrite=self._ignore_existing) | |
| # Core captures the restore point once. Final worker writers must | |
| # restore that same snapshot even if another commit finishes later. | |
| self._base_snapshot_id = self._assigner.restore_snapshot_id | |
| @property | |
| def restore_snapshot_id(self): | |
| if self._assigner is None: | |
| self._prepare() | |
| return self._base_snapshot_id | |
| def assign(self, data, key_hashes=None) -> pa.Int32Array: | |
| batches = input_batches(data) | |
| if isinstance(key_hashes, pa.ChunkedArray): | |
| key_hashes = key_hashes.combine_chunks() | |
| if key_hashes is not None: | |
| if not isinstance(key_hashes, pa.Int32Array): | |
| raise TypeError('Dynamic bucket key hashes must be an Int32Array') | |
| # Validate the whole input before the stateful assigner sees any | |
| # batch, including a null or missing hash in a later chunk. | |
| if len(key_hashes) != data.num_rows or key_hashes.null_count: | |
| raise ValueError('Dynamic bucket key hashes must be non-null and match the input row count') | |
| if self._assigner is None: | |
| self._prepare() | |
| buckets = [] | |
| offset = 0 | |
| for batch in batches: | |
| hashes = None if key_hashes is None else key_hashes.slice(offset, batch.num_rows) | |
| buckets.append(self._assigner.assign(batch, key_hashes=hashes)) | |
| offset += batch.num_rows | |
| return buckets[0] if len(buckets) == 1 else pa.concat_arrays(buckets) |