[python] Fix quadratic cost of writing multiple batches

[yugan95]

Purpose

Linked issue: None

Filling one data file with M batches costs O(M^2). So write_arrow_batch in a loop gets slower with every call, even when the data is nowhere near target-file-size and no file ever rolls.

M calls with 100 rows each, plus prepare_commit and commit, unaware table, default options:

batches	rows	before	after
500	50k	0.42s	0.05s
1000	100k	1.61s	0.08s
2000	200k	6.35s	0.16s

Doubling M used to quadruple the time; now it doubles it — quadratic became linear. Per call: 810 µs at M=500 and 3160 µs at M=2000 before, a flat ~64 µs either way after.

Root Cause

Every write merged the incoming batch into the pending table, then measured that table to decide whether to roll the file:

# DataWriter.write
if self.pending_data is None:
    self.pending_data = processed_data
else:
    self.pending_data = self._merge_data(self.pending_data, processed_data)
self._check_and_roll_if_needed()   # reads pending_data.nbytes

_merge_data is pa.concat_tables in every implementation, and concat_tables does not combine data — the result just references the chunks of both inputs. So after the M-th write the pending table holds M chunks per column.

nbytes adds up the bytes the table's buffers reference, so reading it walks all M chunks. That makes the rolling check O(M) on the M-th write, and it runs on every write: 1 + 2 + ... + M = O(M^2). The concats themselves are cheap — 1% of an isolated 2000-write run, since they only copy chunk pointers — and the nbytes walks are the other 99%.

The check needs two numbers, total bytes and total rows, and it almost always answers "don't roll". The rows only have to be one table when a file actually rolls.

Changes

• Add WriteBuffer (write_buffer.py): appends go into a list with running nbytes / num_rows totals, and concatenation happens only in materialize / take. The totals are exact rather than estimates — concat_tables only collects chunks, so the per-table sizes sum to what the concatenated table reports (a test asserts this).
• _check_and_roll_if_needed gates on the running totals and materializes only once a threshold is crossed. This is the condition the old loop reached by computing split_row and breaking when it landed at num_rows, so which writes roll is unchanged and _find_optimal_split_point still sees the same table.
• pending_data becomes self._buffer across all five writers. That removes four copies of the same "if the buffer is empty assign, else concat" block, since WriteBuffer.append treats the first table like any other. A new DataWriter.pending_row_count reads the buffered row count without materializing.
• BlobWriter still drains on every write: its roll decision depends on external blob bytes, which the buffered descriptors say nothing about.
• A flush clears the buffer only once the write lands, so a failed prepare_commit can be retried. DataWriter.prepare_commit and both composite writers' _close_current_writers reset at the end. KeyValueDataWriter._flush_all spans several files, so it keeps the rows no file has taken yet instead — a retry resumes rather than rewriting.
• WriteBuffer.append rejects a schema mismatch. TableWrite._validate_pyarrow_schema admits differences concat_tables rejects (nullability, binary vs fixed_size_binary), and those used to fail in write, which aborts and cleans up; deferring the fold would move them to prepare_commit, which has no handler, orphaning already-rolled files. Schema.equals has the same tolerance as concat — ignores metadata, enforces nullability — so nothing the old path accepted is newly rejected.

Tests

New pypaimon/tests/write/write_buffer_test.py: 200 writes fold zero times while every rolling trigger still fires, and fail-once writers cover the failure paths — a failed prepare_commit keeps its rows for the retry, and a PK flush that fails on its second file leaves exactly the unwritten remainder.

test_write_merge_buffer.py migrated to WriteBuffer.

API and Format

No format or on-disk change. DataWriter.pending_data is replaced by the private _buffer plus a public pending_row_count; DataWriter is internal to pypaimon.write.writer.

Documentation

N/A

[JingsongLi]

[P2] Preserve buffered rows until the prepare flush succeeds

take() resets the buffer before _write_data_to_file performs fallible file I/O. For a reusable StreamTableWrite, a transient storage error followed by retrying prepare_commit on the same writer silently omits these rows: with a fail-once writer, this head leaves 0 rows buffered and the retry writes nothing, while the base version retains all 3 rows and writes them on retry. Please materialize without draining, perform the write, and reset only after it succeeds (or explicitly poison/abort the writer and reject reuse).

[yugan95]

Good catch, fixed. Neither prepare_commit nor close drains up front any more: prepare_commit resets after the write returns, and close keeps its finally reset — both matching master. Regression test added.

Same pattern in KeyValueDataWriter._flush_all (present on master too), fixed as well but differently: _roll_write spans several files, so keeping every row would make a retry rewrite the ones that already landed. The buffer now keeps only the rows no file has taken. Safe because the sort and the fold are idempotent on their own output — PKs are unique once folded.

BlobWriter.take() is fine as-is: it only runs inside write(), which aborts on failure.

[JingsongLi]

[P2] Keep the normal and sidecar flush failure-atomic

_normal_buffer.take() discards the normal rows before the normal file is written. If that write fails, the vector buffer remains pending; retrying prepare_commit then has no normal_meta, flushes the vector files, and skips the row-count consistency check, so it can return sidecar-only metadata and lose the normal half. Please retain the normal table until the coordinated flush succeeds, or abort/poison the composite writer on any phase failure. The same ordering issue is present in DedicatedFormatWriter._close_current_writers at its _normal_buffer.take() call.

[yugan95]

Fixed the same way: _close_current_writers resets at the end, once the normal file and the sidecars have both landed. Regression test added.

[JingsongLi]

+1

[JingsongLi]

cc @XiaoHongbo-Hope to take a review.

[XiaoHongbo-Hope]

Sidecar failure can duplicate normal data and break normal-sidecar consistency on retry. Please make the composite flush retry-safe.

[XiaoHongbo-Hope]

Changelog failure can commit the same data chunk twice on retry, making data and changelog inconsistent. Please make the data/changelog flush retry-safe.