more cleanup

Author: EmilyMatt

native/core/src/execution/shuffle/metrics.rs

@@ -0,0 +1,44 @@
+use datafusion::physical_plan::metrics::{
+    BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, Time,
+};
+
+pub(super) struct ShuffleRepartitionerMetrics {
+    /// metrics
+    pub(super) baseline: BaselineMetrics,
+
+    /// Time to perform repartitioning
+    pub(super) repart_time: Time,
+
+    /// Time encoding batches to IPC format
+    pub(super) encode_time: Time,
+
+    /// Time spent writing to disk. Maps to "shuffleWriteTime" in Spark SQL Metrics.
+    pub(super) write_time: Time,
+
+    /// Number of input batches
+    pub(super) input_batches: Count,
+
+    /// count of spills during the execution of the operator
+    pub(super) spill_count: Count,
+
+    /// total spilled bytes during the execution of the operator
+    pub(super) spilled_bytes: Count,
+
+    /// The original size of spilled data. Different to `spilled_bytes` because of compression.
+    pub(super) data_size: Count,
+}
+
+impl ShuffleRepartitionerMetrics {
+    pub(super) fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self {
+        Self {
+            baseline: BaselineMetrics::new(metrics, partition),
+            repart_time: MetricBuilder::new(metrics).subset_time("repart_time", partition),
+            encode_time: MetricBuilder::new(metrics).subset_time("encode_time", partition),
+            write_time: MetricBuilder::new(metrics).subset_time("write_time", partition),
+            input_batches: MetricBuilder::new(metrics).counter("input_batches", partition),
+            spill_count: MetricBuilder::new(metrics).spill_count(partition),
+            spilled_bytes: MetricBuilder::new(metrics).spilled_bytes(partition),
+            data_size: MetricBuilder::new(metrics).counter("data_size", partition),
+        }
+    }
+}

native/core/src/execution/shuffle/mod.rs

@@ -17,9 +17,11 @@
 
 pub(crate) mod codec;
 mod comet_partitioning;
-mod repartitioners;
+mod metrics;
+mod partitioners;
 mod shuffle_writer;
 pub mod spark_unsafe;
+mod writers;
 
 pub use codec::{read_ipc_compressed, CompressionCodec, ShuffleBlockWriter};
 pub use comet_partitioning::CometPartitioning;

…/execution/shuffle/repartitioners/mod.rs …rc/execution/shuffle/partitioners/mod.rsnative/core/src/execution/shuffle/repartitioners/mod.rs renamed to native/core/src/execution/shuffle/partitioners/mod.rs native/core/src/execution/shuffle/repartitioners/mod.rs renamed to native/core/src/execution/shuffle/partitioners/mod.rs

@@ -1,9 +1,13 @@
 use arrow::array::RecordBatch;
 
-pub(super) use multi_partition::MultiPartitionShuffleRepartitioner;
+pub(super) use multi_partition::MultiPartitionShufflePartitioner;
+pub(super) use partitioned_batches_producer::{
+    PartitionedBatchIterator, PartitionedBatchesProducer,
+};
 pub(super) use single_partition::SinglePartitionShufflePartitioner;
 
 mod multi_partition;
+mod partitioned_batches_producer;
 mod single_partition;
 
 #[async_trait::async_trait]

…huffle/repartitioners/multi_partition.rs …/shuffle/partitioners/multi_partition.rsnative/core/src/execution/shuffle/repartitioners/multi_partition.rs renamed to native/core/src/execution/shuffle/partitioners/multi_partition.rs native/core/src/execution/shuffle/repartitioners/multi_partition.rs renamed to native/core/src/execution/shuffle/partitioners/multi_partition.rs

