feat: charge native write buffers to the task's memory pool

[andygrove]

Which issue does this PR close?

Closes #5648.
Closes #6115.

Rationale for this change

Neither native writer reserved the memory it buffers, so none of it was visible to Comet's memory pool or to Spark.

• ParquetWriterExec holds its in-progress row group. It uses parquet-rs's default limit of 1Mi rows and no byte limit.
• The native Iceberg writer holds an in-progress row group for every open file. A fanout write keeps a file open for every partition a task writes to. Each partition also holds up to 999 rows back, waiting to fill the 1000-row unit the rolling writer is fed in.

A wide fanout write therefore grew in memory that no budget accounted for, until the container was killed. The JVM writers hold the same buffers on the JVM heap, which executors are sized for.

#5648 asked for an audit of the Iceberg writer's buffers, and a test that a fanout write over many partitions with a small memory budget fails with a Comet out-of-memory error rather than a process-level one. #6115 asked the same of both writers, with the buffers charged to the task's pool. This PR does both.

What changes are included in this PR?

• Both writers register one MemoryConsumer per task, ParquetWriterExec[N] and IcebergWriteExec[N], and try_resize it after every batch. It is one per task rather than one per file because every consumer registered with fair_unified lowers the share of every other consumer in the task.
• ParquetWriterExec reserves ArrowWriter::memory_size(), which counts encoded pages, encoder buffers, dictionaries and Bloom filters. On the remote (HDFS) path it adds the staging buffer's capacity.
• The Iceberg writer wraps iceberg-rust's ParquetWriterBuilder in a MeteredParquetWriterBuilder. That is the only hook, because the rolling and partitioning writers keep their file writers private. iceberg-rust's ParquetWriter exposes current_written_size(), which is every byte the file has written, the in-progress row group's encoded size included. It does not expose parquet-rs's memory_size(). So the builder hands ParquetWriter each file's OutputFile behind a CountedOutput, whose writer counts the bytes that leave memory on their way to storage, and each open file reports what it has written less those. When bytes leave depends on the storage:
  • A local file (file, memory) writes a flushed row group out at once, so it reports its in-progress row group and the few KiB parquet-rs buffers in front of storage.
  • S3 and GCS take a file in parts of at least 5 MiB, and OpenDAL holds each part in memory until the next one is complete or the file closes. A file there also reports what it has flushed since its last upload: at the default 128 MiB row group, the last row group it flushed.
• The Iceberg reservation also includes the rows each RowPacer holds back.
• Neither writer can spill, so a resize the pool refuses fails the task, the way DataFusion's own parquet sink does. The JVM sees CometNativeException: Additional allocation failed for IcebergWriteExec[0] ..., a task failure Spark retries. The Iceberg writer's abort guard deletes the files the task had opened, as on any other task failure.
• Docs:
  • memory_management.md has a new "Native writers" section.
  • The contributor iceberg-writes.md describes the metered builder.
  • The user guide's failure-handling section explains what a fanout write needs and how to make it fit.
  • The Iceberg-write review skill gets a checklist item.

What this does not do:

