You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "NoahKusaba (via GitHub)" <gi...@apache.org> on 2026/10/10 18:09:35 UTC
[PR] refactor(core): drain and size the sort-shuffle buffer through BufferedBatches [datafusion-ballista]
NoahKusaba opened a new pull request, #2579:
URL: https://github.com/apache/datafusion-ballista/pull/2579
## Which issue does this PR close?
Part of #2578 (hybrid sort-shuffle buffering: split rows by partition on
arrival when memory allows). This is the first of three PRs; it changes no
behaviour on its own.
## Rationale for this change
`SortShuffleWriterExec` reaches into `BufferedBatches` from three places:
- the final encode (`encode_buffered_partitions`) and the spill
(`spill_all_partitions`) each call `take()` and rebuild the same
`PartitionedBatchIterator` gather loop;
- the input loop estimates the buffer's memory by hand: the input batch's
size, plus the index lists' capacity growth measured with
`indices_allocated_size()` before and after each push, accumulated into a
local counter.
The buffer's layout therefore leaks into the writer. Moving these two
operations into `BufferedBatches` gives the writer a narrow interface
(`push_batch`, `drain`, `memory_size`, `is_empty`) and keeps the gather loop
and the memory accounting in one place.
That interface is what the follow-up needs. It adds a second layout that
copies each partition's rows out of an input batch as it arrives, while the
batch is still in cache, and falls back to today's layout when the
per-partition buffers would not fit a memory budget. This is the "auto" mode
Comet planned for its immediate-mode shuffle (apache/datafusion-comet#3845)
but set aside over memory at high partition counts, and close to Gluten's
default hash shuffle writer, which falls back to a buffered writer for many
partitions or wide tables (apache/gluten#6475). A single-node TPC-H SF10
prototype is 8.5% faster in total, and 10–23% faster on shuffle-heavy
queries. With this PR in place, the follow-up adds the new layout inside
`BufferedBatches`; its only writer changes are computing the budget and two
metrics, leaving encode, spill and memory accounting untouched.
## What changes are included in this PR?
- `BufferedBatches::drain(f)`: empties the buffer, passing each output batch
of at most `batch_size` rows to `f` with its partition, in partition order
and then arrival order. Encode and spill now just call it.
- `BufferedBatches::memory_size()`: the buffered batches'
`get_array_memory_size` plus the index lists' allocated capacity, kept as a
running total as the lists grow. The writer mirrors it in the pool with
`try_resize` instead of accumulating growth itself.
- `BufferedBatches::new` takes the batch size, so `encode_buffered_partitions`
and `spill_all_partitions` no longer need it passed in.
`memory_size()` is the same number the writer computed before (index lists
grow identically), so the pool reservation, the per-task budget check and
therefore every spill decision are unchanged, and the gather produces the
same batches in the same order. A side effect: the index-capacity total is no
longer recomputed over every partition twice per input batch.
## Are there any user-facing changes?
No behaviour change: same shuffle files, same spills, same metrics.
**API change** (`sort_shuffle::BufferedBatches` is public):
- `new(num_partitions, schema)` → `new(num_partitions, schema, batch_size)`
- `take()` / `BufferedTake`, `indices_for()`, `indices_allocated_size()` and
`num_buffered_rows()` are removed, replaced by `drain()` and
`memory_size()`.
## How was this tested?
- Buffer unit tests rewritten against the new API, comparing whole outputs:
drained batches per partition (order, `batch_size` chunking, emptiness
afterwards), and `memory_size()` against the batches plus index capacity.
- Existing writer tests unchanged and passing, including the spill tests
(per-task budget, pool rejection, zero budget) and the round trips.
- `cargo test -p ballista-core -p ballista-executor -p ballista`, clippy
(`-D warnings`) and fmt are clean.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor(core): drain and size the sort-shuffle buffer through BufferedBatches [datafusion-ballista]
Posted by "NoahKusaba (via GitHub)" <gi...@apache.org>.
NoahKusaba commented on PR #2579:
URL: https://github.com/apache/datafusion-ballista/pull/2579#issuecomment-6138296412
Peak RSS growth during the write, zstd, one 8192-row batch per partition:
| partitions | `main` | before the fix | after the fix |
|---|---|---|---|
| 200 | 36 MB | 139 MB | 36 MB |
| 1000 | 140 MB | 655 MB | 140 MB |
I also wrote the same 37-partition shuffle (40 batches with a string column)
on `main` and on this branch, with lz4 and zstd, with and without spilling.
The data files, index files and reported stats are byte-identical, and the
spill counts match.
`sort_shuffle` bench against `main` (laptop, on AC, criterion's 10 samples):
| benchmark | `main` | this branch | change |
|---|---|---|---|
| no spill, 10 batches, 200 partitions | 71.0 ms | 65.1 ms | −8.8% |
| no spill, 50 batches, 200 partitions | 159.4 ms | 140.4 ms | −10.6% |
| spill, 8 MB limit, 200 partitions | 1.23 s | 1.00 s | −18.6% |
| spill, 2 MB limit, 200 partitions | 2.27 s | 1.77 s | −22.3% |
No regression. I wouldn't read these as a speedup though: the change does the
same work as `main`, and the `main` baseline ran first, so I expect most of
the gap is machine noise.
I skipped the TPC-H run since the plan and the output bytes don't change, but
I'm happy to run it if you'd like.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor(core): drain and size the sort-shuffle buffer through BufferedBatches [datafusion-ballista]
Posted by "NoahKusaba (via GitHub)" <gi...@apache.org>.
NoahKusaba commented on code in PR #2579:
URL: https://github.com/apache/datafusion-ballista/pull/2579#discussion_r4254759397
##########
ballista/core/src/execution_plans/sort_shuffle/buffer.rs:
##########
@@ -99,26 +83,40 @@ impl BufferedBatches {
"BufferedBatches::push_batch schema mismatch"
);
let batch_idx = self.batches.len() as u32;
- for (p, rows) in per_partition_rows.iter().enumerate() {
- let dst = &mut self.indices[p];
- dst.reserve(rows.len());
- for &r in rows {
- dst.push((batch_idx, r));
- self.num_buffered_rows += 1;
- }
+ for (partition_indices, rows) in self.indices.iter_mut().zip(per_partition_rows) {
+ let capacity_before = partition_indices.capacity();
+ partition_indices.extend(rows.iter().map(|&r| (batch_idx, r)));
+ self.indices_bytes += (partition_indices.capacity() - capacity_before)
+ * size_of::<(u32, u32)>();
}
+ self.batches_bytes += batch.get_array_memory_size();
self.batches.push(batch);
}
- /// Drains all state, returning the buffered batches and per-partition
- /// index lists. After this call the buffer is empty.
- pub fn take(&mut self) -> BufferedTake {
- self.num_buffered_rows = 0;
+ /// Empties the buffer, passing each output batch of at most `batch_size`
+ /// rows to `f` with its partition, in partition order and then arrival
+ /// order. Output batches are gathered one at a time, as `f` consumes them.
+ /// The buffer is emptied first, so if `f` or a gather fails, the remaining
+ /// rows are dropped.
+ pub fn drain(
+ &mut self,
+ mut f: impl FnMut(usize, RecordBatch) -> Result<()>,
+ ) -> Result<()> {
+ self.batches_bytes = 0;
+ self.indices_bytes = 0;
let batches = std::mem::take(&mut self.batches);
- // Preserve the partition count by replacing each inner vec with empty
- let indices: Vec<Vec<(u32, u32)>> =
- self.indices.iter_mut().map(std::mem::take).collect();
- (batches, indices)
+ let indices: Vec<_> = self.indices.iter_mut().map(std::mem::take).collect();
+ for (partition, partition_indices) in indices.iter().enumerate() {
Review Comment:
Done. ~~I'll do building the batch refs once per drain in a follow-up PR,
since `main` has the same cost.~~
Actually, your suggested change is pretty trivial, now that I have time to better look at it today. Going to add it in the next commit to, to save time.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor(core): drain and size the sort-shuffle buffer through BufferedBatches [datafusion-ballista]
Posted by "NoahKusaba (via GitHub)" <gi...@apache.org>.
NoahKusaba commented on PR #2579:
URL: https://github.com/apache/datafusion-ballista/pull/2579#issuecomment-6328948009
Note on the "Retrigger CI" commit (d273b17): the failure on 47fc2f9 was a flaky chaos test, not this change.
[`test linux crates`](https://github.com/apache/datafusion-ballista/actions/runs/38662611398/job/116187325468) failed in `ballista-chaos --test ha`, on `baseline_matches_local_datafusion::case_2_aqe_on`, with `cluster must start: timed out waiting for 2 executors to register`. The other 15 `ha` tests passed, and so did the rest of the workflow.
The executor logs show both executors got the same Flight port (`127.0.0.1:45565`). Executor 0 then failed with `AddrInUse` and shut down, so only one executor was left to register:
```
executor-0: Built-in arrow flight server listening on: 127.0.0.1:45565
executor-0: Could not start built-in arrow flight server.
executor-0: executor services stopped with reason Err(TonicError(... AddrInUse ...))
executor-1: Built-in arrow flight server listening on: 127.0.0.1:45565
```
This is a race in the test harness. `spawn_executor` calls `free_ports()` once per executor (`chaos-testing/src/cluster.rs`). That function binds to `:0` to find a free port, then releases it before the child process binds. Within one call the ports are distinct, but nothing stops a later call from getting the same number. Executor 0 starts its Flight server only after gRPC setup and registration, so executor 1's `free_ports()` got 45565 back from the kernel before executor 0 had bound it.
This PR doesn't touch the chaos harness or executor startup, and the same suite passed on the rerun.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor(core): drain and size the sort-shuffle buffer through BufferedBatches [datafusion-ballista]
Posted by "NoahKusaba (via GitHub)" <gi...@apache.org>.
NoahKusaba commented on code in PR #2579:
URL: https://github.com/apache/datafusion-ballista/pull/2579#discussion_r4254759397
##########
ballista/core/src/execution_plans/sort_shuffle/buffer.rs:
##########
@@ -99,26 +83,40 @@ impl BufferedBatches {
"BufferedBatches::push_batch schema mismatch"
);
let batch_idx = self.batches.len() as u32;
- for (p, rows) in per_partition_rows.iter().enumerate() {
- let dst = &mut self.indices[p];
- dst.reserve(rows.len());
- for &r in rows {
- dst.push((batch_idx, r));
- self.num_buffered_rows += 1;
- }
+ for (partition_indices, rows) in self.indices.iter_mut().zip(per_partition_rows) {
+ let capacity_before = partition_indices.capacity();
+ partition_indices.extend(rows.iter().map(|&r| (batch_idx, r)));
+ self.indices_bytes += (partition_indices.capacity() - capacity_before)
+ * size_of::<(u32, u32)>();
}
+ self.batches_bytes += batch.get_array_memory_size();
self.batches.push(batch);
}
- /// Drains all state, returning the buffered batches and per-partition
- /// index lists. After this call the buffer is empty.
- pub fn take(&mut self) -> BufferedTake {
- self.num_buffered_rows = 0;
+ /// Empties the buffer, passing each output batch of at most `batch_size`
+ /// rows to `f` with its partition, in partition order and then arrival
+ /// order. Output batches are gathered one at a time, as `f` consumes them.
+ /// The buffer is emptied first, so if `f` or a gather fails, the remaining
+ /// rows are dropped.
+ pub fn drain(
+ &mut self,
+ mut f: impl FnMut(usize, RecordBatch) -> Result<()>,
+ ) -> Result<()> {
+ self.batches_bytes = 0;
+ self.indices_bytes = 0;
let batches = std::mem::take(&mut self.batches);
- // Preserve the partition count by replacing each inner vec with empty
- let indices: Vec<Vec<(u32, u32)>> =
- self.indices.iter_mut().map(std::mem::take).collect();
- (batches, indices)
+ let indices: Vec<_> = self.indices.iter_mut().map(std::mem::take).collect();
+ for (partition, partition_indices) in indices.iter().enumerate() {
Review Comment:
Done. ~~I'll do building the batch refs once per drain in a follow-up PR,
since `main` has the same cost.~~
Actually, your r suggested change is pretty trivial, now that I have time to better look at it today. Going to add it in the next commit to, to save time.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor(core): drain and size the sort-shuffle buffer through BufferedBatches [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2579:
URL: https://github.com/apache/datafusion-ballista/pull/2579#discussion_r4250112359
##########
ballista/core/src/execution_plans/sort_shuffle/writer.rs:
##########
@@ -124,44 +123,37 @@ struct EncodedPartitions {
fn encode_buffered_partitions(
buffered: &mut BufferedBatches,
schema: &SchemaRef,
- config: &SortShuffleConfig,
opts: &datafusion::arrow::ipc::writer::IpcWriteOptions,
) -> Result<EncodedPartitions> {
let num_partitions = buffered.num_partitions();
- let (batches, indices) = buffered.take();
-
- let mut encoded = Vec::with_capacity(num_partitions);
- let mut stats = Vec::with_capacity(num_partitions);
-
- for partition_indices in indices.iter() {
- let mut buf: Vec<u8> = Vec::new();
- let (mut num_batches, mut num_rows, mut num_bytes) = (0u64, 0u64, 0u64);
-
- if !partition_indices.is_empty() {
- let iter = PartitionedBatchIterator::new(
- &batches,
- partition_indices,
- config.batch_size,
- );
- let mut writer =
- StreamWriter::try_new_with_options(&mut buf, schema, opts.clone())?;
- for result in iter {
- let batch = result?;
- num_rows += batch.num_rows() as u64;
- num_bytes += batch.get_array_memory_size() as u64;
- num_batches += 1;
- writer.write(&batch)?;
- }
- writer.finish()?;
- }
+ let mut writers: Vec<Option<StreamWriter<Vec<u8>>>> =
Review Comment:
This is the part I'd like changed before merging. The old loop created, finished and dropped one `StreamWriter` per partition, so only one was open at a time. Now every partition's writer stays open until `drain` returns. Each arrow-ipc `StreamWriter` caches its own zstd compression context after its first write, so with `ballista.shuffle.compression.codec = zstd` this holds one context per non-empty output partition.
I measured peak RSS with arrow-ipc 59.3, one batch per partition:
| partitions | rows per batch | one writer open | all writers open |
|---|---|---|---|
| 200 | 8192 | 19 MB | 159 MB |
| 1000 | 8192 | 95 MB | 789 MB |
That's per input partition, and a task can run several of these encodes at once. None of it goes through the memory pool, so it can push an executor past its limit. lz4 stays roughly flat, but zstd is a supported codec and we shouldn't regress it.
Since `drain` hands out each partition's batches together, finishing the previous writer as soon as the partition changes keeps it to one open writer. Something like:
```rust
let mut encoded = vec![Vec::new(); num_partitions];
let mut open: Option<(usize, StreamWriter<Vec<u8>>)> = None;
buffered.drain(|partition, batch| {
if open.as_ref().is_some_and(|(p, _)| *p != partition) {
let (p, writer) = open.take().unwrap();
encoded[p] = writer.into_inner()?;
}
let (_, writer) = match &mut open {
Some(current) => current,
slot => slot.insert((
partition,
StreamWriter::try_new_with_options(Vec::new(), schema, opts.clone())?,
)),
};
// stats and write as before
Ok(())
})?;
if let Some((p, writer)) = open {
encoded[p] = writer.into_inner()?;
}
```
Another option would be for `drain` to hand `f` one partition at a time with an iterator over its batches. That would keep the old shape in both callers.
##########
ballista/core/src/execution_plans/sort_shuffle/buffer.rs:
##########
@@ -129,58 +127,67 @@ mod tests {
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use std::sync::Arc;
- fn create_test_schema() -> SchemaRef {
+ fn schema() -> SchemaRef {
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]))
}
- fn create_test_batch(schema: &SchemaRef, values: Vec<i32>) -> RecordBatch {
- let array = Int32Array::from(values);
- RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap()
+ fn batch(values: Vec<i32>) -> RecordBatch {
+ RecordBatch::try_new(schema(), vec![Arc::new(Int32Array::from(values))]).unwrap()
+ }
+
+ fn drain(bb: &mut BufferedBatches) -> Vec<(usize, RecordBatch)> {
+ let mut out = vec![];
+ bb.drain(|p, b| {
+ out.push((p, b));
+ Ok(())
+ })
+ .unwrap();
+ out
}
#[test]
- fn buffered_batches_pushes_and_partitions_indices() {
- let schema = create_test_schema();
- let mut bb = BufferedBatches::new(3, schema.clone());
+ fn drains_each_partition_in_arrival_order_and_batch_size() {
+ let mut bb = BufferedBatches::new(3, schema(), 2);
assert!(bb.is_empty());
assert_eq!(bb.num_partitions(), 3);
- let batch_a = create_test_batch(&schema, vec![10, 20, 30, 40]);
- let batch_b = create_test_batch(&schema, vec![50, 60]);
-
- // Partition 0 gets rows {0, 2} from batch 0; partition 1 gets {1, 3} from
- // batch 0; partition 2 gets {0, 1} from batch 1.
- let per_partition_a: Vec<Vec<u32>> = vec![vec![0, 2], vec![1, 3], vec![]];
- let per_partition_b: Vec<Vec<u32>> = vec![vec![], vec![], vec![0, 1]];
-
- bb.push_batch(batch_a, &per_partition_a);
- bb.push_batch(batch_b, &per_partition_b);
-
+ // Partition 0 gets rows {0, 2} of the first batch and all of the
+ // second; partition 1 gets {3, 1} of the first.
+ bb.push_batch(
+ batch(vec![10, 20, 30, 40]),
+ &[vec![0, 2], vec![3, 1], vec![]],
+ );
+ bb.push_batch(batch(vec![50, 60]), &[vec![0, 1], vec![], vec![]]);
assert!(!bb.is_empty());
- // Total rows referenced by indices: 2 + 2 + 2 = 6
- assert_eq!(bb.num_buffered_rows(), 6);
-
- let p0 = bb.indices_for(0);
- assert_eq!(p0, &[(0u32, 0u32), (0, 2)]);
- let p1 = bb.indices_for(1);
- assert_eq!(p1, &[(0u32, 1u32), (0, 3)]);
- let p2 = bb.indices_for(2);
- assert_eq!(p2, &[(1u32, 0u32), (1, 1)]);
+
+ assert_eq!(
+ drain(&mut bb),
+ vec![
+ (0, batch(vec![10, 30])),
+ (0, batch(vec![50, 60])),
+ (1, batch(vec![40, 20])),
+ ]
+ );
+ assert!(bb.is_empty());
+ assert_eq!(drain(&mut bb), vec![]);
}
#[test]
- fn buffered_batches_take_drains_state() {
- let schema = create_test_schema();
- let mut bb = BufferedBatches::new(2, schema.clone());
- bb.push_batch(create_test_batch(&schema, vec![1, 2]), &[vec![0], vec![1]]);
-
- let (batches, indices) = bb.take();
- assert_eq!(batches.len(), 1);
- assert_eq!(indices.len(), 2);
- assert_eq!(indices[0], vec![(0, 0)]);
- assert_eq!(indices[1], vec![(0, 1)]);
+ fn memory_size_counts_batches_and_index_capacity() {
+ let mut bb = BufferedBatches::new(2, schema(), 8);
+ assert_eq!(bb.memory_size(), 0);
+
+ let input = batch(vec![1, 2, 3]);
+ let batch_bytes = input.get_array_memory_size();
+ bb.push_batch(input, &[vec![0, 2], vec![1]]);
Review Comment:
It might be worth pushing a second batch here that makes one of the index lists reallocate. A few more rows into partition 0 would do it. That would check the running total across a capacity change, which is the part most likely to drift when the follow-up adds the second layout.
##########
ballista/core/src/execution_plans/sort_shuffle/mod.rs:
##########
@@ -38,7 +38,6 @@ mod reader;
mod spill;
mod writer;
-pub use buffer::BufferedBatches;
Review Comment:
Agreed this belongs inside the crate. Nothing outside the writer can do much with it. It has been exported as `ballista_core::execution_plans::sort_shuffle::BufferedBatches` since #1598 though, so it shipped in 54.0.0, 54.1.0 and 55.0.0. Could you add a short note to the upgrade guide so library users aren't caught out? We'll want the `api-change` label on this one too.
##########
ballista/core/src/execution_plans/sort_shuffle/buffer.rs:
##########
@@ -99,26 +83,40 @@ impl BufferedBatches {
"BufferedBatches::push_batch schema mismatch"
);
let batch_idx = self.batches.len() as u32;
- for (p, rows) in per_partition_rows.iter().enumerate() {
- let dst = &mut self.indices[p];
- dst.reserve(rows.len());
- for &r in rows {
- dst.push((batch_idx, r));
- self.num_buffered_rows += 1;
- }
+ for (partition_indices, rows) in self.indices.iter_mut().zip(per_partition_rows) {
+ let capacity_before = partition_indices.capacity();
+ partition_indices.extend(rows.iter().map(|&r| (batch_idx, r)));
+ self.indices_bytes += (partition_indices.capacity() - capacity_before)
+ * size_of::<(u32, u32)>();
}
+ self.batches_bytes += batch.get_array_memory_size();
self.batches.push(batch);
}
- /// Drains all state, returning the buffered batches and per-partition
- /// index lists. After this call the buffer is empty.
- pub fn take(&mut self) -> BufferedTake {
- self.num_buffered_rows = 0;
+ /// Empties the buffer, passing each output batch of at most `batch_size`
+ /// rows to `f` with its partition, in partition order and then arrival
+ /// order. Output batches are gathered one at a time, as `f` consumes them.
+ /// The buffer is emptied first, so if `f` or a gather fails, the remaining
+ /// rows are dropped.
+ pub fn drain(
+ &mut self,
+ mut f: impl FnMut(usize, RecordBatch) -> Result<()>,
+ ) -> Result<()> {
+ self.batches_bytes = 0;
+ self.indices_bytes = 0;
let batches = std::mem::take(&mut self.batches);
- // Preserve the partition count by replacing each inner vec with empty
- let indices: Vec<Vec<(u32, u32)>> =
- self.indices.iter_mut().map(std::mem::take).collect();
- (batches, indices)
+ let indices: Vec<_> = self.indices.iter_mut().map(std::mem::take).collect();
+ for (partition, partition_indices) in indices.iter().enumerate() {
Review Comment:
Small one. The old encode and spill loops skipped empty partitions before building the iterator. `PartitionedBatchIterator::new` collects a `Vec<&RecordBatch>` over every buffered batch even when `indices` is empty, so `drain` now pays that for each empty partition. Could we keep the skip here?
```rust
if partition_indices.is_empty() {
continue;
}
```
As a follow-up, those refs are the same for every partition. Building them once per drain and passing a slice in would save rebuilding them for each partition.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor(core): drain and size the sort-shuffle buffer through BufferedBatches [datafusion-ballista]
Posted by "NoahKusaba (via GitHub)" <gi...@apache.org>.
NoahKusaba commented on code in PR #2579:
URL: https://github.com/apache/datafusion-ballista/pull/2579#discussion_r4277860948
##########
ballista/core/src/execution_plans/sort_shuffle/buffer.rs:
##########
@@ -129,58 +129,78 @@ mod tests {
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use std::sync::Arc;
- fn create_test_schema() -> SchemaRef {
+ fn schema() -> SchemaRef {
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]))
}
- fn create_test_batch(schema: &SchemaRef, values: Vec<i32>) -> RecordBatch {
- let array = Int32Array::from(values);
- RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap()
+ fn batch(values: Vec<i32>) -> RecordBatch {
+ RecordBatch::try_new(schema(), vec![Arc::new(Int32Array::from(values))]).unwrap()
+ }
+
+ /// Each partition `drain` visits, with its batches.
+ fn drain(bb: &mut BufferedBatches) -> Vec<(usize, Vec<RecordBatch>)> {
+ let mut out = vec![];
+ bb.drain(|p, batches| {
+ out.push((p, batches.collect::<Result<_>>()?));
+ Ok(())
+ })
+ .unwrap();
+ out
+ }
+
+ fn index_bytes(bb: &BufferedBatches) -> usize {
+ bb.indices
+ .iter()
+ .map(|v| v.capacity() * size_of::<(u32, u32)>())
+ .sum()
}
#[test]
- fn buffered_batches_pushes_and_partitions_indices() {
- let schema = create_test_schema();
- let mut bb = BufferedBatches::new(3, schema.clone());
+ fn drains_each_partition_in_arrival_order_and_batch_size() {
+ let mut bb = BufferedBatches::new(3, schema(), 2);
Review Comment:
That makes sense, i set batch_size to 3 in the latest commit.
Thanks for the review :smile:
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor(core): drain and size the sort-shuffle buffer through BufferedBatches [datafusion-ballista]
Posted by "NoahKusaba (via GitHub)" <gi...@apache.org>.
NoahKusaba commented on code in PR #2579:
URL: https://github.com/apache/datafusion-ballista/pull/2579#discussion_r4254759241
##########
ballista/core/src/execution_plans/sort_shuffle/writer.rs:
##########
@@ -124,44 +123,37 @@ struct EncodedPartitions {
fn encode_buffered_partitions(
buffered: &mut BufferedBatches,
schema: &SchemaRef,
- config: &SortShuffleConfig,
opts: &datafusion::arrow::ipc::writer::IpcWriteOptions,
) -> Result<EncodedPartitions> {
let num_partitions = buffered.num_partitions();
- let (batches, indices) = buffered.take();
-
- let mut encoded = Vec::with_capacity(num_partitions);
- let mut stats = Vec::with_capacity(num_partitions);
-
- for partition_indices in indices.iter() {
- let mut buf: Vec<u8> = Vec::new();
- let (mut num_batches, mut num_rows, mut num_bytes) = (0u64, 0u64, 0u64);
-
- if !partition_indices.is_empty() {
- let iter = PartitionedBatchIterator::new(
- &batches,
- partition_indices,
- config.batch_size,
- );
- let mut writer =
- StreamWriter::try_new_with_options(&mut buf, schema, opts.clone())?;
- for result in iter {
- let batch = result?;
- num_rows += batch.num_rows() as u64;
- num_bytes += batch.get_array_memory_size() as u64;
- num_batches += 1;
- writer.write(&batch)?;
- }
- writer.finish()?;
- }
+ let mut writers: Vec<Option<StreamWriter<Vec<u8>>>> =
Review Comment:
Fixed with your second option: `drain` now calls `f` once per partition with
an iterator over its batches, so encode keeps one writer open at a time.
Memory numbers are in my PR comment.
##########
ballista/core/src/execution_plans/sort_shuffle/buffer.rs:
##########
@@ -99,26 +83,40 @@ impl BufferedBatches {
"BufferedBatches::push_batch schema mismatch"
);
let batch_idx = self.batches.len() as u32;
- for (p, rows) in per_partition_rows.iter().enumerate() {
- let dst = &mut self.indices[p];
- dst.reserve(rows.len());
- for &r in rows {
- dst.push((batch_idx, r));
- self.num_buffered_rows += 1;
- }
+ for (partition_indices, rows) in self.indices.iter_mut().zip(per_partition_rows) {
+ let capacity_before = partition_indices.capacity();
+ partition_indices.extend(rows.iter().map(|&r| (batch_idx, r)));
+ self.indices_bytes += (partition_indices.capacity() - capacity_before)
+ * size_of::<(u32, u32)>();
}
+ self.batches_bytes += batch.get_array_memory_size();
self.batches.push(batch);
}
- /// Drains all state, returning the buffered batches and per-partition
- /// index lists. After this call the buffer is empty.
- pub fn take(&mut self) -> BufferedTake {
- self.num_buffered_rows = 0;
+ /// Empties the buffer, passing each output batch of at most `batch_size`
+ /// rows to `f` with its partition, in partition order and then arrival
+ /// order. Output batches are gathered one at a time, as `f` consumes them.
+ /// The buffer is emptied first, so if `f` or a gather fails, the remaining
+ /// rows are dropped.
+ pub fn drain(
+ &mut self,
+ mut f: impl FnMut(usize, RecordBatch) -> Result<()>,
+ ) -> Result<()> {
+ self.batches_bytes = 0;
+ self.indices_bytes = 0;
let batches = std::mem::take(&mut self.batches);
- // Preserve the partition count by replacing each inner vec with empty
- let indices: Vec<Vec<(u32, u32)>> =
- self.indices.iter_mut().map(std::mem::take).collect();
- (batches, indices)
+ let indices: Vec<_> = self.indices.iter_mut().map(std::mem::take).collect();
+ for (partition, partition_indices) in indices.iter().enumerate() {
Review Comment:
Done. I'll do building the batch refs once per drain in a follow-up PR,
since `main` has the same cost.
##########
ballista/core/src/execution_plans/sort_shuffle/mod.rs:
##########
@@ -38,7 +38,6 @@ mod reader;
mod spill;
mod writer;
-pub use buffer::BufferedBatches;
Review Comment:
Added a note in `docs/source/upgrading/56.0.0.md`. Could you add the
`api-change` label? I don't have access.
##########
ballista/core/src/execution_plans/sort_shuffle/buffer.rs:
##########
@@ -129,58 +127,67 @@ mod tests {
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use std::sync::Arc;
- fn create_test_schema() -> SchemaRef {
+ fn schema() -> SchemaRef {
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]))
}
- fn create_test_batch(schema: &SchemaRef, values: Vec<i32>) -> RecordBatch {
- let array = Int32Array::from(values);
- RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap()
+ fn batch(values: Vec<i32>) -> RecordBatch {
+ RecordBatch::try_new(schema(), vec![Arc::new(Int32Array::from(values))]).unwrap()
+ }
+
+ fn drain(bb: &mut BufferedBatches) -> Vec<(usize, RecordBatch)> {
+ let mut out = vec![];
+ bb.drain(|p, b| {
+ out.push((p, b));
+ Ok(())
+ })
+ .unwrap();
+ out
}
#[test]
- fn buffered_batches_pushes_and_partitions_indices() {
- let schema = create_test_schema();
- let mut bb = BufferedBatches::new(3, schema.clone());
+ fn drains_each_partition_in_arrival_order_and_batch_size() {
+ let mut bb = BufferedBatches::new(3, schema(), 2);
assert!(bb.is_empty());
assert_eq!(bb.num_partitions(), 3);
- let batch_a = create_test_batch(&schema, vec![10, 20, 30, 40]);
- let batch_b = create_test_batch(&schema, vec![50, 60]);
-
- // Partition 0 gets rows {0, 2} from batch 0; partition 1 gets {1, 3} from
- // batch 0; partition 2 gets {0, 1} from batch 1.
- let per_partition_a: Vec<Vec<u32>> = vec![vec![0, 2], vec![1, 3], vec![]];
- let per_partition_b: Vec<Vec<u32>> = vec![vec![], vec![], vec![0, 1]];
-
- bb.push_batch(batch_a, &per_partition_a);
- bb.push_batch(batch_b, &per_partition_b);
-
+ // Partition 0 gets rows {0, 2} of the first batch and all of the
+ // second; partition 1 gets {3, 1} of the first.
+ bb.push_batch(
+ batch(vec![10, 20, 30, 40]),
+ &[vec![0, 2], vec![3, 1], vec![]],
+ );
+ bb.push_batch(batch(vec![50, 60]), &[vec![0, 1], vec![], vec![]]);
assert!(!bb.is_empty());
- // Total rows referenced by indices: 2 + 2 + 2 = 6
- assert_eq!(bb.num_buffered_rows(), 6);
-
- let p0 = bb.indices_for(0);
- assert_eq!(p0, &[(0u32, 0u32), (0, 2)]);
- let p1 = bb.indices_for(1);
- assert_eq!(p1, &[(0u32, 1u32), (0, 3)]);
- let p2 = bb.indices_for(2);
- assert_eq!(p2, &[(1u32, 0u32), (1, 1)]);
+
+ assert_eq!(
+ drain(&mut bb),
+ vec![
+ (0, batch(vec![10, 30])),
+ (0, batch(vec![50, 60])),
+ (1, batch(vec![40, 20])),
+ ]
+ );
+ assert!(bb.is_empty());
+ assert_eq!(drain(&mut bb), vec![]);
}
#[test]
- fn buffered_batches_take_drains_state() {
- let schema = create_test_schema();
- let mut bb = BufferedBatches::new(2, schema.clone());
- bb.push_batch(create_test_batch(&schema, vec![1, 2]), &[vec![0], vec![1]]);
-
- let (batches, indices) = bb.take();
- assert_eq!(batches.len(), 1);
- assert_eq!(indices.len(), 2);
- assert_eq!(indices[0], vec![(0, 0)]);
- assert_eq!(indices[1], vec![(0, 1)]);
+ fn memory_size_counts_batches_and_index_capacity() {
+ let mut bb = BufferedBatches::new(2, schema(), 8);
+ assert_eq!(bb.memory_size(), 0);
+
+ let input = batch(vec![1, 2, 3]);
+ let batch_bytes = input.get_array_memory_size();
+ bb.push_batch(input, &[vec![0, 2], vec![1]]);
Review Comment:
Added: a second push now reallocates partition 0's index list.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor(core): drain and size the sort-shuffle buffer through BufferedBatches [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2579:
URL: https://github.com/apache/datafusion-ballista/pull/2579#discussion_r4276171535
##########
ballista/core/src/execution_plans/sort_shuffle/buffer.rs:
##########
@@ -129,58 +129,78 @@ mod tests {
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use std::sync::Arc;
- fn create_test_schema() -> SchemaRef {
+ fn schema() -> SchemaRef {
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]))
}
- fn create_test_batch(schema: &SchemaRef, values: Vec<i32>) -> RecordBatch {
- let array = Int32Array::from(values);
- RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap()
+ fn batch(values: Vec<i32>) -> RecordBatch {
+ RecordBatch::try_new(schema(), vec![Arc::new(Int32Array::from(values))]).unwrap()
+ }
+
+ /// Each partition `drain` visits, with its batches.
+ fn drain(bb: &mut BufferedBatches) -> Vec<(usize, Vec<RecordBatch>)> {
+ let mut out = vec![];
+ bb.drain(|p, batches| {
+ out.push((p, batches.collect::<Result<_>>()?));
+ Ok(())
+ })
+ .unwrap();
+ out
+ }
+
+ fn index_bytes(bb: &BufferedBatches) -> usize {
+ bb.indices
+ .iter()
+ .map(|v| v.capacity() * size_of::<(u32, u32)>())
+ .sum()
}
#[test]
- fn buffered_batches_pushes_and_partitions_indices() {
- let schema = create_test_schema();
- let mut bb = BufferedBatches::new(3, schema.clone());
+ fn drains_each_partition_in_arrival_order_and_batch_size() {
+ let mut bb = BufferedBatches::new(3, schema(), 2);
Review Comment:
With `batch_size` 2, partition 0's rows split exactly where the first input batch ends, so this test can't tell "chunk by `batch_size`" from "one output batch per input batch". With 3, the first output batch has to take rows from both inputs:
```rust
let mut bb = BufferedBatches::new(3, schema(), 3);
// ...
(0, vec![batch(vec![10, 30, 50]), batch(vec![60])]),
```
That's what the coalescing layout in #2578 will have to reproduce, so it seems worth pinning here.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor(core): drain and size the sort-shuffle buffer through BufferedBatches [datafusion-ballista]
Posted by "NoahKusaba (via GitHub)" <gi...@apache.org>.
NoahKusaba commented on code in PR #2579:
URL: https://github.com/apache/datafusion-ballista/pull/2579#discussion_r4254509391
##########
ballista/core/src/execution_plans/sort_shuffle/writer.rs:
##########
@@ -124,44 +123,37 @@ struct EncodedPartitions {
fn encode_buffered_partitions(
buffered: &mut BufferedBatches,
schema: &SchemaRef,
- config: &SortShuffleConfig,
opts: &datafusion::arrow::ipc::writer::IpcWriteOptions,
) -> Result<EncodedPartitions> {
let num_partitions = buffered.num_partitions();
- let (batches, indices) = buffered.take();
-
- let mut encoded = Vec::with_capacity(num_partitions);
- let mut stats = Vec::with_capacity(num_partitions);
-
- for partition_indices in indices.iter() {
- let mut buf: Vec<u8> = Vec::new();
- let (mut num_batches, mut num_rows, mut num_bytes) = (0u64, 0u64, 0u64);
-
- if !partition_indices.is_empty() {
- let iter = PartitionedBatchIterator::new(
- &batches,
- partition_indices,
- config.batch_size,
- );
- let mut writer =
- StreamWriter::try_new_with_options(&mut buf, schema, opts.clone())?;
- for result in iter {
- let batch = result?;
- num_rows += batch.num_rows() as u64;
- num_bytes += batch.get_array_memory_size() as u64;
- num_batches += 1;
- writer.write(&batch)?;
- }
- writer.finish()?;
- }
+ let mut writers: Vec<Option<StreamWriter<Vec<u8>>>> =
Review Comment:
🤦♂️ Good catch, I'll go with your second suggestion of handing one partition at-a-time with an iterator over batches.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor(core): drain and size the sort-shuffle buffer through BufferedBatches [datafusion-ballista]
Posted by "NoahKusaba (via GitHub)" <gi...@apache.org>.
NoahKusaba commented on code in PR #2579:
URL: https://github.com/apache/datafusion-ballista/pull/2579#discussion_r4254759397
##########
ballista/core/src/execution_plans/sort_shuffle/buffer.rs:
##########
@@ -99,26 +83,40 @@ impl BufferedBatches {
"BufferedBatches::push_batch schema mismatch"
);
let batch_idx = self.batches.len() as u32;
- for (p, rows) in per_partition_rows.iter().enumerate() {
- let dst = &mut self.indices[p];
- dst.reserve(rows.len());
- for &r in rows {
- dst.push((batch_idx, r));
- self.num_buffered_rows += 1;
- }
+ for (partition_indices, rows) in self.indices.iter_mut().zip(per_partition_rows) {
+ let capacity_before = partition_indices.capacity();
+ partition_indices.extend(rows.iter().map(|&r| (batch_idx, r)));
+ self.indices_bytes += (partition_indices.capacity() - capacity_before)
+ * size_of::<(u32, u32)>();
}
+ self.batches_bytes += batch.get_array_memory_size();
self.batches.push(batch);
}
- /// Drains all state, returning the buffered batches and per-partition
- /// index lists. After this call the buffer is empty.
- pub fn take(&mut self) -> BufferedTake {
- self.num_buffered_rows = 0;
+ /// Empties the buffer, passing each output batch of at most `batch_size`
+ /// rows to `f` with its partition, in partition order and then arrival
+ /// order. Output batches are gathered one at a time, as `f` consumes them.
+ /// The buffer is emptied first, so if `f` or a gather fails, the remaining
+ /// rows are dropped.
+ pub fn drain(
+ &mut self,
+ mut f: impl FnMut(usize, RecordBatch) -> Result<()>,
+ ) -> Result<()> {
+ self.batches_bytes = 0;
+ self.indices_bytes = 0;
let batches = std::mem::take(&mut self.batches);
- // Preserve the partition count by replacing each inner vec with empty
- let indices: Vec<Vec<(u32, u32)>> =
- self.indices.iter_mut().map(std::mem::take).collect();
- (batches, indices)
+ let indices: Vec<_> = self.indices.iter_mut().map(std::mem::take).collect();
+ for (partition, partition_indices) in indices.iter().enumerate() {
Review Comment:
Done. ~~~I'll do building the batch refs once per drain in a follow-up PR,
since `main` has the same cost.~~~
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor(core): drain and size the sort-shuffle buffer through BufferedBatches [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove merged PR #2579:
URL: https://github.com/apache/datafusion-ballista/pull/2579
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor(core): drain and size the sort-shuffle buffer through BufferedBatches [datafusion-ballista]
Posted by "NoahKusaba (via GitHub)" <gi...@apache.org>.
NoahKusaba commented on code in PR #2579:
URL: https://github.com/apache/datafusion-ballista/pull/2579#discussion_r4254759241
##########
ballista/core/src/execution_plans/sort_shuffle/writer.rs:
##########
@@ -124,44 +123,37 @@ struct EncodedPartitions {
fn encode_buffered_partitions(
buffered: &mut BufferedBatches,
schema: &SchemaRef,
- config: &SortShuffleConfig,
opts: &datafusion::arrow::ipc::writer::IpcWriteOptions,
) -> Result<EncodedPartitions> {
let num_partitions = buffered.num_partitions();
- let (batches, indices) = buffered.take();
-
- let mut encoded = Vec::with_capacity(num_partitions);
- let mut stats = Vec::with_capacity(num_partitions);
-
- for partition_indices in indices.iter() {
- let mut buf: Vec<u8> = Vec::new();
- let (mut num_batches, mut num_rows, mut num_bytes) = (0u64, 0u64, 0u64);
-
- if !partition_indices.is_empty() {
- let iter = PartitionedBatchIterator::new(
- &batches,
- partition_indices,
- config.batch_size,
- );
- let mut writer =
- StreamWriter::try_new_with_options(&mut buf, schema, opts.clone())?;
- for result in iter {
- let batch = result?;
- num_rows += batch.num_rows() as u64;
- num_bytes += batch.get_array_memory_size() as u64;
- num_batches += 1;
- writer.write(&batch)?;
- }
- writer.finish()?;
- }
+ let mut writers: Vec<Option<StreamWriter<Vec<u8>>>> =
Review Comment:
Fixed with your second option: `drain` now calls `f` once per partition with
an iterator over its batches, so encode keeps one writer open at a time.
Memory numbers are in my PR comment.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org