| # 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. | |
| from typing import Any, Dict, List, Optional | |
| import pyarrow as pa | |
| from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER | |
| from pypaimon.table.bucket_mode import BucketMode | |
| from pypaimon.write.commit_message import CommitMessage | |
| from pypaimon.write.file_store_write import PostponeFixedBucketFileStoreWrite | |
| from pypaimon.write.postpone_bucket import PostponeBucketPlanner | |
| from pypaimon.write.row_key_extractor import PostponeFixedBucketRowKeyExtractor | |
| from pypaimon.write.row_utils import ( | |
| require_columns, | |
| row_to_named_values, | |
| row_values_to_arrow_table, | |
| ) | |
| from pypaimon.write.table_write import BatchTableWrite | |
| from pypaimon.write.write_builder import BatchWriteBuilder | |
| class PostponeFixedBucketWriteBuilder(BatchWriteBuilder): | |
| """Write builder for fixed-bucket batches on a postpone table.""" | |
| def __init__(self, table): | |
| if table.bucket_mode() != BucketMode.POSTPONE_MODE: | |
| raise ValueError( | |
| "Postpone fixed-bucket write requires a postpone-bucket table" | |
| ) | |
| super().__init__(table) | |
| self._bucket_plan = None | |
| def with_bucket_plan(self, bucket_plan): | |
| """Use a precomputed partition bucket plan.""" | |
| self._bucket_plan = bucket_plan | |
| return self | |
| def new_write(self): | |
| # Automatic size estimation still belongs to the Python coordinator. | |
| # Native direct writes accept a resolved plan or Java's explicit default. | |
| if (self._bucket_plan is not None | |
| or self.table.options.postpone_default_bucket_num() is not None): | |
| native = self._native_write( | |
| self.static_partition, fixed_bucket=True, bucket_plan=self._bucket_plan) | |
| if native is not None: | |
| return native | |
| return PostponeFixedBucketBatchTableWrite( | |
| self.table, | |
| self.commit_user, | |
| self.static_partition, | |
| self._bucket_plan, | |
| ) | |
| class PostponeFixedBucketBatchTableWrite(BatchTableWrite): | |
| """Batch writer which plans postpone rows before routing new partitions.""" | |
| def __init__( | |
| self, | |
| table, | |
| commit_user, | |
| static_partition: Optional[dict] = None, | |
| bucket_plan=None, | |
| ): | |
| self._planner = PostponeBucketPlanner( | |
| table, | |
| known_num_buckets=( | |
| bucket_plan.as_dict() if bucket_plan is not None | |
| else {} if static_partition is not None else None | |
| ), | |
| ) | |
| self._bucket_plan = ( | |
| bucket_plan | |
| if bucket_plan is not None | |
| else self._planner.current_plan() | |
| ) | |
| self._plan_provided = bucket_plan is not None | |
| self._pending_inputs = [] | |
| super().__init__(table, commit_user, static_partition) | |
| def _create_file_store_write(self, commit_user): | |
| return PostponeFixedBucketFileStoreWrite(self.table, commit_user) | |
| def _create_row_key_extractor(self, static_partition): | |
| return PostponeFixedBucketRowKeyExtractor( | |
| self.table, self._bucket_plan) | |
| def _write_partition_bucket_batch(self, partition, bucket, data): | |
| self.file_store_write.write( | |
| partition, | |
| bucket, | |
| data, | |
| self.row_key_extractor.num_buckets(partition), | |
| ) | |
| def _write_partition_bucket_row( | |
| self, partition, bucket, row, values_by_name | |
| ): | |
| self.file_store_write.write_row( | |
| partition, | |
| bucket, | |
| row, | |
| values_by_name, | |
| self.row_key_extractor.num_buckets(partition), | |
| ) | |
| def _buffer_input(self, data) -> bool: | |
| if self._plan_provided: | |
| return False | |
| return any( | |
| not self._bucket_plan.contains(partition) | |
| for partition in self._planner.input_partition_stats(data) | |
| ) | |
| def write_arrow(self, table: pa.Table, bucket=None): | |
| if bucket is not None: | |
| raise ValueError('Precomputed bucket writes require HASH_FIXED or HASH_DYNAMIC tables') | |
| table = self._prepare_arrow_data(table) | |
| if not self._buffer_input(table): | |
| return super().write_arrow(table) | |
| self._pending_inputs.extend( | |
| ("batch", batch) for batch in table.to_batches()) | |
| def write_arrow_batch(self, data: pa.RecordBatch, bucket=None): | |
| if bucket is not None: | |
| raise ValueError('Precomputed bucket writes require HASH_FIXED or HASH_DYNAMIC tables') | |
| data = self._prepare_arrow_data(data) | |
| if not self._buffer_input(data): | |
| return super().write_arrow_batch(data) | |
| self._pending_inputs.append(("batch", data)) | |
| def write_row(self, row): | |
| if self._plan_provided: | |
| return super().write_row(row) | |
| values_by_name = row_to_named_values( | |
| row, self.table.table_schema.fields) | |
| column_names = ( | |
| self.file_store_write.write_cols | |
| if self.file_store_write.write_cols is not None | |
| else list(self.table.field_names) | |
| ) | |
| require_columns(values_by_name, column_names, "write_row") | |
| require_columns(values_by_name, self.table.partition_keys, "write_row") | |
| partition = tuple( | |
| values_by_name[key] for key in self.table.partition_keys) | |
| if self._bucket_plan.contains(partition): | |
| return super().write_row(row) | |
| arrow_row = row_values_to_arrow_table( | |
| values_by_name, self.table.table_schema.fields, column_names) | |
| size = self._planner.input_partition_stats(arrow_row)[partition][1] | |
| self._pending_inputs.append( | |
| ("row", (row, partition, size))) | |
| def _flush_pending_inputs(self): | |
| if not self._pending_inputs: | |
| return | |
| partition_stats = {} | |
| for input_type, value in self._pending_inputs: | |
| if input_type == "batch": | |
| stats_by_partition = ( | |
| self._planner.input_partition_stats(value)) | |
| else: | |
| _, partition, size = value | |
| stats_by_partition = {partition: (1, size)} | |
| for partition, stats in stats_by_partition.items(): | |
| rows, size = partition_stats.get(partition, (0, 0)) | |
| partition_stats[partition] = ( | |
| rows + stats[0], size + stats[1]) | |
| self._bucket_plan = self._planner.plan( | |
| partition_stats, | |
| include_postpone_rows=self.static_partition is None, | |
| ) | |
| self.row_key_extractor.with_bucket_plan(self._bucket_plan) | |
| inputs = self._pending_inputs | |
| self._pending_inputs = [] | |
| for input_type, value in inputs: | |
| if input_type == "batch": | |
| super().write_arrow_batch(value) | |
| else: | |
| super().write_row(value[0]) | |
| def prepare_commit(self) -> List[CommitMessage]: | |
| if self.batch_committed: | |
| raise RuntimeError( | |
| "BatchTableWrite only supports one-time committing.") | |
| self.batch_committed = True | |
| self._flush_pending_inputs() | |
| return self._prepare_commit(BATCH_COMMIT_IDENTIFIER) | |
| def _distributed_write_options(self) -> Dict[str, Any]: | |
| return {"postpone_bucket_planner": self._planner} | |
| def close(self): | |
| self._pending_inputs = [] | |
| super().close() | |
| def abort(self): | |
| self._pending_inputs = [] | |
| super().abort() |