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