Pass execution props to scalar subqueries

datafusion/core/src/physical_planner.rs

@@ -143,6 +143,35 @@ pub trait PhysicalPlanner: Send + Sync {
     ) -> Result<Arc<dyn PhysicalExpr>>;
 }
 
+/// A planner adapter that lowers expressions with the properties for the plan
+/// currently being built rather than the properties stored on the session.
+struct PhysicalPlannerWithProps<'a> {
+    inner: &'a dyn PhysicalPlanner,
+    execution_props: &'a ExecutionProps,
+}
+
+#[async_trait]
+impl PhysicalPlanner for PhysicalPlannerWithProps<'_> {
+    async fn create_physical_plan(
+        &self,
+        logical_plan: &LogicalPlan,
+        session_state: &SessionState,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        self.inner
+            .create_physical_plan(logical_plan, session_state)
+            .await
+    }
+
+    fn create_physical_expr(
+        &self,
+        expr: &Expr,
+        input_dfschema: &DFSchema,
+        _session_state: &SessionState,
+    ) -> Result<Arc<dyn PhysicalExpr>> {
+        create_physical_expr(expr, input_dfschema, self.execution_props)
+    }
+}
+
 /// This trait exposes the ability to plan an [`ExecutionPlan`] out of a [`LogicalPlan`].
 #[async_trait]
 pub trait ExtensionPlanner {
@@ -469,8 +498,8 @@ impl DefaultPhysicalPlanner {
             // Create the shared `ScalarSubqueryResults` container and register
             // it in `ExecutionProps` so that `create_physical_expr` can resolve
             // `Expr::ScalarSubquery` into `ScalarSubqueryExpr` nodes. We clone
-            // the `SessionState` so these are available throughout physical
-            // planning without mutating the caller's state.
+            // the session's `ExecutionProps` so these are available throughout
+            // physical planning without mutating the caller's state.
             //
             // Ideally, the subquery state would live in a dedicated planning
             // context rather than in `ExecutionProps`. It's here because
@@ -694,13 +723,18 @@ impl DefaultPhysicalPlanner {
                     Arc::clone(res.plan())
                 } else {
                     let mut maybe_plan = None;
-                    for planner in &self.extension_planners {
+                    let planner = PhysicalPlannerWithProps {
+                        inner: self,
+                        execution_props,
+                    };
+                    for extension_planner in &self.extension_planners {
                         if maybe_plan.is_some() {
                             break;
                         }
 
-                        maybe_plan =
-                            planner.plan_table_scan(self, scan, session_state).await?;
+                        maybe_plan = extension_planner
+                            .plan_table_scan(&planner, scan, session_state)
+                            .await?;
                     }
 
                     let plan = match maybe_plan {
@@ -1819,15 +1853,19 @@ impl DefaultPhysicalPlanner {
             LogicalPlan::Extension(Extension { node }) => {
                 let mut maybe_plan = None;
                 let children = children.vec();
-                for planner in &self.extension_planners {
+                let planner = PhysicalPlannerWithProps {
+                    inner: self,
+                    execution_props,
+                };
+                for extension_planner in &self.extension_planners {
                     if maybe_plan.is_some() {
                         break;
                     }
 
                     let logical_input = node.inputs();
-                    maybe_plan = planner
+                    maybe_plan = extension_planner
                         .plan_extension(
-                            self,
+                            &planner,
                             node.as_ref(),
                             &logical_input,
                             &children,
@@ -3260,6 +3298,7 @@ mod tests {
         Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, LogicalPlanBuilder,
         Partitioning as LogicalPartitioning, RangePartitioning, Signature, TableSource,
         UserDefinedLogicalNodeCore, Volatility, WindowFunctionDefinition, col, lit,
+        scalar_subquery,
     };
     use datafusion_functions_aggregate::count::{count_all, count_udaf};
     use datafusion_functions_aggregate::expr_fn::sum;
@@ -3887,6 +3926,29 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn scalar_subquery_in_extension_expr_plans() -> Result<()> {
+        let subquery = LogicalPlanBuilder::empty(true)
+            .project(vec![lit(42_i32)])?
+            .build()?;
+        let logical_plan = LogicalPlan::Extension(Extension {
+            node: Arc::new(NoOpExtensionNode {
+                expressions: vec![scalar_subquery(Arc::new(subquery))],
+                ..Default::default()
+            }),
+        });
+        let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(
+            ExpressionExtensionPlanner,
+        )]);
+
+        let plan = planner
+            .create_initial_plan(&logical_plan, &make_session_state())
+            .await?;
+
+        assert_contains!(format!("{plan:?}"), "ScalarSubqueryExec");
+        Ok(())
+    }
+
     #[tokio::test]
     async fn error_during_extension_planning() {
         let session_state = make_session_state();
@@ -4416,6 +4478,7 @@ mod tests {
     #[derive(PartialEq, Eq, Hash)]
     struct NoOpExtensionNode {
         schema: DFSchemaRef,
+        expressions: Vec<Expr>,
     }
 
     impl Default for NoOpExtensionNode {
@@ -4428,6 +4491,7 @@ mod tests {
                     )
                     .unwrap(),
                 ),
+                expressions: vec![],
             }
         }
     }
@@ -4460,7 +4524,7 @@ mod tests {
         }
 
         fn expressions(&self) -> Vec<Expr> {
-            vec![]
+            self.expressions.clone()
         }
 
         fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
@@ -4469,10 +4533,13 @@ mod tests {
 
         fn with_exprs_and_inputs(
             &self,
-            _exprs: Vec<Expr>,
+            exprs: Vec<Expr>,
             _inputs: Vec<LogicalPlan>,
         ) -> Result<Self> {
-            unimplemented!("NoOp");
+            Ok(Self {
+                schema: Arc::clone(&self.schema),
+                expressions: exprs,
+            })
         }
 
         fn supports_limit_pushdown(&self) -> bool {
@@ -4548,6 +4615,27 @@ mod tests {
         }
     }
 
+    struct ExpressionExtensionPlanner;
+
+    #[async_trait]
+    impl ExtensionPlanner for ExpressionExtensionPlanner {
+        async fn plan_extension(
+            &self,
+            planner: &dyn PhysicalPlanner,
+            node: &dyn UserDefinedLogicalNode,
+            _logical_inputs: &[&LogicalPlan],
+            _physical_inputs: &[Arc<dyn ExecutionPlan>],
+            session_state: &SessionState,
+        ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+            for expr in node.expressions() {
+                planner.create_physical_expr(&expr, node.schema(), session_state)?;
+            }
+            Ok(Some(Arc::new(NoOpExecutionPlan::new(Arc::clone(
+                node.schema().inner(),
+            )))))
+        }
+    }
+
     //  Produces an execution plan where the schema is mismatched from
     //  the logical plan node.
     struct BadExtensionPlanner {}