[python] Reduce memory usage for Parquet row-id updates

[XiaoHongbo-Hope]

Background

Data-evolution column updates rewrite the affected file-group range, even when only a few rows change. Materializing the projected columns and buffering the merged output can cause high memory usage for large string or nested columns.

Changes

Stream ordinary Parquet updates with bounded row-group buffering independent of read batch size. Preserve row-id ranges, statistics, metadata, and failure cleanup. Deduplicate update projections consistently. Specialized format paths remain unchanged.

This reduces whole-group buffering, but does not eliminate read/write amplification or impose a hard process-memory limit.

Tests

Cover batch/stream upserts, duplicate update columns, nested values, nulls, row-group boundaries, and write/close failure cleanup, including Python 3.6 / PyArrow 6 compatibility.

[JingsongLi]

Reviewed by pulling the branch locally (head d8c9955) and running the affected suites.

Verification notes

• table_upsert_by_key_test.py, table_update_test.py, table_update_by_row_id_chunked_test.py and the four data_evolution_* suites: 301 passed (excluding vortex/lance/mosaic; the test_vortex_basic segfault on my machine is a local native-lib problem, unrelated to this PR).
• I instrumented _write_group to count which branch is taken: the streaming path is hit 263 times and the legacy path 0 times across those suites, while the blob/vector suites still exercise the legacy path 23 times. So the new path has broad regression coverage and the old one is still guarded.
• _create_data_file_meta is behaviour-preserving. The base revision had external_path_str = file_path if is_external_path else None and value_stats_enabled = self.options.metadata_stats_enabled(), which the helper reproduces exactly; the extra_files / creation_time None-fallbacks are unreachable from the existing call site.
• The fast-path predicate lines up field-by-field with the writer dispatch in FileStoreWrite._create_writer: _has_blob_columns() uses the same is_blob_file_field, the sidecar and shredding conditions match _should_write_row_sidecar() and DataWriter._variant_shredding, and the vector condition is strictly stronger, which is the safe direction.
• file_io.write_parquet is just pq.write_table(compression=..., compression_level=...), so the pq.ParquetWriter usage keeps the same layout and store_schema metadata; new_output_stream creates parent directories in both the local and pyarrow FileIO implementations.
• _row_groups cannot spin or drop rows: the buffer is flushed as soon as num_rows >= _ROW_GROUP_MAX_ROWS, so count >= 1 always holds, and the degenerate low == 0 case either flushes a non-empty buffer or forces a single oversized row.
• Cross-row-group stats aggregation equals the single-shot computation (_get_column_stats does no truncation), and _collect_value_stats(None, fields, stats) never touches data when column_stats is supplied, so passing None is safe.
• Making merged_schema explicit is a fix rather than a regression: _merge_chunked_column returns pa.chunked_array(..., type=original_col.type), so the type always matches original_data.schema.field(name) and only the nullability flag is newly carried over.
• Sequence numbers: hardcoding 0/0 agrees with the legacy path, because FileStoreWrite also passes seq_number = 0 for BUCKET_UNAWARE append tables and AppendOnlyDataWriter never calls sequence_generator.next().

No correctness or data-loss issue found; the new tests are solid (injecting AssertionError into to_arrow to prove nothing is materialized, the interleaved read == written assertion to prove real streaming, and injected write/close failures to check overlay cleanup). A few non-blocking suggestions inline.

Suggested follow-up tests

1. Guard tests asserting _write_group_streaming is not called for blob / variant-shredding / row-sidecar / vector / changelog-producer != none tables, so the predicate cannot be loosened by accident later.
2. A schema-evolution case: update a column that does not exist in the original file yet (the _ROW_ID anchor plus all-null baseline path) through the streaming writer.
3. Boundary row ids: update exactly the first and the last row of a group.

[JingsongLi]

_original_file_read deliberately keeps SpecialFields.ROW_ID in read_type ("Keep _ROW_ID as a row-count anchor"), but the actual values are then dropped by select(column_names) and the row mapping is re-derived arithmetically from offset = first_row_id and end = offset + batch.num_rows.

