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