zhuqi-lucas opened a new issue, #25271:
URL: https://github.com/apache/datafusion/issues/25271
### 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.
<details>
<summary>Full repro tests</summary>
```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<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
);
}
}
```
</details>
### 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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
