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