Uncomment tests in favor of just #[ignore]-ing them

Author: gabotechs

tests/common/mod.rs

@@ -1,3 +1,4 @@
 pub mod insta;
 pub mod localhost;
 pub mod parquet;
+pub mod plan;

tests/common/plan.rs

@@ -0,0 +1,68 @@
+use datafusion::common::plan_err;
+use datafusion::common::tree_node::{Transformed, TreeNode};
+use datafusion::error::DataFusionError;
+use datafusion::physical_expr::Partitioning;
+use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode};
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion_distributed::ArrowFlightReadExec;
+use std::sync::Arc;
+
+pub fn distribute_aggregate(
+    plan: Arc<dyn ExecutionPlan>,
+) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
+    let mut aggregate_partial_found = false;
+    Ok(plan
+        .transform_up(|node| {
+            let Some(agg) = node.as_any().downcast_ref::<AggregateExec>() else {
+                return Ok(Transformed::no(node));
+            };
+
+            match agg.mode() {
+                AggregateMode::Partial => {
+                    if aggregate_partial_found {
+                        return plan_err!("Two consecutive partial aggregations found");
+                    }
+                    aggregate_partial_found = true;
+                    let expr = agg
+                        .group_expr()
+                        .expr()
+                        .iter()
+                        .map(|(v, _)| Arc::clone(v))
+                        .collect::<Vec<_>>();
+
+                    if node.children().len() != 1 {
+                        return plan_err!("Aggregate must have exactly one child");
+                    }
+                    let child = node.children()[0].clone();
+
+                    let node = node.with_new_children(vec![Arc::new(ArrowFlightReadExec::new(
+                        Partitioning::Hash(expr, 1),
+                        child.schema(),
+                        0,
+                    ))])?;
+                    Ok(Transformed::yes(node))
+                }
+                AggregateMode::Final
+                | AggregateMode::FinalPartitioned
+                | AggregateMode::Single
+                | AggregateMode::SinglePartitioned => {
+                    if !aggregate_partial_found {
+                        return plan_err!("No partial aggregate found before the final one");
+                    }
+
+                    if node.children().len() != 1 {
+                        return plan_err!("Aggregate must have exactly one child");
+                    }
+                    let child = node.children()[0].clone();
+
+                    let node = node.with_new_children(vec![Arc::new(ArrowFlightReadExec::new(
+                        Partitioning::RoundRobinBatch(8),
+                        child.schema(),
+                        1,
+                    ))])?;
+                    Ok(Transformed::yes(node))
+                }
+            }
+        })?
+        .data)
+}

tests/custom_extension_codec.rs

@@ -1,6 +1,6 @@
 #[allow(dead_code)]
 mod common;
-/*
+
 #[cfg(test)]
 mod tests {
     use crate::assert_snapshot;
@@ -27,7 +27,7 @@ mod tests {
     use datafusion::physical_plan::{
         displayable, execute_stream, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
     };
-    use datafusion_distributed::{assign_stages, ArrowFlightReadExec, SessionBuilder};
+    use datafusion_distributed::{ArrowFlightReadExec, SessionBuilder};
     use datafusion_proto::physical_plan::PhysicalExtensionCodec;
     use datafusion_proto::protobuf::proto_error;
     use futures::{stream, TryStreamExt};
@@ -37,6 +37,7 @@ mod tests {
     use std::sync::Arc;
 
     #[tokio::test]
+    #[ignore]
     async fn custom_extension_codec() -> Result<(), Box<dyn std::error::Error>> {
         #[derive(Clone)]
         struct CustomSessionBuilder;
@@ -66,7 +67,6 @@ mod tests {
         ");
 
         let distributed_plan = build_plan(true)?;
-        let distributed_plan = assign_stages(distributed_plan, &ctx)?;
 
         assert_snapshot!(displayable(distributed_plan.as_ref()).indent(true).to_string(), @r"
         SortExec: expr=[numbers@0 DESC NULLS LAST], preserve_partitioning=[false]
@@ -124,8 +124,9 @@ mod tests {
 
         if distributed {
             plan = Arc::new(ArrowFlightReadExec::new(
-                plan.clone(),
                 Partitioning::Hash(vec![col("numbers", &plan.schema())?], 1),
+                plan.clone().schema(),
+                0, // TODO: stage num should be assigned by someone else
             ));
         }
 
@@ -139,8 +140,9 @@ mod tests {
 
         if distributed {
             plan = Arc::new(ArrowFlightReadExec::new(
-                plan.clone(),
                 Partitioning::RoundRobinBatch(10),
+                plan.clone().schema(),
+                1, // TODO: stage num should be assigned by someone else
             ));
 
             plan = Arc::new(RepartitionExec::try_new(
@@ -266,4 +268,4 @@ mod tests {
             .map_err(|err| proto_error(format!("{err}")))
         }
     }
-}*/
+}

