SortPreservingMerge over-reserves memory by orders of magnitude for Utf8View sort keys: reservation measured before gc(), causing spurious ResourcesExhausted

[zhuqi-lucas]

Describe the bug

FieldCursorStream::convert_batch (single-sort-key merges) reserves memory for the sort-key array before ArrayValues::new compacts it, and never adjusts afterwards. For Utf8View keys the two figures can differ by orders of magnitude:

• sorts/stream.rs: let size_in_mem = array.get_buffer_memory_size(); → array_reservation.try_grow(size_in_mem)? — measured on the pre-gc array. For a StringViewArray this returns the size of all referenced data buffers, not the bytes the views actually use.
• sorts/cursor.rs: impl CursorArray for StringViewArray { fn values(&self) -> Self { self.gc() } } — the cursor then keeps only the gc'd compact copy…
• …while ArrayValues::new stores the untouched inflated reservation for the cursor's whole lifetime. Nothing resizes it.

A selective filter over view arrays is zero-copy on the data buffers (arrow-select's filter_byte_view), so a 10-row batch can keep multi-MB buffers referenced. Measured: 100k rows of ~100-byte strings filtered down to 10 rows → get_buffer_memory_size() = 10,469,568 B reserved, gc'd cursor actually retains 1,180 B — an 8,872× over-reservation, held until the cursor drops. The merge then spuriously fails with ResourcesExhausted (or spuriously spills) on data that fits comfortably.

Two adjacent problems in the same path (secondary):

• Double-counting even for compact arrays: BatchBuilder::push_batch separately reserves get_record_batch_memory_size(&batch), which already counts the same key-column buffers the cursor reservation counts.
• gc() copies unconditionally: every single-Utf8View-key merge re-copies the entire sort-key column once per batch per stream even when the array is already compact (gc() only short-circuits on empty data buffers).

To Reproduce

Tests appended to sorts/sort_preserving_merge.rs (view_gc_accounting_tests):

1. Size-discrepancy assertion: build 100k × ~100 B StringViewArray, arrow::compute::filter to 10 rows, compare pre-gc vs post-gc get_buffer_memory_size() → 10,469,568 B vs 1,180 B.
2. End-to-end: SortPreservingMergeExec over 2 partitions (one such batch each) with a TrackConsumersPool limit = real retained bytes + 5 MB slack (25.9 MB total):

Resources exhausted: Additional allocation failed for SortPreservingMergeExec[0] ...
  SortPreservingMergeExec[0]#0(can spill: false) consumed 20.0 MB, peak 20.0 MB.
Error: Failed to allocate additional 10.0 MB ... 4.8 MB remain available for the total pool

Two cursor reservations of ~10 MB each, for 20 rows of actual key data.

Counterfactual: patching convert_batch to reserve the post-values() retained size makes the identical merge pass under the same limit.

Full repro tests

mod view_gc_accounting_tests {
    use super::*;
    use crate::collect;
    use crate::test::TestMemoryExec;
    use arrow::array::{Array, ArrayRef, BooleanArray, RecordBatch, StringViewArray};
    use arrow::datatypes::{DataType, Field, Schema};
    use datafusion_common::utils::memory::get_record_batch_memory_size;
    use datafusion_execution::TaskContext;
    use datafusion_execution::memory_pool::{GreedyMemoryPool, TrackConsumersPool};
    use datafusion_execution::runtime_env::RuntimeEnvBuilder;
    use datafusion_physical_expr::expressions::Column;
    use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
    use std::num::NonZeroUsize;
    use std::sync::Arc;

    /// Build a batch with a single sorted Utf8View column of `keep` rows whose
    /// views still reference the large (~`rows * ~100B`) source data buffers,
    /// as produced by `arrow::compute::filter` (zero-copy on the data buffers).
    fn filtered_view_batch(seed: u8, rows: usize, keep: usize) -> RecordBatch {
        // > 12 bytes so strings are not inlined in the view
        let strings: Vec<String> = (0..rows)
            .map(|i| format!("{seed:02}-{i:010}-{}", "x".repeat(88)))
            .collect();
        let array = StringViewArray::from_iter_values(strings.iter().map(|s| s.as_str()));
        let step = rows / keep;
        let mask = BooleanArray::from_iter((0..rows).map(|i| Some(i % step == 0)));
        let filtered = arrow::compute::filter(&array, &mask).unwrap();
        let schema = Arc::new(Schema::new(vec![Field::new(
            "c",
            DataType::Utf8View,
            false,
        )]));
        RecordBatch::try_new(schema, vec![Arc::clone(&filtered)]).unwrap()
    }

    /// Demonstrates the size discrepancy: a filtered view array reports the full
    /// source buffers via get_buffer_memory_size(), while gc() (what ArrayValues
    /// actually retains) is tiny.
    #[test]
    fn filtered_view_array_size_discrepancy() {
        let batch = filtered_view_batch(0, 100_000, 10);
        let col = batch.column(0);
        let pre_gc = col.get_buffer_memory_size();
        let gcd = col
            .as_any()
            .downcast_ref::<StringViewArray>()
            .unwrap()
            .gc();
        let post_gc = gcd.get_buffer_memory_size();
        println!("rows kept          : {}", col.len());
        println!("pre-gc buffer size : {pre_gc}");
        println!("post-gc buffer size: {post_gc}");
        assert!(pre_gc > 8_000_000, "expected inflated size, got {pre_gc}");
        assert!(post_gc < 100_000, "expected compact size, got {post_gc}");
    }

    /// If cursor memory accounting were based on what the cursor actually
    /// retains (the gc'd copy), this merge would fit comfortably in the pool:
    /// the pool limit is set to the real retained bytes (batches, counted the
    /// same way BatchBuilder counts them, plus gc'd cursors) plus 4MB slack.
    ///
    /// It fails with ResourcesExhausted because FieldCursorStream::convert_batch
    /// reserves array.get_buffer_memory_size() (pre-gc, includes the full
    /// filtered-away source buffers) and never adjusts after ArrayValues::new
    /// gc's the array.
    #[tokio::test]
    async fn spm_utf8view_filtered_batches_should_fit_in_pool() {
        let b0 = filtered_view_batch(0, 100_000, 10);
        let b1 = filtered_view_batch(1, 100_000, 10);
        let schema = b0.schema();

        // Real retained memory while merging:
        //  - both batches, as accounted by BatchBuilder::push_batch
        //  - the gc'd cursor copies, generously overestimated at 1MB total
        let real_batches =
            get_record_batch_memory_size(&b0) + get_record_batch_memory_size(&b1);
        let limit = real_batches + 1_000_000 + 4_000_000; // real + slack
        println!("real batch bytes  : {real_batches}");
        println!("pool limit        : {limit}");
        println!(
            "pre-gc key sizes  : {} + {}",
            b0.column(0).get_buffer_memory_size(),
            b1.column(0).get_buffer_memory_size()
        );

        let pool = Arc::new(TrackConsumersPool::new(
            GreedyMemoryPool::new(limit),
            NonZeroUsize::new(5).unwrap(),
        ));
        let runtime = RuntimeEnvBuilder::new()
            .with_memory_pool(pool)
            .build_arc()
            .unwrap();
        let task_ctx = Arc::new(TaskContext::default().with_runtime(runtime));

        let source =
            TestMemoryExec::try_new_exec(&[vec![b0], vec![b1]], schema.clone(), None)
                .unwrap();
        let spm = SortPreservingMergeExec::new(
            [PhysicalSortExpr::new_default(Arc::new(Column::new("c", 0)))].into(),
            source,
        );

        let result = collect(Arc::new(spm), task_ctx).await;
        match result {
            Ok(batches) => {
                let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
                assert_eq!(rows, 20);
            }
            Err(e) => panic!(
                "merge should fit: real retained ~{real_batches} bytes, limit {limit}, but got: {e}"
            ),
        }
    }
}

#[cfg(test)]
mod view_gc_perf_probe {
    use arrow::array::{Array, StringViewArray};

    /// gc() copies all string bytes per batch per stream even when the array
    /// is already compact. Quantify bytes copied + wall time (debug build,
    /// order-of-magnitude only).
    #[test]
    fn gc_copies_even_when_compact() {
        let strings: Vec<String> =
            (0..80_000).map(|i| format!("{i:010}-{}", "x".repeat(88))).collect();
        let compact =
            StringViewArray::from_iter_values(strings.iter().map(|s| s.as_str()));
        let bytes = compact.get_buffer_memory_size();
        let t = std::time::Instant::now();
        let mut sink = 0usize;
        const N: usize = 20;
        for _ in 0..N {
            sink += compact.gc().get_buffer_memory_size();
        }
        let el = t.elapsed();
        println!(
            "compact array {} bytes; gc() x{N} took {:?} ({:?}/call), sink={}",
            bytes,
            el,
            el / N as u32,
            sink
        );
    }
}

Expected behavior

Reserve what the cursor retains: convert first (gc once), then try_grow on the compacted array's size. That fixes the spurious ResourcesExhausted and half of the double-count in one move. A follow-up can gate gc() on an "already compact" heuristic to remove the unconditional copy.

Additional context

• Only the single-sort-key Utf8View path is affected (streaming_merge.rs routes single keys to FieldCursorStream; multi-key sorts use RowCursorStream, which sizes via rows.size() and is correct).
• Reachability: with default config, CoalesceBatchesExec usually compacts post-filter views before the merge. It becomes realistic with pushdown_filters = true (common perf setting) + schema_force_view_types (default true) + a sorted multi-partition table (WITH ORDER): the plan is DataSourceExec → SortPreservingMergeExec with no coalesce in between, and parquet late materialization emits view arrays referencing full decoded pages. Order-preserving RepartitionExec likewise passes sparse views straight through.
• Verified present on current main (same code at tip: stream.rs measuring pre-gc, cursor.rs gc-ing after). Found during a sorts audit; repro + counterfactual patch available. I plan to follow up with a fix PR.