Refactor partition groups and add tests

This change adds a new abstraction called a `Grouper` in `distribution_strategy.rs`. This is responsible for grouping 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 grouper 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,96 @@
+// A Grouper is used to group partitions.
+pub trait Grouper {
+    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 PartitionGrouper {
+    partition_group_size: usize,
+}
+
+impl PartitionGrouper {
+    pub fn new(partition_group_size: usize) -> Self {
+        assert!(
+            partition_group_size > 0,
+            "partition groups cannot be size 0"
+        );
+        PartitionGrouper {
+            partition_group_size,
+        }
+    }
+}
+
+impl Grouper for PartitionGrouper {
+    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_grouper_basic() {
+        let grouper = PartitionGrouper::new(4);
+        let groups = grouper.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_grouper_uneven() {
+        let grouper = PartitionGrouper::new(2);
+        let groups = grouper.split(5);
+
+        let expected = vec![
+            PartitionGroup { start: 0, end: 2 },
+            PartitionGroup { start: 2, end: 4 },
+            PartitionGroup { start: 4, end: 5 },
+        ];
+
+        assert_eq!(groups, expected);
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_invalid_group_size() {
+        PartitionGrouper::new(0);
+    }
+}

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 crate::{record_batch_exec::RecordBatchExec, vocab::CtxPartitionGroup};
+    use arrow::array::{Int32Array, RecordBatch};
+    use datafusion::{
+        arrow::datatypes::{DataType, Field, Schema},
+        prelude::SessionContext,
+    };
+    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() {
+        let input = create_test_record_batch_exec();
+        let partition_count = 3;
+        let isolator = PartitionIsolatorExec::new(input, partition_count);
+
+        // Test success case: valid partition with partition group
+        let ctx = SessionContext::new();
+        let partition_group = vec![0u64, 1u64, 2u64];
+        {
+            let state = ctx.state_ref();
+            let mut guard = state.write();
+            let config = guard.config_mut();
+            config.set_extension(Arc::new(CtxPartitionGroup(partition_group)));
+        }
+
+        let task_context = ctx.task_ctx();
+
+        // Success case: execute valid partition
+        let result = isolator.execute(0, task_context.clone());
+        assert!(result.is_ok());
+
+        // Error case: try to execute partition beyond partition_count
+        let result = isolator.execute(4, task_context.clone());
+        assert!(result.is_err());
+        if let Err(e) = result {
+            assert!(e
+                .to_string()
+                .contains("Invalid partition 4 for PartitionIsolatorExec"));
+        }
+
+        // Error case: test without partition group extension
+        let empty_ctx = SessionContext::new();
+        let empty_task_context = empty_ctx.task_ctx();
+
+        let result = isolator.execute(0, empty_task_context);
+        assert!(result.is_err());
+        if let Err(e) = result {
+            assert!(e
+                .to_string()
+                .contains("PartitionGroup not set in session config"));
+        }
+    }
+}

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;