tests/distributed_aggregation.rs

@@ -6,12 +6,14 @@ mod tests {
     use crate::assert_snapshot;
     use crate::common::localhost::{start_localhost_context, NoopSessionBuilder};
     use crate::common::parquet::register_parquet_tables;
+    use crate::common::plan::distribute_aggregate;
     use datafusion::arrow::util::pretty::pretty_format_batches;
     use datafusion::physical_plan::{displayable, execute_stream};
     use futures::TryStreamExt;
     use std::error::Error;
 
     #[tokio::test]
+    #[ignore]
     async fn distributed_aggregation() -> Result<(), Box<dyn Error>> {
         // FIXME these ports are in use on my machine, we should find unused ports
         // Changed them for now
@@ -26,9 +28,13 @@ mod tests {
 
         let physical_str = displayable(physical.as_ref()).indent(true).to_string();
 
-        println!("\n\nPhysical Plan:\n{}", physical_str);
+        let physical_distributed = distribute_aggregate(physical.clone())?;
 
-        /*assert_snapshot!(physical_str,
+        let physical_distributed_str = displayable(physical_distributed.as_ref())
+            .indent(true)
+            .to_string();
+
+        assert_snapshot!(physical_str,
             @r"
         ProjectionExec: expr=[count(*)@0 as count(*), RainToday@1 as RainToday]
           SortPreservingMergeExec: [count(Int64(1))@2 ASC NULLS LAST]
@@ -41,7 +47,24 @@ mod tests {
                         AggregateExec: mode=Partial, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))]
                           DataSourceExec: file_groups={1 group: [[/testdata/weather.parquet]]}, projection=[RainToday], file_type=parquet
         ",
-        );*/
+        );
+
+        assert_snapshot!(physical_distributed_str,
+            @r"
+        ProjectionExec: expr=[count(*)@0 as count(*), RainToday@1 as RainToday]
+          SortPreservingMergeExec: [count(Int64(1))@2 ASC NULLS LAST]
+            SortExec: expr=[count(Int64(1))@2 ASC NULLS LAST], preserve_partitioning=[true]
+              ProjectionExec: expr=[count(Int64(1))@1 as count(*), RainToday@0 as RainToday, count(Int64(1))@1 as count(Int64(1))]
+                AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))]
+                  ArrowFlightReadExec: input_tasks=8 hash_expr=[] stage_id=UUID input_stage_id=UUID input_hosts=[http://localhost:50050/, http://localhost:50051/, http://localhost:50052/, http://localhost:50050/, http://localhost:50051/, http://localhost:50052/, http://localhost:50050/, http://localhost:50051/]
+                    CoalesceBatchesExec: target_batch_size=8192
+                      RepartitionExec: partitioning=Hash([RainToday@0], CPUs), input_partitions=CPUs
+                        RepartitionExec: partitioning=RoundRobinBatch(CPUs), input_partitions=1
+                          AggregateExec: mode=Partial, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))]
+                            ArrowFlightReadExec: input_tasks=1 hash_expr=[RainToday@0] stage_id=UUID input_stage_id=UUID input_hosts=[http://localhost:50052/]
+                              DataSourceExec: file_groups={1 group: [[/testdata/weather.parquet]]}, projection=[RainToday], file_type=parquet
+        ",
+        );
 
         let batches = pretty_format_batches(
             &execute_stream(physical, ctx.task_ctx())?
@@ -58,6 +81,20 @@ mod tests {
         +----------+-----------+
         ");
 
+        let batches_distributed = pretty_format_batches(
+            &execute_stream(physical_distributed, ctx.task_ctx())?
+                .try_collect::<Vec<_>>()
+                .await?,
+        )?;
+        assert_snapshot!(batches_distributed, @r"
+        +----------+-----------+
+        | count(*) | RainToday |
+        +----------+-----------+
+        | 66       | Yes       |
+        | 300      | No        |
+        +----------+-----------+
+        ");
+
         Ok(())
     }
 }

tests/error_propagation.rs

@@ -1,6 +1,6 @@
 #[allow(dead_code)]
 mod common;
-/*
+
 #[cfg(test)]
 mod tests {
     use crate::common::localhost::start_localhost_context;
@@ -26,6 +26,7 @@ mod tests {
     use std::sync::Arc;
 
     #[tokio::test]
+    #[ignore]
     async fn test_error_propagation() -> Result<(), Box<dyn Error>> {
         #[derive(Clone)]
         struct CustomSessionBuilder;
@@ -48,14 +49,13 @@ mod tests {
 
         let mut plan: Arc<dyn ExecutionPlan> = Arc::new(ErrorExec::new("something failed"));
 
-        for size in [1, 2, 3] {
+        for (i, size) in [1, 2, 3].iter().enumerate() {
             plan = Arc::new(ArrowFlightReadExec::new(
-                Partitioning::RoundRobinBatch(size),
+                Partitioning::RoundRobinBatch(*size as usize),
                 plan.schema(),
-                0,
+                i,
             ));
         }
-        let plan = assign_stages(plan, &ctx)?;
         let stream = execute_stream(plan, ctx.task_ctx())?;
 
         let Err(err) = stream.try_collect::<Vec<_>>().await else {
@@ -170,4 +170,4 @@ mod tests {
             .map_err(|err| proto_error(format!("{err}")))
         }
     }
-}*/
+}

