| # 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. |
| |
| import logging |
| import random |
| from typing import Dict, List, Tuple |
| |
| import pyarrow as pa |
| |
| |
| logger = logging.getLogger(__name__) |
| |
| from pypaimon.common.options.core_options import CoreOptions |
| from pypaimon.schema.data_types import is_blob_file_field |
| from pypaimon.write.commit_message import CommitMessage |
| from pypaimon.write.row_utils import row_values_to_arrow_table |
| from pypaimon.write.writer.append_only_data_writer import AppendOnlyDataWriter |
| from pypaimon.write.writer.dedicated_format_writer import DedicatedFormatWriter |
| from pypaimon.write.writer.data_vector_writer import DataVectorWriter |
| from pypaimon.write.writer.data_writer import DataWriter |
| from pypaimon.write.writer.key_value_data_writer import KeyValueDataWriter |
| from pypaimon.table.bucket_mode import BucketMode |
| |
| |
| class FileStoreWrite: |
| """Base class for file store write operations.""" |
| |
| def __init__(self, table, commit_user): |
| from pypaimon.table.file_store_table import FileStoreTable |
| |
| self.table: FileStoreTable = table |
| self.data_writers: Dict[Tuple, DataWriter] = {} |
| self._runtime_total_buckets: Dict[Tuple, int] = {} |
| self.max_seq_numbers: dict = {} |
| self.write_cols = None |
| self.blob_consumer = None |
| self.commit_identifier = 0 |
| self.options = CoreOptions.copy(table.options) |
| self.changelog_producer = self.options.changelog_producer() |
| self._configure_data_file_prefix(commit_user) |
| |
| def _configure_data_file_prefix(self, commit_user): |
| if self.table.bucket_mode() == BucketMode.POSTPONE_MODE: |
| self.options.set(CoreOptions.DATA_FILE_PREFIX, |
| (f"{self.options.data_file_prefix()}-u-{commit_user}" |
| f"-s-{random.randint(0, 2 ** 31 - 2)}-w-")) |
| |
| def disable_rolling(self): |
| """Disable size- and row-based file rolling.""" |
| max_value = CoreOptions.TARGET_FILE_ROW_NUM.default_value() |
| self.options.set( |
| CoreOptions.TARGET_FILE_SIZE, str(max_value)) |
| self.options.set( |
| CoreOptions.TARGET_FILE_ROW_NUM, str(max_value)) |
| |
| def write( |
| self, |
| partition: Tuple, |
| bucket: int, |
| data: pa.RecordBatch, |
| total_buckets=None, |
| ): |
| self._check_runtime_bucket(partition, bucket, total_buckets) |
| key = (partition, bucket) |
| if key not in self.data_writers: |
| self.data_writers[key] = self._create_data_writer(partition, bucket, self.options) |
| writer = self.data_writers[key] |
| writer.write(data) |
| |
| def write_row( |
| self, |
| partition: Tuple, |
| bucket: int, |
| row, |
| values_by_name: dict, |
| total_buckets=None, |
| ): |
| self._check_runtime_bucket(partition, bucket, total_buckets) |
| key = (partition, bucket) |
| if key not in self.data_writers: |
| self.data_writers[key] = self._create_data_writer(partition, bucket, self.options) |
| writer = self.data_writers[key] |
| if hasattr(writer, 'write_row'): |
| writer.write_row(row) |
| return |
| |
| column_names = ( |
| self.write_cols |
| if self.write_cols is not None |
| else list(self.table.field_names) |
| ) |
| data = row_values_to_arrow_table( |
| values_by_name, |
| self.table.table_schema.fields, |
| column_names, |
| ) |
| writer.write(data.to_batches()[0]) |
| |
| def _check_runtime_bucket(self, partition, bucket, total_buckets): |
| if total_buckets is None: |
| return |
| if (isinstance(total_buckets, bool) |
| or not isinstance(total_buckets, int) |
| or total_buckets <= 0): |
| raise ValueError("Total number of buckets must be positive") |
| if bucket < 0 or bucket >= total_buckets: |
| raise ValueError( |
| "Bucket {} is out of range [0, {})".format( |
| bucket, total_buckets |
| ) |
| ) |
| |
| partition = tuple(partition) |
| previous = self._runtime_total_buckets.get(partition) |
| if previous is not None and previous != total_buckets: |
| raise RuntimeError( |
| "Try to write partition {} with a new bucket num {}, but " |
| "the previous bucket num is {}.".format( |
| partition, total_buckets, previous |
| ) |
| ) |
| self._runtime_total_buckets[partition] = total_buckets |
| |
| def _create_data_writer(self, partition: Tuple, bucket: int, options: CoreOptions) -> DataWriter: |
| row_limit = options.target_file_row_num() |
| max_value = CoreOptions.TARGET_FILE_ROW_NUM.default_value() |
| if row_limit < 1: |
| raise ValueError( |
| "target-file-row-num should be at least 1") |
| if row_limit > max_value: |
| raise ValueError( |
| f"target-file-row-num should be at most {max_value}") |
| if row_limit != max_value: |
| # Row-count rolling is implemented in the base append writer only. |
| # DE (data-evolution) append tables are the target; primary-key, |
| # blob and vector writers override rolling and are not supported yet. |
| row_rolling_supported = ( |
| self.table.options.data_evolution_enabled() |
| and not self.table.is_primary_key_table |
| and not self._has_blob_columns() |
| and not (self._has_vector_columns() |
| and options.with_vector_format())) |
| if not row_rolling_supported: |
| raise NotImplementedError( |
| "target-file-row-num is set on this table but pypaimon supports row-count " |
| "based file rolling only for data-evolution append tables (no primary key, " |
| "blob or vector columns); unset it or write with Java/Flink/Spark.") |
| |
| def max_seq_number(): |
| return self._seq_number_stats(partition).get(bucket, 1) |
| |
| # Check if table has blob columns |
| if self._has_blob_columns(): |
| return DedicatedFormatWriter( |
| table=self.table, |
| partition=partition, |
| bucket=bucket, |
| max_seq_number=0, |
| options=options, |
| write_cols=self.write_cols, |
| blob_consumer=self.blob_consumer, |
| changelog_producer=self.changelog_producer, |
| ) |
| elif self._has_vector_columns() and options.with_vector_format(): |
| return DataVectorWriter( |
| table=self.table, |
| partition=partition, |
| bucket=bucket, |
| max_seq_number=0, |
| options=options, |
| write_cols=self.write_cols, |
| ) |
| elif self.table.is_primary_key_table: |
| return KeyValueDataWriter( |
| table=self.table, |
| partition=partition, |
| bucket=bucket, |
| max_seq_number=max_seq_number(), |
| options=options, |
| merge_function=self._build_pk_merge_function(), |
| changelog_producer=self.changelog_producer) |
| else: |
| seq_number = 0 if self.table.bucket_mode() == BucketMode.BUCKET_UNAWARE else max_seq_number() |
| return AppendOnlyDataWriter( |
| table=self.table, |
| partition=partition, |
| bucket=bucket, |
| max_seq_number=seq_number, |
| options=options, |
| write_cols=self.write_cols, |
| changelog_producer=self.changelog_producer |
| ) |
| |
| def _build_pk_merge_function(self): |
| """Build the merge function for the in-memory write buffer. |
| |
| Shares ``merge_engine_dispatch.build_merge_function`` with the |
| read path so the supported engines (deduplicate, first-row, |
| partial-update with no out-of-scope options) cannot drift |
| between sides. |
| |
| For wholly unsupported engines (``aggregation``) the writer |
| falls back to ``DeduplicateMergeFunction`` so the flushed file |
| still maintains the LSM "PK unique within a file" invariant. |
| The read path's dispatch still raises ``NotImplementedError``, |
| so the user gets an explicit error before they observe |
| wrong-engine data; the fallback only narrows the damage to |
| "file is deduped, not aggregated" rather than the silent |
| multi-row-per-PK corruption that existed pre-PR. |
| |
| Partial-update with out-of-scope options (sequence-group, |
| per-field aggregator, ignore-delete, remove-record-on-*) does |
| **not** fall back: ``partial_update_unsupported_options`` sees |
| the configured keys and re-raises, so the first |
| ``write_arrow`` call (where ``_create_data_writer`` first runs) |
| surfaces the error. Silently degrading to dedupe there is the |
| same live corruption pattern this PR exists to close. |
| |
| ``with_write_type`` (column-subset writes) on a PK table is |
| also rejected here. The buffer layout |
| ``_add_system_fields`` produces would carry only the subset |
| on the value side, while a ``MergeFunction`` such as |
| ``PartialUpdateMergeFunction`` is built against the full table |
| arity -- the two sides would mismatch on |
| ``KeyValue.value.get_field`` and raise ``IndexError`` at |
| flush time. Refusing it explicitly avoids that obscure failure |
| and keeps the supported surface narrow. |
| |
| The value-side schema must match the layout |
| ``KeyValueDataWriter`` flushes -- ``_add_system_fields`` keeps |
| every original user column on the value side (the primary keys |
| are duplicated as ``_KEY_<pk>`` columns to the left of the |
| value side). So ``value_arity`` here is ``len(table.fields)``, |
| not ``len(table.fields) - len(primary_keys)``. |
| """ |
| from pypaimon.common.merge_engine_dispatch import ( |
| build_merge_function, partial_update_unsupported_options) |
| from pypaimon.common.options.core_options import MergeEngine |
| from pypaimon.read.reader.deduplicate_merge_function import \ |
| DeduplicateMergeFunction |
| |
| engine = self.options.merge_engine() |
| raw_options = self.options.options.to_map() |
| |
| if self.write_cols is not None: |
| raise NotImplementedError( |
| "with_write_type is not yet supported on primary-key " |
| "tables: the writer-side merge buffer assumes the " |
| "input batch carries the full table schema. Drop the " |
| "with_write_type call or write the missing columns as " |
| "nulls in the input batch." |
| ) |
| |
| # PARTIAL_UPDATE + out-of-scope option: never silently fall |
| # back -- forward the read-side error verbatim so writes fail |
| # before the first flush rather than corrupt the file. |
| if engine == MergeEngine.PARTIAL_UPDATE \ |
| and partial_update_unsupported_options(raw_options): |
| return build_merge_function( |
| engine=engine, raw_options=raw_options, |
| key_arity=len(self.table.trimmed_primary_keys), |
| value_arity=len(self.table.table_schema.fields), |
| value_field_nullables=[ |
| f.type.nullable for f in self.table.table_schema.fields], |
| value_field_names=[ |
| f.name for f in self.table.table_schema.fields], |
| ) |
| |
| # Catch the dispatch's "wholly unsupported engine" raise only |
| # for the engines we know are out of scope today; any other |
| # NotImplementedError is a bug we want to surface, not swallow. |
| if engine == MergeEngine.AGGREGATE: |
| # Surface the silent semantic mismatch in logs: the file |
| # will be PK-unique (better than the pre-PR multi-row |
| # corruption), but any reader that honours the declared |
| # engine will see wrong values. Users sharing tables |
| # across writers especially need to see this. |
| logger.warning( |
| "merge-engine '%s' is not implemented on the pypaimon " |
| "write path; falling back to deduplicate so the flushed " |
| "file stays PK-unique. The file contents reflect " |
| "deduplicate semantics (latest writer wins), not %s " |
| "semantics. Any reader that interprets the file under " |
| "the declared engine will return incorrect results. " |
| "Avoid the pypaimon writer for tables on this engine.", |
| engine.value, engine.value) |
| return DeduplicateMergeFunction() |
| |
| all_value_fields = self.table.table_schema.fields |
| return build_merge_function( |
| engine=engine, raw_options=raw_options, |
| key_arity=len(self.table.trimmed_primary_keys), |
| value_arity=len(all_value_fields), |
| value_field_nullables=[ |
| f.type.nullable for f in all_value_fields], |
| value_field_names=[f.name for f in all_value_fields], |
| ) |
| |
| def _has_blob_columns(self) -> bool: |
| """Check if the table schema contains blob columns.""" |
| return any(is_blob_file_field(field) for field in self.table.table_schema.fields) |
| |
| def _has_vector_columns(self) -> bool: |
| from pypaimon.schema.data_types import VectorType |
| return any(isinstance(f.type, VectorType) for f in self.table.table_schema.fields) |
| |
| def prepare_commit(self, commit_identifier) -> List[CommitMessage]: |
| self.commit_identifier = commit_identifier |
| commit_messages = [] |
| for (partition, bucket), writer in self.data_writers.items(): |
| committed_files = writer.prepare_commit() |
| changelog_files = writer.prepare_changelog_commit() |
| if committed_files or changelog_files: |
| commit_message = CommitMessage( |
| partition=partition, |
| bucket=bucket, |
| new_files=committed_files, |
| changelog_files=changelog_files, |
| total_buckets=self._runtime_total_buckets.get(partition), |
| ) |
| commit_messages.append(commit_message) |
| return commit_messages |
| |
| def close(self): |
| """Close all data writers and clean up resources.""" |
| for writer in self.data_writers.values(): |
| writer.close() |
| self.data_writers.clear() |
| self._runtime_total_buckets.clear() |
| |
| def abort(self): |
| """Abort all data writers and clean up files produced by this write.""" |
| for writer in self.data_writers.values(): |
| try: |
| writer.abort() |
| except Exception as e: |
| logger.warning("Failed to abort data writer.", exc_info=e) |
| self.data_writers.clear() |
| self._runtime_total_buckets.clear() |
| |
| def _seq_number_stats(self, partition: Tuple) -> Dict[int, int]: |
| buckets = self.max_seq_numbers.get(partition) |
| if buckets is None: |
| buckets = self._load_seq_number_stats(partition) |
| self.max_seq_numbers[partition] = buckets |
| return buckets |
| |
| def _load_seq_number_stats(self, partition: Tuple) -> dict: |
| read_builder = self.table.new_read_builder() |
| predicate_builder = read_builder.new_predicate_builder() |
| sub_predicates = [] |
| for key, value in zip(self.table.partition_keys, partition): |
| sub_predicates.append(predicate_builder.equal(key, value)) |
| partition_filter = predicate_builder.and_predicates(sub_predicates) |
| |
| scan = read_builder.with_filter(partition_filter).new_scan() |
| splits = scan.plan_for_write().splits() |
| |
| max_seq_numbers = {} |
| for split in splits: |
| current_seq_num = max([file.max_sequence_number for file in split.files]) |
| existing_max = max_seq_numbers.get(split.bucket, -1) |
| if current_seq_num > existing_max: |
| max_seq_numbers[split.bucket] = current_seq_num |
| return max_seq_numbers |
| |
| |
| class PostponeFixedBucketFileStoreWrite(FileStoreWrite): |
| """File store write with runtime bucket counts for postpone tables.""" |
| |
| def _configure_data_file_prefix(self, commit_user): |
| pass |