@@ -1,8 +1,8 @@
-use crate::execution::shuffle::repartitioners::ShufflePartitioner;
-use crate::execution::shuffle::shuffle_writer::{
-    BufBatchWriter, PartitionWriter, PartitionedBatchIterator, PartitionedBatchesProducer,
-    ShuffleRepartitionerMetrics,
+use crate::execution::shuffle::metrics::ShuffleRepartitionerMetrics;
+use crate::execution::shuffle::partitioners::{
+    PartitionedBatchIterator, PartitionedBatchesProducer, ShufflePartitioner,
 };
+use crate::execution::shuffle::writers::{BufBatchWriter, PartitionWriter};
 use crate::execution::shuffle::{
     shuffle_writer, CometPartitioning, CompressionCodec, ShuffleBlockWriter,
 };
@@ -24,7 +24,7 @@ use std::sync::Arc;
 use tokio::time::Instant;
 
 /// A partitioner that uses a hash function to partition data into multiple partitions
-pub(crate) struct MultiPartitionShuffleRepartitioner {
+pub(crate) struct MultiPartitionShufflePartitioner {
     output_data_file: String,
     output_index_file: String,
     buffered_batches: Vec<RecordBatch>,
@@ -62,7 +62,7 @@ struct ScratchSpace {
     partition_starts: Vec<u32>,
 }
 
-impl MultiPartitionShuffleRepartitioner {
+impl MultiPartitionShufflePartitioner {
     #[allow(clippy::too_many_arguments)]
     pub(crate) fn try_new(
         partition: usize,
@@ -509,14 +509,14 @@ impl MultiPartitionShuffleRepartitioner {
 }
 
 #[cfg(test)]
-impl MultiPartitionShuffleRepartitioner {
+impl MultiPartitionShufflePartitioner {
     pub(crate) fn partition_writers(&self) -> &[PartitionWriter] {
         &self.partition_writers
     }
 }
 
 #[async_trait::async_trait]
-impl ShufflePartitioner for MultiPartitionShuffleRepartitioner {
+impl ShufflePartitioner for MultiPartitionShufflePartitioner {
     /// Shuffles rows in input batch into corresponding partition buffer.
     /// This function will slice input batch according to configured batch size and then
     /// shuffle rows into corresponding partition buffer.
@@ -567,10 +567,9 @@ impl ShufflePartitioner for MultiPartitionShuffleRepartitioner {
 
                 // if we wrote a spill file for this partition then copy the
                 // contents into the shuffle file
-                if let Some(spill_data) = self.partition_writers[i].spill_file.as_ref() {
+                if let Some(spill_data) = self.partition_writers[i].spill_file() {
                     let mut spill_file = BufReader::new(
-                        File::open(spill_data.temp_file.path())
-                            .map_err(shuffle_writer::to_df_err)?,
+                        File::open(spill_data.path()).map_err(shuffle_writer::to_df_err)?,
                     );
                     let mut write_timer = self.metrics.write_time.timer();
                     std::io::copy(&mut spill_file, &mut output_data)
@@ -622,7 +621,7 @@ impl ShufflePartitioner for MultiPartitionShuffleRepartitioner {
     }
 }
 
-impl Debug for MultiPartitionShuffleRepartitioner {
+impl Debug for MultiPartitionShufflePartitioner {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         f.debug_struct("ShuffleRepartitioner")
             .field("memory_used", &self.used())

native/core/src/execution/shuffle/partitioners/partitioned_batches_producer.rs

@@ -0,0 +1,93 @@
+use arrow::array::RecordBatch;
+use arrow::compute::interleave_record_batch;
+use datafusion::common::DataFusionError;
+
+/// A helper struct to produce shuffled batches.
+/// This struct takes ownership of the buffered batches and partition indices from the
+/// ShuffleRepartitioner, and provides an iterator over the batches in the specified partitions.
+pub(crate) struct PartitionedBatchesProducer {
+    buffered_batches: Vec<RecordBatch>,
+    partition_indices: Vec<Vec<(u32, u32)>>,
+    batch_size: usize,
+}
+
+impl PartitionedBatchesProducer {
+    pub(crate) fn new(
+        buffered_batches: Vec<RecordBatch>,
+        indices: Vec<Vec<(u32, u32)>>,
+        batch_size: usize,
+    ) -> Self {
+        Self {
+            partition_indices: indices,
+            buffered_batches,
+            batch_size,
+        }
+    }
+
+    pub(crate) fn produce(&mut self, partition_id: usize) -> PartitionedBatchIterator<'_> {
+        PartitionedBatchIterator::new(
+            &self.partition_indices[partition_id],
+            &self.buffered_batches,
+            self.batch_size,
+        )
+    }
+}
+
+pub(crate) struct PartitionedBatchIterator<'a> {
+    record_batches: Vec<&'a RecordBatch>,
+    batch_size: usize,
+    indices: Vec<(usize, usize)>,
+    pos: usize,
+}
+
+impl<'a> PartitionedBatchIterator<'a> {
+    fn new(
+        indices: &'a [(u32, u32)],
+        buffered_batches: &'a [RecordBatch],
+        batch_size: usize,
+    ) -> Self {
+        if indices.is_empty() {
+            // Avoid unnecessary allocations when the partition is empty
+            return Self {
+                record_batches: vec![],
+                batch_size,
+                indices: vec![],
+                pos: 0,
+            };
+        }
+        let record_batches = buffered_batches.iter().collect::<Vec<_>>();
+        let current_indices = indices
+            .iter()
+            .map(|(i_batch, i_row)| (*i_batch as usize, *i_row as usize))
+            .collect::<Vec<_>>();
+        Self {
+            record_batches,
+            batch_size,
+            indices: current_indices,
+            pos: 0,
+        }
+    }
+}
+
+impl Iterator for PartitionedBatchIterator<'_> {
+    type Item = datafusion::common::Result<RecordBatch>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.pos >= self.indices.len() {
+            return None;
+        }
+
+        let indices_end = std::cmp::min(self.pos + self.batch_size, self.indices.len());
+        let indices = &self.indices[self.pos..indices_end];
+        match interleave_record_batch(&self.record_batches, indices) {
+            Ok(batch) => {
+                self.pos = indices_end;
+                Some(Ok(batch))
+            }
+            Err(e) => Some(Err(DataFusionError::ArrowError(
+                Box::from(e),
+                Some(DataFusionError::get_back_trace()),
+            ))),
+        }
+    }
+}

…uffle/repartitioners/single_partition.rs …shuffle/partitioners/single_partition.rsnative/core/src/execution/shuffle/repartitioners/single_partition.rs renamed to native/core/src/execution/shuffle/partitioners/single_partition.rs native/core/src/execution/shuffle/repartitioners/single_partition.rs renamed to native/core/src/execution/shuffle/partitioners/single_partition.rs

@@ -1,11 +1,12 @@
-use crate::execution::shuffle::repartitioners::ShufflePartitioner;
-use crate::execution::shuffle::shuffle_writer::{BufBatchWriter, ShuffleRepartitionerMetrics};
+use crate::execution::shuffle::metrics::ShuffleRepartitionerMetrics;
+use crate::execution::shuffle::partitioners::ShufflePartitioner;
+use crate::execution::shuffle::writers::BufBatchWriter;
 use crate::execution::shuffle::{shuffle_writer, CompressionCodec, ShuffleBlockWriter};
 use arrow::array::RecordBatch;
 use arrow::datatypes::SchemaRef;
 use datafusion::common::DataFusionError;
 use std::fs::{File, OpenOptions};
-use std::io::{BufWriter, Seek, Write};
+use std::io::{BufWriter, Write};
 use tokio::time::Instant;
 
 /// A partitioner that writes all shuffle data to a single file and a single index file
@@ -155,11 +156,7 @@ impl ShufflePartitioner for SinglePartitionShufflePartitioner {
             .open(self.output_index_path.clone())
             .map_err(|e| DataFusionError::Execution(format!("shuffle write error: {e:?}")))?;
         let mut index_buf_writer = BufWriter::new(index_file);
-        let data_file_length = self
-            .output_data_writer
-            .writer
-            .stream_position()
-            .map_err(shuffle_writer::to_df_err)?;
+        let data_file_length = self.output_data_writer.writer_position()?;
         for offset in [0, data_file_length] {
             index_buf_writer
                 .write_all(&(offset as i64).to_le_bytes()[..])