| # 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 List, Optional, Tuple |
| |
| import pyarrow as pa |
| |
| from pypaimon.common.options.core_options import CoreOptions |
| 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.generic_row import GenericRow |
| from pypaimon.write.writer.data_writer import DataWriter |
| |
| logger = logging.getLogger(__name__) |
| |
| |
| class DataVectorWriter(DataWriter): |
| """A rolling file writer that stores vector columns separately from normal columns. |
| |
| All vector columns are written to a single `.vector.<format>` file (matching |
| Java behavior), while normal columns go to standard data files. |
| |
| Metadata organization: |
| committed_files = [ |
| normal_file_meta, # data.parquet (id, label columns) |
| vector_file_meta, # data.vector.lance (all vector columns) |
| ... |
| ] |
| """ |
| |
| 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): |
| super().__init__(table, partition, bucket, max_seq_number, options, write_cols=write_cols) |
| |
| self.vector_column_names = self._get_vector_columns_from_schema() |
| self.vector_file_format = options.vector_file_format() |
| |
| all_column_names = self.table.field_names |
| vector_set = set(self.vector_column_names) |
| |
| if write_cols is not None: |
| write_col_set = set(write_cols) |
| self.vector_write_columns = [ |
| col for col in self.vector_column_names if col in write_col_set |
| ] |
| self.normal_column_names = [ |
| col for col in write_cols if col not in vector_set |
| ] |
| else: |
| self.vector_write_columns = list(self.vector_column_names) |
| self.normal_column_names = [ |
| col for col in all_column_names if col not in vector_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 |
| |
| self.record_count = 0 |
| self.closed = False |
| self.pending_normal_data: Optional[pa.Table] = None |
| |
| 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=self.vector_file_format, |
| options=options, |
| ) |
| |
| 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 pa.concat_tables([existing_data, new_data]) |
| |
| def write(self, data: pa.RecordBatch): |
| try: |
| normal_data, vector_data = self._split_data(data) |
| |
| processed_normal = pa.Table.from_batches([normal_data]) if normal_data is not None else None |
| if self.pending_normal_data is None: |
| self.pending_normal_data = processed_normal |
| elif processed_normal is not None: |
| self.pending_normal_data = pa.concat_tables([self.pending_normal_data, processed_normal]) |
| |
| 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 |
| |
| if self._should_roll_normal(): |
| 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]: |
| self._close_current_writers() |
| return self.committed_files.copy() |
| |
| def close(self): |
| if self.closed: |
| return |
| try: |
| self._close_current_writers() |
| 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): |
| if self.vector_writer is not None: |
| self.vector_writer.abort() |
| self.pending_normal_data = None |
| super().abort() |
| |
| def _split_data(self, data: pa.RecordBatch) -> Tuple[pa.RecordBatch, pa.RecordBatch]: |
| normal_data = ( |
| pa.RecordBatch.from_arrays( |
| [data.column(name) for name in self.normal_column_names], |
| names=self.normal_column_names, |
| ) |
| if self.normal_column_names else None |
| ) |
| 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, vector_data |
| |
| def _should_roll_normal(self) -> bool: |
| if self.pending_normal_data is None: |
| return False |
| if self.record_count % self.CHECK_ROLLING_RECORD_CNT != 0: |
| return False |
| return self.pending_normal_data.nbytes > self.target_file_size |
| |
| def _close_current_writers(self): |
| has_normal = self.pending_normal_data is not None and self.pending_normal_data.num_rows > 0 |
| |
| normal_meta = None |
| if has_normal: |
| normal_meta = self._write_normal_data_to_file(self.pending_normal_data) |
| self.committed_files.append(normal_meta) |
| |
| if self.vector_writer is not None: |
| vector_metas = self.vector_writer.prepare_commit() |
| if vector_metas: |
| if normal_meta is not None: |
| self._validate_consistency(normal_meta, vector_metas) |
| self.committed_files.extend(vector_metas) |
| self.vector_writer.committed_files.clear() |
| |
| self.pending_normal_data = None |
| |
| 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) |
| |
| 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, options=self.mosaic_writer_options) |
| 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}") |
| |
| is_external_path = self.external_path_provider is not None |
| external_path_str = file_path if is_external_path else None |
| |
| 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_str, |
| file_path=file_path, |
| write_cols=self.write_cols, |
| ) |
| |
| def _validate_consistency( |
| self, normal_meta: DataFileMeta, vector_metas: List[DataFileMeta]): |
| if normal_meta is None: |
| return |
| normal_row_count = normal_meta.row_count |
| vector_row_count = sum(meta.row_count for meta in vector_metas) |
| if normal_row_count != vector_row_count: |
| raise RuntimeError( |
| f"Row count mismatch between main file and vector files. " |
| f"Main file: {normal_meta.file_name} (rows: {normal_row_count}), " |
| f"vector files: {[m.file_name for m in vector_metas]} (rows: {vector_row_count})" |
| ) |
