perf: charge the local shuffle writer's buffers to the memory pool

docs/source/contributor-guide/memory_management.md

@@ -582,14 +582,21 @@ diverge for several structural reasons:
   retained/dirty page cache all add resident bytes that no layer above the allocator can see.
   Freeing memory does not necessarily return pages to the OS.
 - **Non-Rust allocations.** Memory allocated by C dependencies through libc `malloc`, and anything
-  `mmap`ed, never passes through Rust's `GlobalAlloc`, so neither the memory pool nor the
-  allocation counters (`native_allocated`, `jemalloc_allocated`) see it. In a default build the C
-  dependencies are libzstd (`zstd-sys`, behind the Parquet `zstd` codec), libhdfs (`hdfs-sys`,
-  pulled in by the default `hdfs-opendal` feature), and the TLS stack used for cloud object stores
-  (`aws-lc-sys`). With the `python-udf` feature, the embedded Python interpreter and PyArrow also
-  allocate outside `GlobalAlloc`; their memory is absent from the executor's native `allocated`
-  figure, and Python working allocations are not reserved in Comet's pool. Building with the
-  `jemalloc` or `mimalloc` feature adds the allocator itself (`tikv-jemalloc-sys`,
+  `mmap`ed, never passes through Rust's `GlobalAlloc`, so the allocation counters
+  (`native_allocated`, `jemalloc_allocated`) never see it, and the memory pool tracks it only where
+  an operator charges it by hand. The one place that happens today is the native shuffle writer's
+  zstd context when a task writes to more than one local partition: the writer reserves the
+  context's size for as long as it holds one, so those bytes are in `reserved` but not in
+  `allocated`. It holds one between blocks only up to level 8. A larger context exceeds the 8 MiB
+  the writer keeps, so `release_zstd_if_oversized` frees it after each block, and that workspace,
+  about 15 MiB at level 9 and 90 MiB at level 19, is in neither figure. In a default build the C
+  dependencies are libzstd (`zstd-sys`, behind the Parquet
+  `zstd` codec and the shuffle writer's `zstd` codec), libhdfs (`hdfs-sys`, pulled in by the default
+  `hdfs-opendal` feature), and the TLS stack used for cloud object stores (`aws-lc-sys`). With the
+  `python-udf` feature, the embedded Python interpreter and PyArrow also allocate outside
+  `GlobalAlloc`; their memory is absent from the executor's native `allocated` figure, and Python
+  working allocations are not reserved in Comet's pool. Building with the `jemalloc` or `mimalloc`
+  feature adds the allocator itself (`tikv-jemalloc-sys`,
   `libmimalloc-sys`). It is worth knowing which dependencies are _not_ C, because several names
   suggest otherwise: the other Parquet codecs are pure Rust in this build, `snap` for Snappy,
   `lz4_flex` for LZ4 and `zlib-rs` for gzip, as is `libbz2-rs-sys` despite its name, so those

docs/source/contributor-guide/native_shuffle.md

@@ -483,6 +483,26 @@ The `MultiPartitionShuffleRepartitioner` holds:
 - `max_buffer_bytes`, the optional fixed spill threshold described above. `None` leaves pool
   pressure as the only trigger.
 - `scratch`, reusable buffers for partition ID computation.
+- Through the writer, a second reservation for the local writer's own buffers: the data file
+  and spill file write buffers, the spill copy buffer, the encode scratch and the cached zstd
+  context. `try_new` hands `LocalPartitionWriter` a `reservation.new_empty()` sibling, so the
+  bytes land on the same consumer without a new one, and the writer keeps it equal to what
+  those buffers hold. A spill frees `reservation` but not this one, and it counts toward
+  neither `memory_spilled_bytes` nor `spark.comet.shuffle.native.maxBufferBytes`. Each file
+  and copy buffer is allocated when the phase that uses it starts: the spill file buffer when
+  the first spill opens the file, the other two at the first `finish_partition`. That call
+  flushes the spill file and releases its buffer, then sizes the copy buffer, which serves
+  every spilled range, and then the data file buffer, which blocks of the write buffer size
+  bypass anyway. Each asks the pool for its full size, and when the pool refuses the writer
+  uses an 8 KiB buffer instead and charges that. Any other pool error, such as an exception
+  Spark threw while spilling another consumer, fails the task. A spill forced by pool pressure
+  opens the spill file while its batches are still charged, so that buffer often starts
+  small. After each spill has freed its batches, the writer drops the drained encode scratch,
+  which one wide row can leave larger than the pool, and grows the spill file buffer toward
+  the full size, but to no more than half of what the spill freed, so the next batches have
+  room. It never shrinks that buffer, and moves the buffered bytes into the larger one. The
+  scratch and the zstd context already exist when they are charged, so the writer records
+  them without asking.
 
 The spill file is owned by `PartitionedSpill` in `writers/local/spill.rs`. It holds one DataFusion
 `SpillFile`, created lazily on the first spill and shared by every output partition, and tracks per

docs/source/user-guide/latest/tuning/memory.md

@@ -44,10 +44,16 @@ needs. A native Iceberg fanout write is the exception: it closes some of its ope
 
 Comet's memory pool only tracks memory that an operator explicitly reserves, which in practice means the batches
 an operator deliberately accumulates: the sort buffer, the build side of a hash join, hash aggregation state, the
-shuffle writer's buffered partitions, and the files a native writer holds open. Memory that is not reserved is
-invisible to the pool no matter how much of it there is. That includes:
+shuffle writer's buffered partitions, the files a native writer holds open, and the write buffers, encode scratch
+and zstd context of the native shuffle writer when it writes to more than one local partition (the scratch is
+charged as each write ends, so its transient peak within a write, the write buffer size plus one block, is not, and
+the zstd context only up to level 8, as [Sizing the Overhead from the Memory Usage Log] describes).
+Memory that is not reserved is invisible to the pool no matter how much of it there is. That includes:
 
 - per-batch working memory in expression kernels and Arrow array builders,
+- the write buffers and encode scratch of the single-partition, empty-schema and remote shuffle writers
+  (`spark.comet.shuffle.native.writeBufferSize` sizes the local ones), and the zstd context they keep while
+  encoding when `spark.comet.shuffle.compression.codec` is `zstd`,
 - decompression buffers and Parquet reader structures,
 - object store request buffers and the async runtime's own machinery,
 - Arrow buffers allocated on the JVM side, which no budget covers at all,
@@ -182,11 +188,19 @@ alongside the JVM's own non-heap memory. To size the overhead from it:
    non-heap memory, and add the most untracked memory seen on any executor.
 3. Add a margin on top. The log can miss the true peak between samples, and none of the figures
    includes the allocator's fragmentation and retained pages, or memory allocated by native C
-   libraries such as zstd. With native Arrow UDFs, budget the embedded Python interpreter and
-   PyArrow allocations in addition to the untracked memory calculated from the log. A PyArrow
-   result buffer later imported into the JVM can appear in the `imported from native` figure that
-   the estimate subtracts, even though Rust's `allocated` does not count it. Use the executor's
-   peak resident memory to size this additional overhead.
+   libraries such as zstd. The zstd context of the native shuffle writer is the exception at levels
+   up to 8: when `spark.comet.shuffle.compression.codec` is `zstd`, a shuffle task writing to more
+   than one local partition reserves the context it keeps between blocks while spilling or writing
+   its output, so that memory is in `reserved` but never in `allocated`, and the untracked memory
+   worked out from the log understates what has to fit in the overhead by one context per such task
+   (about 1.3 MiB at the default level, close to 8 MiB at levels 7 and 8). Above level 8 the context
+   is too large to keep, so each block creates and frees its own and it appears in neither figure:
+   budget one per task writing shuffle output, about 15 MiB at level 9 and 90 MiB at level 19, on
+   top of the untracked memory. With native Arrow UDFs, budget the embedded
+   Python interpreter and PyArrow allocations in addition to the untracked memory calculated from
+   the log. A PyArrow result buffer later imported into the JVM can appear in the
+   `imported from native` figure that the estimate subtracts, even though Rust's `allocated` does
+   not count it. Use the executor's peak resident memory to size this additional overhead.
 
 For example, a 16 GiB executor derives an overhead of 1638 MiB. If the line above has the most
 untracked memory in its log, that is 5412.3 + (310.4 - 96.2) - 3890.0 = 1736.5 MiB. The overhead

native/shuffle/src/codec_context.rs

@@ -112,6 +112,10 @@ impl ShuffleCodecContext {
         }
     }
 
+    pub(crate) fn retained_bytes(&self) -> usize {
+        self.zstd.as_ref().map_or(0, |cctx| cctx.sizeof())
+    }
+
     /// Test hook for the release-vs-retain contract of the two encode paths.
     #[cfg(test)]
     pub(crate) fn holds_zstd_cctx(&self) -> bool {