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: [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]
