| # 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 uuid |
| from typing import Dict, List, Optional, Tuple |
| |
| import pyarrow as pa |
| |
| from pypaimon.common.options.core_options import CoreOptions, ChangelogProducer |
| from pypaimon.data.timestamp import Timestamp |
| from pypaimon.manifest.schema.data_file_meta import DataFileMeta |
| from pypaimon.manifest.schema.simple_stats import SimpleStats |
| from pypaimon.schema.data_types import VectorType |
| from pypaimon.table.row.blob import BlobConsumer |
| from pypaimon.table.row.generic_row import GenericRow |
| from pypaimon.write.writer.data_writer import DataWriter |
| |
| logger = logging.getLogger(__name__) |
| |
| |
| class DedicatedFormatWriter(DataWriter): |
| """A rolling file writer that writes normal, blob, and vector columns to dedicated files. |
| |
| Splits incoming data three ways: |
| - Normal columns → standard data files (.parquet / .orc / .vortex / …) |
| - Blob columns (large_binary) → .blob files |
| - Vector columns (when vector.file.format is configured) → .vector.<format> files |
| |
| This mirrors Java's DedicatedFormatRollingFileWriter. |
| |
| Metadata order in committed_files: |
| [normal_meta, blob_meta1, …, vector_meta1, …] |
| """ |
| |
| # Constant for checking rolling condition periodically |
| CHECK_ROLLING_RECORD_CNT = 1000 |
| |
| def __init__(self, table, partition: Tuple, bucket: int, max_seq_number: int, options: CoreOptions = None, |
| write_cols: Optional[List[str]] = None, blob_consumer: Optional[BlobConsumer] = None, |
| changelog_producer: ChangelogProducer = ChangelogProducer.NONE): |
| super().__init__(table, partition, bucket, max_seq_number, options, write_cols=write_cols, |
| changelog_producer=changelog_producer) |
| |
| # Determine blob columns from table schema |
| self.blob_column_names = self._get_blob_columns_from_schema() |
| self.blob_descriptor_fields = CoreOptions.blob_descriptor_fields(self.options) |
| self.blob_view_fields = CoreOptions.blob_view_fields(self.options) |
| self.blob_inline_fields = self.blob_descriptor_fields.union(self.blob_view_fields) |
| |
| unknown_descriptor_fields = self.blob_descriptor_fields.difference( |
| set(self.blob_column_names) |
| ) |
| if unknown_descriptor_fields: |
| raise ValueError( |
| "Fields in 'blob-descriptor-field' must be blob fields in schema. " |
| f"Unknown fields: {sorted(unknown_descriptor_fields)}" |
| ) |
| |
| # Blob fields that should still be written to `.blob` files. |
| self.blob_file_column_names = [ |
| col for col in self.blob_column_names if col not in self.blob_inline_fields |
| ] |
| full_blob_file_set = set(self.blob_file_column_names) |
| all_column_names = self.table.field_names |
| |
| # Detect vector columns that should be written to dedicated files. |
| full_vector_column_names = self._get_vector_columns_from_schema() |
| full_vector_set = set(full_vector_column_names) |
| # Only split vector columns when vector.file.format is configured. |
| has_dedicated_vector = bool(full_vector_column_names) and options.with_vector_format() |
| dedicated_set = full_blob_file_set | (full_vector_set if has_dedicated_vector else set()) |
| |
| # Narrow columns when TableWrite.with_write_type(...) supplies a partial column list. |
| # Incoming RecordBatches only contain those columns; selecting full normal/blob lists |
| # would raise KeyError. |
| if write_cols is not None: |
| write_col_set = set(write_cols) |
| self.blob_file_column_names = [ |
| col for col in self.blob_file_column_names if col in write_col_set |
| ] |
| self.vector_write_columns = [ |
| col for col in full_vector_column_names if col in write_col_set |
| ] if has_dedicated_vector else [] |
| self.normal_column_names = [ |
| col for col in write_cols if col not in dedicated_set |
| ] |
| else: |
| self.vector_write_columns = list(full_vector_column_names) if has_dedicated_vector else [] |
| self.normal_column_names = [ |
| col for col in all_column_names if col not in dedicated_set |
| ] |
| normal_name_set = set(self.normal_column_names) |
| self.normal_columns = [ |
| field for field in self.table.table_schema.fields if field.name in normal_name_set |
| ] |
| self.write_cols = self.normal_column_names |
| |
| # State management for blob writer |
| self.record_count = 0 |
| self.closed = False |
| |
| # Track pending data for normal data only |
| self.pending_normal_data: Optional[pa.Table] = None |
| |
| # Initialize blob writers for each blob-file column. |
| from pypaimon.write.writer.blob_writer import BlobWriter |
| self.blob_writers: Dict[str, BlobWriter] = {} |
| for blob_column in self.blob_file_column_names: |
| self.blob_writers[blob_column] = BlobWriter( |
| table=self.table, |
| partition=self.partition, |
| bucket=self.bucket, |
| max_seq_number=max_seq_number, |
| blob_column=blob_column, |
| options=options, |
| blob_consumer=blob_consumer, |
| ) |
| |
| # Initialize vector writer when vector.file.format is configured. |
| from pypaimon.write.writer.vector_writer import VectorWriter |
| self.vector_writer: Optional[VectorWriter] = None |
| if self.vector_write_columns: |
| self.vector_writer = VectorWriter( |
| table=self.table, |
| partition=self.partition, |
| bucket=self.bucket, |
| max_seq_number=max_seq_number, |
| vector_columns=self.vector_write_columns, |
| vector_file_format=options.vector_file_format(), |
| options=options, |
| ) |
| |
| # Initialize ExternalStorageBlobWriter if configured |
| self._external_storage_writer = None |
| external_storage_fields = self.options.blob_external_storage_fields() |
| external_storage_path = self.options.blob_external_storage_path() |
| if external_storage_fields and external_storage_path: |
| from pypaimon.write.writer.external_storage_blob_writer import \ |
| ExternalStorageBlobWriter |
| self._external_storage_writer = ExternalStorageBlobWriter( |
| file_io=self.file_io, |
| external_storage_path=external_storage_path, |
| external_storage_fields=external_storage_fields, |
| blob_target_file_size=self.options.blob_target_file_size(), |
| data_file_prefix=CoreOptions.data_file_prefix(self.options), |
| ) |
| |
| logger.info( |
| "Initialized DedicatedFormatWriter with blob columns: %s, blob file columns: %s, " |
| "vector columns: %s, descriptor stored columns: %s, external storage fields: %s, view stored columns: %s", |
| self.blob_column_names, |
| self.blob_file_column_names, |
| self.vector_write_columns, |
| sorted(self.blob_descriptor_fields), |
| sorted(external_storage_fields) if external_storage_fields else [], |
| sorted(self.blob_view_fields) |
| ) |
| |
| def _get_blob_columns_from_schema(self) -> List[str]: |
| blob_columns = [] |
| for field in self.table.table_schema.fields: |
| type_str = str(field.type).lower() |
| if 'blob' in type_str: |
| blob_columns.append(field.name) |
| |
| if len(blob_columns) == 0: |
| raise ValueError("No blob field found in table schema.") |
| return blob_columns |
| |
| def _get_vector_columns_from_schema(self) -> List[str]: |
| return [ |
| field.name for field in self.table.table_schema.fields |
| if isinstance(field.type, VectorType) |
| ] |
| |
| def _process_data(self, data: pa.RecordBatch) -> pa.RecordBatch: |
| normal_data, _, _ = self._split_data(data) |
| return normal_data |
| |
| def _merge_data(self, existing_data: pa.Table, new_data: pa.Table) -> pa.Table: |
| return self._merge_normal_data(existing_data, new_data) |
| |
| def write(self, data: pa.RecordBatch): |
| try: |
| # Transform external-storage fields: write raw blob to external storage, |
| # replace with serialized BlobDescriptor |
| if self._external_storage_writer: |
| data = self._external_storage_writer.transform_batch(data) |
| |
| # Split data into normal, blob, and vector parts |
| normal_data, blob_data_map, vector_data = self._split_data(data) |
| self._validate_inline_stored_fields_input(data) |
| |
| # Process and accumulate normal data (may be None for partial writes) |
| processed_normal = self._process_normal_data(normal_data) |
| if processed_normal is not None: |
| if self.pending_normal_data is None: |
| self.pending_normal_data = processed_normal |
| else: |
| self.pending_normal_data = self._merge_normal_data(self.pending_normal_data, processed_normal) |
| |
| # Write blob-file columns to dedicated blob writers. |
| for blob_column, blob_data in blob_data_map.items(): |
| if blob_data is not None and blob_data.num_rows > 0: |
| self.blob_writers[blob_column].write(blob_data) |
| |
| # Write vector columns to dedicated vector writer. |
| if self.vector_writer is not None and vector_data is not None and vector_data.num_rows > 0: |
| self.vector_writer.write(vector_data) |
| |
| self.record_count += data.num_rows |
| |
| # Check if normal data rolling is needed |
| if self._should_roll_normal(): |
| # When normal data rolls, close both writers and fetch blob metadata |
| self._close_current_writers() |
| |
| except Exception as e: |
| logger.error("Exception occurs when writing data. Cleaning up.", exc_info=e) |
| self.abort() |
| raise e |
| |
| def prepare_commit(self) -> List[DataFileMeta]: |
| # Close any remaining data |
| self._close_current_writers() |
| |
| return self.committed_files.copy() |
| |
| def close(self): |
| if self.closed: |
| return |
| |
| try: |
| self._close_current_writers() |
| if self._external_storage_writer: |
| self._external_storage_writer.close() |
| except Exception as e: |
| logger.error("Exception occurs when closing writer. Cleaning up.", exc_info=e) |
| self.abort() |
| raise |
| finally: |
| self.closed = True |
| self.pending_normal_data = None |
| |
| def abort(self): |
| """Abort all writers and clean up resources.""" |
| for blob_writer in self.blob_writers.values(): |
| blob_writer.abort() |
| if self.vector_writer is not None: |
| self.vector_writer.abort() |
| if self._external_storage_writer: |
| self._external_storage_writer.abort() |
| committed_non_blob_files = [ |
| file_meta for file_meta in self.committed_files |
| if not DataFileMeta.is_blob_file(file_meta.file_name) |
| ] |
| self._delete_committed_files(committed_non_blob_files) |
| self.pending_normal_data = None |
| self.pending_data = None |
| self.committed_files.clear() |
| |
| def _split_data(self, data: pa.RecordBatch) -> Tuple[ |
| pa.RecordBatch, Dict[str, pa.RecordBatch], Optional[pa.RecordBatch]]: |
| """Split data into normal, blob, and vector parts based on column names.""" |
| normal_data = data.select(self.normal_column_names) if self.normal_column_names else None |
| blob_data_map = { |
| blob_column: data.select([blob_column]) for blob_column in self.blob_file_column_names |
| } |
| vector_data = ( |
| pa.RecordBatch.from_arrays( |
| [data.column(name) for name in self.vector_write_columns], |
| names=self.vector_write_columns, |
| ) if self.vector_write_columns else None |
| ) |
| return normal_data, blob_data_map, vector_data |
| |
| def _validate_inline_stored_fields_input(self, data: pa.RecordBatch): |
| if not self.blob_inline_fields: |
| return |
| |
| from pypaimon.table.row.blob import BlobDescriptor, BlobViewStruct |
| |
| for field_name in self.blob_descriptor_fields: |
| if field_name not in data.schema.names: |
| continue |
| values = data.column(data.schema.get_field_index(field_name)).to_pylist() |
| for value in values: |
| if value is None: |
| continue |
| if hasattr(value, 'as_py'): |
| value = value.as_py() |
| if isinstance(value, str): |
| value = value.encode('utf-8') |
| if not isinstance(value, (bytes, bytearray)): |
| raise ValueError( |
| "blob-descriptor-field requires blob field value to be a serialized " |
| "BlobDescriptor." |
| ) |
| try: |
| descriptor_bytes = bytes(value) |
| descriptor = BlobDescriptor.deserialize(descriptor_bytes) |
| if descriptor.serialize() != descriptor_bytes: |
| raise ValueError("Descriptor payload contains trailing bytes.") |
| except Exception as e: |
| raise ValueError( |
| "blob-descriptor-field requires blob field value to be a serialized " |
| "BlobDescriptor." |
| ) from e |
| |
| for field_name in self.blob_view_fields: |
| if field_name not in data.schema.names: |
| continue |
| values = data.column(data.schema.get_field_index(field_name)).to_pylist() |
| for value in values: |
| if value is None: |
| continue |
| if hasattr(value, 'as_py'): |
| value = value.as_py() |
| if isinstance(value, str): |
| value = value.encode('utf-8') |
| if not isinstance(value, (bytes, bytearray)): |
| raise ValueError( |
| "blob-view-field requires blob field value to be a serialized " |
| "BlobViewStruct." |
| ) |
| try: |
| view_bytes = bytes(value) |
| view_struct = BlobViewStruct.deserialize(view_bytes) |
| if view_struct.serialize() != view_bytes: |
| raise ValueError("BlobViewStruct payload contains trailing bytes.") |
| except Exception as e: |
| raise ValueError( |
| "blob-view-field requires blob field value to be a serialized " |
| "BlobViewStruct." |
| ) from e |
| |
| @staticmethod |
| def _process_normal_data(data: pa.RecordBatch) -> Optional[pa.Table]: |
| """Process normal data (similar to base DataWriter).""" |
| if data is None or data.num_rows == 0: |
| return None |
| return pa.Table.from_batches([data]) |
| |
| @staticmethod |
| def _merge_normal_data(existing_data: pa.Table, new_data: pa.Table) -> pa.Table: |
| return pa.concat_tables([existing_data, new_data]) |
| |
| def _should_roll_normal(self) -> bool: |
| if self.pending_normal_data is None: |
| return False |
| |
| # Check rolling condition periodically (every CHECK_ROLLING_RECORD_CNT records) |
| if self.record_count % self.CHECK_ROLLING_RECORD_CNT != 0: |
| return False |
| |
| # Check if normal data exceeds target size |
| current_size = self.pending_normal_data.nbytes |
| return current_size > self.target_file_size |
| |
| def _close_current_writers(self): |
| """Close normal, blob, and vector writers; add metadata in order: normal, blob, vector.""" |
| normal_meta = None |
| if self.pending_normal_data is not None and self.pending_normal_data.num_rows > 0: |
| normal_meta = self._write_normal_data_to_file(self.pending_normal_data) |
| self.committed_files.append(normal_meta) |
| |
| blob_metas = [] |
| for blob_column in self.blob_file_column_names: |
| writer_metas = self.blob_writers[blob_column].prepare_commit() |
| if normal_meta is not None: |
| self._validate_consistency(normal_meta, writer_metas, blob_column) |
| blob_metas.extend(writer_metas) |
| self.committed_files.extend(blob_metas) |
| |
| vector_metas = [] |
| if self.vector_writer is not None: |
| vector_metas = self.vector_writer.prepare_commit() |
| if vector_metas and normal_meta is not None: |
| self._validate_consistency(normal_meta, vector_metas, 'vector') |
| self.committed_files.extend(vector_metas) |
| self.vector_writer.committed_files.clear() |
| |
| self.pending_normal_data = None |
| |
| if normal_meta is not None or blob_metas or vector_metas: |
| normal_name = normal_meta.file_name if normal_meta is not None else '<none>' |
| logger.info(f"Closed writers - normal: {normal_name}, " |
| f"{len(blob_metas)} blob metas, {len(vector_metas)} vector metas") |
| |
| def _write_normal_data_to_file(self, data: pa.Table) -> Optional[DataFileMeta]: |
| if data.num_rows == 0: |
| return None |
| |
| file_name = f"{CoreOptions.data_file_prefix(self.options)}{uuid.uuid4()}-0.{self.file_format}" |
| file_path = self._generate_file_path(file_name) |
| |
| # Write file based on format |
| if self.file_format == CoreOptions.FILE_FORMAT_PARQUET: |
| self.file_io.write_parquet(file_path, data, compression=self.compression, zstd_level=self.zstd_level) |
| elif self.file_format == CoreOptions.FILE_FORMAT_ORC: |
| self.file_io.write_orc(file_path, data, compression=self.compression, zstd_level=self.zstd_level) |
| elif self.file_format == CoreOptions.FILE_FORMAT_AVRO: |
| self.file_io.write_avro(file_path, data, compression=self.compression, zstd_level=self.zstd_level) |
| elif self.file_format == CoreOptions.FILE_FORMAT_LANCE: |
| self.file_io.write_lance(file_path, data) |
| elif self.file_format == CoreOptions.FILE_FORMAT_VORTEX: |
| self.file_io.write_vortex(file_path, data) |
| elif self.file_format == CoreOptions.FILE_FORMAT_MOSAIC: |
| self.file_io.write_mosaic(file_path, data) |
| elif self.file_format == CoreOptions.FILE_FORMAT_ROW: |
| self.file_io.write_row(file_path, data, zstd_level=self.zstd_level) |
| else: |
| raise ValueError(f"Unsupported file format: {self.file_format}") |
| |
| # Determine if this is an external path |
| is_external_path = self.external_path_provider is not None |
| external_path_str = file_path if is_external_path else None |
| |
| return self._create_data_file_meta(file_name, file_path, data, external_path_str) |
| |
| def _create_data_file_meta(self, file_name: str, file_path: str, data: pa.Table, |
| external_path: Optional[str] = None) -> DataFileMeta: |
| # Column stats (only for normal columns) |
| metadata_stats_enabled = self.options.metadata_stats_enabled() |
| stats_columns = self.normal_columns if metadata_stats_enabled else [] |
| value_stats = self._collect_value_stats(data, stats_columns) |
| |
| min_seq, max_seq = self._append_file_sequence_range(data.num_rows) |
| |
| return DataFileMeta.create( |
| file_name=file_name, |
| file_size=self.file_io.get_file_size(file_path), |
| row_count=data.num_rows, |
| min_key=GenericRow([], []), |
| max_key=GenericRow([], []), |
| key_stats=SimpleStats.empty_stats(), |
| value_stats=value_stats, |
| min_sequence_number=min_seq, |
| max_sequence_number=max_seq, |
| schema_id=self.table.table_schema.id, |
| level=0, |
| extra_files=[], |
| creation_time=Timestamp.now(), |
| delete_row_count=0, |
| file_source=0, |
| value_stats_cols=[column.name for column in stats_columns], |
| external_path=external_path, |
| file_path=file_path, |
| write_cols=self.write_cols) |
| |
| def _validate_consistency( |
| self, normal_meta: DataFileMeta, blob_metas: List[DataFileMeta], blob_column: str): |
| if normal_meta is None: |
| return |
| |
| normal_row_count = normal_meta.row_count |
| blob_row_count = sum(meta.row_count for meta in blob_metas) |
| |
| if normal_row_count != blob_row_count: |
| raise RuntimeError( |
| f"This is a bug: The row count of main file and blob files does not match. " |
| f"Main file: {normal_meta.file_name} (row count: {normal_row_count}), " |
| f"blob field: {blob_column}, " |
| f"blob files: {[meta.file_name for meta in blob_metas]} (total row count: {blob_row_count})" |
| ) |
