Skip distribution for plans with ScalarSubqueryExec

Physical scalar subqueries (apache/datafusion#21240) cannot be
distributed yet because ScalarSubqueryExpr serialization requires
context from the surrounding ScalarSubqueryExec.

Author: marc-pydantic

src/distributed_planner/distributed_physical_optimizer_rule.rs

@@ -11,6 +11,7 @@ use datafusion::config::ConfigOptions;
 use datafusion::error::DataFusionError;
 use datafusion::physical_optimizer::PhysicalOptimizerRule;
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion::physical_plan::scalar_subquery::ScalarSubqueryExec;
 use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 use std::fmt::Debug;
 use std::ops::AddAssign;
@@ -39,6 +40,16 @@ use super::insert_broadcast::insert_broadcast_execs;
 #[derive(Debug, Default)]
 pub struct DistributedPhysicalOptimizerRule;
 
+/// Checks if the plan tree contains a [ScalarSubqueryExec] node.
+fn contains_scalar_subquery_exec(plan: &Arc<dyn ExecutionPlan>) -> bool {
+    if plan.downcast_ref::<ScalarSubqueryExec>().is_some() {
+        return true;
+    }
+    plan.children()
+        .iter()
+        .any(|c| contains_scalar_subquery_exec(c))
+}
+
 impl PhysicalOptimizerRule for DistributedPhysicalOptimizerRule {
     fn optimize(
         &self,
@@ -49,6 +60,11 @@ impl PhysicalOptimizerRule for DistributedPhysicalOptimizerRule {
             return Ok(original);
         }
 
+        // Skip distribution for plans with scalar subqueries; serialization not yet supported (https://github.com/apache/datafusion/pull/21240)
+        if contains_scalar_subquery_exec(&original) {
+            return Ok(original);
+        }
+
         let mut plan = Arc::clone(&original);
         if original.output_partitioning().partition_count() > 1 {
             plan = Arc::new(CoalescePartitionsExec::new(plan))

src/metrics/task_metrics_collector.rs

@@ -310,6 +310,8 @@ mod tests {
         .await;
     }
 
+    /// Skipped: scalar subquery distribution not yet supported (https://github.com/apache/datafusion/pull/21240)
+    #[ignore]
     #[tokio::test]
     async fn test_metrics_collection_e2e_3() {
         run_metrics_collection_e2e_test(