address lints and add back stream_partitioner_registry

Author: robtandy

src/flight_service/mod.rs

@@ -1,6 +1,7 @@
 mod do_get;
 mod service;
 mod session_builder;
+mod stream_partitioner_registry;
 
 pub(crate) use do_get::DoGet;
 

src/flight_service/stream_partitioner_registry.rs

@@ -0,0 +1,166 @@
+use dashmap::{DashMap, Entry};
+use datafusion::error::DataFusionError;
+use datafusion::physical_plan::repartition::RepartitionExec;
+use datafusion::physical_plan::{ExecutionPlan, Partitioning};
+use std::sync::Arc;
+use uuid::Uuid;
+// TODO: find some way of cleaning up abandoned partitioners
+
+/// Keeps track of all the [StreamPartitioner] currently running in the program, identifying them
+/// by stage id.
+#[derive(Default)]
+pub struct StreamPartitionerRegistry {
+    map: DashMap<(Uuid, usize), Arc<RepartitionExec>>,
+}
+
+impl StreamPartitionerRegistry {
+    /// Builds a new [StreamPartitioner] if there was not one for this specific stage id.
+    /// If there was already one, return a reference to it.
+    pub fn get_or_create_stream_partitioner(
+        &self,
+        id: Uuid,
+        actor_idx: usize,
+        plan: Arc<dyn ExecutionPlan>,
+        partitioning: Partitioning,
+    ) -> Result<Arc<RepartitionExec>, DataFusionError> {
+        match self.map.entry((id, actor_idx)) {
+            Entry::Occupied(entry) => Ok(Arc::clone(entry.get())),
+            Entry::Vacant(entry) => Ok(Arc::clone(
+                &entry.insert(Arc::new(RepartitionExec::try_new(plan, partitioning)?)),
+            )),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::test_utils::mock_exec::MockExec;
+    use datafusion::arrow::array::{RecordBatch, UInt32Array};
+    use datafusion::arrow::datatypes::{DataType, Field, Schema};
+    use datafusion::execution::TaskContext;
+    use datafusion::physical_expr::expressions::col;
+    use futures::StreamExt;
+
+    #[tokio::test]
+    async fn round_robin_1() -> Result<(), Box<dyn std::error::Error>> {
+        const PARTITIONS: usize = 10;
+
+        let registry = StreamPartitionerRegistry::default();
+        let partitioner = registry.get_or_create_stream_partitioner(
+            Uuid::new_v4(),
+            0,
+            mock_exec(15, 10),
+            Partitioning::RoundRobinBatch(PARTITIONS),
+        )?;
+
+        let rows_per_partition = gather_rows_per_partition(&partitioner).await;
+
+        assert_eq!(
+            rows_per_partition,
+            vec![20, 20, 20, 20, 20, 10, 10, 10, 10, 10]
+        );
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn round_robin_2() -> Result<(), Box<dyn std::error::Error>> {
+        const PARTITIONS: usize = 10;
+
+        let registry = StreamPartitionerRegistry::default();
+        let partitioner = registry.get_or_create_stream_partitioner(
+            Uuid::new_v4(),
+            0,
+            mock_exec(5, 10),
+            Partitioning::RoundRobinBatch(PARTITIONS),
+        )?;
+
+        let rows_per_partition = gather_rows_per_partition(&partitioner).await;
+
+        assert_eq!(rows_per_partition, vec![10, 10, 10, 10, 10, 0, 0, 0, 0, 0]);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn hash_1() -> Result<(), Box<dyn std::error::Error>> {
+        const PARTITIONS: usize = 10;
+
+        let registry = StreamPartitionerRegistry::default();
+        let partitioner = registry.get_or_create_stream_partitioner(
+            Uuid::new_v4(),
+            0,
+            mock_exec(15, 10),
+            Partitioning::Hash(vec![col("c0", &test_schema())?], PARTITIONS),
+        )?;
+
+        let rows_per_partition = gather_rows_per_partition(&partitioner).await;
+
+        assert_eq!(
+            rows_per_partition,
+            vec![30, 15, 0, 45, 0, 15, 0, 15, 15, 15]
+        );
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn hash_2() -> Result<(), Box<dyn std::error::Error>> {
+        const PARTITIONS: usize = 10;
+
+        let registry = StreamPartitionerRegistry::default();
+        let partitioner = registry.get_or_create_stream_partitioner(
+            Uuid::new_v4(),
+            0,
+            mock_exec(5, 10),
+            Partitioning::Hash(vec![col("c0", &test_schema())?], PARTITIONS),
+        )?;
+
+        let rows_per_partition = gather_rows_per_partition(&partitioner).await;
+
+        assert_eq!(rows_per_partition, vec![10, 5, 0, 15, 0, 5, 0, 5, 5, 5]);
+        Ok(())
+    }
+
+    async fn gather_rows_per_partition(partitioner: &RepartitionExec) -> Vec<usize> {
+        let mut data = vec![];
+        let n_partitions = partitioner.partitioning().partition_count();
+        let ctx = Arc::new(TaskContext::default());
+        for i in 0..n_partitions {
+            let mut stream = partitioner.execute(i, ctx.clone()).unwrap();
+            data.push(0);
+            while let Some(msg) = stream.next().await {
+                data[i] += msg.unwrap().num_rows();
+            }
+        }
+        data
+    }
+
+    fn test_schema() -> Arc<Schema> {
+        Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]))
+    }
+
+    fn mock_exec(n_batches: usize, n_rows: usize) -> Arc<dyn ExecutionPlan> {
+        Arc::new(MockExec::new(
+            create_vec_batches(n_batches, n_rows),
+            test_schema(),
+        ))
+    }
+
+    /// Create vector batches
+    fn create_vec_batches(
+        n_batches: usize,
+        n_rows: usize,
+    ) -> Vec<Result<RecordBatch, DataFusionError>> {
+        let batch = create_batch(n_rows);
+        (0..n_batches).map(|_| Ok(batch.clone())).collect()
+    }
+
+    /// Create batch
+    fn create_batch(n_rows: usize) -> RecordBatch {
+        let schema = test_schema();
+        let mut data = vec![];
+        for i in 0..n_rows {
+            data.push(i as u32)
+        }
+        RecordBatch::try_new(schema, vec![Arc::new(UInt32Array::from(data))]).unwrap()
+    }
+}

src/stage/display.rs

@@ -26,19 +26,13 @@ use super::ExecutionStage;
 
 // Unicode box-drawing characters for creating borders and connections.
 const LTCORNER: &str = "┌"; // Left top corner
-const RTCORNER: &str = "┐"; // Right top corner
 const LDCORNER: &str = "└"; // Left bottom corner
-const RDCORNER: &str = "┘"; // Right bottom corner
-
-const TMIDDLE: &str = "┬"; // Top T-junction (connects down)
-const LMIDDLE: &str = "├"; // Left T-junction (connects right)
-const DMIDDLE: &str = "┴"; // Bottom T-junction (connects up)
-
 const VERTICAL: &str = "│"; // Vertical line
 const HORIZONTAL: &str = "─"; // Horizontal line
 
 impl DisplayAs for ExecutionStage {
     fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        #[allow(clippy::format_in_format_args)]
         match t {
             DisplayFormatType::Default => {
                 write!(f, "{}", self.name)

src/stage/mod.rs

@@ -1,7 +1,7 @@
 mod display;
+mod execution_stage;
 mod proto;
-mod stage;
 
 pub use display::display_stage_graphviz;
+pub use execution_stage::ExecutionStage;
 pub use proto::{proto_from_stage, stage_from_proto, ExecutionStageProto};
-pub use stage::ExecutionStage;

src/stage/proto.rs

@@ -31,7 +31,7 @@ pub struct ExecutionStageProto {
     pub plan: Option<Box<PhysicalPlanNode>>,
     /// The input stages to this stage
     #[prost(repeated, message, tag = "5")]
-    pub inputs: Vec<Box<ExecutionStageProto>>,
+    pub inputs: Vec<ExecutionStageProto>,
     /// Our tasks which tell us how finely grained to execute the partitions in
     /// the plan
     #[prost(message, repeated, tag = "6")]
@@ -45,7 +45,7 @@ pub fn proto_from_stage(
     let proto_plan = PhysicalPlanNode::try_from_physical_plan(stage.plan.clone(), codec)?;
     let inputs = stage
         .child_stages_iter()
-        .map(|s| Ok(Box::new(proto_from_stage(s, codec)?)))
+        .map(|s| proto_from_stage(s, codec))
         .collect::<Result<Vec<_>>>()?;
 
     Ok(ExecutionStageProto {
@@ -74,7 +74,7 @@ pub fn stage_from_proto(
         .inputs
         .into_iter()
         .map(|s| {
-            stage_from_proto(*s, registry, runtime, codec)
+            stage_from_proto(s, registry, runtime, codec)
                 .map(|s| Arc::new(s) as Arc<dyn ExecutionPlan>)
         })
         .collect::<Result<Vec<_>>>()?;