This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 02c79da229 [python] Consolidate composite writer lifecycle and file
output (#9948)
02c79da229 is described below
commit 02c79da229a031d20f9f136ae8e6d01b0472d285
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):
