chore: Cleanup shuffle structure

Author: EmilyMatt

native/core/benches/row_columnar.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use arrow::datatypes::DataType as ArrowDataType;
-use comet::execution::shuffle::row::{
+use comet::execution::shuffle::spark_unsafe::row::{
     process_sorted_row_partition, SparkUnsafeObject, SparkUnsafeRow,
 };
 use comet::execution::shuffle::CompressionCodec;

native/core/src/execution/jni_api.rs

@@ -22,7 +22,7 @@ use crate::{
     errors::{try_unwrap_or_throw, CometError, CometResult},
     execution::{
         metrics::utils::update_comet_metric, planner::PhysicalPlanner, serde::to_arrow_datatype,
-        shuffle::row::process_sorted_row_partition, sort::RdxSort,
+        shuffle::spark_unsafe::row::process_sorted_row_partition, sort::RdxSort,
     },
     jvm_bridge::{jni_new_global_ref, JVMClasses},
 };

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

@@ -17,10 +17,9 @@
 
 pub(crate) mod codec;
 mod comet_partitioning;
-mod list;
-mod map;
-pub mod row;
+mod repartitioners;
 mod shuffle_writer;
+pub mod spark_unsafe;
 
 pub use codec::{read_ipc_compressed, CompressionCodec, ShuffleBlockWriter};
 pub use comet_partitioning::CometPartitioning;

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

@@ -0,0 +1,15 @@
+use arrow::array::RecordBatch;
+
+pub(super) use multi_partition::MultiPartitionShuffleRepartitioner;
+pub(super) use single_partition::SinglePartitionShufflePartitioner;
+
+mod multi_partition;
+mod single_partition;
+
+#[async_trait::async_trait]
+pub(super) trait ShufflePartitioner: Send + Sync {
+    /// Insert a batch into the partitioner
+    async fn insert_batch(&mut self, batch: RecordBatch) -> datafusion::common::Result<()>;
+    /// Write shuffle data and shuffle index file to disk
+    fn shuffle_write(&mut self) -> datafusion::common::Result<()>;
+}

native/core/src/execution/shuffle/repartitioners/multi_partition.rs

@@ -0,0 +1,176 @@
+use crate::execution::shuffle::repartitioners::ShufflePartitioner;
+use crate::execution::shuffle::shuffle_writer::{BufBatchWriter, ShuffleRepartitionerMetrics};
+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 tokio::time::Instant;
+
+/// A partitioner that writes all shuffle data to a single file and a single index file
+pub(crate) struct SinglePartitionShufflePartitioner {
+    // output_data_file: File,
+    output_data_writer: BufBatchWriter<ShuffleBlockWriter, File>,
+    output_index_path: String,
+    /// Batches that are smaller than the batch size and to be concatenated
+    buffered_batches: Vec<RecordBatch>,
+    /// Number of rows in the concatenating batches
+    num_buffered_rows: usize,
+    /// Metrics for the repartitioner
+    metrics: ShuffleRepartitionerMetrics,
+    /// The configured batch size
+    batch_size: usize,
+}
+
+impl SinglePartitionShufflePartitioner {
+    pub(crate) fn try_new(
+        output_data_path: String,
+        output_index_path: String,
+        schema: SchemaRef,
+        metrics: ShuffleRepartitionerMetrics,
+        batch_size: usize,
+        codec: CompressionCodec,
+        write_buffer_size: usize,
+    ) -> datafusion::common::Result<Self> {
+        let shuffle_block_writer = ShuffleBlockWriter::try_new(schema.as_ref(), codec.clone())?;
+
+        let output_data_file = OpenOptions::new()
+            .write(true)
+            .create(true)
+            .truncate(true)
+            .open(output_data_path)
+            .map_err(shuffle_writer::to_df_err)?;
+
+        let output_data_writer =
+            BufBatchWriter::new(shuffle_block_writer, output_data_file, write_buffer_size);
+
+        Ok(Self {
+            output_data_writer,
+            output_index_path,
+            buffered_batches: vec![],
+            num_buffered_rows: 0,
+            metrics,
+            batch_size,
+        })
+    }
+
+    /// Add a batch to the buffer of the partitioner, these buffered batches will be concatenated
+    /// and written to the output data file when the number of rows in the buffer reaches the batch size.
+    fn add_buffered_batch(&mut self, batch: RecordBatch) {
+        self.num_buffered_rows += batch.num_rows();
+        self.buffered_batches.push(batch);
+    }
+
+    /// Consumes buffered batches and return a concatenated batch if successful
+    fn concat_buffered_batches(&mut self) -> datafusion::common::Result<Option<RecordBatch>> {
+        if self.buffered_batches.is_empty() {
+            Ok(None)
+        } else if self.buffered_batches.len() == 1 {
+            let batch = self.buffered_batches.remove(0);
+            self.num_buffered_rows = 0;
+            Ok(Some(batch))
+        } else {
+            let schema = &self.buffered_batches[0].schema();
+            match arrow::compute::concat_batches(schema, self.buffered_batches.iter()) {
+                Ok(concatenated) => {
+                    self.buffered_batches.clear();
+                    self.num_buffered_rows = 0;
+                    Ok(Some(concatenated))
+                }
+                Err(e) => Err(DataFusionError::ArrowError(
+                    Box::from(e),
+                    Some(DataFusionError::get_back_trace()),
+                )),
+            }
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl ShufflePartitioner for SinglePartitionShufflePartitioner {
+    async fn insert_batch(&mut self, batch: RecordBatch) -> datafusion::common::Result<()> {
+        let start_time = Instant::now();
+        let num_rows = batch.num_rows();
+
+        if num_rows > 0 {
+            self.metrics.data_size.add(batch.get_array_memory_size());
+            self.metrics.baseline.record_output(num_rows);
+
+            if num_rows >= self.batch_size || num_rows + self.num_buffered_rows > self.batch_size {
+                let concatenated_batch = self.concat_buffered_batches()?;
+
+                // Write the concatenated buffered batch
+                if let Some(batch) = concatenated_batch {
+                    self.output_data_writer.write(
+                        &batch,
+                        &self.metrics.encode_time,
+                        &self.metrics.write_time,
+                    )?;
+                }
+
+                if num_rows >= self.batch_size {
+                    // Write the new batch
+                    self.output_data_writer.write(
+                        &batch,
+                        &self.metrics.encode_time,
+                        &self.metrics.write_time,
+                    )?;
+                } else {
+                    // Add the new batch to the buffer
+                    self.add_buffered_batch(batch);
+                }
+            } else {
+                self.add_buffered_batch(batch);
+            }
+        }
+
+        self.metrics.input_batches.add(1);
+        self.metrics
+            .baseline
+            .elapsed_compute()
+            .add_duration(start_time.elapsed());
+        Ok(())
+    }
+
+    fn shuffle_write(&mut self) -> datafusion::common::Result<()> {
+        let start_time = Instant::now();
+        let concatenated_batch = self.concat_buffered_batches()?;
+
+        // Write the concatenated buffered batch
+        if let Some(batch) = concatenated_batch {
+            self.output_data_writer.write(
+                &batch,
+                &self.metrics.encode_time,
+                &self.metrics.write_time,
+            )?;
+        }
+        self.output_data_writer.flush(&self.metrics.write_time)?;
+
+        // Write index file. It should only contain 2 entries: 0 and the total number of bytes written
+        let index_file = OpenOptions::new()
+            .write(true)
+            .create(true)
+            .truncate(true)
+            .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)?;
+        for offset in [0, data_file_length] {
+            index_buf_writer
+                .write_all(&(offset as i64).to_le_bytes()[..])
+                .map_err(shuffle_writer::to_df_err)?;
+        }
+        index_buf_writer.flush()?;
+
+        self.metrics
+            .baseline
+            .elapsed_compute()
+            .add_duration(start_time.elapsed());
+        Ok(())
+    }
+}