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

@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use arrow::datatypes::Schema;
+use arrow::datatypes::{Field, Schema};
 use datafusion::{
     common::{internal_datafusion_err, internal_err, Result},
     execution::FunctionRegistry,
@@ -243,10 +243,10 @@ mod test {
         isolator::PartitionIsolatorExec, max_rows::MaxRowsExec, stage_reader::DDStageReaderExec,
     };
 
-    fn create_test_schema() -> Arc<arrow::datatypes::Schema> {
-        Arc::new(arrow::datatypes::Schema::new(vec![
-            arrow::datatypes::Field::new("a", DataType::Int32, false),
-            arrow::datatypes::Field::new("b", DataType::Int32, false),
+    fn create_test_schema() -> Arc<Schema> {
+        Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int32, false),
+            Field::new("b", DataType::Int32, false),
         ]))
     }
 

src/distribution_strategy.rs

@@ -0,0 +1,133 @@
+// A Grouper is used to group partitions.
+pub trait Grouper {
+    // group groups the number of partitions into a vec of groups.
+    fn group(&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 {
+    // new creates a new PartitionGroup containing partitions in the range [start..end).
+    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
+    }
+}
+
+// PartitionGrouper groups a number partitions together depending on a partition_group_size.
+// Ex. 10 partitions with a group size of 3 will yield groups [(0..3), (3..6), (6..9), (9)].
+// - A partition_group_size of 0 will panic
+// - Grouping 0 partitions will result an empty vec
+// - It's possible for bad groupings to exist. Ex. if the group size is 99 and there are 100
+//   partitions, then you will get unbalanced partitions [(0..99), (99)]
+#[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 {
+    // group implements the Grouper trait
+    fn group(&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.group(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.group(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);
+    }
+
+    #[test]
+    fn test_num_partitions_smaller_than_group_size() {
+        let g =  PartitionGrouper::new(2);
+        let groups = g.group(1);
+        let expected = vec![
+            PartitionGroup { start: 0, end: 1 },
+        ];
+        assert_eq!(groups, expected);
+    }
+
+    #[test]
+    fn test_num_partitions_equal_to_group_size() {
+        let g =  PartitionGrouper::new(2);
+        let groups = g.group(2);
+        let expected = vec![
+            PartitionGroup { start: 0, end: 2 },
+        ];
+        assert_eq!(groups, expected);
+    }
+
+    #[test]
+    fn test_zero_partitions_to_group() {
+        let g =  PartitionGrouper::new(2);
+        let groups = g.group(0);
+        let expected = vec![];
+        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,75 @@ 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());
+        assert!(result.err().unwrap().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());
+        assert!(result.err().unwrap().to_string()
+            .contains("PartitionGroup not set in session config"));
+
+        // Empty task ctx
+        let result = isolator.execute(1, SessionContext::new().task_ctx());
+        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;