refactor: pass SubqueryContext explicitly through planner traits (apache#23649)

## Which issue does this PR close?

- Alternative to apache#22340

## Rationale for this change

When we turn logical plans into physical plans, we have a work around to
get the context for scalar subqueries. Specifically we clone the session
and mutate the `execution_props` where the subquery context currently
sits. This is even called out in the code with comments:

```rust
            // Ideally, the subquery state would live in a dedicated planning
            // context rather than in `ExecutionProps`. It's here because
            // `create_physical_expr` only receives `&ExecutionProps`.
```

The real reason for this is that we want to expose the `QueryPlanner`
via FFI and currently it depends upon downcasts from `Session` to
`SessionState`. That does not work for the FFI crate where we have a
different `Session` implementation. Since we cannot do the downcast, we
cannot use the work around in the current code.

This PR is designed to plumb through the required context properly
rather than rely on manipulating the execution properties.

## What changes are included in this PR?

The major change here is to remove `subquery_indexes` and
`subquery_results` from `ExecutionProps` and put them into their own
struct `SubqueryContext`. Everything else is just refactoring and
plumbing to align around this split of the data.

## Are these changes tested?

- Existing unit tests all pass.
- Added additional unit test to cover user defined expressions.

## Are there any user-facing changes?

Yes, this adds a single parameter into `create_physical_expr`, the
subquery context. The migration guide demonstrates how to add a default
context in for users, if necessary.

Author: timsaucer

datafusion-examples/examples/dataframe/cache_factory.rs

@@ -29,6 +29,7 @@ use datafusion::error::Result;
 use datafusion::execution::context::QueryPlanner;
 use datafusion::execution::session_state::CacheFactory;
 use datafusion::execution::{SessionState, SessionStateBuilder};
+use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion::logical_expr::{
     Extension, LogicalPlan, UserDefinedLogicalNode, UserDefinedLogicalNodeCore,
 };
@@ -146,6 +147,7 @@ impl ExtensionPlanner for CacheNodePlanner {
         logical_inputs: &[&LogicalPlan],
         physical_inputs: &[Arc<dyn ExecutionPlan>],
         session_state: &SessionState,
+        _planning_ctx: &PhysicalPlanningContext,
     ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
         if let Some(cache_node) = node.as_any().downcast_ref::<CacheNode>() {
             assert_eq!(logical_inputs.len(), 1, "Inconsistent number of inputs");

datafusion-examples/examples/query_planning/expr_api.rs

@@ -33,6 +33,7 @@ use datafusion::functions_aggregate::first_last::first_value_udaf;
 use datafusion::logical_expr::execution_props::ExecutionProps;
 use datafusion::logical_expr::expr::BinaryExpr;
 use datafusion::logical_expr::interval_arithmetic::Interval;
+use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion::logical_expr::simplify::SimplifyContext;
 use datafusion::logical_expr::{ColumnarValue, ExprFunctionExt, ExprSchemable, Operator};
 use datafusion::optimizer::analyzer::type_coercion::TypeCoercionRewriter;
@@ -541,8 +542,12 @@ fn type_coercion_demo() -> Result<()> {
 
     // Evaluation with an expression that has not been type coerced cannot succeed.
     let props = ExecutionProps::default();
-    let physical_expr =
-        datafusion::physical_expr::create_physical_expr(&expr, &df_schema, &props)?;
+    let physical_expr = datafusion::physical_expr::create_physical_expr(
+        &expr,
+        &df_schema,
+        &props,
+        &PhysicalPlanningContext::default(),
+    )?;
     let e = physical_expr.evaluate(&batch).unwrap_err();
     assert!(
         e.find_root()
@@ -566,6 +571,7 @@ fn type_coercion_demo() -> Result<()> {
         &coerced_expr,
         &df_schema,
         &props,
+        &PhysicalPlanningContext::default(),
     )?;
     assert!(physical_expr.evaluate(&batch).is_ok());
 
@@ -578,6 +584,7 @@ fn type_coercion_demo() -> Result<()> {
         &coerced_expr,
         &df_schema,
         &props,
+        &PhysicalPlanningContext::default(),
     )?;
     assert!(physical_expr.evaluate(&batch).is_ok());
 
@@ -606,6 +613,7 @@ fn type_coercion_demo() -> Result<()> {
         &coerced_expr,
         &df_schema,
         &props,
+        &PhysicalPlanningContext::default(),
     )?;
     assert!(physical_expr.evaluate(&batch).is_ok());
 

datafusion-examples/examples/query_planning/pruning.rs

@@ -26,6 +26,7 @@ use datafusion::common::pruning::PruningStatistics;
 use datafusion::common::{DFSchema, ScalarValue};
 use datafusion::error::Result;
 use datafusion::execution::context::ExecutionProps;
+use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion::physical_expr::create_physical_expr;
 use datafusion::physical_optimizer::pruning::PruningPredicate;
 use datafusion::prelude::*;
@@ -194,7 +195,13 @@ impl PruningStatistics for MyCatalog {
 fn create_pruning_predicate(expr: Expr, schema: &SchemaRef) -> PruningPredicate {
     let df_schema = DFSchema::try_from(Arc::clone(schema)).unwrap();
     let props = ExecutionProps::new();
-    let physical_expr = create_physical_expr(&expr, &df_schema, &props).unwrap();
+    let physical_expr = create_physical_expr(
+        &expr,
+        &df_schema,
+        &props,
+        &PhysicalPlanningContext::default(),
+    )
+    .unwrap();
     PruningPredicate::try_new(physical_expr, Arc::clone(schema)).unwrap()
 }
 

datafusion-examples/examples/relation_planner/table_sample.rs

@@ -119,6 +119,7 @@ use datafusion_common::{
     DFSchemaRef, DataFusionError, Result, Statistics, internal_err, not_impl_err,
     plan_datafusion_err, plan_err,
 };
+use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion_expr::{
     UserDefinedLogicalNode, UserDefinedLogicalNodeCore,
     logical_plan::{Extension, LogicalPlan, LogicalPlanBuilder},
@@ -587,6 +588,7 @@ impl ExtensionPlanner for TableSampleExtensionPlanner {
         _logical_inputs: &[&LogicalPlan],
         physical_inputs: &[Arc<dyn ExecutionPlan>],
         _session_state: &SessionState,
+        _planning_ctx: &PhysicalPlanningContext,
     ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
         let Some(sample_node) = node.as_any().downcast_ref::<TableSamplePlanNode>()
         else {

datafusion/catalog-listing/src/helpers.rs

@@ -34,6 +34,7 @@ use arrow::{
     record_batch::RecordBatch,
 };
 use datafusion_expr::execution_props::ExecutionProps;
+use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
 use futures::stream::FuturesUnordered;
 use futures::{StreamExt, TryStreamExt, stream::BoxStream};
 use log::{debug, trace};
@@ -328,7 +329,12 @@ pub fn filter_partitioned_file(
 
     let filter = utils::conjunction(filters.iter().cloned()).unwrap_or_else(|| lit(true));
     let props = ExecutionProps::new();
-    let expr = create_physical_expr(&filter, df_schema, &props)?;
+    let expr = create_physical_expr(
+        &filter,
+        df_schema,
+        &props,
+        &PhysicalPlanningContext::default(),
+    )?;
 
     // Since we're only operating on a single file, our batch and resulting "array" holds only one
     // value indicating if the input file matches the provided filters

datafusion/catalog-listing/src/table.rs

@@ -44,6 +44,7 @@ use datafusion_execution::cache::cache_manager::{
 };
 use datafusion_expr::dml::InsertOp;
 use datafusion_expr::execution_props::ExecutionProps;
+use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion_expr::{
     Expr, Partitioning as LogicalPartitioning, TableProviderFilterPushDown, TableType,
 };
@@ -614,6 +615,7 @@ impl TableProvider for ListingTable {
                         output_partitioning,
                         &df_schema,
                         state.execution_props(),
+                        &PhysicalPlanningContext::default(),
                     )?
                 }
             };

datafusion/catalog/src/memory/table.rs

@@ -36,6 +36,7 @@ use datafusion_datasource::memory::{MemSink, MemorySourceConfig};
 use datafusion_datasource::sink::DataSinkExec;
 use datafusion_datasource::source::DataSourceExec;
 use datafusion_expr::dml::InsertOp;
+use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion_expr::{Expr, SortExpr, TableType};
 use datafusion_physical_expr::{
     LexOrdering, PhysicalExpr, create_physical_expr, create_physical_sort_exprs,
@@ -209,8 +210,12 @@ impl TableProvider for MemTable {
             let eqp = state.execution_props();
             let mut file_sort_order = vec![];
             for sort_exprs in sort_order.iter() {
-                let physical_exprs =
-                    create_physical_sort_exprs(sort_exprs, &df_schema, eqp)?;
+                let physical_exprs = create_physical_sort_exprs(
+                    sort_exprs,
+                    &df_schema,
+                    eqp,
+                    &PhysicalPlanningContext::default(),
+                )?;
                 file_sort_order.extend(LexOrdering::new(physical_exprs));
             }
             source = source.try_with_sort_information(file_sort_order)?;
@@ -356,8 +361,12 @@ impl TableProvider for MemTable {
         let physical_assignments: HashMap<String, Arc<dyn PhysicalExpr>> = assignments
             .iter()
             .map(|(name, expr)| {
-                let physical_expr =
-                    create_physical_expr(expr, &df_schema, state.execution_props())?;
+                let physical_expr = create_physical_expr(
+                    expr,
+                    &df_schema,
+                    state.execution_props(),
+                    &PhysicalPlanningContext::default(),
+                )?;
                 Ok((name.clone(), physical_expr))
             })
             .collect::<Result<_>>()?;
@@ -470,8 +479,12 @@ fn evaluate_filters_to_mask(
     let mut combined_mask: Option<BooleanArray> = None;
 
     for filter_expr in filters {
-        let physical_expr =
-            create_physical_expr(filter_expr, df_schema, execution_props)?;
+        let physical_expr = create_physical_expr(
+            filter_expr,
+            df_schema,
+            execution_props,
+            &PhysicalPlanningContext::default(),
+        )?;
 
         let result = physical_expr.evaluate(batch)?;
         let array = result.into_array(batch.num_rows())?;

datafusion/catalog/src/streaming.rs

@@ -22,6 +22,7 @@ use std::sync::Arc;
 use arrow::datatypes::SchemaRef;
 use async_trait::async_trait;
 use datafusion_common::{DFSchema, Result, plan_err};
+use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion_expr::{Expr, SortExpr, TableType};
 use datafusion_physical_expr::equivalence::project_ordering;
 use datafusion_physical_expr::projection::ProjectionMapping;
@@ -131,8 +132,12 @@ impl TableProvider for StreamingTable {
             let df_schema = DFSchema::try_from(Arc::clone(&self.schema))?;
             let eqp = state.execution_props();
 
-            let original_sort_exprs =
-                create_physical_sort_exprs(&self.sort_order, &df_schema, eqp)?;
+            let original_sort_exprs = create_physical_sort_exprs(
+                &self.sort_order,
+                &df_schema,
+                eqp,
+                &PhysicalPlanningContext::default(),
+            )?;
 
             if let Some(p) = projection {
                 // When performing a projection, the output columns will not match

datafusion/core/src/execution/context/mod.rs

@@ -2383,6 +2383,7 @@ mod tests {
     use arrow_schema::FieldRef;
     use datafusion_common::DataFusionError;
     use datafusion_common::datatype::DataTypeExt;
+    use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
     use std::error::Error;
     use std::path::PathBuf;
 
@@ -2851,6 +2852,7 @@ mod tests {
             _expr: &Expr,
             _input_dfschema: &DFSchema,
             _session_state: &SessionState,
+            _planning_ctx: &PhysicalPlanningContext,
         ) -> Result<Arc<dyn PhysicalExpr>> {
             unimplemented!()
         }

datafusion/core/src/execution/session_state.rs

@@ -55,6 +55,7 @@ use datafusion_execution::runtime_env::RuntimeEnv;
 use datafusion_expr::TableSource;
 use datafusion_expr::execution_props::ExecutionProps;
 use datafusion_expr::expr_rewriter::FunctionRewrite;
+use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion_expr::planner::ExprPlanner;
 #[cfg(feature = "sql")]
 use datafusion_expr::planner::{RelationPlanner, TypePlanner};
@@ -799,7 +800,12 @@ impl SessionState {
                 .transform_up(|expr| rewrite.rewrite(expr, df_schema, config_options))?
                 .data;
         }
-        create_physical_expr(&expr, df_schema, self.execution_props())
+        create_physical_expr(
+            &expr,
+            df_schema,
+            self.execution_props(),
+            &PhysicalPlanningContext::default(),
+        )
     }
 
     /// Return the session ID