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 8100277d75 [python] Fix quadratic cost of writing multiple batches
(#9375)
8100277d75 is described below
commit 8100277d75f6744013fab6cbedbf1d9c758fef66
Author: Yu Gan <[email protected]>
AuthorDate: Wed Aug 26 11:42:09 2026 +0800
[python] Fix quadratic cost of writing multiple batches (#9375)
---
.../pypaimon/tests/test_write_merge_buffer.py | 116 ++-
.../tests/write/changelog_producer_test.py | 66 ++
.../pypaimon/tests/write/table_write_test.py | 6 +-
.../pypaimon/tests/write/write_buffer_test.py | 801 +++++++++++++++++++++
paimon-python/pypaimon/write/writer/blob_writer.py | 18 +-
.../pypaimon/write/writer/data_vector_writer.py | 84 ++-
paimon-python/pypaimon/write/writer/data_writer.py | 292 +++++---
.../write/writer/dedicated_format_writer.py | 107 ++-
.../pypaimon/write/writer/key_value_data_writer.py | 63 +-
.../pypaimon/write/writer/write_buffer.py | 101 +++
10 files changed, 1419 insertions(+), 235 deletions(-)
diff --git a/paimon-python/pypaimon/tests/test_write_merge_buffer.py
b/paimon-python/pypaimon/tests/test_write_merge_buffer.py
index 031d1a3a02..b3ee177a39 100644
--- a/paimon-python/pypaimon/tests/test_write_merge_buffer.py
+++ b/paimon-python/pypaimon/tests/test_write_merge_buffer.py
@@ -19,9 +19,9 @@
"""Unit tests for ``KeyValueDataWriter`` buffer behaviour.
Covers the fold algorithm (`_merge_pending_by_pk`), the flush lifecycle
-(`_flush_all` empties the buffer + clears pending_data), and the
-roll-write helper (`_roll_write` splits oversized buffers across
-multiple files). Drives a thin harness that bypasses
+(`_flush_all` drains the buffer), and the roll-write helper
+(`_roll_write` splits oversized buffers across multiple files).
+Drives a thin harness that bypasses
``DataWriter.__init__`` so tests can exercise these paths without
spinning up the real catalog/write stack.
"""
@@ -36,6 +36,7 @@ from pypaimon.read.reader.deduplicate_merge_function import \
from pypaimon.read.reader.partial_update_merge_function import \
PartialUpdateMergeFunction
from pypaimon.write.writer.key_value_data_writer import KeyValueDataWriter
+from pypaimon.write.writer.write_buffer import WriteBuffer
# Layout matches what ``KeyValueDataWriter._add_system_fields`` emits:
@@ -87,7 +88,7 @@ class _Harness(KeyValueDataWriter):
# Large enough that ``_check_and_roll_if_needed`` does not
# trigger on its own in tests that don't care about rolling.
self.target_file_size = target_file_size
- self.pending_data = None
+ self._buffer = WriteBuffer(self._merge_data)
self.committed_files = []
self.written_chunks = []
@@ -133,13 +134,13 @@ class WriteMergeBufferTest(unittest.TestCase):
def test_dedupe_nullable_pk_uses_null_safe_equality_and_nulls_first(self):
writer = _Harness(DeduplicateMergeFunction())
- writer.pending_data = pa.Table.from_pylist(
+ writer._buffer.append(pa.Table.from_pylist(
[_row(2, 1, 'two', None),
_row(None, 2, 'null-old', None),
_row(1, 3, 'one', None),
_row(None, 4, 'null-new', None)],
schema=_NULLABLE_PK_SCHEMA,
- )
+ ))
writer._flush_all()
@@ -290,18 +291,18 @@ class WriteMergeBufferTest(unittest.TestCase):
# responsible for sorting before folding, so unsorted input is
# the right stress case.
writer = _Harness(DeduplicateMergeFunction())
- writer.pending_data = pa.Table.from_pylist(
+ writer._buffer.append(pa.Table.from_pylist(
[_row(2, 5, 'B2-new', None),
_row(1, 2, 'A1-mid', None),
_row(1, 1, 'A1-old', None),
_row(2, 4, 'B2-old', None),
_row(1, 3, 'A1-new', None)],
schema=_SCHEMA,
- )
+ ))
writer._flush_all()
# Buffer cleared.
- self.assertIsNone(writer.pending_data)
+ self.assertTrue(writer._buffer.is_empty)
# Exactly one file written (size well under target).
self.assertEqual(len(writer.written_chunks), 1)
flushed = writer.written_chunks[0]
@@ -314,15 +315,14 @@ class WriteMergeBufferTest(unittest.TestCase):
def test_flush_all_on_empty_buffer_is_noop(self):
writer = _Harness(DeduplicateMergeFunction())
- writer.pending_data = None
writer._flush_all()
- self.assertIsNone(writer.pending_data)
+ self.assertTrue(writer._buffer.is_empty)
self.assertEqual(writer.written_chunks, [])
def test_flush_all_clears_buffer_even_when_fold_drops_everything(self):
# MergeFunction that returns None for every group; verifies
- # ``_flush_all`` still resets ``pending_data`` so a subsequent
- # write starts from a clean slate.
+ # ``_flush_all`` still drains the buffer so a subsequent write
+ # starts from a clean slate.
class DropAll:
def reset(self):
pass
@@ -334,12 +334,12 @@ class WriteMergeBufferTest(unittest.TestCase):
return None
writer = _Harness(DropAll())
- writer.pending_data = pa.Table.from_pylist(
+ writer._buffer.append(pa.Table.from_pylist(
[_row(1, 1, 'A', None), _row(1, 2, 'B', None)],
schema=_SCHEMA,
- )
+ ))
writer._flush_all()
- self.assertIsNone(writer.pending_data)
+ self.assertTrue(writer._buffer.is_empty)
self.assertEqual(writer.written_chunks, [])
# -- _roll_write ------------------------------------------------------
@@ -377,6 +377,90 @@ class WriteMergeBufferTest(unittest.TestCase):
for chunk in writer.written_chunks[:-1]:
self.assertLessEqual(chunk.nbytes, target)
+ def test_roll_write_leaves_only_unwritten_rows_when_a_file_fails(self):
+ # A flush spans several files, so it cannot be all-or-nothing: keeping
+ # every row would make the retry rewrite what the first files already
+ # took. The buffer has to hold the remainder and nothing else.
+ rows = [_row(i, i, 'x' * 64, 'y' * 64) for i in range(1, 401)]
+ data = pa.Table.from_pylist(rows, schema=_SCHEMA)
+
+ class _FailOnSecondFile(_Harness):
+ """Fails once, on the second file, then writes normally."""
+
+ failed = False
+
+ def _write_data_to_file(self, chunk):
+ if not self.failed and len(self.written_chunks) == 1:
+ self.failed = True
+ raise IOError('transient storage failure')
+ super()._write_data_to_file(chunk)
+
+ writer = _FailOnSecondFile(DeduplicateMergeFunction(),
+ target_file_size=data.nbytes // 4)
+ writer._buffer.append(data)
+ with self.assertRaises(IOError):
+ writer._flush_all()
+
+ written = sum(c.num_rows for c in writer.written_chunks)
+ self.assertEqual(len(writer.written_chunks), 1)
+ self.assertEqual(written + writer._buffer.num_rows, data.num_rows)
+
+ # The retry picks up exactly where the failure left off.
+ writer._flush_all()
+ self.assertTrue(writer._buffer.is_empty)
+ self.assertEqual(sum(c.num_rows for c in writer.written_chunks),
+ data.num_rows)
+ flushed_ids = [r['id'] for c in writer.written_chunks
+ for r in c.to_pylist()]
+ self.assertEqual(sorted(flushed_ids), [r['id'] for r in rows])
+
+
+class KeyValueWriteBufferTest(unittest.TestCase):
+ """The write path must not fold the buffer on every write.
+
+ ``_check_and_roll_if_needed`` flushes on size alone, and it reads that size
+ from the buffer's running total, so a write that stays under
+ ``target_file_size`` never needs the buffered batches as one table. Folding
+ there anyway would make writing N batches O(N^2): each fold leaves one more
+ chunk per column for the next ``pa.Table.nbytes`` walk to visit.
+ """
+
+ def test_writes_under_target_size_do_not_fold(self):
+ writer = _Harness(DeduplicateMergeFunction())
+ writer.sequence_generator = _StubSeqGen()
+ batch = pa.RecordBatch.from_pylist(
+ [{'id': 1, 'a': 'A', 'b': None}],
+ schema=pa.schema([
+ pa.field('id', pa.int64(), nullable=False),
+ pa.field('a', pa.string()),
+ pa.field('b', pa.string()),
+ ]),
+ )
+
+ folds = []
+ original = pa.concat_tables
+
+ def counting(*args, **kwargs):
+ folds.append(1)
+ return original(*args, **kwargs)
+
+ pa.concat_tables = counting
+ try:
+ for _ in range(100):
+ writer.write(batch)
+ finally:
+ pa.concat_tables = original
+
+ self.assertEqual(folds, [])
+ self.assertEqual(writer.written_chunks, [])
+ self.assertEqual(writer.pending_row_count, 100)
+
+ # The rows are all still there; flushing folds them exactly once.
+ writer._flush_all()
+ self.assertEqual(len(writer.written_chunks), 1)
+ # Dedup collapses the 100 same-PK rows to one.
+ self.assertEqual(writer.written_chunks[0].num_rows, 1)
+
class _StubSeqGen:
"""Stand-in for ``SequenceGenerator`` so the harness can call
diff --git a/paimon-python/pypaimon/tests/write/changelog_producer_test.py
b/paimon-python/pypaimon/tests/write/changelog_producer_test.py
index b789593038..2e2b2e20da 100644
--- a/paimon-python/pypaimon/tests/write/changelog_producer_test.py
+++ b/paimon-python/pypaimon/tests/write/changelog_producer_test.py
@@ -267,6 +267,72 @@ class ChangelogProducerTest(unittest.TestCase):
table_write.close()
table_commit.close()
+ def test_failed_changelog_write_leaves_nothing_to_commit(self):
+ """A data file and its changelog are committed together or not at all.
+
+ The data file used to be recorded before its changelog was written, so
a
+ changelog failure left the meta committed while the rows stayed
buffered
+ for the retry -- and the retry then wrote a second data file covering
+ rows the first meta already claimed. Committing both metas would double
+ every row in the snapshot.
+ """
+ table = self._create_table(
+ 'test_changelog_atomic',
+ options={'changelog-producer': 'input', 'bucket': '1'}
+ )
+ # Streaming write: ``BatchTableWrite`` refuses a second prepare_commit,
+ # and the retry is the whole point here.
+ write_builder = table.new_stream_write_builder()
+ table_write = write_builder.new_write()
+ table_commit = write_builder.new_commit()
+
+ original_write_parquet = table.file_io.write_parquet
+ state = {'failed': False}
+
+ def failing_write_parquet(path, data, **kwargs):
+ if not state['failed'] and '/changelog-' in str(path):
+ state['failed'] = True
+ raise IOError('transient storage failure')
+ return original_write_parquet(path, data, **kwargs)
+
+ table.file_io.write_parquet = failing_write_parquet
+ try:
+ table_write.write_arrow(self._sample_data())
+ with self.assertRaises(IOError):
+ table_write.prepare_commit(0)
+
+ bucket_dir = os.path.join(
+ self.warehouse, 'default.db', 'test_changelog_atomic',
+ 'dt=p1', 'bucket-0')
+ self.assertEqual(glob.glob(os.path.join(bucket_dir, 'data-*')), [],
+ "The data file must not outlive its failed
changelog")
+ self.assertEqual(glob.glob(os.path.join(bucket_dir,
'changelog-*')), [],
+ "A half-written changelog must not be left
behind")
+
+ messages = table_write.prepare_commit(0)
+ finally:
+ table.file_io.write_parquet = original_write_parquet
+
+ self.assertTrue(state['failed'], "the changelog write never failed")
+ # One data file and one changelog for the 3 rows, not two of each.
+ self.assertEqual(len(glob.glob(os.path.join(bucket_dir, 'data-*'))), 1)
+ self.assertEqual(len(glob.glob(os.path.join(bucket_dir,
'changelog-*'))), 1)
+ new_files = [meta for msg in messages for meta in msg.new_files]
+ changelog_files = [meta for msg in messages for meta in
msg.changelog_files]
+ self.assertEqual(len(new_files), 1)
+ self.assertEqual(len(changelog_files), 1)
+ self.assertEqual(sum(meta.row_count for meta in new_files), 3)
+
+ table_commit.commit(messages, 0)
+ read_builder = table.new_read_builder()
+ actual = read_builder.new_read().to_arrow(
+ read_builder.new_scan().plan().splits())
+ self.assertEqual(actual.num_rows, 3, "the retry must not duplicate
rows")
+ self.assertEqual(sorted(actual.column('user_id').to_pylist()), [1, 2,
3])
+
+ table_write.close()
+ table_commit.close()
+
def test_reject_changelog_producer_on_append_only_table(self):
append_schema = pa.schema([
('user_id', pa.int32()),
diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py
b/paimon-python/pypaimon/tests/write/table_write_test.py
index d7c12735e4..aa46603d7d 100644
--- a/paimon-python/pypaimon/tests/write/table_write_test.py
+++ b/paimon-python/pypaimon/tests/write/table_write_test.py
@@ -1958,9 +1958,9 @@ class TableWriteTest(unittest.TestCase):
)
writer.write(big_batch)
- pending_rows = writer.pending_data.num_rows if writer.pending_data is
not None else 0
+ pending_rows = writer.pending_row_count
committed_rows = sum(f.row_count for f in writer.committed_files)
self.assertEqual(committed_rows + pending_rows, num_rows)
self.assertGreater(len(writer.committed_files), 0)
- if writer.pending_data is not None:
- self.assertLessEqual(writer.pending_data.nbytes, target)
+ if pending_rows > 0:
+ self.assertLessEqual(writer._buffer.materialize().nbytes, target)
diff --git a/paimon-python/pypaimon/tests/write/write_buffer_test.py
b/paimon-python/pypaimon/tests/write/write_buffer_test.py
new file mode 100644
index 0000000000..434f328897
--- /dev/null
+++ b/paimon-python/pypaimon/tests/write/write_buffer_test.py
@@ -0,0 +1,801 @@
+# 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.
+
+"""Tests for ``WriteBuffer`` and the deferred-fold write path.
+
+The regression these guard: ``DataWriter.write`` used to fold every incoming
+batch into its buffer and then measure the result, which made writing N batches
+O(N^2) -- ``pa.concat_tables`` leaves N chunks per column and the
+``pa.Table.nbytes`` walk behind each rolling decision re-visits all of them.
+Folding is now deferred, so the tests below assert both halves of that: the
+fold count no longer scales with the number of writes, and every rolling
+trigger the eager path had still fires.
+"""
+
+import contextlib
+import unittest
+
+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.dedicated_format_writer import DedicatedFormatWriter
+from pypaimon.write.writer.write_buffer import WriteBuffer
+
+_SCHEMA = pa.schema([
+ pa.field('id', pa.int64(), nullable=False),
+ pa.field('name', pa.string()),
+])
+
+# Anything past this is effectively "no rolling"; matches how
+# ``target_file_row_num`` defaults to the max long when the option is unset.
+_NO_LIMIT = 2 ** 63 - 1
+
+
+def _batch(start: int, num_rows: int) -> pa.RecordBatch:
+ return pa.RecordBatch.from_pydict(
+ {
+ 'id': list(range(start, start + num_rows)),
+ 'name': ['n%d' % i for i in range(start, start + num_rows)],
+ },
+ schema=_SCHEMA,
+ )
+
+
+def _table(start: int, num_rows: int) -> pa.Table:
+ return pa.Table.from_batches([_batch(start, num_rows)])
+
+
+# Same columns and types as ``_SCHEMA``, but ``id`` is nullable. This is one of
+# the differences ``TableWrite._validate_pyarrow_schema`` lets through (it only
+# compares field types) while ``pa.concat_tables`` rejects it.
+_NULLABLE_SCHEMA = pa.schema([
+ pa.field('id', pa.int64()),
+ pa.field('name', pa.string()),
+])
+
+# ``_SCHEMA`` carrying metadata, which neither ``Schema.equals`` nor
+# ``concat_tables`` looks at.
+_ANNOTATED_SCHEMA = pa.schema(
+ [
+ pa.field('id', pa.int64(), nullable=False, metadata={b'k': b'v'}),
+ pa.field('name', pa.string()),
+ ],
+ metadata={b'origin': b'test'},
+)
+
+
+def _batch_with(schema: pa.Schema, start: int, num_rows: int) ->
pa.RecordBatch:
+ return pa.RecordBatch.from_pydict(
+ {
+ 'id': list(range(start, start + num_rows)),
+ 'name': ['n%d' % i for i in range(start, start + num_rows)],
+ },
+ schema=schema,
+ )
+
+
+def _table_with(schema: pa.Schema, start: int, num_rows: int) -> pa.Table:
+ return pa.Table.from_batches([_batch_with(schema, start, num_rows)])
+
+
[email protected]
+def _count_folds():
+ """Count the concat calls a fold performs, as a list of one entry each.
+
+ ``pa.concat_tables`` is the only way the buffer or a writer's
+ ``_merge_data`` collapses tables, so its call count is the fold count. The
+ eager path called it once per write; the deferred path calls it once per
+ ``materialize`` that has something to fold.
+ """
+ calls = []
+ original = pa.concat_tables
+
+ def counting(*args, **kwargs):
+ calls.append(1)
+ return original(*args, **kwargs)
+
+ pa.concat_tables = counting
+ try:
+ yield calls
+ finally:
+ pa.concat_tables = original
+
+
+class _Harness(AppendOnlyDataWriter):
+ """Append-only writer with the file layer stubbed out.
+
+ Bypasses ``DataWriter.__init__`` -- which needs a real table, catalog and
+ file IO -- and sets up only what the write/roll path reads.
+ """
+
+ def __init__(self, target_file_size: int = _NO_LIMIT,
+ target_file_row_num: int = _NO_LIMIT):
+ self.target_file_size = target_file_size
+ self.target_file_row_num = target_file_row_num
+ self._buffer = WriteBuffer(self._merge_data)
+ self.committed_files = []
+ self.written_chunks = []
+ self.aborted = False
+
+ def _write_data_to_file(self, data: pa.Table):
+ self.written_chunks.append(data)
+
+ def abort(self):
+ # The real ``abort`` deletes the committed files through file IO this
+ # harness has none of; record that it ran and do the rest.
+ self.aborted = True
+ self._buffer.reset()
+ self.committed_files.clear()
+
+
+class WriteBufferTest(unittest.TestCase):
+
+ def _buffer(self):
+ return WriteBuffer(lambda a, b: pa.concat_tables([a, b]))
+
+ def test_append_tracks_counts_without_folding(self):
+ buffer = self._buffer()
+ with _count_folds() as folds:
+ for i in range(50):
+ buffer.append(_table(i * 10, 10))
+ self.assertEqual(folds, [])
+ self.assertEqual(buffer.num_rows, 500)
+ self.assertFalse(buffer.is_empty)
+
+ def test_running_nbytes_matches_the_folded_table(self):
+ # The gate on the write path trusts the running sum, so it has to agree
+ # with what the folded table reports -- separate tables share no Arrow
+ # buffers, so the per-table sizes add up exactly.
+ buffer = self._buffer()
+ for i in range(10):
+ buffer.append(_table(i * 10, 10))
+ running = buffer.nbytes
+ self.assertEqual(running, buffer.materialize().nbytes)
+ self.assertEqual(running, buffer.nbytes)
+
+ def test_materialize_folds_once_and_is_idempotent(self):
+ buffer = self._buffer()
+ for i in range(20):
+ buffer.append(_table(i, 1))
+ with _count_folds() as folds:
+ first = buffer.materialize()
+ second = buffer.materialize()
+ # One concat for the 20 appended tables, and nothing on re-read.
+ self.assertEqual(len(folds), 1)
+ self.assertIs(first, second)
+ self.assertEqual(first.num_rows, 20)
+ self.assertEqual(first.column('id').to_pylist(), list(range(20)))
+
+ def test_single_append_skips_the_concat_entirely(self):
+ buffer = self._buffer()
+ buffer.append(_table(0, 5))
+ with _count_folds() as folds:
+ self.assertEqual(buffer.materialize().num_rows, 5)
+ self.assertEqual(folds, [])
+
+ def test_append_after_materialize_goes_through_merge(self):
+ merged = []
+
+ def merge(existing, new):
+ merged.append((existing.num_rows, new.num_rows))
+ return pa.concat_tables([existing, new])
+
+ buffer = WriteBuffer(merge)
+ buffer.append(_table(0, 3))
+ buffer.materialize()
+ buffer.append(_table(3, 4))
+ self.assertEqual(buffer.materialize().num_rows, 7)
+ self.assertEqual(merged, [(3, 4)])
+
+ def test_empty_buffer_materializes_to_none(self):
+ buffer = self._buffer()
+ self.assertTrue(buffer.is_empty)
+ self.assertIsNone(buffer.materialize())
+ self.assertEqual(buffer.nbytes, 0)
+ self.assertEqual(buffer.num_rows, 0)
+
+ def test_reset_replaces_contents_and_remeasures(self):
+ buffer = self._buffer()
+ buffer.append(_table(0, 100))
+ replacement = _table(0, 7)
+ buffer.reset(replacement)
+ self.assertIs(buffer.materialize(), replacement)
+ self.assertEqual(buffer.num_rows, 7)
+ self.assertEqual(buffer.nbytes, replacement.nbytes)
+
+ def test_reset_to_none_empties_the_buffer(self):
+ buffer = self._buffer()
+ buffer.append(_table(0, 100))
+ buffer.reset()
+ self.assertTrue(buffer.is_empty)
+ self.assertIsNone(buffer.materialize())
+ self.assertEqual(buffer.num_rows, 0)
+ self.assertEqual(buffer.nbytes, 0)
+
+ def test_zero_row_table_is_not_reported_empty(self):
+ # ``is_empty`` has to distinguish "nothing set" from "an empty table
was
+ # set", because the writers use it to decide whether there is anything
+ # to roll at all.
+ buffer = self._buffer()
+ buffer.reset(_table(0, 0))
+ self.assertFalse(buffer.is_empty)
+ self.assertEqual(buffer.num_rows, 0)
+ self.assertIsNotNone(buffer.materialize())
+
+
+class SchemaGuardTest(unittest.TestCase):
+ """The guard keeps a mismatch failing where the eager fold failed.
+
+ Folding on append meant a batch ``concat_tables`` could not accept raised
+ inside ``DataWriter.write``, which aborts and cleans up. Deferring the fold
+ would otherwise push that failure out to ``prepare_commit``, which has no
+ such handler, so ``append`` rejects up front exactly what concat rejects.
+ """
+
+ def _buffer(self):
+ return WriteBuffer(lambda a, b: pa.concat_tables([a, b]))
+
+ def test_append_rejects_a_schema_concat_would_reject(self):
+ buffer = self._buffer()
+ buffer.append(_table(0, 3))
+ with self.assertRaises(ValueError) as caught:
+ buffer.append(_table_with(_NULLABLE_SCHEMA, 3, 3))
+ self.assertIn('schema differs', str(caught.exception))
+ # And a rejected batch leaves the running counts describing the rows
+ # that are actually buffered.
+ self.assertEqual(buffer.num_rows, 3)
+ self.assertEqual(buffer.materialize().num_rows, 3)
+
+ def test_append_rejects_a_mismatch_after_a_materialize(self):
+ # ``materialize`` rebases the schema onto the folded table, so the
+ # second half of a buffer's life is guarded too.
+ buffer = self._buffer()
+ buffer.append(_table(0, 3))
+ buffer.materialize()
+ with self.assertRaises(ValueError):
+ buffer.append(_table_with(_NULLABLE_SCHEMA, 3, 3))
+
+ def test_append_accepts_a_metadata_only_difference(self):
+ # ``concat_tables`` ignores schema and field metadata, so the guard has
+ # to as well or it would reject batches the old path folded fine.
+ buffer = self._buffer()
+ buffer.append(_table(0, 3))
+ buffer.append(_table_with(_ANNOTATED_SCHEMA, 3, 4))
+ self.assertEqual(buffer.materialize().num_rows, 7)
+
+ def test_reset_rebases_the_schema(self):
+ # A writer that resets to a new table starts a new file; batches
+ # matching that table's schema have to be accepted afterwards.
+ buffer = self._buffer()
+ buffer.append(_table(0, 3))
+ buffer.reset(_table_with(_NULLABLE_SCHEMA, 0, 2))
+ buffer.append(_table_with(_NULLABLE_SCHEMA, 2, 2))
+ self.assertEqual(buffer.materialize().num_rows, 4)
+ with self.assertRaises(ValueError):
+ buffer.append(_table(4, 1))
+
+ def test_first_append_after_an_empty_reset_sets_the_schema(self):
+ buffer = self._buffer()
+ buffer.append(_table(0, 3))
+ buffer.reset()
+ buffer.append(_table_with(_NULLABLE_SCHEMA, 0, 3))
+ self.assertEqual(buffer.materialize().num_rows, 3)
+
+
+class DeferredFoldWritePathTest(unittest.TestCase):
+
+ def test_many_small_writes_fold_a_constant_number_of_times(self):
+ # The regression: with an eager fold this was 199 concats for 200
+ # writes, each walking a buffer one chunk longer than the last.
+ writer = _Harness()
+ with _count_folds() as folds:
+ for i in range(200):
+ writer.write(_batch(i * 5, 5))
+ self.assertEqual(folds, [])
+ # Reading the buffer folds it -- once, not once per write.
+ with _count_folds() as folds:
+ pending = writer._buffer.materialize()
+ self.assertEqual(len(folds), 1)
+ self.assertEqual(pending.num_rows, 1000)
+ self.assertEqual(pending.column('id').to_pylist(), list(range(1000)))
+
+ def test_row_num_rolling_still_fires(self):
+ # ``target_file_row_num`` postdates the eager fold, so a byte-only gate
+ # would leave row-based rolling silently dead for tables that set it
+ # while staying well under target_file_size.
+ writer = _Harness(target_file_row_num=10)
+ for i in range(10):
+ writer.write(_batch(i * 3, 3))
+ self.assertEqual([c.num_rows for c in writer.written_chunks],
+ [10, 10])
+ self.assertEqual(writer.pending_row_count, 10)
+ self.assertEqual(
+ [row for c in writer.written_chunks
+ for row in c.column('id').to_pylist()],
+ list(range(20)),
+ )
+
+ def test_row_num_rolling_splits_a_single_oversized_write(self):
+ writer = _Harness(target_file_row_num=4)
+ writer.write(_batch(0, 14))
+ self.assertEqual([c.num_rows for c in writer.written_chunks],
+ [4, 4, 4])
+ self.assertEqual(writer.pending_row_count, 2)
+
+ def test_size_rolling_still_bounds_written_files(self):
+ target = _table(0, 100).nbytes // 4
+ writer = _Harness(target_file_size=target)
+ for i in range(20):
+ writer.write(_batch(i * 5, 5))
+ self.assertGreaterEqual(len(writer.written_chunks), 3)
+ for chunk in writer.written_chunks:
+ self.assertLessEqual(chunk.nbytes, target)
+ written = sum(c.num_rows for c in writer.written_chunks)
+ self.assertEqual(written + writer.pending_row_count, 100)
+
+ def test_row_that_alone_exceeds_target_size_is_rolled_by_itself(self):
+ writer = _Harness(target_file_size=1)
+ writer.write(_batch(0, 3))
+ # Rolls the first two rows one at a time; the last stays buffered
+ # because the loop stops once the buffer is down to a single row.
+ self.assertEqual([c.num_rows for c in writer.written_chunks], [1, 1])
+ self.assertEqual(writer.pending_row_count, 1)
+
+ def test_rolling_does_not_refold_on_every_subsequent_write(self):
+ # After a roll the remainder goes back into the buffer, which
+ # re-measures it. If that left the roll condition stuck open, every
+ # later write would fold again and the quadratic would be back.
+ writer = _Harness(target_file_row_num=8)
+ writer.write(_batch(0, 12))
+ # Rolled 8 rows, 4 left buffered.
+ self.assertEqual(len(writer.written_chunks), 1)
+ with _count_folds() as folds:
+ for i in range(20):
+ writer.write(_batch(100 + i, 1))
+ # 4 buffered + 20 single-row writes crosses 8 rows exactly twice, so
+ # only two of those 20 writes fold anything. Asserting a bound rather
+ # than an exact count: how many concats one fold takes is up to
+ # ``WriteBuffer``, but it must not be one per write.
+ self.assertLess(len(folds), 20)
+ self.assertEqual(len(writer.written_chunks), 3)
+
+ def test_pending_row_count_does_not_fold(self):
+ # The composite writers read this on every write to size the next
+ # slice; folding here would bring the quadratic back through the side
+ # door.
+ writer = _Harness()
+ with _count_folds() as folds:
+ for i in range(30):
+ writer.write(_batch(i, 1))
+ self.assertEqual(writer.pending_row_count, i + 1)
+ self.assertEqual(folds, [])
+
+ def test_prepare_commit_flushes_the_deferred_buffer(self):
+ writer = _Harness()
+ for i in range(10):
+ writer.write(_batch(i * 2, 2))
+ writer.prepare_commit()
+ self.assertEqual([c.num_rows for c in writer.written_chunks], [20])
+ self.assertTrue(writer._buffer.is_empty)
+
+ def test_close_flushes_the_deferred_buffer(self):
+ writer = _Harness()
+ for i in range(10):
+ writer.write(_batch(i * 2, 2))
+ writer.close()
+ self.assertEqual([c.num_rows for c in writer.written_chunks], [20])
+ self.assertTrue(writer._buffer.is_empty)
+
+ def test_mismatched_schema_fails_the_write_that_carries_it(self):
+ # Not ``prepare_commit``: only ``write`` aborts, so a failure deferred
+ # to commit time would leave the files rolled so far orphaned.
+ writer = _Harness()
+ writer.write(_batch(0, 3))
+ with self.assertRaises(ValueError):
+ writer.write(_batch_with(_NULLABLE_SCHEMA, 3, 3))
+ self.assertTrue(writer.aborted)
+
+
+class FlushFailureTest(unittest.TestCase):
+ """A failed flush leaves the rows buffered for the retry.
+
+ ``StreamTableWrite`` is reusable, so a transient storage error followed by
+ another ``prepare_commit`` on the same writer has to write the same rows,
+ not silently skip them. Draining the buffer before the write would lose
+ them.
+ """
+
+ class _FailOnceHarness(_Harness):
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.fail_next = True
+
+ def _write_data_to_file(self, data: pa.Table):
+ if self.fail_next:
+ self.fail_next = False
+ raise IOError('transient storage failure')
+ super()._write_data_to_file(data)
+
+ def test_failed_prepare_commit_keeps_the_rows_for_the_retry(self):
+ writer = self._FailOnceHarness()
+ for i in range(3):
+ writer.write(_batch(i, 1))
+ with self.assertRaises(IOError):
+ writer.prepare_commit()
+ self.assertEqual(writer.pending_row_count, 3)
+ writer.prepare_commit()
+ self.assertEqual([c.num_rows for c in writer.written_chunks], [3])
+ self.assertEqual(writer.pending_row_count, 0)
+
+ class _FailOnceVectorHarness(DataVectorWriter):
+ """``_close_current_writers`` with the file layer stubbed out.
+
+ No ``vector_writer``, so this covers the normal half on its own; the
+ point is only where the buffer is cleared relative to the write.
+ """
+
+ def __init__(self):
+ self.target_file_size = _NO_LIMIT
+ self.target_file_row_num = _NO_LIMIT
+ self.record_count = 0
+ self.vector_writer = None
+ self._normal_buffer = WriteBuffer(self._merge_data)
+ self.committed_files = []
+ self.written = []
+ self.fail_next = True
+
+ def _write_normal_data_to_file(self, data: pa.Table):
+ if self.fail_next:
+ self.fail_next = False
+ raise IOError('transient storage failure')
+ self.written.append(data)
+ return object()
+
+ def test_failed_normal_flush_keeps_the_rows_for_the_retry(self):
+ # Otherwise the retry finds no normal_meta, flushes the sidecars alone
+ # and skips the row-count check, committing sidecar-only metadata.
+ writer = self._FailOnceVectorHarness()
+ writer._normal_buffer.append(_table(0, 3))
+ with self.assertRaises(IOError):
+ writer._close_current_writers()
+ self.assertEqual(writer._normal_buffer.num_rows, 3)
+ writer._close_current_writers()
+ self.assertEqual([t.num_rows for t in writer.written], [3])
+ self.assertEqual(writer._normal_buffer.num_rows, 0)
+
+
+class _StubMeta:
+ """The handful of ``DataFileMeta`` fields the flush and abort paths
read."""
+
+ def __init__(self, row_count: int, file_name: str):
+ self.row_count = row_count
+ self.file_name = file_name
+ self.file_path = '/warehouse/%s' % file_name
+ self.external_path = None
+ self.extra_files = []
+
+
+class _RecordingFileIO:
+ def __init__(self):
+ self.deleted = []
+
+ def delete_quietly(self, path):
+ self.deleted.append(path)
+
+
+class _StubSidecarWriter:
+ """A blob/vector writer that fails its first ``prepare_commit``.
+
+ Models the real ones in the way that matters here: a failure produces no
+ metadata, and because the sub-writer drains its own buffer as it writes, a
+ later call returns whatever has landed so far -- so the parent must be able
+ to harvest the same metas twice without double-counting them.
+ """
+
+ def __init__(self, row_count: int, file_name: str, fail_times: int = 0,
+ delete_on_abort: bool = True):
+ self.committed_files = []
+ self.pending_row_count = 0
+ self.prepare_commit_calls = 0
+ self.aborted = False
+ self._row_count = row_count
+ self._file_name = file_name
+ self._fail_times = fail_times
+ self._delete_on_abort = delete_on_abort
+
+ def prepare_commit(self):
+ self.prepare_commit_calls += 1
+ if self._fail_times > 0:
+ self._fail_times -= 1
+ raise IOError('transient sidecar failure')
+ if not self.committed_files:
+ self.committed_files.append(
+ _StubMeta(self._row_count, self._file_name))
+ return self.committed_files.copy()
+
+ def delete_file_upon_abort(self):
+ return self._delete_on_abort
+
+ def abort(self):
+ self.aborted = True
+ self.committed_files.clear()
+
+
+class CompositeFlushResumeTest(unittest.TestCase):
+ """A composite flush publishes all of its files or none of them.
+
+ One flush writes the normal data file and then the blob/vector sidecars.
The
+ sidecar writers drain their own buffers as they go, so a failure part way
+ through cannot be rolled back -- deleting the sidecars that already landed
+ would lose rows nothing can replay. So the flush resumes instead: the
normal
+ rows stay buffered until their file lands, the landed file is remembered so
+ the retry skips it, and no metadata is published until every phase is done.
+
+ Publishing the normal meta before the sidecars ran, as the code used to,
+ left the retry writing a second copy of rows the first meta already covered
+ -- and ``_validate_consistency`` then checked the sidecars against that
+ second copy only.
+ """
+
+ class _VectorHarness(DataVectorWriter):
+ def __init__(self, vector_writer, fail_normal_times: int = 0):
+ self.target_file_size = _NO_LIMIT
+ self.target_file_row_num = _NO_LIMIT
+ self.record_count = 0
+ self.vector_writer = vector_writer
+ self.normal_column_names = ['id', 'name']
+ self.vector_write_columns = []
+ self._normal_buffer = WriteBuffer(self._merge_data)
+ self._buffer = WriteBuffer(self._merge_data)
+ self.committed_files = []
+ self.committed_changelog_files = []
+ self.file_io = _RecordingFileIO()
+ self.written = []
+ self._fail_normal_times = fail_normal_times
+
+ def _write_normal_data_to_file(self, data: pa.Table):
+ if self._fail_normal_times > 0:
+ self._fail_normal_times -= 1
+ raise IOError('transient storage failure')
+ self.written.append(data)
+ return _StubMeta(data.num_rows, 'data-%d' % len(self.written))
+
+ class _DedicatedHarness(DedicatedFormatWriter):
+ def __init__(self, blob_writers, vector_writer=None):
+ self.target_file_size = _NO_LIMIT
+ self.target_file_row_num = _NO_LIMIT
+ self.record_count = 0
+ self.blob_writers = blob_writers
+ self.blob_file_column_names = list(blob_writers)
+ self.vector_writer = vector_writer
+ self._normal_buffer = WriteBuffer(self._merge_normal_data)
+ self._buffer = WriteBuffer(self._merge_normal_data)
+ self.committed_files = []
+ self._committed_files_to_delete_on_abort = []
+ self.file_io = _RecordingFileIO()
+ self.written = []
+
+ def _write_normal_data_to_file(self, data: pa.Table):
+ self.written.append(data)
+ return _StubMeta(data.num_rows, 'data-%d' % len(self.written))
+
+ def test_failed_sidecar_publishes_nothing_and_the_retry_resumes(self):
+ vector = _StubSidecarWriter(3, 'vector-0', fail_times=1)
+ writer = self._VectorHarness(vector)
+ writer._normal_buffer.append(_table(0, 3))
+
+ with self.assertRaises(IOError):
+ writer._close_current_writers()
+ # The normal file landed, so the rows are no longer buffered -- but
+ # nothing is committed and the file is remembered for the retry.
+ self.assertEqual([t.num_rows for t in writer.written], [3])
+ self.assertEqual(writer.committed_files, [])
+ self.assertEqual(writer._normal_buffer.num_rows, 0)
+ self.assertIsNotNone(writer._pending_normal_meta)
+
+ writer._close_current_writers()
+ # Still one normal file: the retry resumed at the sidecar phase rather
+ # than writing the same 3 rows again.
+ self.assertEqual([t.num_rows for t in writer.written], [3])
+ self.assertEqual([m.file_name for m in writer.committed_files],
+ ['data-1', 'vector-0'])
+ self.assertIsNone(writer._pending_normal_meta)
+ # Harvested once the flush completed, and cleared only then.
+ self.assertEqual(vector.committed_files, [])
+
+ def test_successful_flush_publishes_normal_then_sidecars(self):
+ vector = _StubSidecarWriter(3, 'vector-0')
+ writer = self._VectorHarness(vector)
+ writer._normal_buffer.append(_table(0, 3))
+ writer._close_current_writers()
+ self.assertEqual([m.file_name for m in writer.committed_files],
+ ['data-1', 'vector-0'])
+ self.assertEqual(vector.committed_files, [])
+ self.assertIsNone(writer._pending_normal_meta)
+ self.assertEqual(writer.record_count, 0)
+
+ def test_failed_normal_write_keeps_the_rows_and_publishes_nothing(self):
+ # The other half of the same rule: the sidecars are never reached, so
+ # the rows have to stay where a retry can find them.
+ vector = _StubSidecarWriter(3, 'vector-0')
+ writer = self._VectorHarness(vector, fail_normal_times=1)
+ writer._normal_buffer.append(_table(0, 3))
+ with self.assertRaises(IOError):
+ writer._close_current_writers()
+ self.assertEqual(writer._normal_buffer.num_rows, 3)
+ self.assertEqual(writer.committed_files, [])
+ self.assertIsNone(writer._pending_normal_meta)
+ self.assertEqual(vector.prepare_commit_calls, 0)
+
+ writer._close_current_writers()
+ self.assertEqual([t.num_rows for t in writer.written], [3])
+ self.assertEqual([m.file_name for m in writer.committed_files],
+ ['data-1', 'vector-0'])
+
+ def test_write_is_rejected_while_a_flush_is_unfinished(self):
+ # Rows appended between a failed flush and its retry would belong to no
+ # file: the resumed flush skips the normal write, while the sidecar
+ # writer would drain them -- breaking the row-count check.
+ vector = _StubSidecarWriter(3, 'vector-0', fail_times=1)
+ writer = self._VectorHarness(vector)
+ writer._normal_buffer.append(_table(0, 3))
+ with self.assertRaises(IOError):
+ writer._close_current_writers()
+
+ with self.assertRaises(RuntimeError) as caught:
+ writer.write(_batch(3, 1))
+ self.assertIn('Cannot write', str(caught.exception))
+ # Rejecting does not abort, so the flush is still resumable.
+ self.assertIsNotNone(writer._pending_normal_meta)
+ writer._close_current_writers()
+ self.assertEqual([t.num_rows for t in writer.written], [3])
+
+ def test_abort_deletes_the_unpublished_normal_file(self):
+ # It is in no committed list, so abort has to know about it separately
+ # or it leaks a data file no snapshot references.
+ vector = _StubSidecarWriter(3, 'vector-0', fail_times=1)
+ writer = self._VectorHarness(vector)
+ writer._normal_buffer.append(_table(0, 3))
+ with self.assertRaises(IOError):
+ writer._close_current_writers()
+
+ writer.abort()
+ self.assertEqual(writer.file_io.deleted, ['/warehouse/data-1'])
+ self.assertIsNone(writer._pending_normal_meta)
+ self.assertTrue(vector.aborted)
+
+ def test_dedicated_writer_failed_blob_phase_publishes_nothing(self):
+ blob = _StubSidecarWriter(3, 'blob-0', fail_times=1)
+ writer = self._DedicatedHarness({'payload': blob})
+ writer._normal_buffer.append(_table(0, 3))
+
+ with self.assertRaises(IOError):
+ writer._close_current_writers()
+ self.assertEqual([t.num_rows for t in writer.written], [3])
+ self.assertEqual(writer.committed_files, [])
+ # Already tracked for abort, since no committed list holds it yet.
+ self.assertEqual(
+ [m.file_name for m in writer._committed_files_to_delete_on_abort],
+ ['data-1'])
+
+ writer._close_current_writers()
+ self.assertEqual([t.num_rows for t in writer.written], [3])
+ self.assertEqual([m.file_name for m in writer.committed_files],
+ ['data-1', 'blob-0'])
+ self.assertEqual(blob.committed_files, [])
+ self.assertIsNone(writer._pending_normal_meta)
+
+ def test_dedicated_writer_keeps_the_documented_meta_order(self):
+ blob = _StubSidecarWriter(3, 'blob-0')
+ vector = _StubSidecarWriter(3, 'vector-0')
+ writer = self._DedicatedHarness({'payload': blob}, vector)
+ writer._normal_buffer.append(_table(0, 3))
+ writer._close_current_writers()
+ self.assertEqual([m.file_name for m in writer.committed_files],
+ ['data-1', 'blob-0', 'vector-0'])
+ self.assertEqual(
+ [m.file_name for m in writer._committed_files_to_delete_on_abort],
+ ['data-1', 'blob-0', 'vector-0'])
+
+ def test_dedicated_writer_respects_the_blob_delete_policy(self):
+ # Externally managed blob files are not the writer's to delete, so they
+ # must stay out of the abort list even now that it is filled at publish
+ # time rather than as each sidecar lands.
+ blob = _StubSidecarWriter(3, 'blob-0', delete_on_abort=False)
+ writer = self._DedicatedHarness({'payload': blob})
+ writer._normal_buffer.append(_table(0, 3))
+ writer._close_current_writers()
+ self.assertEqual([m.file_name for m in writer.committed_files],
+ ['data-1', 'blob-0'])
+ self.assertEqual(
+ [m.file_name for m in writer._committed_files_to_delete_on_abort],
+ ['data-1'])
+
+
+class VectorNormalBufferTest(unittest.TestCase):
+ """``DataVectorWriter`` keeps its own buffer for the normal columns."""
+
+ class _VectorHarness(DataVectorWriter):
+ """Only the fields ``_should_roll_normal`` reads.
+
+ ``CHECK_ROLLING_RECORD_CNT`` is 1 rather than the real 1000 so the size
+ branch is reached on every write instead of being short-circuited by
the
+ periodic check.
+ """
+
+ CHECK_ROLLING_RECORD_CNT = 1
+
+ def __init__(self, target_file_size: int = _NO_LIMIT):
+ self.target_file_size = target_file_size
+ self.target_file_row_num = _NO_LIMIT
+ self.record_count = 0
+ self.vector_writer = None
+ self._normal_buffer = WriteBuffer(self._merge_data)
+ self._buffer = WriteBuffer(self._merge_data)
+
+ def test_should_roll_normal_does_not_fold(self):
+ writer = self._VectorHarness()
+ with _count_folds() as folds:
+ for i in range(50):
+ writer._normal_buffer.append(_table(i * 5, 5))
+ writer.record_count += 5
+ self.assertFalse(writer._should_roll_normal())
+ self.assertEqual(folds, [])
+ self.assertEqual(writer._normal_buffer.num_rows, 250)
+
+ def test_should_roll_normal_fires_off_the_running_size(self):
+ writer = self._VectorHarness(target_file_size=_table(0, 20).nbytes)
+ writer.record_count = 1
+ for i in range(10):
+ writer._normal_buffer.append(_table(i * 5, 5))
+ if writer._should_roll_normal():
+ break
+ self.assertTrue(writer._should_roll_normal())
+ # Fires before the buffer grows far past the target, i.e. off the
+ # accumulated size rather than at some arbitrary later point.
+ self.assertLessEqual(writer._normal_buffer.num_rows, 30)
+
+ def test_pending_row_count_reports_the_normal_buffer(self):
+ # The inherited property reads the base ``_buffer``, which this writer
+ # never fills, so it has to be overridden or it always answers zero.
+ writer = self._VectorHarness()
+ self.assertEqual(writer.pending_row_count, 0)
+ writer._normal_buffer.append(_table(0, 7))
+ self.assertEqual(writer.pending_row_count, 7)
+ self.assertTrue(writer._buffer.is_empty)
+
+ def test_pending_row_count_falls_back_to_the_vector_writer(self):
+ # A table whose columns are all vectors buffers nothing normal, so the
+ # count has to come from the sidecar instead of reading as zero.
+ writer = self._VectorHarness()
+ writer.vector_writer = _StubSidecarWriter(0, 'vector-0')
+ writer.vector_writer.pending_row_count = 4
+ self.assertEqual(writer.pending_row_count, 4)
+ writer._normal_buffer.append(_table(0, 7))
+ self.assertEqual(writer.pending_row_count, 7)
+
+ def
test_dedicated_writer_pending_row_count_reports_the_normal_buffer(self):
+ writer = CompositeFlushResumeTest._DedicatedHarness({})
+ self.assertEqual(writer.pending_row_count, 0)
+ writer._normal_buffer.append(_table(0, 3))
+ self.assertEqual(writer.pending_row_count, 3)
+ self.assertTrue(writer._buffer.is_empty)
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/paimon-python/pypaimon/write/writer/blob_writer.py
b/paimon-python/pypaimon/write/writer/blob_writer.py
index 40d301c2f1..276968eee2 100644
--- a/paimon-python/pypaimon/write/writer/blob_writer.py
+++ b/paimon-python/pypaimon/write/writer/blob_writer.py
@@ -58,21 +58,23 @@ class BlobWriter(AppendOnlyDataWriter):
logger.info(f"Initialized BlobWriter with blob file format,
blob_target_file_size={self.blob_target_file_size}")
def _check_and_roll_if_needed(self):
- if self.pending_data is None:
+ # Rolling here is driven by the size of the external blobs, which the
+ # buffered descriptors say nothing about, so there is no cheap count to
+ # check first: every write drains the buffer.
+ pending = self._buffer.take()
+ if pending is None:
return
# Always write blob rows one-by-one so rolling uses actual blob bytes
size rather than
# in-memory serialized descriptor size.
- for i in range(self.pending_data.num_rows):
- row_data = self.pending_data.slice(i, 1)
+ for i in range(pending.num_rows):
+ row_data = pending.slice(i, 1)
self._write_row_to_file(row_data)
self.record_count += 1
if self.rolling_file():
self.close_current_writer()
- self.pending_data = None
-
def _write_row_to_file(self, row_data: pa.Table):
"""Write a single row to the current blob file. Opens a new file if
needed."""
if row_data.num_rows == 0:
@@ -228,7 +230,7 @@ class BlobWriter(AppendOnlyDataWriter):
if self.current_writer is not None:
self.close_current_writer()
- # Call parent to handle pending_data fallback.
+ # Call parent to flush anything left in the buffer.
return super().prepare_commit()
def close(self):
@@ -237,7 +239,7 @@ class BlobWriter(AppendOnlyDataWriter):
if self.current_writer is not None:
self.close_current_writer()
- # Call parent to handle pending_data fallback.
+ # Call parent to flush anything left in the buffer.
super().close()
def delete_file_upon_abort(self) -> bool:
@@ -252,7 +254,7 @@ class BlobWriter(AppendOnlyDataWriter):
self.current_writer = None
self.current_file_path = None
if not self.delete_file_upon_abort():
- self.pending_data = None
+ self._buffer.reset()
self.committed_files.clear()
else:
super().abort()
diff --git a/paimon-python/pypaimon/write/writer/data_vector_writer.py
b/paimon-python/pypaimon/write/writer/data_vector_writer.py
index 72291cb738..0aa4c51810 100644
--- a/paimon-python/pypaimon/write/writer/data_vector_writer.py
+++ b/paimon-python/pypaimon/write/writer/data_vector_writer.py
@@ -28,6 +28,7 @@ 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
logger = logging.getLogger(__name__)
@@ -80,7 +81,13 @@ class DataVectorWriter(DataWriter):
self.record_count = 0
self.closed = False
- self.pending_normal_data: Optional[pa.Table] = None
+ # 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
@@ -109,12 +116,13 @@ class DataVectorWriter(DataWriter):
return pa.concat_tables([existing_data, new_data])
def write(self, data: pa.RecordBatch):
+ self._require_finished_flush()
try:
offset = 0
# _write_batch keeps normal and vector pending rows in lockstep
# and closes both writers when the shared row limit is reached.
while offset < data.num_rows:
- capacity = self.target_file_row_num - self._current_row_count()
+ capacity = self.target_file_row_num - self.pending_row_count
if capacity <= 0:
self._close_current_writers()
capacity = self.target_file_row_num
@@ -133,11 +141,8 @@ class DataVectorWriter(DataWriter):
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 normal_data is not None:
+ self._normal_buffer.append(pa.Table.from_batches([normal_data]))
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)
@@ -162,12 +167,12 @@ class DataVectorWriter(DataWriter):
raise
finally:
self.closed = True
- self.pending_normal_data = None
+ self._normal_buffer.reset()
def abort(self):
if self.vector_writer is not None:
self.vector_writer.abort()
- self.pending_normal_data = None
+ self._normal_buffer.reset()
super().abort()
def _split_data(self, data: pa.RecordBatch) -> Tuple[pa.RecordBatch,
pa.RecordBatch]:
@@ -188,38 +193,59 @@ class DataVectorWriter(DataWriter):
return normal_data, vector_data
def _should_roll_normal(self) -> bool:
- if self.pending_normal_data is None:
+ # Runs on every write, so it answers from the running counts only.
+ if self._normal_buffer.is_empty:
return False
- if self.pending_normal_data.num_rows >= self.target_file_row_num:
+ 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.pending_normal_data.nbytes > self.target_file_size
-
- def _current_row_count(self) -> int:
- if self.pending_normal_data is not None:
- return self.pending_normal_data.num_rows
- if self.vector_writer is not None and self.vector_writer.pending_data
is not None:
- return self.vector_writer.pending_data.num_rows
+ 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
+ # fills. Normal and vector rows are kept in lockstep, so either half
+ # answers for the pair; the vector writer is asked only when the table
+ # has no normal columns at all.
+ if not self._normal_buffer.is_empty:
+ return self._normal_buffer.num_rows
+ if self.vector_writer is not None:
+ # Running count, not a folded buffer: this runs on every write.
+ return self.vector_writer.pending_row_count
return 0
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)
+ # 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:
- if normal_meta is not None:
- self._validate_consistency(normal_meta, vector_metas)
- self.committed_files.extend(vector_metas)
+ 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_data = None
+ self._pending_normal_meta = None
self.record_count = 0
def _write_normal_data_to_file(self, data: pa.Table) ->
Optional[DataFileMeta]:
diff --git a/paimon-python/pypaimon/write/writer/data_writer.py
b/paimon-python/pypaimon/write/writer/data_writer.py
index e34127ddaf..24a67186c3 100644
--- a/paimon-python/pypaimon/write/writer/data_writer.py
+++ b/paimon-python/pypaimon/write/writer/data_writer.py
@@ -30,6 +30,7 @@ from pypaimon.schema.data_types import PyarrowFieldParser
from pypaimon.table.bucket_mode import BucketMode
from pypaimon.table.row.generic_row import GenericRow
from pypaimon.write.writer.mosaic_writer_options import
create_mosaic_writer_options
+from pypaimon.write.writer.write_buffer import WriteBuffer
class DataWriter(ABC):
@@ -72,7 +73,7 @@ class DataWriter(ABC):
)
self.sequence_generator = SequenceGenerator(max_seq_number)
- self.pending_data: Optional[pa.Table] = None
+ self._buffer = WriteBuffer(self._merge_data)
self.committed_files: List[DataFileMeta] = []
self.committed_changelog_files: List[DataFileMeta] = []
self.changelog_producer = changelog_producer
@@ -104,15 +105,38 @@ class DataWriter(ABC):
# the table schema is fixed for the lifetime of this writer.
self._paimon_field_id: Dict[str, int] = {pf.name: pf.id for pf in
self.table.fields}
+ # 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.
+
+ The composite writers override this: they keep their normal rows in
+ ``_normal_buffer`` and hand the rest to their sub-writers, so the base
+ ``_buffer`` stays empty for them.
+ """
+ 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)
-
- if self.pending_data is None:
- self.pending_data = processed_data
- else:
- self.pending_data = self._merge_data(self.pending_data,
processed_data)
-
+ self._buffer.append(processed_data)
self._check_and_roll_if_needed()
except Exception as e:
import logging
@@ -122,9 +146,11 @@ class DataWriter(ABC):
raise e
def prepare_commit(self) -> List[DataFileMeta]:
- if self.pending_data is not None and self.pending_data.num_rows > 0:
- self._write_data_to_file(self.pending_data)
- self.pending_data = None
+ if self._buffer.num_rows > 0:
+ # Clear only once the write lands: a caller that retries
+ # prepare_commit after a failed write has to still find its rows.
+ self._write_data_to_file(self._buffer.materialize())
+ self._buffer.reset()
return self.committed_files.copy()
@@ -133,8 +159,8 @@ class DataWriter(ABC):
def close(self):
try:
- if self.pending_data is not None and self.pending_data.num_rows >
0:
- self._write_data_to_file(self.pending_data)
+ if self._buffer.num_rows > 0:
+ self._write_data_to_file(self._buffer.materialize())
except Exception as e:
import logging
logger = logging.getLogger(__name__)
@@ -142,7 +168,7 @@ class DataWriter(ABC):
self.abort()
raise e
finally:
- self.pending_data = None
+ self._buffer.reset()
# Note: Don't clear committed_files in close() - they should be
returned by prepare_commit()
def abort(self):
@@ -150,10 +176,15 @@ class DataWriter(ABC):
Abort all writers and clean up resources. This method should be called
when an error occurs
during writing. It deletes any files that were written and cleans up
resources.
"""
- self._delete_committed_files(self.committed_files +
self.committed_changelog_files)
+ 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
- self.pending_data = None
+ self._buffer.reset()
self.committed_files.clear()
self.committed_changelog_files.clear()
@@ -191,16 +222,20 @@ class DataWriter(ABC):
return -1, -1
def _check_and_roll_if_needed(self):
- while self.pending_data is not None:
- num_rows = self.pending_data.num_rows
+ # Neither trigger can fire below these thresholds, so the running
counts
+ # rule out rolling -- the common case -- without concatenating
anything.
+ while (self._buffer.nbytes > self.target_file_size
+ or self._buffer.num_rows > self.target_file_row_num):
+ pending = self._buffer.materialize()
+ num_rows = pending.num_rows
# Row-count trigger: keep at most target_file_row_num rows per
file.
split_row = num_rows
if num_rows > self.target_file_row_num:
split_row = self.target_file_row_num
# Size trigger: roll earlier if the size split point comes first.
- if self.pending_data.nbytes > self.target_file_size:
+ if pending.nbytes > self.target_file_size:
size_split = self._find_optimal_split_point(
- self.pending_data, self.target_file_size)
+ pending, self.target_file_size)
# First row alone exceeds target_file_size: roll it by itself.
if size_split <= 0:
size_split = 1
@@ -208,8 +243,8 @@ class DataWriter(ABC):
split_row = size_split
if split_row <= 0 or split_row >= num_rows:
break
- self._write_data_to_file(self.pending_data.slice(0, split_row))
- self.pending_data = self.pending_data.slice(split_row)
+ self._write_data_to_file(pending.slice(0, split_row))
+ self._buffer.reset(pending.slice(split_row))
def _write_data_to_file(self, data: pa.Table):
if data.num_rows == 0:
@@ -223,9 +258,16 @@ class DataWriter(ABC):
logical_data = data
extra_files = []
row_sidecar_path = None
+ changelog_meta = None
if self._variant_shredding:
data = self._apply_variant_shredding(data)
+ # One data file means up to three files on disk -- the data file, its
row
+ # sidecar and its changelog -- and none of them is committed until all
of
+ # them have landed. A caller that retries the flush still holds these
+ # rows in its buffer, so publishing the data file before the changelog
+ # 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:
self.file_io.write_parquet(file_path, data,
compression=self.compression, zstd_level=self.zstd_level)
@@ -255,74 +297,80 @@ class DataWriter(ABC):
fields=self._row_sidecar_fields(logical_data),
zstd_level=self.zstd_level)
extra_files.append(row_sidecar_name)
+
+ # min key & max key
+
+ selected_table = data.select(self.trimmed_primary_keys)
+ key_columns_batch = selected_table.to_batches()[0]
+ min_key_row_batch = key_columns_batch.slice(0, 1)
+ max_key_row_batch =
key_columns_batch.slice(key_columns_batch.num_rows - 1, 1)
+ min_key = [col.to_pylist()[0] for col in min_key_row_batch.columns]
+ max_key = [col.to_pylist()[0] for col in max_key_row_batch.columns]
+
+ # key stats & value stats
+ value_stats_enabled = self.options.metadata_stats_enabled()
+ if value_stats_enabled:
+ stats_fields = self.table.fields if
self.table.is_primary_key_table \
+ else PyarrowFieldParser.to_paimon_schema(data.schema)
+ else:
+ stats_fields = self.table.trimmed_primary_keys_fields
+ column_stats = {
+ field.name: self._get_column_stats(data, field.name)
+ for field in stats_fields
+ }
+ key_fields = self.trimmed_primary_keys_fields
+ key_stats = self._collect_value_stats(data, key_fields,
column_stats)
+ if not self.options.primary_key_nullable() and not all(
+ count == 0 for count in key_stats.null_counts):
+ raise RuntimeError("Primary key should not be null")
+
+ value_fields = stats_fields if value_stats_enabled else []
+ value_stats = self._collect_value_stats(data, value_fields,
column_stats)
+
+ # Read the range without advancing it: the advance belongs with the
+ # append below, so a retried flush derives the same range.
+ min_seq = self.sequence_generator.start
+ max_seq = self.sequence_generator.current
+ creation_time = Timestamp.now()
+ data_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(min_key, self.trimmed_primary_keys_fields),
+ max_key=GenericRow(max_key, self.trimmed_primary_keys_fields),
+ key_stats=key_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=extra_files,
+ creation_time=creation_time,
+ delete_row_count=0,
+ file_source=0,
+ value_stats_cols=None if value_stats_enabled else [],
+ external_path=external_path_str,
+ first_row_id=None,
+ write_cols=self.write_cols,
+ file_path=file_path,
+ )
+
+ if self.changelog_producer == ChangelogProducer.INPUT:
+ changelog_meta = self._write_changelog_file(
+ data, min_key, max_key, key_stats, value_stats,
+ min_seq, max_seq, creation_time,
+ value_stats_enabled, external_path_str is not None,
+ )
except Exception:
self.file_io.delete_quietly(file_path)
if row_sidecar_path is not None:
self.file_io.delete_quietly(row_sidecar_path)
raise
- # min key & max key
-
- selected_table = data.select(self.trimmed_primary_keys)
- key_columns_batch = selected_table.to_batches()[0]
- min_key_row_batch = key_columns_batch.slice(0, 1)
- max_key_row_batch = key_columns_batch.slice(key_columns_batch.num_rows
- 1, 1)
- min_key = [col.to_pylist()[0] for col in min_key_row_batch.columns]
- max_key = [col.to_pylist()[0] for col in max_key_row_batch.columns]
-
- # key stats & value stats
- value_stats_enabled = self.options.metadata_stats_enabled()
- if value_stats_enabled:
- stats_fields = self.table.fields if
self.table.is_primary_key_table \
- else PyarrowFieldParser.to_paimon_schema(data.schema)
- else:
- stats_fields = self.table.trimmed_primary_keys_fields
- column_stats = {
- field.name: self._get_column_stats(data, field.name)
- for field in stats_fields
- }
- key_fields = self.trimmed_primary_keys_fields
- key_stats = self._collect_value_stats(data, key_fields, column_stats)
- if not self.options.primary_key_nullable() and not all(
- count == 0 for count in key_stats.null_counts):
- raise RuntimeError("Primary key should not be null")
-
- value_fields = stats_fields if value_stats_enabled else []
- value_stats = self._collect_value_stats(data, value_fields,
column_stats)
-
- min_seq = self.sequence_generator.start
- max_seq = self.sequence_generator.current
self.sequence_generator.start = self.sequence_generator.current
- creation_time = Timestamp.now()
- self.committed_files.append(DataFileMeta.create(
- file_name=file_name,
- file_size=self.file_io.get_file_size(file_path),
- row_count=data.num_rows,
- min_key=GenericRow(min_key, self.trimmed_primary_keys_fields),
- max_key=GenericRow(max_key, self.trimmed_primary_keys_fields),
- key_stats=key_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=extra_files,
- creation_time=creation_time,
- delete_row_count=0,
- file_source=0,
- value_stats_cols=None if value_stats_enabled else [],
- external_path=external_path_str,
- first_row_id=None,
- write_cols=self.write_cols,
- file_path=file_path,
- ))
-
- if self.changelog_producer == ChangelogProducer.INPUT:
- self._write_changelog_file(
- data, min_key, max_key, key_stats, value_stats,
- min_seq, max_seq, creation_time,
- value_stats_enabled, external_path_str is not None,
- )
+ self.committed_files.append(data_meta)
+ if changelog_meta is not None:
+ self.committed_changelog_files.append(changelog_meta)
def _apply_variant_shredding(self, data: pa.Table) -> pa.Table:
"""Transform VARIANT columns into shredded Parquet format.
@@ -351,48 +399,58 @@ class DataWriter(ABC):
def _write_changelog_file(self, data, min_key, max_key, key_stats,
value_stats,
min_seq, max_seq, creation_time,
- value_stats_enabled, is_external):
+ value_stats_enabled, is_external) ->
DataFileMeta:
+ """Write the changelog file for one data file and return its meta.
+
+ The caller appends the returned meta only once the whole data file has
+ landed, so a failure here leaves nothing behind: no meta to commit, and
+ no file on disk either.
+ """
cl_fmt = self.changelog_file_format
changelog_file_name = f"changelog-{uuid.uuid4()}-0.{cl_fmt}"
changelog_file_path = self._generate_file_path(changelog_file_name)
changelog_external_path = changelog_file_path if is_external else None
- if cl_fmt == CoreOptions.FILE_FORMAT_PARQUET:
- self.file_io.write_parquet(changelog_file_path, data,
compression=self.compression,
+ try:
+ if cl_fmt == CoreOptions.FILE_FORMAT_PARQUET:
+ self.file_io.write_parquet(changelog_file_path, data,
compression=self.compression,
+ zstd_level=self.zstd_level)
+ 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_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:
- raise ValueError(f"Unsupported changelog file format: {cl_fmt}. "
- f"Supported formats: parquet, orc, avro.")
-
- self.committed_changelog_files.append(DataFileMeta.create(
- file_name=changelog_file_name,
- file_size=self.file_io.get_file_size(changelog_file_path),
- row_count=data.num_rows,
- min_key=GenericRow(min_key, self.trimmed_primary_keys_fields),
- max_key=GenericRow(max_key, self.trimmed_primary_keys_fields),
- key_stats=key_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=creation_time,
- delete_row_count=0,
- file_source=0,
- value_stats_cols=None if value_stats_enabled else [],
- external_path=changelog_external_path,
- first_row_id=None,
- write_cols=self.write_cols,
- file_path=changelog_file_path,
- ))
+ 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:
+ raise ValueError(f"Unsupported changelog file format:
{cl_fmt}. "
+ f"Supported formats: parquet, orc, avro.")
+
+ return DataFileMeta.create(
+ file_name=changelog_file_name,
+ file_size=self.file_io.get_file_size(changelog_file_path),
+ row_count=data.num_rows,
+ min_key=GenericRow(min_key, self.trimmed_primary_keys_fields),
+ max_key=GenericRow(max_key, self.trimmed_primary_keys_fields),
+ key_stats=key_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=creation_time,
+ delete_row_count=0,
+ file_source=0,
+ value_stats_cols=None if value_stats_enabled else [],
+ external_path=changelog_external_path,
+ first_row_id=None,
+ write_cols=self.write_cols,
+ file_path=changelog_file_path,
+ )
+ except Exception:
+ self.file_io.delete_quietly(changelog_file_path)
+ raise
def _generate_file_path(self, file_name: str) -> str:
if self.external_path_provider:
diff --git a/paimon-python/pypaimon/write/writer/dedicated_format_writer.py
b/paimon-python/pypaimon/write/writer/dedicated_format_writer.py
index ae73880fd0..001b0f9f69 100644
--- a/paimon-python/pypaimon/write/writer/dedicated_format_writer.py
+++ b/paimon-python/pypaimon/write/writer/dedicated_format_writer.py
@@ -39,6 +39,7 @@ from pypaimon.write.row_utils import (
row_values_to_arrow_table,
)
from pypaimon.write.writer.data_writer import DataWriter
+from pypaimon.write.writer.write_buffer import WriteBuffer
logger = logging.getLogger(__name__)
@@ -134,9 +135,14 @@ class DedicatedFormatWriter(DataWriter):
self.record_count = 0
self.closed = False
- # Track pending data for normal data only
- self.pending_normal_data: Optional[pa.Table] = 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
@@ -200,12 +206,15 @@ class DedicatedFormatWriter(DataWriter):
return self._merge_normal_data(existing_data, new_data)
def write(self, data: pa.RecordBatch):
+ # Outside the try on purpose: rejecting the write must not abort the
+ # writer, or the unfinished flush would lose its chance to be retried.
+ self._require_finished_flush()
try:
offset = 0
# _write_batch keeps normal/blob/vector pending rows in lockstep
# and closes all writers when the shared row limit is reached.
while offset < data.num_rows:
- capacity = self.target_file_row_num - self._current_row_count()
+ capacity = self.target_file_row_num - self.pending_row_count
if capacity <= 0:
self._close_current_writers()
capacity = self.target_file_row_num
@@ -229,10 +238,7 @@ class DedicatedFormatWriter(DataWriter):
# 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)
+ self._normal_buffer.append(processed_normal)
# Write blob-file columns to dedicated blob writers.
for blob_column, blob_data in blob_data_map.items():
@@ -251,6 +257,7 @@ class DedicatedFormatWriter(DataWriter):
self._close_current_writers()
def write_row(self, row):
+ self._require_finished_flush()
try:
values_by_name = row_to_named_values(
row, self.table.table_schema.fields)
@@ -274,11 +281,7 @@ class DedicatedFormatWriter(DataWriter):
).to_batches()[0]
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)
+ self._normal_buffer.append(processed_normal)
for blob_column in self.blob_file_column_names:
arrow_type = PyarrowFieldParser.from_paimon_type(
@@ -344,7 +347,7 @@ class DedicatedFormatWriter(DataWriter):
raise
finally:
self.closed = True
- self.pending_normal_data = None
+ self._normal_buffer.reset()
def abort(self):
"""Abort all writers and clean up resources."""
@@ -352,9 +355,12 @@ class DedicatedFormatWriter(DataWriter):
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.pending_normal_data = None
- self.pending_data = None
+ self._normal_buffer.reset()
+ self._buffer.reset()
self.committed_files.clear()
self._committed_files_to_delete_on_abort.clear()
@@ -465,10 +471,11 @@ class DedicatedFormatWriter(DataWriter):
return pa.concat_tables([existing_data, new_data])
def _should_roll_normal(self) -> bool:
- if self.pending_normal_data is None:
+ # Runs on every write, so it answers from the running counts only.
+ if self._normal_buffer.is_empty:
return False
- if self.pending_normal_data.num_rows >= self.target_file_row_num:
+ if self._normal_buffer.num_rows >= self.target_file_row_num:
return True
# Check rolling condition periodically (every CHECK_ROLLING_RECORD_CNT
records)
@@ -476,28 +483,46 @@ class DedicatedFormatWriter(DataWriter):
return False
# Check if normal data exceeds target size
- current_size = self.pending_normal_data.nbytes
- return current_size > self.target_file_size
-
- def _current_row_count(self) -> int:
- if self.pending_normal_data is not None:
- return self.pending_normal_data.num_rows
+ 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
+ # fills. Normal, blob and vector rows are kept in lockstep, so any half
+ # answers for all of them; the sidecars are asked only when the table
+ # has no normal columns at all.
+ if not self._normal_buffer.is_empty:
+ return self._normal_buffer.num_rows
for blob_writer in self.blob_writers.values():
if blob_writer.current_writer is not None:
return blob_writer.current_writer.row_count
- if self.vector_writer is not None and self.vector_writer.pending_data
is not None:
- return self.vector_writer.pending_data.num_rows
+ if self.vector_writer is not None:
+ # Running count, not a folded buffer: this runs on every write.
+ return self.vector_writer.pending_row_count
return 0
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)
- self._committed_files_to_delete_on_abort.append(normal_meta)
+ # 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 = []
for blob_column in self.blob_file_column_names:
blob_writer = self.blob_writers[blob_column]
writer_metas = blob_writer.prepare_commit()
@@ -505,20 +530,30 @@ class DedicatedFormatWriter(DataWriter):
self._validate_consistency(normal_meta, writer_metas,
blob_column)
blob_metas.extend(writer_metas)
if blob_writer.delete_file_upon_abort():
- self._committed_files_to_delete_on_abort.extend(writer_metas)
- blob_writer.committed_files.clear()
- self.committed_files.extend(blob_metas)
+ deletable_blob_metas.extend(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')
- self.committed_files.extend(vector_metas)
- self._committed_files_to_delete_on_abort.extend(vector_metas)
+
+ # 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_data = None
+ self._pending_normal_meta = None
self.record_count = 0
if normal_meta is not None or blob_metas or vector_metas:
diff --git a/paimon-python/pypaimon/write/writer/key_value_data_writer.py
b/paimon-python/pypaimon/write/writer/key_value_data_writer.py
index 86b6ad365a..1f1537f6d1 100644
--- a/paimon-python/pypaimon/write/writer/key_value_data_writer.py
+++ b/paimon-python/pypaimon/write/writer/key_value_data_writer.py
@@ -31,7 +31,7 @@ from pypaimon.write.writer.data_writer import DataWriter
class KeyValueDataWriter(DataWriter):
"""Data writer for primary key tables with system fields and sorting.
- Accumulates incoming batches in ``pending_data`` without sorting or
+ Accumulates incoming batches in the write buffer without sorting or
folding on the write path. Sort and ``MergeFunction``-based fold
are deferred to flush time (``_flush_all``), where the result is
roll-written into one or more data files. This enforces the LSM
@@ -52,9 +52,9 @@ class KeyValueDataWriter(DataWriter):
def _process_data(self, data: pa.RecordBatch) -> pa.Table:
# No sort here: sorting once at flush is strictly cheaper than
- # per-batch sort + a final global sort. ``pending_data`` ends
- # up as a concat of unsorted batches; ``_flush_all`` sorts it
- # exactly once before folding.
+ # per-batch sort + a final global sort. The buffer ends up as a
+ # concat of unsorted batches; ``_flush_all`` sorts it exactly
+ # once before folding.
enhanced_data = self._add_system_fields(data)
return pa.Table.from_batches([enhanced_data])
@@ -64,31 +64,30 @@ class KeyValueDataWriter(DataWriter):
return pa.concat_tables([existing_data, new_data])
def prepare_commit(self) -> List[DataFileMeta]:
- if self.pending_data is not None and self.pending_data.num_rows > 0:
+ if self._buffer.num_rows > 0:
self._flush_all()
- # ``_flush_all`` leaves ``pending_data = None``, so super's
- # prepare_commit just returns ``committed_files``.
+ # ``_flush_all`` empties the buffer, so super's prepare_commit just
+ # returns ``committed_files``.
return super().prepare_commit()
def _check_and_roll_if_needed(self):
- # Buffer overflowed target_file_size: sort + fold + roll-write
- # the whole buffer as multiple files in one pass. Unlike the
- # base class's slice loop, we never keep a slice remainder in
- # ``pending_data`` -- flush empties the buffer outright.
- if (self.pending_data is not None
- and self.pending_data.num_rows > 0
- and self.pending_data.nbytes > self.target_file_size):
+ # Buffer overflowed target_file_size: sort + fold + roll-write the
whole
+ # buffer as multiple files in one pass. Unlike the base class's slice
+ # loop, a successful flush leaves no remainder -- it empties the buffer
+ # outright. There is no row-count trigger here, and none is needed:
+ # ``FileStoreWrite`` rejects target-file-row-num on primary key tables.
+ if self._buffer.num_rows > 0 and self._buffer.nbytes >
self.target_file_size:
self._flush_all()
def close(self):
# Override the base ``close`` because its straight
- # ``_write_data_to_file(pending_data)`` would land an unsorted,
- # un-folded buffer on disk -- violating the file-internal
+ # ``_write_data_to_file`` of the whole buffer would land an unsorted,
+ # un-folded table on disk -- violating the file-internal
# PK-unique invariant. Route the final flush through
# ``_flush_all`` so the contract holds even on the
# close-without-prepare_commit path.
try:
- if self.pending_data is not None and self.pending_data.num_rows >
0:
+ if self._buffer.num_rows > 0:
self._flush_all()
except Exception as e:
import logging
@@ -99,24 +98,29 @@ class KeyValueDataWriter(DataWriter):
self.abort()
raise e
finally:
- self.pending_data = None
+ self._buffer.reset()
def _flush_all(self) -> None:
"""Sort + fold the entire buffer, then roll-write as files.
- On return, ``pending_data is None`` and every flushed chunk
- has been recorded in ``committed_files``. The buffer is
- always fully drained per flush: no slice remainder is
- carried back into ``pending_data``.
+ On return, the buffer is empty and every flushed chunk has been
+ recorded in ``committed_files``. If a file write fails, the buffer is
+ left holding exactly the rows no file has taken yet, so a retried
+ flush neither loses nor duplicates them.
"""
- if self.pending_data is None or self.pending_data.num_rows == 0:
- self.pending_data = None
+ pending = self._buffer.materialize()
+ if pending is None or pending.num_rows == 0:
+ self._buffer.reset()
return
- sorted_data = self._sort_by_primary_key(self.pending_data)
+ sorted_data = self._sort_by_primary_key(pending)
folded = self._merge_pending_by_pk(sorted_data)
- self.pending_data = None
if folded.num_rows == 0:
+ self._buffer.reset()
return
+ # Park the folded rows in the buffer for ``_roll_write`` to drain. Both
+ # the sort and the fold are idempotent on their own output -- PKs are
+ # unique once folded -- so a retry over what is left is still correct.
+ self._buffer.reset(folded)
self._roll_write(folded)
def _roll_write(self, data: pa.Table) -> None:
@@ -127,10 +131,15 @@ class KeyValueDataWriter(DataWriter):
size does not violate the LSM file-internal invariant.
Reuses ``_find_optimal_split_point`` / ``_write_data_to_file``
from the base class.
+
+ The buffer is narrowed to the rows still unwritten after each file, so
+ a failure part way through leaves the remainder -- and only the
+ remainder -- for whoever flushes next.
"""
while data.num_rows > 0:
if data.nbytes <= self.target_file_size:
self._write_data_to_file(data)
+ self._buffer.reset()
return
split_row = self._find_optimal_split_point(
data, self.target_file_size)
@@ -138,9 +147,11 @@ class KeyValueDataWriter(DataWriter):
# Single row already exceeds target_file_size; nothing
# to gain from further slicing, write it as-is.
self._write_data_to_file(data)
+ self._buffer.reset()
return
self._write_data_to_file(data.slice(0, split_row))
data = data.slice(split_row)
+ self._buffer.reset(data)
def _merge_pending_by_pk(self, data: pa.Table) -> pa.Table:
"""Fold same-PK runs in ``data`` using ``self._merge_function``.
diff --git a/paimon-python/pypaimon/write/writer/write_buffer.py
b/paimon-python/pypaimon/write/writer/write_buffer.py
new file mode 100644
index 0000000000..b07a899fba
--- /dev/null
+++ b/paimon-python/pypaimon/write/writer/write_buffer.py
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import Callable, List, Optional
+
+import pyarrow as pa
+
+
+class WriteBuffer:
+ """Accumulates Arrow tables, concatenating them only when asked.
+
+ ``append`` adds a table's row count and size to running totals, so a writer
+ can answer its rolling check without touching Arrow. ``materialize`` and
+ ``take`` are the operations that concatenate.
+ """
+
+ def __init__(self, merge: Callable[[pa.Table, pa.Table], pa.Table]):
+ # ``merge(existing, new)`` combines two tables the way the owning
writer
+ # wants, and runs once per ``materialize`` rather than once per append.
+ self._merge = merge
+ self._table: Optional[pa.Table] = None
+ self._appended: List[pa.Table] = []
+ self._schema: Optional[pa.Schema] = None
+ # ``concat_tables`` only collects chunks and ``nbytes`` sums the
buffers
+ # they reference, so this running total is what the fold will report.
+ self.nbytes = 0
+ self.num_rows = 0
+
+ @property
+ def is_empty(self) -> bool:
+ """True when nothing has been appended and no table has been set.
+
+ Distinct from ``num_rows == 0``, which a zero-row table also satisfies.
+ """
+ return self._table is None and not self._appended
+
+ def append(self, data: pa.Table) -> None:
+ # ``concat_tables`` rejects any schema difference while ``TableWrite``
+ # admits a few (differing nullability, ``binary`` vs
+ # ``fixed_size_binary``). Reject here so those keep failing in
+ # ``write``, which aborts, instead of in ``prepare_commit``, which does
+ # not. ``Schema.equals`` ignores metadata, and so does concat.
+ if self._schema is None:
+ self._schema = data.schema
+ elif not data.schema.equals(self._schema):
+ raise ValueError(
+ "Cannot buffer a batch whose schema differs from the batches "
+ f"already buffered.\nBuffered schema is: {self._schema}\n"
+ f"Incoming schema is: {data.schema}")
+ self._appended.append(data)
+ self.nbytes += data.nbytes
+ self.num_rows += data.num_rows
+
+ def materialize(self) -> Optional[pa.Table]:
+ """Concatenate everything appended so far into one table and return it.
+
+ Returns None while the buffer is empty, and is a no-op when called
again
+ with nothing appended since.
+ """
+ if self._appended:
+ folded = (self._appended[0] if len(self._appended) == 1
+ else pa.concat_tables(self._appended))
+ self._appended = []
+ self._table = (folded if self._table is None
+ else self._merge(self._table, folded))
+ self._schema = self._table.schema
+ self.nbytes = self._table.nbytes
+ self.num_rows = self._table.num_rows
+ return self._table
+
+ def take(self) -> Optional[pa.Table]:
+ """Return everything buffered as one table and empty the buffer."""
+ table = self.materialize()
+ self.reset()
+ return table
+
+ def reset(self, table: Optional[pa.Table] = None) -> None:
+ """Replace the contents with ``table``, or empty the buffer.
+
+ ``table`` is measured from scratch because the usual caller passes a
+ slice of the table the running totals were describing.
+ """
+ self._appended = []
+ self._table = table
+ self._schema = None if table is None else table.schema
+ self.nbytes = 0 if table is None else table.nbytes
+ self.num_rows = 0 if table is None else table.num_rows
