| # 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 List, Optional | |
| import pyarrow as pa | |
| from pypaimon.common.options.core_options import ChangelogProducer | |
| from pypaimon.schema.arrow_schema import ( | |
| arrow_schemas_compatible, normalize_arrow_strings, normalize_arrow_write_nullability, project_arrow_batch) | |
| from pypaimon.schema.data_types import PyarrowFieldParser | |
| from pypaimon.table.row.blob import BlobConsumer | |
| from pypaimon.write.row_utils import ( | |
| require_columns, | |
| row_to_named_values, | |
| validate_write_columns, | |
| ) | |
| from pypaimon.write.commit_message import CommitMessage | |
| from pypaimon.write.commit_identifier import validate_commit_identifier | |
| from pypaimon.execution.python.write.file_store_write import FileStoreWrite | |
| from pypaimon.write.input_adapter import _pandas_write_batch | |
| class PythonTableWrite: | |
| """Bound Python writer implementing the common write execution protocol.""" | |
| def __init__(self, public_writer): | |
| self.public_writer = public_writer | |
| def initialize(self, table, commit_user, static_partition: Optional[dict] = None, | |
| *, restore_snapshot_id=None): | |
| from pypaimon.table.file_store_table import FileStoreTable | |
| self.table: FileStoreTable = table | |
| self.table_pyarrow_schema = PyarrowFieldParser.from_paimon_schema(self.table.table_schema.fields) | |
| self.commit_user = commit_user | |
| self.static_partition = static_partition | |
| self.restore_snapshot_id = restore_snapshot_id | |
| self.file_store_write = self._create_file_store_write(commit_user) | |
| # Java overwrite writers ignore previous files even without an explicit | |
| # restore snapshot. HASH restoration is separately disabled by the row | |
| # key extractor; sequence restoration must also start from empty state. | |
| self.file_store_write.restore_snapshot_id = ( | |
| 0 if static_partition is not None else restore_snapshot_id) | |
| if static_partition is not None: | |
| # An overwrite replaces state, not an input changelog. Java's | |
| # overwrite commit does not publish changelog manifests; avoid | |
| # writing unreferenced changelog files in the first place. | |
| self.file_store_write.changelog_producer = ChangelogProducer.NONE | |
| self.row_key_extractor = self._create_row_key_extractor(static_partition) | |
| def _create_file_store_write(self, commit_user): | |
| return FileStoreWrite(self.table, commit_user) | |
| def _create_row_key_extractor(self, static_partition): | |
| if self.restore_snapshot_id is not None: | |
| from pypaimon.execution.python.write.row_key_extractor import DynamicBucketRowKeyExtractor | |
| return DynamicBucketRowKeyExtractor( | |
| self.table.table_schema, table=self.table, | |
| ignore_existing=static_partition is not None, | |
| base_snapshot_id=self.restore_snapshot_id) | |
| from pypaimon.execution.python.write.row_key_extractor import create_row_key_extractor | |
| return create_row_key_extractor(self.table, ignore_existing=static_partition is not None) | |
| def write_arrow(self, data, bucket: Optional[int] = None): | |
| """Write Arrow data, optionally to a bucket assigned upstream.""" | |
| data = self._prepare_arrow_schema(data) | |
| # Public batch overrides remain observable for Arrow and pandas input. | |
| # Schema, routing and preparation helpers stay within this executor. | |
| if isinstance(data, pa.RecordBatch): | |
| return self.public_writer.write_arrow_batch(data, bucket) | |
| for batch in data.to_batches(): | |
| self.public_writer.write_arrow_batch(batch, bucket) | |
| def write_arrow_batch(self, data: pa.RecordBatch, bucket: Optional[int] = None): | |
| if bucket is not None: | |
| return self._write_precomputed_bucket(data, bucket) | |
| data = self._prepare_arrow_data(data) | |
| if data.num_rows == 0: | |
| return | |
| groups = self.row_key_extractor.extract_partition_bucket_groups(data) | |
| for partition, bucket, row_indices in groups: | |
| if row_indices is None: | |
| # Every input row belongs to the same partition/bucket. Passing the | |
| # original batch through avoids copying large BLOB values through | |
| # Arrow take before the dedicated BLOB writer consumes them. | |
| sub_table = data | |
| else: | |
| # row_indices is an int64 array of this group's rows in | |
| # ascending input order (the extractor sorts grouped indices so | |
| # sequence-number assignment stays latest-wins correct), so the | |
| # span is just first..last. | |
| lo = row_indices[0].as_py() | |
| hi = row_indices[-1].as_py() | |
| count = len(row_indices) | |
| if hi - lo + 1 == count: | |
| # Distinct row indices spanning exactly `count` values are | |
| # contiguous, so share the original Arrow buffers instead of | |
| # gathering their rows into newly allocated buffers with take. | |
| sub_table = data.slice(lo, count) | |
| else: | |
| sub_table = pa.compute.take(data, row_indices) | |
| self._write_partition_bucket_batch(partition, bucket, sub_table) | |
| def roll_before_group_if_needed(self, row_count: int): | |
| """Roll current files before the next logical write group if needed.""" | |
| if isinstance(row_count, bool) or not isinstance(row_count, int) \ | |
| or row_count <= 0: | |
| raise ValueError("Group row count must be a positive integer.") | |
| self.file_store_write.roll_before_group_if_needed(row_count) | |
| def _write_partition_bucket_batch(self, partition, bucket, data): | |
| self.file_store_write.write(partition, bucket, data) | |
| def _write_precomputed_bucket(self, data, bucket): | |
| """Write one partition/bucket group using normal index maintenance.""" | |
| from pypaimon.table.bucket_mode import BucketMode | |
| from pypaimon.execution.python.write.row_key_extractor import DynamicBucketRowKeyExtractor | |
| bucket_mode = self.table.bucket_mode() | |
| if bucket_mode not in (BucketMode.HASH_FIXED, BucketMode.HASH_DYNAMIC): | |
| raise ValueError( | |
| "Precomputed bucket writes are only valid for HASH_FIXED or " | |
| "HASH_DYNAMIC tables" | |
| ) | |
| if not isinstance(self.row_key_extractor, DynamicBucketRowKeyExtractor): | |
| if bucket_mode == BucketMode.HASH_DYNAMIC: | |
| raise RuntimeError("Dynamic bucket extractor is not configured") | |
| from pypaimon.index.dynamic_bucket import MAX_DYNAMIC_BUCKETS | |
| upper_bound = MAX_DYNAMIC_BUCKETS if bucket_mode == BucketMode.HASH_DYNAMIC else self.table.options.bucket() | |
| if not 0 <= bucket < upper_bound: | |
| raise ValueError('Bucket id must be between 0 and {}'.format(upper_bound - 1)) | |
| data = self._prepare_arrow_data(data) | |
| if bucket_mode == BucketMode.HASH_DYNAMIC: | |
| partition = self.row_key_extractor.notify_precomputed_bucket_batch(data, bucket) | |
| else: | |
| partitions = self.row_key_extractor.extract_partitions_batch(data) | |
| if not partitions: | |
| return | |
| partition = tuple(partitions[0]) | |
| if any(tuple(actual) != partition for actual in partitions): | |
| raise ValueError('A precomputed bucket group contained multiple partitions') | |
| if partition is None: | |
| return | |
| self._write_partition_bucket_batch(partition, bucket, data) | |
| def write_row(self, row): | |
| values_by_name = row_to_named_values(row, self.table.table_schema.fields) | |
| require_columns(values_by_name, self.table.partition_keys, "write_row") | |
| from pypaimon.write.row_kind import skip_write_row | |
| write_cols = (self.file_store_write.write_cols | |
| if set(values_by_name) == set(self.table.field_names) else list(values_by_name)) | |
| if skip_write_row(self.table, values_by_name, row.get_row_kind(), write_cols): | |
| return | |
| partition, bucket = ( | |
| self.row_key_extractor.extract_partition_bucket_row(values_by_name) | |
| ) | |
| self._write_partition_bucket_row( | |
| partition, bucket, row, values_by_name | |
| ) | |
| def _write_partition_bucket_row( | |
| self, partition, bucket, row, values_by_name | |
| ): | |
| self.file_store_write.write_row( | |
| partition, bucket, row, values_by_name | |
| ) | |
| def write_pandas(self, dataframe): | |
| write_cols = self.file_store_write.write_cols | |
| if write_cols is not None: | |
| # Column-subset write (append-only ``with_write_type``): build the | |
| # RecordBatch against the subset schema so the input only needs the | |
| # written columns, mirroring the ``write_arrow`` path. | |
| pa_schema = self._write_cols_pyarrow_schema(write_cols) | |
| else: | |
| pa_schema = self.table_pyarrow_schema | |
| record_batch = _pandas_write_batch( | |
| getattr(self, 'table', None), dataframe, pa_schema, self.table_pyarrow_schema) | |
| return self.public_writer.write_arrow_batch(record_batch) | |
| def with_write_type(self, write_cols: List[str]): | |
| write_cols = validate_write_columns(write_cols) | |
| for col in write_cols: | |
| if col not in self.table_pyarrow_schema.names: | |
| raise ValueError(f"Column {col} is not in table schema.") | |
| full_names = self.table_pyarrow_schema.names | |
| # Java preserves the declared append column order in writeCols. PK | |
| # writes still require the complete table layout in the Python buffer. | |
| full_pk_type = (self.table.is_primary_key_table and len(write_cols) == len(full_names) | |
| and set(write_cols) == set(full_names)) | |
| if write_cols == full_names or full_pk_type: | |
| write_cols = None | |
| if self.file_store_write._ever_written: | |
| if write_cols == self.file_store_write.write_cols: | |
| return self | |
| raise RuntimeError('with_write_type must be called before any write operation') | |
| self.file_store_write.write_cols = write_cols | |
| return self | |
| def with_blob_consumer(self, blob_consumer: BlobConsumer): | |
| if blob_consumer is not None and not callable(blob_consumer): | |
| raise TypeError('blob_consumer must be callable') | |
| self.file_store_write.blob_consumer = blob_consumer | |
| return self | |
| def with_blob_uri_reader_factory(self, uri_reader_factory): | |
| if uri_reader_factory is not None and not callable(getattr(uri_reader_factory, 'create', None)): | |
| raise TypeError('URI reader factory must provide callable create') | |
| if self.file_store_write._ever_written: | |
| raise RuntimeError( | |
| "with_blob_uri_reader_factory must be called before any " | |
| "write operation." | |
| ) | |
| self.file_store_write.blob_uri_reader_factory = uri_reader_factory | |
| return self | |
| def close(self): | |
| try: | |
| self.file_store_write.close() | |
| finally: | |
| self._release_prepared_indexes() | |
| def abort(self): | |
| try: | |
| self.file_store_write.abort() | |
| finally: | |
| abort = getattr(self.row_key_extractor, "abort", None) | |
| if abort is not None: | |
| abort() | |
| def prepare_commit(self, commit_identifier) -> List[CommitMessage]: | |
| commit_identifier = validate_commit_identifier(commit_identifier) | |
| commit_messages = self.file_store_write._prepare_commit_messages(commit_identifier) | |
| prepare_indexes = getattr(self.row_key_extractor, "prepare_commit", None) | |
| index_changes = prepare_indexes() if prepare_indexes is not None else {} | |
| messages_by_bucket = { | |
| (tuple(message.partition), message.bucket): message | |
| for message in commit_messages | |
| } | |
| for (partition, bucket), changes in index_changes.items(): | |
| message = messages_by_bucket.get((partition, bucket)) | |
| if message is None: | |
| message = CommitMessage( | |
| partition=partition, | |
| bucket=bucket, | |
| new_files=[], | |
| ) | |
| commit_messages.append(message) | |
| messages_by_bucket[(partition, bucket)] = message | |
| message.index_adds.extend(changes.additions) | |
| message.index_deletes.extend(changes.deletions) | |
| # Hand off data, changelog and HASH indexes together. Until every | |
| # preparation succeeds, retain the metadata for retry and abort. | |
| self.file_store_write._release_prepared_files() | |
| self._release_prepared_indexes() | |
| return commit_messages | |
| def _release_prepared_indexes(self) -> None: | |
| release = getattr(self.row_key_extractor, "release_prepared", None) | |
| if release is not None: | |
| release() | |
| def _prepare_arrow_data(self, data): | |
| data = self._prepare_arrow_schema(data) | |
| # Schema-only writers can normalize input without table options. | |
| table = getattr(self, 'table', None) | |
| if table is None: | |
| return data | |
| data = self._normalize_arrow_write_schema(data) | |
| from pypaimon.write.row_kind import filter_write_batch | |
| selected = (self.file_store_write.write_cols if not table.is_primary_key_table else None) | |
| write_cols = selected | |
| if selected is None: | |
| write_cols = (self.file_store_write.write_cols | |
| if data.schema.names == self.table_pyarrow_schema.names else data.schema.names) | |
| data = filter_write_batch(table, data, write_cols) | |
| if not data.num_rows: | |
| return data | |
| if selected is not None: | |
| routing = self._write_routing_columns() | |
| columns = selected + [name for name in data.schema.names | |
| if name in routing and name not in selected] | |
| data = self._project_write_input(data, columns) | |
| # Routing fields are not written values. Arrow's whole-table | |
| # string cast must not reapply their declared top-level NOT NULL. | |
| schema = pa.schema([field if field.name in selected else field.with_nullable(True) | |
| for field in data.schema], metadata=data.schema.metadata) | |
| data = (pa.RecordBatch.from_arrays(data.columns, schema=schema) if isinstance(data, pa.RecordBatch) | |
| else pa.Table.from_arrays(data.columns, schema=schema)) | |
| data = normalize_arrow_strings(data) | |
| physical = self._project_write_input(data, selected) if selected is not None else data | |
| target = self.table_pyarrow_schema | |
| if physical.schema.names != target.names: | |
| target = self._write_cols_pyarrow_schema(physical.schema.names) | |
| # Validate only surviving values. Keep the input transport metadata: | |
| # PyArrow otherwise rejects hidden NULL children of NULL ROW parents. | |
| validated = normalize_arrow_write_nullability(physical, target) | |
| self._validate_pyarrow_schema(validated.schema) | |
| return data | |
| @staticmethod | |
| def _project_write_input(data, columns): | |
| return (project_arrow_batch(data, columns) if isinstance(data, pa.RecordBatch) | |
| else data.select(columns)) | |
| def _write_routing_columns(self): | |
| table = self.table | |
| routing = set(table.partition_keys) | |
| if table.options.bucket() > 0: | |
| routing.update(table.table_schema.bucket_keys) | |
| row_kind = table.options.options.to_map().get('rowkind.field') | |
| if row_kind is not None: | |
| routing.add(row_kind) | |
| return routing | |
| def _prepare_arrow_schema(self, data): | |
| # Legacy schema-only callers have no row-kind options to evaluate. | |
| # A real writer must defer casts until ignored events have been removed. | |
| if getattr(self, 'table', None) is None: | |
| data = normalize_arrow_strings(data) | |
| self._validate_pyarrow_schema( | |
| data.schema, check_nested_nullability=getattr(self, 'table', None) is None) | |
| return data | |
| def _normalize_arrow_write_schema(self, data): | |
| # Compatible input batches may differ in top-level nullability or | |
| # metadata. Buffered writes and row-to-Arrow conversion must use the | |
| # same declared fields, while keeping normalized Arrow array types. | |
| fields = [] | |
| for field in data.schema: | |
| declared = self.table_pyarrow_schema.field(field.name) | |
| fields.append(pa.field(field.name, field.type, nullable=declared.nullable, | |
| metadata=declared.metadata)) | |
| schema = pa.schema(fields, metadata=self.table_pyarrow_schema.metadata) | |
| from pypaimon.write.row_kind import write_value_schema | |
| schema = write_value_schema(self.table, schema) | |
| if data.schema.equals(schema, check_metadata=True): | |
| return data | |
| if isinstance(data, pa.RecordBatch): | |
| return pa.RecordBatch.from_arrays(data.columns, schema=schema) | |
| return pa.Table.from_arrays(data.columns, schema=schema) | |
| def _validate_pyarrow_schema(self, data_schema: pa.Schema, *, check_nested_nullability=True): | |
| if self._is_compatible_pyarrow_schema( | |
| data_schema, self.table_pyarrow_schema, check_nested_nullability=check_nested_nullability): | |
| return | |
| write_cols = self.file_store_write.write_cols | |
| if write_cols is not None: | |
| write_cols_schema = self._write_cols_pyarrow_schema(write_cols) | |
| if self._is_compatible_pyarrow_schema( | |
| data_schema, write_cols_schema, check_nested_nullability=check_nested_nullability): | |
| return | |
| table = getattr(self, 'table', None) | |
| if table is not None and not table.is_primary_key_table: | |
| names = data_schema.names | |
| selected = set(write_cols) | |
| allowed = selected | self._write_routing_columns() | |
| if (len(set(names)) == len(names) and set(names).issubset(allowed) | |
| and [name for name in names if name in selected] == write_cols): | |
| transport_schema = self._write_cols_pyarrow_schema(names) | |
| if self._is_compatible_pyarrow_schema( | |
| data_schema, transport_schema, check_nested_nullability=check_nested_nullability): | |
| return | |
| self._raise_inconsistent_schema(data_schema) | |
| def _is_compatible_pyarrow_schema( | |
| self, data_schema: pa.Schema, expected_schema: pa.Schema, *, check_nested_nullability=True) -> bool: | |
| return arrow_schemas_compatible( | |
| data_schema, expected_schema, check_top_level_nullability=False, | |
| allow_binary_compatibility=True, check_nested_nullability=check_nested_nullability) | |
| def _write_cols_pyarrow_schema(self, write_cols: List[str]) -> pa.Schema: | |
| table_fields = { | |
| field.name: field for field in self.table_pyarrow_schema | |
| } | |
| return pa.schema([table_fields[col] for col in write_cols]) | |
| def _raise_inconsistent_schema(self, data_schema: pa.Schema): | |
| raise ValueError(f"Input schema isn't consistent with table schema and write cols. " | |
| f"Input schema is: {data_schema} " | |
| f"Table schema is: {self.table_pyarrow_schema} " | |
| f"Write cols is: {self.file_store_write.write_cols}") |