comphead commented on code in PR #26931:
URL: https://github.com/apache/datafusion/pull/26931#discussion_r4473956438
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -852,60 +851,93 @@ impl ExternalSorter {
Box::pin(RecordBatchStreamAdapter::new(schema, batches))
}
- /// Concatenates `batches` into fewer, larger runs, each bounded by
- /// `sort_in_place_threshold_bytes`, to reduce merge fan-in.
`self.reservation`
- /// is resized to the coalesced footprint so the caller's per-run splits
stay
- /// exact.
+ /// Tries to concatenate `batches` into runs of at least `batch_size` rows,
+ /// except for the final run. With `coalesce_large_runs`,
+ /// runs also grow toward `sort_in_place_threshold_bytes` to reduce merge
fan-in.
+ /// `self.reservation` is resized to the coalesced footprint so the
caller's
+ /// per-run splits stay exact.
+ /// Keeps the original batches if concatenation fails or the memory pool
+ /// cannot accommodate the combined footprint.
fn coalesce_in_mem_batches_into_runs(
&mut self,
batches: Vec<RecordBatch>,
+ coalesce_large_runs: bool,
) -> Result<Vec<RecordBatch>> {
+ if !coalesce_large_runs
+ && batches
+ .iter()
+ .all(|batch| batch.num_rows() >= self.batch_size)
+ {
+ return Ok(batches);
+ }
let target = self.sort_in_place_threshold_bytes.max(1);
let mut runs: Vec<RecordBatch> = Vec::new();
- let mut group: Vec<RecordBatch> = Vec::new();
+ let mut group: Vec<&RecordBatch> = Vec::new();
Review Comment:
`group` now borrows from `batches` so that the whole input can be returned
on failure. That keeps every original batch alive until the last group has been
copied. On `main`, each group's originals were dropped right after its
`concat_batches`, so the transient peak was one group, not the whole buffer.
The copies are also not reserved before they are made. On the spill path the
pool is full by construction. Any growth in footprint (for example the aliased
columns in `test_small_run_coalescing_falls_back_when_memory_is_full`) is only
caught by the final `try_resize`. By then the whole buffer has been copied, and
all of it is thrown away.
Owning the group and falling back per group bounds both:
```rust
let flush = |group: &mut Vec<RecordBatch>, runs: &mut Vec<RecordBatch>| {
if group.len() > 1
&& let Ok(run) = concat_batches(&schema, group.iter())
{
group.clear();
runs.push(run);
} else {
runs.append(group);
}
};
```
If the group's copy is reserved with `try_grow` before concatenating, and
the group is kept as is when that fails, the final realign becomes a shrink.
This also removes both `return Ok(batches)` paths and the
`pop().cloned().unwrap()`.
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -852,60 +851,93 @@ impl ExternalSorter {
Box::pin(RecordBatchStreamAdapter::new(schema, batches))
}
- /// Concatenates `batches` into fewer, larger runs, each bounded by
- /// `sort_in_place_threshold_bytes`, to reduce merge fan-in.
`self.reservation`
- /// is resized to the coalesced footprint so the caller's per-run splits
stay
- /// exact.
+ /// Tries to concatenate `batches` into runs of at least `batch_size` rows,
+ /// except for the final run. With `coalesce_large_runs`,
+ /// runs also grow toward `sort_in_place_threshold_bytes` to reduce merge
fan-in.
+ /// `self.reservation` is resized to the coalesced footprint so the
caller's
+ /// per-run splits stay exact.
+ /// Keeps the original batches if concatenation fails or the memory pool
+ /// cannot accommodate the combined footprint.
fn coalesce_in_mem_batches_into_runs(
&mut self,
batches: Vec<RecordBatch>,
+ coalesce_large_runs: bool,
) -> Result<Vec<RecordBatch>> {
+ if !coalesce_large_runs
+ && batches
+ .iter()
+ .all(|batch| batch.num_rows() >= self.batch_size)
+ {
+ return Ok(batches);
+ }
let target = self.sort_in_place_threshold_bytes.max(1);
let mut runs: Vec<RecordBatch> = Vec::new();
- let mut group: Vec<RecordBatch> = Vec::new();
+ let mut group: Vec<&RecordBatch> = Vec::new();
+ let mut group_rows = 0usize;
let mut group_bytes = 0usize;
// Count the buffers the batches share as when they were buffered
let mut counter = self.new_buffered_batches_counter();
// Flush a group into a run, skipping the copy for a single-batch
group.
- let flush = |group: &mut Vec<RecordBatch>,
+ let flush = |group: &mut Vec<&RecordBatch>,
runs: &mut Vec<RecordBatch>,
schema: &SchemaRef|
-> Result<()> {
match group.len() {
0 => {}
- 1 => runs.push(group.pop().unwrap()),
+ 1 => runs.push(group.pop().cloned().unwrap()),
_ => {
- runs.push(concat_batches(schema, group.iter())?);
+ runs.push(concat_batches(schema, group.iter().copied())?);
group.clear();
}
}
Ok(())
};
- for batch in batches {
- let bytes =
- get_reserved_bytes_for_next_record_batch(&batch,
counter.as_mut())?;
- if !group.is_empty() && group_bytes.saturating_add(bytes) > target
{
- flush(&mut group, &mut runs, &self.schema)?;
+ // Keep the input batches until their replacement fits the reservation.
+ // Concatenation can need more memory when it breaks shared input
buffers.
+ // Array offset and dictionary-key limits can also reject a combined
run.
+ for batch in &batches {
+ let bytes = if coalesce_large_runs {
+ get_reserved_bytes_for_next_record_batch(batch,
counter.as_mut())?
+ } else {
+ 0
+ };
+ if group_rows >= self.batch_size
+ && (!coalesce_large_runs || group_bytes.saturating_add(bytes)
> target)
+ {
+ if flush(&mut group, &mut runs, &self.schema).is_err() {
Review Comment:
Two points on this fallback:
1. `is_err()` swallows every error, not only the offset and dictionary-key
limits named in the comment above. Matching the Arrow overflow variants
(`DictionaryKeyOverflowError`, `OffsetOverflowError`) would keep real failures
visible.
2. The in-place path in `in_mem_sort_stream` still calls `concat_batches` on
all buffered batches. With the default `sort_in_place_threshold_bytes` (1 MiB),
the input from `test_small_run_coalescing_falls_back_on_dictionary_overflow`
takes that path and fails with `Arrow error: Dictionary key bigger than the key
type`. I checked this on `main`, and this PR doesn't change that path. The test
only passes because it sets the threshold to 0. Either fall through to the run
path when that concat fails, or narrow the dictionary claim in the description
and leave the rest as a follow-up.
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -852,60 +851,93 @@ impl ExternalSorter {
Box::pin(RecordBatchStreamAdapter::new(schema, batches))
}
- /// Concatenates `batches` into fewer, larger runs, each bounded by
- /// `sort_in_place_threshold_bytes`, to reduce merge fan-in.
`self.reservation`
- /// is resized to the coalesced footprint so the caller's per-run splits
stay
- /// exact.
+ /// Tries to concatenate `batches` into runs of at least `batch_size` rows,
+ /// except for the final run. With `coalesce_large_runs`,
+ /// runs also grow toward `sort_in_place_threshold_bytes` to reduce merge
fan-in.
+ /// `self.reservation` is resized to the coalesced footprint so the
caller's
+ /// per-run splits stay exact.
+ /// Keeps the original batches if concatenation fails or the memory pool
+ /// cannot accommodate the combined footprint.
fn coalesce_in_mem_batches_into_runs(
&mut self,
batches: Vec<RecordBatch>,
+ coalesce_large_runs: bool,
) -> Result<Vec<RecordBatch>> {
+ if !coalesce_large_runs
+ && batches
+ .iter()
+ .all(|batch| batch.num_rows() >= self.batch_size)
+ {
+ return Ok(batches);
+ }
let target = self.sort_in_place_threshold_bytes.max(1);
let mut runs: Vec<RecordBatch> = Vec::new();
- let mut group: Vec<RecordBatch> = Vec::new();
+ let mut group: Vec<&RecordBatch> = Vec::new();
+ let mut group_rows = 0usize;
let mut group_bytes = 0usize;
// Count the buffers the batches share as when they were buffered
let mut counter = self.new_buffered_batches_counter();
// Flush a group into a run, skipping the copy for a single-batch
group.
- let flush = |group: &mut Vec<RecordBatch>,
+ let flush = |group: &mut Vec<&RecordBatch>,
runs: &mut Vec<RecordBatch>,
schema: &SchemaRef|
-> Result<()> {
match group.len() {
0 => {}
- 1 => runs.push(group.pop().unwrap()),
+ 1 => runs.push(group.pop().cloned().unwrap()),
_ => {
- runs.push(concat_batches(schema, group.iter())?);
+ runs.push(concat_batches(schema, group.iter().copied())?);
group.clear();
}
}
Ok(())
};
- for batch in batches {
- let bytes =
- get_reserved_bytes_for_next_record_batch(&batch,
counter.as_mut())?;
- if !group.is_empty() && group_bytes.saturating_add(bytes) > target
{
- flush(&mut group, &mut runs, &self.schema)?;
+ // Keep the input batches until their replacement fits the reservation.
+ // Concatenation can need more memory when it breaks shared input
buffers.
+ // Array offset and dictionary-key limits can also reject a combined
run.
+ for batch in &batches {
+ let bytes = if coalesce_large_runs {
+ get_reserved_bytes_for_next_record_batch(batch,
counter.as_mut())?
+ } else {
+ 0
+ };
+ if group_rows >= self.batch_size
+ && (!coalesce_large_runs || group_bytes.saturating_add(bytes)
> target)
Review Comment:
A full batch that follows a short group gets absorbed and copied. With
`batch_size = 4`, `[2] [4]` becomes one 6-row run, built by copying the 4-row
batch. Outside large-run mode, pushing batches with `num_rows() >= batch_size`
straight to `runs` avoids that copy. Run order doesn't matter, since equal keys
have no ordering guarantee.
With that change, the `all(..)` pre-scan at the top is no longer needed
either. An all-full input copies nothing, so a `copied` flag that skips the
realign pass at the end covers the same fast path in one traversal.
--
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]