That silently assumes the reader hands back every row of the group, contiguously, in physical order. It holds today, but any future row-level filtering on the read path (predicate push-down, an authorization filter, deletion vectors, a reordering or parallel reader) would make updates land on the wrong rows with no error at all. Since the column is already being read, one check turns that class of regression into a loud failure:

row_ids = batch[SpecialFields.ROW_ID.name]
if row_ids[0].as_py() != offset:
    raise ValueError(
        f'Original file group is not contiguous at {offset}')

The previous implementation made the same assumption, but it materialized the whole group in one shot, so the assumption was less load-bearing; streaming relies on it per batch.

[JingsongLi]

This only catches update row ids past the end of the group. A row id below first_row_id is picked up by the < end test in the very first batch and then surfaces from _merge_update_with_original as IndexError: Update position -N is outside column range [0, M), which does not point at the real cause.

_calculate_first_row_id's valid_row_id_ranges check normally rejects such input, so this is about diagnosability rather than correctness. Making the selection two-sided would cover both directions:

while update_index < len(updates) and updates[update_index][1] < end:
    if updates[update_index][1] < offset:
        raise ValueError('Update row IDs precede the original file group')

[JingsongLi]

_write_batches is a private name carrying a public contract: it names the file, owns the output stream, aggregates stats, appends to committed_files and cleans up on failure. In other words it re-implements a good part of DataWriter._write_data_to_file, yet it is reached only from here and never from DataWriter.write / prepare_commit.

The risk is drift: whoever later adds a step to _write_data_to_file (a new sidecar, an extra metadata field, a validation) will silently miss this second file-producing path. Extracting _create_data_file_meta is a good start; consider either dropping the underscore to signal this is a real entry point, or moving it into a small dedicated writer so both paths share the finalization.

Minor, same area: this is the first call site that uses _to_managed_arrow_batch_reader as a context manager (multimodal/temporal.py and multimodal/query.py do not). _ClosableArrowBatchReader implements __enter__ / __exit__ so it works, and RecordBatchReader.from_stream only exists on new enough pyarrow, but it is worth saying so in that method's docstring.

[XiaoHongbo-Hope]

_write_batches is a private name carrying a public contract: it names the file, owns the output stream, aggregates stats, appends to committed_files and cleans up on failure. In other words it re-implements a good part of DataWriter._write_data_to_file, yet it is reached only from here and never from DataWriter.write / prepare_commit.

The risk is drift: whoever later adds a step to _write_data_to_file (a new sidecar, an extra metadata field, a validation) will silently miss this second file-producing path. Extracting _create_data_file_meta is a good start; consider either dropping the underscore to signal this is a real entry point, or moving it into a small dedicated writer so both paths share the finalization.

Minor, same area: this is the first call site that uses _to_managed_arrow_batch_reader as a context manager (multimodal/temporal.py and multimodal/query.py do not). _ClosableArrowBatchReader implements __enter__ / __exit__ so it works, and RecordBatchReader.from_stream only exists on new enough pyarrow, but it is worth saying so in that method's docstring.

I prefer to move it into a small dedicated writer.

[JingsongLi]

pa.__version__ is re-parsed on every slice, and _row_group_slice is called O(log n) times per row group from the binary search. Worth hoisting to a module-level constant, e.g. _ARROW_MAJOR = int(pa.__version__.split('.')[0]).

[JingsongLi]

Two small things here.

• 128 * 1024 * 1024 duplicates the Parquet block-size default that lives on the Java side, and file.block-size is declared no_default_value(), so this magic number is the effective default for Python writers. A named constant with a short comment would make that coupling explicit.
• The budget is measured in Arrow in-memory bytes, not encoded Parquet bytes, so the resulting row groups end up substantially smaller than file.block-size suggests. The docstring already says the byte count is an estimate; spelling out that it is an Arrow-side estimate would set expectations better.

Also, the file.block-size must be positive check on the next line lives in a generator body, so it only fires once the first item is pulled. By then _write_batches has already created the output file (it is removed by delete_quietly, so nothing leaks). Config validation like this fits better at writer construction time.