Refactor partition groups and add tests

This change adds a new abstraction called a `Splitter` in `distribution_strategy.rs`. This is responsible for splitting partitions into partition groups. The new `distribution_strategy` module can be used to store other splitters and, in the future, distribution strategies like round robin etc.

A "partition group" is now a range [starting_partition, ending_partition) instead of a Vec, so we improve space efficiency. In the proto layer, the `DDTask` proto still uses a Vec, so we simply expand the range into a Vec when creating them in `planning.rs::assign_to_workers`.

Testing
- unit tests in isolator.rs
- unit tests for `build_replacement` in `planning.rs`
- unit tests for the splitter in `distribution_strategy.rs`

Author: jayshrivastava

src/codec.rs

@@ -305,7 +305,7 @@ mod test {
         let schema = create_test_schema();
         let part = Partitioning::UnknownPartitioning(2);
         let reader_exec = Arc::new(DDStageReaderExec::try_new(part, schema, 1).unwrap());
-        let exec = Arc::new(PartitionIsolatorExec::new(reader_exec, 4));
+        let exec = Arc::new(PartitionIsolatorExec::new(reader_exec, 2));
 
         verify_round_trip(exec);
     }

src/distribution_strategy.rs

@@ -0,0 +1,88 @@
+// A Splitter is used to group partitions.
+pub trait Splitter {
+    fn split(&self, num_partitions: usize) -> Vec<PartitionGroup>;
+}
+
+// PartitionGroup is a struct that represents a range of partitions from [start, end). This is
+// more space efficient than a vector of u64s.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct PartitionGroup {
+    start: usize,
+    end: usize,
+}
+
+impl PartitionGroup {
+    pub fn new(start: usize, end: usize) -> Self {
+        Self { start, end }
+    }
+
+    // start is the first in the range
+    pub fn start(&self) -> usize {
+        self.start
+    }
+
+    // end is the exclusive end partition in the range
+    pub fn end(&self) -> usize {
+        self.end
+    }
+}
+
+#[derive(Debug, Copy, Clone, PartialEq)]
+pub struct PartitionSplitter{
+    partition_group_size: usize,
+}
+
+impl PartitionSplitter {
+    pub fn new(partition_group_size: usize) -> Self {
+        assert!(partition_group_size > 0, "partition groups cannot be size 0");
+        PartitionSplitter {
+            partition_group_size,
+        }
+    }
+}
+
+impl Splitter for PartitionSplitter {
+    fn split(&self, num_partitions: usize) -> Vec<PartitionGroup> {
+        (0..num_partitions)
+            .step_by(self.partition_group_size)
+            .map(|start| {
+                let end = std::cmp::min(start + self.partition_group_size, num_partitions);
+                PartitionGroup { start, end }
+            })
+            .collect()
+    }
+}
+
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_partition_splitter_basic() {
+        let splitter = PartitionSplitter::new(4);
+        let groups = splitter.split(10);
+
+        let expected = vec![
+            PartitionGroup { start: 0, end: 4 },
+            PartitionGroup { start: 4, end: 8 },
+            PartitionGroup { start: 8, end: 10 },
+        ];
+
+        assert_eq!(groups, expected);
+    }
+
+    #[test]
+    fn test_partition_splitter_uneven() {
+        let splitter = PartitionSplitter::new(2);
+        let groups = splitter.split(5);
+
+        let expected = vec![
+            PartitionGroup { start: 0, end: 2 },
+            PartitionGroup { start: 2, end: 4 },
+            PartitionGroup { start: 4, end: 5 },
+        ];
+
+        assert_eq!(groups, expected);
+    }
+}

src/isolator.rs

@@ -73,9 +73,9 @@ impl ExecutionPlan for PartitionIsolatorExec {
     }
 
     fn with_new_children(
-        self: std::sync::Arc<Self>,
-        children: Vec<std::sync::Arc<dyn ExecutionPlan>>,
-    ) -> Result<std::sync::Arc<dyn ExecutionPlan>> {
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
         // TODO: generalize this
         assert_eq!(children.len(), 1);
         Ok(Arc::new(Self::new(
@@ -87,7 +87,7 @@ impl ExecutionPlan for PartitionIsolatorExec {
     fn execute(
         &self,
         partition: usize,
-        context: std::sync::Arc<datafusion::execution::TaskContext>,
+        context: Arc<datafusion::execution::TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
         let config = context.session_config();
         let partition_group = &config
@@ -143,3 +143,73 @@ impl ExecutionPlan for PartitionIsolatorExec {
         output_stream
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use datafusion::{
+        arrow::datatypes::{DataType, Field, Schema},
+        physical_plan::Partitioning,
+    };
+    use crate::record_batch_exec::RecordBatchExec;
+    use arrow::array::{Int32Array, RecordBatch};
+    use std::sync::Arc;
+
+    fn create_test_record_batch_exec() -> Arc<dyn ExecutionPlan> {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("col1", DataType::Int32, false),
+        ]));
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
+        ).unwrap();
+        Arc::new(RecordBatchExec::new(batch))
+    }
+
+    #[test]
+    fn test_partition_isolator_exec_new() {
+        let input = create_test_record_batch_exec();
+        let partition_count = 4;
+        let isolator = PartitionIsolatorExec::new(input.clone(), partition_count);
+
+        assert_eq!(isolator.partition_count, partition_count);
+        assert!(Arc::ptr_eq(&isolator.input, &input));
+        assert_eq!(
+            isolator.properties().output_partitioning().partition_count(),
+            partition_count
+        );
+    }
+
+    #[test]
+    fn test_partition_isolator_exec_properties() {
+        let input = create_test_record_batch_exec();
+        let partition_count = 3;
+        let isolator = PartitionIsolatorExec::new(input, partition_count);
+
+        let properties = isolator.properties();
+        match properties.output_partitioning() {
+            Partitioning::UnknownPartitioning(count) => {
+                assert_eq!(*count, partition_count);
+            }
+            _ => panic!("Expected UnknownPartitioning"),
+        }
+    }
+
+    #[test]
+    fn test_partition_isolator_exec_with_new_children() {
+        let input1 = create_test_record_batch_exec();
+        let input2 = create_test_record_batch_exec();
+        let partition_count = 2;
+        
+        let isolator = Arc::new(PartitionIsolatorExec::new(input1, partition_count));
+        let new_isolator = isolator.with_new_children(vec![input2.clone()]).unwrap();
+        
+        let new_isolator_typed = new_isolator
+            .as_any()
+            .downcast_ref::<PartitionIsolatorExec>()
+            .unwrap();
+        
+        assert_eq!(new_isolator_typed.partition_count, partition_count);
+        assert!(Arc::ptr_eq(&new_isolator_typed.input, &input2));
+    }
+}

src/lib.rs

@@ -25,6 +25,7 @@ pub use proto::generated::protobuf;
 pub mod analyze;
 pub mod codec;
 pub mod customizer;
+pub mod distribution_strategy;
 pub mod explain;
 pub mod flight;
 pub mod friendly;