apache / apache/datafusion

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

Open
#25,271 0 comments 0 reactions 1 assignee Claimed by @zhuqi-lucas View on GitHub
bug
Dominant language
Rust
Stars
9.3k
Forks
2.4k
Avg merge
3d 7h
Merged PRs (30d)
344

Description

### 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

```rust
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 = (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::()
.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 =
(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.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.