This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit 23f5644052b0a8d09a3fc142a9da4e4e2bc382e3 Author: zhigang <[email protected]> AuthorDate: Sat Sep 19 22:05:31 2026 +0800 [python] Consolidate composite writer lifecycle and file output (#9948) --- .../pypaimon/tests/write/write_buffer_test.py | 17 +- .../pypaimon/write/writer/composite_data_writer.py | 186 +++++++++++++++++++ .../pypaimon/write/writer/data_vector_writer.py | 152 +--------------- paimon-python/pypaimon/write/writer/data_writer.py | 89 +++++----- .../write/writer/dedicated_format_writer.py | 197 ++------------------- 5 files changed, 263 insertions(+), 378 deletions(-) diff --git a/paimon-python/pypaimon/tests/write/write_buffer_test.py b/paimon-python/pypaimon/tests/write/write_buffer_test.py index 39f49b3d64..0c539d6282 100644 --- a/paimon-python/pypaimon/tests/write/write_buffer_test.py +++ b/paimon-python/pypaimon/tests/write/write_buffer_test.py @@ -33,6 +33,7 @@ import pyarrow as pa from pypaimon.write.writer.append_only_data_writer import AppendOnlyDataWriter from pypaimon.write.writer.data_vector_writer import DataVectorWriter +from pypaimon.write.writer.data_writer import DataWriter from pypaimon.write.writer.dedicated_format_writer import DedicatedFormatWriter from pypaimon.write.writer.write_buffer import WriteBuffer @@ -456,6 +457,8 @@ class FlushFailureTest(unittest.TestCase): self.vector_writer = None self._normal_buffer = WriteBuffer(self._merge_data) self.committed_files = [] + self._committed_files_to_delete_on_abort = [] + self._pending_normal_meta = None self.written = [] self.fail_next = True @@ -528,6 +531,9 @@ class _StubSidecarWriter: _StubMeta(self._row_count, self._file_name)) return self.committed_files.copy() + def _release_prepared_files(self): + return DataWriter._release_prepared_files(self) + def delete_file_upon_abort(self): return self._delete_on_abort @@ -563,6 +569,8 @@ class CompositeFlushResumeTest(unittest.TestCase): self._normal_buffer = WriteBuffer(self._merge_data) self._buffer = WriteBuffer(self._merge_data) self.committed_files = [] + self._committed_files_to_delete_on_abort = [] + self._pending_normal_meta = None self.committed_changelog_files = [] self.file_io = _RecordingFileIO() self.written = [] @@ -586,7 +594,9 @@ class CompositeFlushResumeTest(unittest.TestCase): self._normal_buffer = WriteBuffer(self._merge_normal_data) self._buffer = WriteBuffer(self._merge_normal_data) self.committed_files = [] + self.committed_changelog_files = [] self._committed_files_to_delete_on_abort = [] + self._pending_normal_meta = None self.file_io = _RecordingFileIO() self.written = [] self._video_group_policy = None @@ -703,9 +713,14 @@ class CompositeFlushResumeTest(unittest.TestCase): def test_dedicated_writer_keeps_the_documented_meta_order(self): blob = _StubSidecarWriter(3, 'blob-0') - vector = _StubSidecarWriter(3, 'vector-0') + vector = _StubSidecarWriter(3, 'vector-0', fail_times=1) writer = self._DedicatedHarness({'payload': blob}, vector) writer._normal_buffer.append(_table(0, 3)) + with self.assertRaises(IOError): + writer._close_current_writers() + self.assertEqual(writer.committed_files, []) + self.assertEqual([m.file_name for m in blob.committed_files], ['blob-0']) + writer._close_current_writers() self.assertEqual([m.file_name for m in writer.committed_files], ['data-1', 'blob-0', 'vector-0']) diff --git a/paimon-python/pypaimon/write/writer/composite_data_writer.py b/paimon-python/pypaimon/write/writer/composite_data_writer.py new file mode 100644 index 0000000000..9cf8abf29c --- /dev/null +++ b/paimon-python/pypaimon/write/writer/composite_data_writer.py @@ -0,0 +1,186 @@ +# 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 abc import abstractmethod +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.table.row.generic_row import GenericRow +from pypaimon.write.writer.data_writer import DataWriter +from pypaimon.write.writer.write_buffer import WriteBuffer + + +logger = logging.getLogger(__name__) + + +class CompositeDataWriter(DataWriter): + """Shared metadata handoff for normal files and blob/vector sidecars.""" + + CHECK_ROLLING_RECORD_CNT = 1000 + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.record_count = 0 + self.closed = False + self._normal_buffer = WriteBuffer(self._merge_data) + self._pending_normal_meta: Optional[DataFileMeta] = None + self._committed_files_to_delete_on_abort: List[DataFileMeta] = [] + + 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._normal_buffer.reset() + + def _should_roll_normal(self) -> bool: + # Runs on every write, so it answers from the running counts only. + if self._normal_buffer.is_empty: + return False + if self._normal_buffer.num_rows >= self.target_file_row_num: + return True + if self.record_count % self.CHECK_ROLLING_RECORD_CNT != 0: + return False + return self._normal_buffer.nbytes > self.target_file_size + + def _require_finished_flush(self): + """Refuse to buffer more rows while a flush is only half done. + + A composite flush writes the normal data file first and the sidecars + after. Once that file is on disk it covers exactly the rows flushed so + far, so rows appended before the retry finishes would belong to no file + the resumed flush writes. + """ + if self._pending_normal_meta is not None: + raise RuntimeError( + "Cannot write: a previous flush left a data file that no commit " + "has taken yet. Retry prepare_commit() to finish that flush, or " + "abort() this writer.") + + def _close_current_writers(self): + normal_meta = self._pending_normal_meta + if normal_meta is None: + normal_data = self._normal_buffer.materialize() + if normal_data is not None and normal_data.num_rows > 0: + normal_meta = self._write_normal_data_to_file(normal_data) + self._pending_normal_meta = normal_meta + # The parent owns this file as soon as it lands, even if a + # sidecar later fails and nothing can be published yet. + self._committed_files_to_delete_on_abort.append(normal_meta) + self._normal_buffer.reset() + + # Prepare and validate all sidecars before publishing any metadata. + # On failure children retain their metadata for retry and abort; the + # normal file is remembered so a retry never writes those rows twice. + prepared = self._prepare_sidecar_commits(normal_meta) + if normal_meta is not None: + self.committed_files.append(normal_meta) + for _, metas in prepared: + self.committed_files.extend(metas) + for writer, _ in prepared: + self._committed_files_to_delete_on_abort.extend(writer._release_prepared_files()) + + self._pending_normal_meta = None + self.record_count = 0 + + @abstractmethod + def _prepare_sidecar_commits( + self, normal_meta: Optional[DataFileMeta] + ) -> List[Tuple[DataWriter, List[DataFileMeta]]]: + """Return all children and metadata in order after preparation and validation. + + Children retain their metadata and cleanup responsibility until the parent + accepts the complete result. + """ + + 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) + + try: + if self.file_format == CoreOptions.FILE_FORMAT_BLOB: + raise ValueError(f"Unsupported file format: {self.file_format}") + shredding_stats = self._write_file(file_path, data, 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) + + meta = 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, + ) + self._map_shared_shredding.file_completed(shredding_stats) + return meta + except Exception: + self.file_io.delete_quietly(file_path) + raise + + def abort(self): + # Subclasses abort children first to clean up files not yet handed off. + # This list holds normal files and sidecars whose cleanup was transferred + # to the parent; externally owned BlobConsumer files are excluded. + self._delete_committed_files( + self._committed_files_to_delete_on_abort + self.committed_changelog_files) + self._pending_normal_meta = None + self._normal_buffer.reset() + self._buffer.reset() + self.committed_files.clear() + self.committed_changelog_files.clear() + self._committed_files_to_delete_on_abort.clear() diff --git a/paimon-python/pypaimon/write/writer/data_vector_writer.py b/paimon-python/pypaimon/write/writer/data_vector_writer.py index 63f5cbc136..5d9fa08225 100644 --- a/paimon-python/pypaimon/write/writer/data_vector_writer.py +++ b/paimon-python/pypaimon/write/writer/data_vector_writer.py @@ -16,24 +16,19 @@ # 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 -from pypaimon.write.writer.write_buffer import WriteBuffer +from pypaimon.write.writer.composite_data_writer import CompositeDataWriter logger = logging.getLogger(__name__) -class DataVectorWriter(DataWriter): +class DataVectorWriter(CompositeDataWriter): """A rolling file writer that stores vector columns separately from normal columns. All vector columns are written to a single `.vector.<format>` file (matching @@ -47,8 +42,6 @@ class DataVectorWriter(DataWriter): ] """ - 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) @@ -88,16 +81,6 @@ class DataVectorWriter(DataWriter): else self.normal_column_names ) - self.record_count = 0 - self.closed = False - # Normal columns are buffered separately from the vector columns, which - # the vector writer owns. - self._normal_buffer = WriteBuffer(self._merge_data) - # A normal data file that landed while a later phase of the same flush - # failed. Held so the retry resumes at that phase instead of writing the - # rows a second time. - self._pending_normal_meta: Optional[DataFileMeta] = None - from pypaimon.write.writer.vector_writer import VectorWriter self.vector_writer: Optional[VectorWriter] = None if self.vector_write_columns: @@ -161,27 +144,9 @@ class DataVectorWriter(DataWriter): if self._should_roll_normal(): self._close_current_writers() - 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._normal_buffer.reset() - def abort(self): if self.vector_writer is not None: self.vector_writer.abort() - self._normal_buffer.reset() super().abort() def _split_data(self, data: pa.RecordBatch) -> Tuple[pa.RecordBatch, pa.RecordBatch]: @@ -201,16 +166,6 @@ class DataVectorWriter(DataWriter): ) return normal_data, vector_data - def _should_roll_normal(self) -> bool: - # Runs on every write, so it answers from the running counts only. - if self._normal_buffer.is_empty: - return False - if self._normal_buffer.num_rows >= self.target_file_row_num: - return True - if self.record_count % self.CHECK_ROLLING_RECORD_CNT != 0: - return False - return self._normal_buffer.nbytes > self.target_file_size - @property def pending_row_count(self) -> int: # Overrides the base property, which reads a buffer this writer never @@ -224,102 +179,13 @@ class DataVectorWriter(DataWriter): return self.vector_writer.pending_row_count return 0 - def _close_current_writers(self): - # A flush spans the normal file and the vector sidecars, and the vector - # writer drains its own buffer as it goes, so its half cannot be replayed - # from scratch. Two rules make a retry resume rather than restart: the - # normal rows stay buffered until their file lands, and once it has - # landed the file is remembered instead of the rows. Nothing reaches - # ``committed_files`` until every phase has succeeded, so a retry never - # finds a half-published flush. - normal_meta = self._pending_normal_meta - if normal_meta is None: - normal_data = self._normal_buffer.materialize() - if normal_data is not None and normal_data.num_rows > 0: - normal_meta = self._write_normal_data_to_file(normal_data) - self._pending_normal_meta = normal_meta - self._normal_buffer.reset() - - 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) - - if normal_meta is not None: - self.committed_files.append(normal_meta) - self.committed_files.extend(vector_metas) - if self.vector_writer is not None: - # Cleared only now: until the flush completes, a retry has to be able - # to harvest the same metas again. - self.vector_writer.committed_files.clear() - - self._pending_normal_meta = None - self.record_count = 0 - - def _write_normal_data_to_file(self, data: pa.Table) -> Optional[DataFileMeta]: - if data.num_rows == 0: - return None - - shredding_stats = {} - - file_name = f"{CoreOptions.data_file_prefix(self.options)}{uuid.uuid4()}-0.{self.file_format}" - file_path = self._generate_file_path(file_name) - - # Until metadata is returned, no caller can track this file for abort. - try: - if self.file_format == CoreOptions.FILE_FORMAT_PARQUET: - shredding_stats = self._write_parquet_data(file_path, data) - 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) - - meta = 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, - ) - self._map_shared_shredding.file_completed(shredding_stats) - return meta - except Exception: - self.file_io.delete_quietly(file_path) - raise + def _prepare_sidecar_commits(self, normal_meta): + if self.vector_writer is None: + return [] + vector_metas = self.vector_writer.prepare_commit() + if vector_metas and normal_meta is not None: + self._validate_consistency(normal_meta, vector_metas) + return [(self.vector_writer, vector_metas)] def _validate_consistency( self, normal_meta: DataFileMeta, vector_metas: List[DataFileMeta]): diff --git a/paimon-python/pypaimon/write/writer/data_writer.py b/paimon-python/pypaimon/write/writer/data_writer.py index 71f5004f7e..79de18f606 100644 --- a/paimon-python/pypaimon/write/writer/data_writer.py +++ b/paimon-python/pypaimon/write/writer/data_writer.py @@ -113,10 +113,6 @@ class DataWriter(ABC): if self.changelog_producer == ChangelogProducer.INPUT else None, ) - # Set by the composite writers when a flush landed its normal data file but a - # later phase of the same flush failed; see their ``_close_current_writers``. - _pending_normal_meta: Optional[DataFileMeta] = None - @property def pending_row_count(self) -> int: """Rows held for the file being written, not yet in a finished file. @@ -127,20 +123,6 @@ class DataWriter(ABC): """ return self._buffer.num_rows - def _require_finished_flush(self): - """Refuse to buffer more rows while a flush is only half done. - - A composite flush writes the normal data file first and the sidecars - after. Once that file is on disk it covers exactly the rows flushed so - far, so rows appended before the retry finishes would belong to no file - the resumed flush writes. - """ - if self._pending_normal_meta is not None: - raise RuntimeError( - "Cannot write: a previous flush left a data file that no commit " - "has taken yet. Retry prepare_commit() to finish that flush, or " - "abort() this writer.") - def write(self, data: pa.RecordBatch): try: processed_data = self._process_data(data) @@ -162,6 +144,19 @@ class DataWriter(ABC): return self.committed_files.copy() + def _release_prepared_files(self) -> List[DataFileMeta]: + """Clear metadata accepted by the parent and return files it must delete on abort. + + Call only after all sidecars are prepared and validated. Until then, this + writer retains its metadata for retry and its responsibility for cleanup. + """ + owned_files = self.committed_files.copy() if self.delete_file_upon_abort() else [] + self.committed_files.clear() + return owned_files + + def delete_file_upon_abort(self) -> bool: + return True + def prepare_changelog_commit(self) -> List[DataFileMeta]: return self.committed_changelog_files.copy() @@ -185,10 +180,6 @@ class DataWriter(ABC): during writing. It deletes any files that were written and cleans up resources. """ to_delete = self.committed_files + self.committed_changelog_files - if self._pending_normal_meta is not None: - # No list tracks this one: it landed but its flush never published it. - to_delete.append(self._pending_normal_meta) - self._pending_normal_meta = None self._delete_committed_files(to_delete) # Clean up resources @@ -254,6 +245,28 @@ class DataWriter(ABC): self._write_data_to_file(pending.slice(0, split_row)) self._buffer.reset(pending.slice(split_row)) + def _write_file(self, file_path: str, data: pa.Table, file_format: str) -> Dict: + """Write contents without publishing metadata; return shredding statistics.""" + if file_format == CoreOptions.FILE_FORMAT_PARQUET: + return self._write_parquet_data(file_path, data) + elif file_format == CoreOptions.FILE_FORMAT_ORC: + self.file_io.write_orc(file_path, data, compression=self.compression, zstd_level=self.zstd_level) + elif file_format == CoreOptions.FILE_FORMAT_AVRO: + self.file_io.write_avro(file_path, data, compression=self.compression, zstd_level=self.zstd_level) + elif file_format == CoreOptions.FILE_FORMAT_BLOB: + self.file_io.write_blob(file_path, data) + elif file_format == CoreOptions.FILE_FORMAT_LANCE: + self.file_io.write_lance(file_path, data) + elif file_format == CoreOptions.FILE_FORMAT_VORTEX: + self.file_io.write_vortex(file_path, data) + elif file_format == CoreOptions.FILE_FORMAT_MOSAIC: + self.file_io.write_mosaic(file_path, data, options=self.mosaic_writer_options) + elif 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: {file_format}") + return {} + def _write_data_to_file(self, data: pa.Table): if data.num_rows == 0: return @@ -278,24 +291,7 @@ class DataWriter(ABC): # exists would make the retry write a second copy of rows the first meta # already covers. try: - if self.file_format == CoreOptions.FILE_FORMAT_PARQUET: - shared_shredding_stats = self._write_parquet_data(file_path, data) - 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_BLOB: - self.file_io.write_blob(file_path, data) - 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}") + shared_shredding_stats = self._write_file(file_path, data, self.file_format) if self._should_write_row_sidecar(): row_sidecar_name = f"{file_name}{self.ROW_SIDECAR_SUFFIX}" @@ -451,17 +447,12 @@ class DataWriter(ABC): changelog_external_path = changelog_file_path if is_external else None try: - if cl_fmt == CoreOptions.FILE_FORMAT_PARQUET: - self._write_parquet_data(changelog_file_path, data) - elif cl_fmt == CoreOptions.FILE_FORMAT_ORC: - self.file_io.write_orc(changelog_file_path, data, compression=self.compression, - zstd_level=self.zstd_level) - elif cl_fmt == CoreOptions.FILE_FORMAT_AVRO: - self.file_io.write_avro(changelog_file_path, data, compression=self.compression, - zstd_level=self.zstd_level) - else: + if cl_fmt not in (CoreOptions.FILE_FORMAT_PARQUET, + CoreOptions.FILE_FORMAT_ORC, + CoreOptions.FILE_FORMAT_AVRO): raise ValueError(f"Unsupported changelog file format: {cl_fmt}. " f"Supported formats: parquet, orc, avro.") + self._write_file(changelog_file_path, data, cl_fmt) return DataFileMeta.create( file_name=changelog_file_name, diff --git a/paimon-python/pypaimon/write/writer/dedicated_format_writer.py b/paimon-python/pypaimon/write/writer/dedicated_format_writer.py index 22dc87be0e..98fc208bef 100644 --- a/paimon-python/pypaimon/write/writer/dedicated_format_writer.py +++ b/paimon-python/pypaimon/write/writer/dedicated_format_writer.py @@ -16,15 +16,12 @@ # 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 ( PyarrowFieldParser, VectorType, @@ -36,20 +33,18 @@ from pypaimon.table.row.blob import ( BlobConsumer, video_payload_descriptor, ) -from pypaimon.table.row.generic_row import GenericRow from pypaimon.write.row_utils import ( require_columns, row_to_named_values, row_values_to_arrow_table, ) -from pypaimon.write.writer.data_writer import DataWriter +from pypaimon.write.writer.composite_data_writer import CompositeDataWriter from pypaimon.write.writer.video_group import VideoGroupRollingPolicy -from pypaimon.write.writer.write_buffer import WriteBuffer logger = logging.getLogger(__name__) -class DedicatedFormatWriter(DataWriter): +class DedicatedFormatWriter(CompositeDataWriter): """A rolling file writer that writes normal, blob, and vector columns to dedicated files. Splits incoming data three ways: @@ -63,9 +58,6 @@ class DedicatedFormatWriter(DataWriter): [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, @@ -151,24 +143,12 @@ class DedicatedFormatWriter(DataWriter): else self.normal_column_names ) - # State management for blob writer - self.record_count = 0 - self.closed = False self._video_group_policy = ( VideoGroupRollingPolicy() if self.video_frame_columns else None ) - # Normal columns are buffered separately from the blob and vector - # columns, which their own writers own. - self._normal_buffer = WriteBuffer(self._merge_normal_data) - self._committed_files_to_delete_on_abort: List[DataFileMeta] = [] - # A normal data file that landed while a later phase of the same flush - # failed. Held so the retry resumes at that phase instead of writing the - # rows a second time. - self._pending_normal_meta: Optional[DataFileMeta] = None - # Initialize blob writers for each blob-file column. from pypaimon.write.writer.blob_writer import BlobWriter self.blob_writers: Dict[str, BlobWriter] = {} @@ -357,40 +337,13 @@ class DedicatedFormatWriter(DataWriter): return value - 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() - except Exception as e: - logger.error("Exception occurs when closing writer. Cleaning up.", exc_info=e) - self.abort() - raise - finally: - self.closed = True - self._normal_buffer.reset() - 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() - # An unpublished normal file is already in the delete list, added when it - # landed, so there is nothing left to resume. - self._pending_normal_meta = None - self._delete_committed_files(self._committed_files_to_delete_on_abort) - self._normal_buffer.reset() - self._buffer.reset() - self.committed_files.clear() - self._committed_files_to_delete_on_abort.clear() + super().abort() def _split_data(self, data: pa.RecordBatch) -> Tuple[ Optional[pa.RecordBatch], Dict[str, pa.RecordBatch], Optional[pa.RecordBatch]]: @@ -508,21 +461,6 @@ class DedicatedFormatWriter(DataWriter): 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: - # Runs on every write, so it answers from the running counts only. - if self._normal_buffer.is_empty: - return False - - if self._normal_buffer.num_rows >= self.target_file_row_num: - return True - - # 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 - return self._normal_buffer.nbytes > self.target_file_size - def roll_before_group_if_needed(self, row_count: int): """Roll current files before the next logical write group if needed.""" self._require_finished_flush() @@ -638,136 +576,25 @@ class DedicatedFormatWriter(DataWriter): return 0 def _close_current_writers(self): - """Close normal, blob, and vector writers; add metadata in order: normal, blob, vector.""" - # A flush spans the normal file and every blob/vector sidecar, and the - # sidecar writers drain their own buffers as they go, so their half cannot - # be replayed from scratch. Two rules make a retry resume rather than - # restart: the normal rows stay buffered until their file lands, and once - # it has landed the file is remembered instead of the rows. Nothing - # reaches ``committed_files`` until every phase has succeeded, so a retry - # never finds a half-published flush. - normal_meta = self._pending_normal_meta - if normal_meta is None: - normal_data = self._normal_buffer.materialize() - if normal_data is not None and normal_data.num_rows > 0: - normal_meta = self._write_normal_data_to_file(normal_data) - self._pending_normal_meta = normal_meta - # Tracked for abort right away: until the flush publishes it, this - # file is in no other list. - self._committed_files_to_delete_on_abort.append(normal_meta) - self._normal_buffer.reset() - - blob_metas = [] - deletable_blob_metas = [] + super()._close_current_writers() + if self._video_group_policy is not None: + self._video_group_policy.reset() + + def _prepare_sidecar_commits(self, normal_meta): + prepared = [] for blob_column in self.blob_file_column_names: blob_writer = self.blob_writers[blob_column] writer_metas = blob_writer.prepare_commit() if normal_meta is not None: self._validate_consistency(normal_meta, writer_metas, blob_column) - blob_metas.extend(writer_metas) - if blob_writer.delete_file_upon_abort(): - deletable_blob_metas.extend(writer_metas) + prepared.append((blob_writer, writer_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') - - # Every phase landed; publish in order: normal, blob, vector. - if normal_meta is not None: - self.committed_files.append(normal_meta) - self.committed_files.extend(blob_metas) - self.committed_files.extend(vector_metas) - self._committed_files_to_delete_on_abort.extend(deletable_blob_metas) - self._committed_files_to_delete_on_abort.extend(vector_metas) - # The sub-writers' metas are cleared only now: before the flush completes, - # a retry has to be able to harvest the same ones again, and an abort has - # to find them so each sub-writer can apply its own delete policy. - for blob_column in self.blob_file_column_names: - self.blob_writers[blob_column].committed_files.clear() - if self.vector_writer is not None: - self.vector_writer.committed_files.clear() - - self._pending_normal_meta = None - self.record_count = 0 - if self._video_group_policy is not None: - self._video_group_policy.reset() - - 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 - - shredding_stats = {} - - file_name = f"{CoreOptions.data_file_prefix(self.options)}{uuid.uuid4()}-0.{self.file_format}" - file_path = self._generate_file_path(file_name) - - # Until metadata is returned, no caller can track this file for abort. - try: - # Write file based on format - if self.file_format == CoreOptions.FILE_FORMAT_PARQUET: - shredding_stats = self._write_parquet_data(file_path, data) - 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}") - - # 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 - - meta = self._create_data_file_meta(file_name, file_path, data, external_path_str) - self._map_shared_shredding.file_completed(shredding_stats) - return meta - except Exception: - self.file_io.delete_quietly(file_path) - raise - - 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) + prepared.append((self.vector_writer, vector_metas)) + return prepared def _validate_consistency( self, normal_meta: DataFileMeta, blob_metas: List[DataFileMeta], blob_column: str):
