andygrove opened a new issue, #27069:
URL: https://github.com/apache/datafusion/issues/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`:
```rust
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`.
--
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]