tests/highly_distributed_query.rs

@@ -1,19 +1,20 @@
 #[allow(dead_code)]
 mod common;
-/*
+
 #[cfg(test)]
 mod tests {
     use crate::assert_snapshot;
     use crate::common::localhost::{start_localhost_context, NoopSessionBuilder};
     use crate::common::parquet::register_parquet_tables;
     use datafusion::physical_expr::Partitioning;
     use datafusion::physical_plan::{displayable, execute_stream};
-    use datafusion_distributed::{assign_stages, ArrowFlightReadExec};
+    use datafusion_distributed::ArrowFlightReadExec;
     use futures::TryStreamExt;
     use std::error::Error;
     use std::sync::Arc;
 
     #[tokio::test]
+    #[ignore]
     async fn highly_distributed_query() -> Result<(), Box<dyn Error>> {
         let (ctx, _guard) = start_localhost_context(
             [
@@ -29,13 +30,13 @@ mod tests {
         let physical_str = displayable(physical.as_ref()).indent(true).to_string();
 
         let mut physical_distributed = physical.clone();
-        for size in [1, 10, 5] {
+        for (i, size) in [1, 10, 5].iter().enumerate() {
             physical_distributed = Arc::new(ArrowFlightReadExec::new(
-                physical_distributed.clone(),
-                Partitioning::RoundRobinBatch(size),
+                Partitioning::RoundRobinBatch(*size as usize),
+                physical_distributed.schema(),
+                i,
             ));
         }
-        let physical_distributed = assign_stages(physical_distributed, &ctx)?;
         let physical_distributed_str = displayable(physical_distributed.as_ref())
             .indent(true)
             .to_string();
@@ -75,4 +76,4 @@ mod tests {
 
         Ok(())
     }
-}*/
+}

tests/stage_planning.rs

@@ -1,8 +1,6 @@
 mod common;
 mod tpch;
 
-// FIXME: commented out until we figure out how to integrate best with tpch
-/*
 #[cfg(test)]
 mod tests {
     use crate::tpch::tpch_query;
@@ -17,7 +15,9 @@ mod tests {
     use std::error::Error;
     use std::sync::Arc;
 
+    // FIXME: ignored out until we figure out how to integrate best with tpch
     #[tokio::test]
+    #[ignore]
     async fn stage_planning() -> Result<(), Box<dyn Error>> {
         let config = SessionConfig::new().with_target_partitions(3);
 
@@ -86,4 +86,4 @@ mod tests {
 
         Ok(())
     }
-}*/
+}