• Nothing sheds memory when the pool refuses. Flushing the row group early in ParquetWriterExec, or closing the largest open file in a fanout write, would let the write continue with a different file layout. Those can be follow-ups.
• The Iceberg reservation misses what parquet-rs holds beyond the encoded estimate: dictionary hash tables, unencoded dictionary indices and buffer capacity. An exact figure needs an upstream accessor on iceberg-rust's ParquetWriter. Native Iceberg writes decline Bloom filters today, and feat: support bloom filters in native Iceberg writes #5724 would need to reserve them itself.
• The S3 and GCS figure follows OpenDAL 0.58's writer: a 5 MiB minimum part, and one complete part held until the next. If OpenDAL changes how it buffers parts, HeldByStorage has to follow.
• There is no cap on open fanout writers. Accounting for them keeps the shape of iceberg-java's FanoutDataWriter, which also keeps every partition's file open.
• ParquetWriterExec still writes 1Mi-row row groups with no byte limit (Native Parquet writer ignores Spark's Parquet writer properties (block size, page size, dictionary, writer version, statistics) #5304), so a wide write needs a correspondingly large share of the pool.

Behavior change: both native writers are off by default. A write whose buffers exceed the task's share of the pool now fails the task. Before, it succeeded in untracked memory, or got the executor killed.

How are these changes tested?

• Rust, iceberg_write.rs:
  • a_fanout_write_reserves_the_rows_its_partitions_hold_back: 900 rows held back in each of 1 and 8 partitions reserve 9.5 KB and 74 KB.
  • the_reservation_follows_the_files_that_are_open: eight one-unit partitions reserve 32 KB fanout, 4.1 KB clustered, and 4.5 KB when an unpartitioned write rolls through eight files.
  • a_file_reserves_no_more_than_one_row_group: a 225 KB file written with 16 KiB row groups reserves at most 17.8 KB, its in-progress row group and parquet-rs's write buffer.
  • files_that_flush_give_back_what_they_reserved: twelve fanout partitions, each a unit of 1,000 distinct 256-byte strings that flushes as soon as it is written into 128 KiB row groups, fit in a pool of eight row groups (1 MiB), with a peak reservation of 0. Charging each flushed file a whole row group failed this write on its ninth partition.
  • a_files_share_is_what_it_still_holds: a local file's share falls back each time a row group flushes and stays under a row group plus parquet-rs's 8 KiB write buffer. A file taken in parts reports everything it has written until that makes a part.
  • storage_taken_in_parts_holds_the_last_part_until_the_next and only_local_storage_writes_through: the part model and which schemes write through.
  • a_write_the_pool_cannot_hold_fails_and_deletes_its_files: fails with ResourcesExhausted naming the consumer, leaves no files, and gives back its reservation.
  • an_open_file_gives_back_its_share_when_it_closes_or_is_dropped and pacer_counts_the_memory_of_the_rows_it_holds_back.
• Rust, parquet_writer.rs: the row group is reserved until the file closes, and a pool too small for it fails the write.
• Mutations, each caught by at least one test: no resize; a resize without the held-back rows; a resize without the open files; a file's share leaked on close; a share not given back on drop; the pacer's count not reset after a hand-over; only one fanout partition counted; a writer that does not count what it hands on; a share that ignores what has left memory; the share capped at the row group again; local storage taken as uploading in parts; a builder that ignores the storage; the part model releasing the new part instead of the one it held.
• JVM. Both tests shrink the pool to about 4 MiB with spark.comet.exec.memoryPool.fraction=0.002.
  • CometIcebergWriteActionSuite "native acceleration: a fanout write that outgrows the memory pool fails its task" writes 64 partitions from one task. It checks the error, that nothing was committed, and that no data files were left (the writer deleted the 12 it had opened). Then the same write with the whole pool succeeds natively.
  • CometParquetWriterSuite "a row group the pool cannot hold fails its task with a native out-of-memory error" does the same for a single 10 MB row group.
  • Both fail with the resizes removed, which is main's behaviour.
• OpenDAL's part holding was checked against MinIO through iceberg-storage-opendal's S3 FileWrite at the pinned rev. Writing 6, 6, 1, 1 and 4 MiB, MinIO received part 1 only during the second write, part 2 only during the fifth, and part 3 at close.
• Runs:
  • cargo test -p datafusion-comet --lib passes (707 tests), and clippy with -D warnings and cargo fmt are clean.
  • On Spark 4.1, CometIcebergWriteActionSuite, CometIcebergRewriteActionSuite and CometIcebergSystemFunctionSuite pass (110 tests), and so does CometParquetWriterSuite (50 tests).
  • The two new JVM tests also passed on Spark 3.5 at the first commit.
  • I have not run Spark 3.4, 4.0 or the Iceberg Spark tests locally.

[sunchao]

Summary

• Prior state and problem: Both native writers retained buffers outside Comet’s memory accounting.
• Design approach: Add one task-level reservation covering writer buffers and Iceberg’s pending rows.
• Correctness / compatibility analysis: Found one introduced P2 issue: flushed Iceberg data remains charged as live memory, causing false out-of-memory failures. Compared Spark’s memory and task-failure handling across supported versions and Iceberg’s Java row-group lifecycle.
• Key design decisions: One consumer avoids multiplying fair-pool participants with partition count. File rolling and cleanup ownership remain unchanged.
• Implementation sketch: Parquet uses ArrowWriter::memory_size(). Iceberg adds a metered writer wrapper, per-file shares, and pending-row accounting. Accounting runs after each input batch, including a walk over fanout pacers.
• Behavioral changes worth calling out: Pool refusals now fail the task. Neither writer spills. Failure cleanup worked in focused validation.
• Suggested improvements: Charge live Iceberg buffers and release reservations after row-group flushes, with the regression described below.

Reviewed the entire 10-file diff from 14f0f59f74dbcf273ff3cbdb1721a5ee535bbe84 to d077400f307cad8e0df66d4ac6850a7433332639. PR remains open and non-draft. No existing reviews, comments, or threads were present.

Routed skills: review-comet-pr, review-comet-memory-pr, and review-comet-iceberg-write-pr.

Exact-head CI: 24 successful checks, 14 skipped, no failures or pending checks. Logs confirm both new JVM failure tests passed. Spark SQL suites, upstream Iceberg suites, and macOS were skipped.

Local validation: 50 Iceberg writer tests, six Parquet writer tests, and the disposable task-path reproduction passed using --no-default-features. JVM and HDFS integration tests were not run locally. Disposable repository changes were removed, and the checkout is clean.

[sunchao]

[P2] Release the row-group charge after a flush. current_written_size() includes bytes already flushed to storage, so this cap leaves every flushed fanout file permanently charged for a full row group until it closes. With 128 KiB row groups, uncompressed 1,000-row batches containing distinct 256-byte strings, a 512 MiB file target, and a 1 MiB pool, the ninth partition requests 1,179,648 bytes and fails with ResourcesExhausted, although the equivalent parquet-rs writers report zero live row-group bytes after every batch. Flushed buffers should release their reservation so subsequent partitions can reuse the budget. This introduces repeatable task failures for writes whose row-group buffers do not accumulate. Please use a live-buffer accessor from iceberg-rust and add a regression covering reservation reuse after flushing.

Evidence: Reproduced through exact-head run_write_task with disposable test review_flushed_fanout_files_false_oom. The unbounded control wrote all 12 partitions. The 1 MiB pool failed requesting another 128 KiB after reserving 1 MiB. Parallel ArrowWriter controls asserted one flushed row group and memory_size() == 0 for every batch. Failure cleanup deleted all files and returned the reservation. Reproduction patch and output are preserved at /tmp/comet-6247-current-validation/reproduction.patch and /tmp/comet-6247-current-validation/task-probe.log. Pinned iceberg-rust implements current_written_size() as bytes_written() + in_progress_size().

[andygrove]

Fixed in e238121, though not with an accessor. iceberg-rust's ParquetWriter has no in-progress size accessor, at the pinned rev or on main; it only has current_written_size(). So each file's OutputFile is now wrapped (CountedOutput), its writer counts the bytes that leave memory on their way to storage, and a file's share is current_written_size() less those. For a local file that is the in-progress row group plus whatever is still in parquet-rs's 8 KiB BufWriter.

Releasing on flush is only right for local files, though. On S3 and GCS the flushed bytes stay in memory. OpenDAL collects what a file hands it into parts of at least 5 MiB (write_multi_min_size), and its MultipartWriter keeps each part until the next one is complete or the file closes, because only then does it know which part is the last (opendal-core 0.58.2, raw/oio/write/multipart_write.rs). I checked this against MinIO through iceberg-storage-opendal's S3 FileWrite at the pinned rev. Writing 6, 6, 1, 1 and 4 MiB, part 1 was uploaded only during the second write, part 2 only during the fifth, and part 3 at close. So at the default 128 MiB row group, a fanout file on S3 holds the last row group it flushed until it flushes another or closes, and releasing on flush would undercount S3 by up to a row group per open file. The share now follows the storage (StorageWrites, HeldByStorage): file and memory release a row group as it is handed on, and s3, s3a and gs release a part only once the next one is complete.

Regression tests:

• files_that_flush_give_back_what_they_reserved is your scenario: 12 fanout partitions, 128 KiB row groups, 1,000 distinct 256-byte strings per partition and a 1 MiB pool. It failed on the ninth partition before the change and now passes, with a peak reservation of 0 bytes.
• a_files_share_is_what_it_still_holds checks both storage kinds at the file level.
• storage_taken_in_parts_holds_the_last_part_until_the_next pins the part model.

[sunchao]

Summary

• Prior state and problem: Both native writers buffered data outside Comet’s memory accounting, allowing writes to exceed the task’s budget.
• Design approach: Register one task-level consumer and resize its reservation after each batch.
• Correctness / compatibility analysis: No additional introduced P1/P2 issues found within this review. The existing P2 about charging flushed Iceberg data remains unresolved and reproduces at this head. Checked Spark’s memory and task-failure handling across supported versions and Iceberg Java’s row-group lifecycle.
• Key design decisions: One consumer avoids reducing fair-pool shares for every open file. Per-file charges use drop-based cleanup. Eligibility and configuration defaults remain unchanged.
• Implementation sketch: Parquet uses ArrowWriter::memory_size() plus remote staging capacity. Iceberg wraps its file writers and includes pending rows. Accounting also scans the fanout feeds after each input batch. No additional substantiated performance or abstraction issue found.
• Behavioral changes worth calling out: Compared with branch-1.1, pool refusals now intentionally fail native writes. Neither writer spills. The existing Iceberg overcharge also causes unintended failures when row-group buffers have already been flushed.
• Suggested improvements: Address the existing thread by charging live Iceberg buffers and releasing their charge after flushes, preserving the demonstrated regression as a test.

Reviewed the entire 10-file diff from 9dc8c3ca89962506078831f7064a25c75da83883 to 91839920807e42e4a7d3ad0fbc23eb2fff1251d1. The PR remains open and non-draft. Read the existing review, inline comment and unresolved thread. Routed skills: review-comet-pr, review-comet-memory-pr, and review-comet-iceberg-write-pr.

Existing blocker: discussion_r4124645199, at native/core/src/execution/operators/iceberg_write.rs:293. With 128 KiB row groups, the unlimited-pool control writes all 12 fanout partitions. A 1 MiB pool fails on the ninth because flushed files remain charged. Equivalent parquet-rs controls report zero live row-group bytes after each batch. Failure cleanup removes the files and releases the reservation. This is existing feedback, so it is not duplicated in findings.

Exact-head CI: 7 successful checks, 14 skipped and 11 in progress, with no reported failures. Native build, Rust tests and lint checks remain pending. Spark SQL, upstream Iceberg and macOS suites were skipped.

Local validation: 65 Iceberg writer tests, six Parquet writer tests and the existing-blocker reproduction passed with --no-default-features. JVM, HDFS, real object-store and full upstream suites were not run locally. Evidence is in /tmp/comet-6247-918399-validation/. Disposable source changes were removed and the checkout is clean.

[sunchao]

Summary

• Prior state and problem: Both native writers buffered data outside Comet’s memory accounting.
• Design approach: Register one task-level consumer and resize its reservation after each batch.
• Correctness / compatibility analysis: No introduced P1/P2 issues found within this review. The previously reported Iceberg flush overcharge is fixed: the original 12-partition reproduction now succeeds with a 1 MiB pool. Checked Spark’s memory and task-failure handling across supported versions and Iceberg Java’s row-group lifecycle.
• Key design decisions: One consumer preserves task-wide fair-pool accounting. Per-file shares release on close or drop. The storage adapter leaves encoding, file rolling and cleanup ownership unchanged.
• Implementation sketch: Parquet charges ArrowWriter::memory_size() plus remote staging capacity. Iceberg subtracts bytes released to storage and adds pending rows. Accounting scans fanout feeds after each batch. No substantiated P1/P2 performance or abstraction issue found.
• Behavioral changes worth calling out: Compared with branch-1.1, pool refusals now intentionally fail native writes. Neither writer spills. Iceberg’s estimate still excludes dictionary overhead and excess buffer capacity, as documented.
• Suggested improvements: No further P1/P2 changes requested.

Reviewed the full 11-file diff from 9dc8c3ca89962506078831f7064a25c75da83883 to e23812130021a658d749488f3867f1a68bdb9bf0. Confirmed the PR is open and non-draft. Read both existing reviews, the inline comment and its thread.

Routed skills: review-comet-pr, review-comet-memory-pr, review-comet-iceberg-write-pr.

Exact-head CI: 17 successful checks, 14 skipped and 7 running, with no reported failures. Native build and Java lint passed. Rust tests, JVM suites and query-result checks remain pending. Spark SQL, upstream Iceberg and macOS suites were skipped.

Local validation: 69 Iceberg writer tests, six Parquet writer tests and two disposable probes passed using --no-default-features. The multipart probe exercised pinned OpenDAL against a local S3 stub and matched released bytes to observed uploads. JVM, HDFS, real object-store and full upstream suites were not run locally. No throughput benchmark was run. Evidence is preserved in /tmp/comet-6247-e238121-validation/. Disposable source changes were removed and the checkout is clean.