SortExec panics in sort_run on FFI-imported sliced batches: assertion reservation.size() + shared.size() >= reserved_bytes #27069

Describe the bug

On branch-56 (1635ed947), SortExec panics when its input batches were imported through the Arrow C Data Interface and are zero-copy slices of one larger batch:

thread 'main' panicked at datafusion/physical-plan/src/sorts/sort.rs:1022:29:
assertion failed: reservation.size() + shared.size() >= reserved_bytes

The check is an assert!, so release builds panic too. 55.2.0 does not have this code, so this is a regression in 56. The assertion and the shared-buffer accounting behind it came from #25820, and #26242 changed the same code.

RecordBatchMemoryCounter identifies a buffer by Buffer::data_ptr() (the start of the allocation) and sizes it with Buffer::capacity(). For a buffer created by Buffer::from_custom_allocation, which is what arrow::ffi::from_ffi uses, capacity() is the length declared at import, not the size of the allocation. When sliced arrays go through FFI, every slice's string values buffer starts at the same address, but each import declares a different length (up to that slice's last offset). The counter keeps the first length it sees. So sort_run_sizes computes a shared input size smaller than a later run's get_record_batch_memory_size, and the assertion in sort_run fails.

A run is only exposed when a batch passes through coalesce_in_mem_batches_into_runs uncopied, i.e. when a single batch is about sort_in_place_threshold_bytes (1 MB by default) or larger. Smaller batches are concatenated into new buffers first.

DataFusion Comet hits this in TPC-DS q67a while testing 56 (apache/datafusion-comet#7249), where sort input arrives from the JVM over FFI. I'd also expect any FFI consumer that sorts sliced input to be exposed, for example datafusion-python with sliced pyarrow data, but I have not tested that.

To Reproduce

Default configuration apart from one target partition. With Cargo.toml dependencies datafusion (branch-56, default-features = false), arrow = { version = "60.1.0", features = ["ffi"] } and tokio:

use std::sync::Arc;
use arrow::array::{Array, ArrayRef, Int32Array, StringArray, make_array};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::ffi::{from_ffi, to_ffi};
use arrow::record_batch::RecordBatch;
use datafusion::datasource::MemTable;
use datafusion::prelude::*;

/// Round-trip an array through the Arrow C Data Interface.
fn ffi_roundtrip(a: &ArrayRef) -> ArrayRef {
    let (array, schema) = to_ffi(&a.to_data()).unwrap();
    make_array(unsafe { from_ffi(array, &schema) }.unwrap())
}

#[tokio::main]
async fn main() -> datafusion::error::Result<()> {
    let n = 65_536;
    let k: ArrayRef = Arc::new(Int32Array::from_iter_values((0..n).rev()));
    let s: ArrayRef = Arc::new(StringArray::from_iter_values(
        (0..n).map(|i| format!("{i:0>200}")),
    ));
    let schema = Arc::new(Schema::new(vec![
        Field::new("k", DataType::Int32, false),
        Field::new("s", DataType::Utf8, false),
    ]));
    let big = RecordBatch::try_new(Arc::clone(&schema), vec![k, s])?;
    // Zero-copy slices of one batch (about 1.6 MB each), each imported over FFI.
    let batches: Vec<RecordBatch> = (0..n as usize)
        .step_by(8192)
        .map(|off| {
            let b = big.slice(off, 8192);
            let cols = b.columns().iter().map(ffi_roundtrip).collect();
            RecordBatch::try_new(Arc::clone(&schema), cols).unwrap()
        })
        .collect();
    let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1));
    ctx.register_table("t", Arc::new(MemTable::try_new(schema, vec![batches])?))?;
    let out = ctx.table("t").await?.sort(vec![col("k").sort(true, false)])?.collect().await?;
    println!("rows: {}", out.iter().map(|b| b.num_rows()).sum::<usize>());
    Ok(())
}

cargo run --release panics with the assertion above.

A smaller unit test in datafusion/physical-plan/src/sorts/sort.rs shows the same thing without FFI. It builds two batches whose Int32 columns are Buffer::from_custom_allocation buffers over the same pointer, declaring 16 and 1024 values, then sorts them with sort_in_place_threshold_bytes = 0.

Expected behavior

The sort returns all 65,536 rows sorted, as it does on 55.2.0. Memory accounting may be approximate for buffers whose real allocation size is unknown, but it should not panic.

Additional context

Possible directions: key the counter on (data_ptr, capacity), or take the larger capacity seen for an address; or turn the assertion into a resize that covers any shortfall. Comet would need the fix backported to branch-56.