perf: Implement physical execution of uncorrelated scalar subqueries

.github/workflows/large_files.yml

@@ -34,9 +34,9 @@ jobs:
           fetch-depth: 0
       - name: Check size of new Git objects
         env:
-          # 1 MB ought to be enough for anybody.
+          # 2 MB ought to be enough for anybody.
           # TODO in case we may want to consciously commit a bigger file to the repo without using Git LFS we may disable the check e.g. with a label
-          MAX_FILE_SIZE_BYTES: 1048576
+          MAX_FILE_SIZE_BYTES: 2097152
         shell: bash
         run: |
           if [ "${{ github.event_name }}" = "merge_group" ]; then

datafusion/core/Cargo.toml

@@ -144,6 +144,7 @@ datafusion-session = { workspace = true }
 datafusion-sql = { workspace = true, optional = true }
 flate2 = { workspace = true, optional = true }
 futures = { workspace = true }
+indexmap = { workspace = true }
 itertools = { workspace = true }
 liblzma = { workspace = true, optional = true }
 log = { workspace = true }

datafusion/expr/src/execution_props.rs

@@ -18,9 +18,13 @@
 use crate::var_provider::{VarProvider, VarType};
 use chrono::{DateTime, Utc};
 use datafusion_common::HashMap;
+use datafusion_common::ScalarValue;
 use datafusion_common::alias::AliasGenerator;
 use datafusion_common::config::ConfigOptions;
-use std::sync::Arc;
+use datafusion_common::{Result, internal_err};
+use std::fmt;
+use std::hash::{Hash, Hasher};
+use std::sync::{Arc, Mutex};
 
 /// Holds per-query execution properties and data (such as statement
 /// starting timestamps).
@@ -42,6 +46,12 @@ pub struct ExecutionProps {
     pub config_options: Option<Arc<ConfigOptions>>,
     /// Providers for scalar variables
     pub var_providers: Option<HashMap<VarType, Arc<dyn VarProvider + Send + Sync>>>,
+    /// Maps each logical `Subquery` to its index in `subquery_results`.
+    /// Populated by the physical planner before calling `create_physical_expr`.
+    pub subquery_indexes: HashMap<crate::logical_plan::Subquery, SubqueryIndex>,
+    /// Shared results container for uncorrelated scalar subquery values.
+    /// Populated at execution time by `ScalarSubqueryExec`.
+    pub subquery_results: ScalarSubqueryResults,
 }
 
 impl Default for ExecutionProps {
@@ -58,6 +68,8 @@ impl ExecutionProps {
             alias_generator: Arc::new(AliasGenerator::new()),
             config_options: None,
             var_providers: None,
+            subquery_indexes: HashMap::new(),
+            subquery_results: ScalarSubqueryResults::default(),
         }
     }
 
@@ -85,8 +97,7 @@ impl ExecutionProps {
         &*self
     }
 
-    /// Registers a variable provider, returning the existing
-    /// provider, if any
+    /// Registers a variable provider, returning the existing provider, if any
     pub fn add_var_provider(
         &mut self,
         var_type: VarType,
@@ -119,15 +130,149 @@ impl ExecutionProps {
     }
 }
 
+/// Index of a scalar subquery within a [`ScalarSubqueryResults`] container.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct SubqueryIndex(usize);
+
+impl SubqueryIndex {
+    /// Creates a new subquery index.
+    pub const fn new(index: usize) -> Self {
+        Self(index)
+    }
+
+    /// Returns the underlying slot index.
+    pub const fn as_usize(self) -> usize {
+        self.0
+    }
+}
+
+/// Shared results container for uncorrelated scalar subqueries.
+///
+/// Each entry corresponds to one scalar subquery, identified by its index.
+/// Each slot is populated at execution time by `ScalarSubqueryExec`, read by
+/// `ScalarSubqueryExpr` instances that share this container, and cleared when
+/// the plan is reset for re-execution.
+#[derive(Clone, Default)]
+pub struct ScalarSubqueryResults {
+    slots: Arc<Vec<Mutex<Option<ScalarValue>>>>,
+}
+
+impl ScalarSubqueryResults {
+    /// Creates a new shared results container with `n` empty slots.
+    pub fn new(n: usize) -> Self {
+        Self {
+            slots: Arc::new((0..n).map(|_| Mutex::new(None)).collect()),
+        }
+    }
+
+    /// Returns the scalar value stored at `index`, if it has been populated.
+    pub fn get(&self, index: SubqueryIndex) -> Option<ScalarValue> {
+        let slot = self.slots.get(index.as_usize())?;
+        slot.lock().unwrap().clone()
+    }
+
+    /// Stores `value` in the slot at `index`.
+    pub fn set(&self, index: SubqueryIndex, value: ScalarValue) -> Result<()> {
+        let Some(slot) = self.slots.get(index.as_usize()) else {
+            return internal_err!(
+                "ScalarSubqueryResults: result index {} is out of bounds",
+                index.as_usize()
+            );
+        };
+
+        let mut slot = slot.lock().unwrap();
+        if slot.is_some() {
+            return internal_err!(
+                "ScalarSubqueryResults: result for index {} was already populated",
+                index.as_usize()
+            );
+        }
+        *slot = Some(value);
+
+        Ok(())
+    }
+
+    /// Clears all populated results so the container can be reused.
+    pub fn clear(&self) {
+        for slot in self.slots.iter() {
+            *slot.lock().unwrap() = None;
+        }
+    }
+
+    /// Returns true if `this` and `other` point to the same shared container.
+    pub fn ptr_eq(this: &Self, other: &Self) -> bool {
+        Arc::ptr_eq(&this.slots, &other.slots)
+    }
+}
+
+impl fmt::Debug for ScalarSubqueryResults {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_list()
+            .entries(self.slots.iter().map(|slot| slot.lock().unwrap().clone()))
+            .finish()
+    }
+}
+
+impl PartialEq for ScalarSubqueryResults {
+    fn eq(&self, other: &Self) -> bool {
+        Self::ptr_eq(self, other)
+    }
+}
+
+impl Eq for ScalarSubqueryResults {}
+
+impl Hash for ScalarSubqueryResults {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        Arc::as_ptr(&self.slots).hash(state);
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
+
     #[test]
     fn debug() {
         let props = ExecutionProps::new();
         assert_eq!(
-            "ExecutionProps { query_execution_start_time: None, alias_generator: AliasGenerator { next_id: 1 }, config_options: None, var_providers: None }",
+            "ExecutionProps { query_execution_start_time: None, alias_generator: AliasGenerator { next_id: 1 }, config_options: None, var_providers: None, subquery_indexes: {}, subquery_results: [] }",
             format!("{props:?}")
         );
     }
+
+    #[test]
+    fn scalar_subquery_results_set_and_get() -> Result<()> {
+        let results = ScalarSubqueryResults::new(1);
+        assert_eq!(results.get(SubqueryIndex::new(0)), None);
+
+        results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(42)))?;
+        assert_eq!(
+            results.get(SubqueryIndex::new(0)),
+            Some(ScalarValue::Int32(Some(42)))
+        );
+        assert!(
+            results
+                .set(SubqueryIndex::new(0), ScalarValue::Int32(Some(7)))
+                .is_err()
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn scalar_subquery_results_clear() -> Result<()> {
+        let results = ScalarSubqueryResults::new(1);
+        results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(42)))?;
+
+        results.clear();
+
+        assert_eq!(results.get(SubqueryIndex::new(0)), None);
+        results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(7)))?;
+        assert_eq!(
+            results.get(SubqueryIndex::new(0)),
+            Some(ScalarValue::Int32(Some(7)))
+        );
+
+        Ok(())
+    }
 }

datafusion/expr/src/expr.rs

@@ -2059,6 +2059,12 @@ impl Expr {
             .expect("exists closure is infallible")
     }
 
+    /// Returns true if the expression contains a scalar subquery.
+    pub fn contains_scalar_subquery(&self) -> bool {
+        self.exists(|expr| Ok(matches!(expr, Expr::ScalarSubquery(_))))
+            .expect("exists closure is infallible")
+    }
+
     /// Returns true if the expression node is volatile, i.e. whether it can return
     /// different results when evaluated multiple times with the same input.
     /// Note: unlike [`Self::is_volatile`], this function does not consider inputs:

datafusion/expr/src/logical_plan/plan.rs

@@ -3768,6 +3768,7 @@ impl PartialOrd for Aggregate {
 /// index among identical entries. For example, if the same set appears three
 /// times, the ordinals are 0, 1, 2 and this function returns 2.
 /// Returns 0 when no grouping set is duplicated.
+#[allow(clippy::allow_attributes, clippy::mutable_key_type)] // Expr contains Arc with interior mutability but is intentionally used as hash key
 fn max_grouping_set_duplicate_ordinal(group_expr: &[Expr]) -> usize {
     if let Some(Expr::GroupingSet(GroupingSet::GroupingSets(sets))) = group_expr.first() {
         let mut counts: HashMap<&[Expr], usize> = HashMap::new();

datafusion/expr/src/logical_plan/tree_node.rs

@@ -808,7 +808,7 @@ impl LogicalPlan {
         transform_down_up_with_subqueries_impl(self, &mut f_down, &mut f_up)
     }
 
-    /// Similarly to [`Self::apply`], calls `f` on  this node and its inputs
+    /// Similarly to [`Self::apply`], calls `f` on this node and its inputs,
     /// including subqueries that may appear in expressions such as `IN (SELECT
     /// ...)`.
     pub fn apply_subqueries<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(
@@ -821,9 +821,7 @@ impl LogicalPlan {
                 | Expr::InSubquery(InSubquery { subquery, .. })
                 | Expr::SetComparison(SetComparison { subquery, .. })
                 | Expr::ScalarSubquery(subquery) => {
-                    // use a synthetic plan so the collector sees a
-                    // LogicalPlan::Subquery (even though it is
-                    // actually a Subquery alias)
+                    // Wrap in LogicalPlan::Subquery to match f's signature
                     f(&LogicalPlan::Subquery(subquery.clone()))
                 }
                 _ => Ok(TreeNodeRecursion::Continue),
@@ -888,4 +886,18 @@ impl LogicalPlan {
             })
         })
     }
+
+    /// Similar to [`Self::map_subqueries`], but only applies `f` to
+    /// uncorrelated subqueries (those with no outer reference columns).
+    pub fn map_uncorrelated_subqueries<F: FnMut(Self) -> Result<Transformed<Self>>>(
+        self,
+        mut f: F,
+    ) -> Result<Transformed<Self>> {
+        self.map_subqueries(|subquery_plan| match &subquery_plan {
+            LogicalPlan::Subquery(sq) if sq.outer_ref_columns.is_empty() => {
+                f(subquery_plan)
+            }
+            _ => Ok(Transformed::no(subquery_plan)),
+        })
+    }
 }

datafusion/optimizer/src/common_subexpr_eliminate.rs

@@ -586,8 +586,12 @@ impl OptimizerRule for CommonSubexprEliminate {
             | LogicalPlan::Unnest(_)
             | LogicalPlan::RecursiveQuery(_) => {
                 // This rule handles recursion itself in a `ApplyOrder::TopDown` like
-                // manner.
-                plan.map_children(|c| self.rewrite(c, config))?
+                // manner. Process uncorrelated subqueries in expressions
+                // (e.g., Expr::ScalarSubquery), then direct children.
+                plan.map_uncorrelated_subqueries(|c| self.rewrite(c, config))?
+                    .transform_sibling(|plan| {
+                        plan.map_children(|c| self.rewrite(c, config))
+                    })?
             }
         };
 

datafusion/optimizer/src/eliminate_cross_join.rs

@@ -212,7 +212,12 @@ fn rewrite_children(
     plan: LogicalPlan,
     config: &dyn OptimizerConfig,
 ) -> Result<Transformed<LogicalPlan>> {
-    let transformed_plan = plan.map_children(|input| optimizer.rewrite(input, config))?;
+    // Process uncorrelated subqueries in expressions, then direct children.
+    let transformed_plan = plan
+        .map_uncorrelated_subqueries(|input| optimizer.rewrite(input, config))?
+        .transform_sibling(|plan| {
+            plan.map_children(|input| optimizer.rewrite(input, config))
+        })?;
 
     // recompute schema if the plan was transformed
     if transformed_plan.transformed {

datafusion/optimizer/src/optimize_projections/mod.rs

@@ -136,9 +136,11 @@ fn optimize_projections(
     // their parents' required indices.
     match plan {
         LogicalPlan::Projection(proj) => {
-            return merge_consecutive_projections(proj)?.transform_data(|proj| {
-                rewrite_projection_given_requirements(proj, config, &indices)
-            });
+            return merge_consecutive_projections(proj)?
+                .transform_data(|proj| {
+                    rewrite_projection_given_requirements(proj, config, &indices)
+                })?
+                .transform_data(|plan| optimize_subqueries(plan, config));
         }
         LogicalPlan::Aggregate(aggregate) => {
             // Split parent requirements to GROUP BY and aggregate sections:
@@ -210,7 +212,8 @@ fn optimize_projections(
                     new_aggr_expr,
                 )
                 .map(LogicalPlan::Aggregate)
-            });
+            })?
+            .transform_data(|plan| optimize_subqueries(plan, config));
         }
         LogicalPlan::Window(window) => {
             let input_schema = Arc::clone(window.input.schema());
@@ -250,7 +253,8 @@ fn optimize_projections(
                         .map(LogicalPlan::Window)
                         .map(Transformed::yes)
                 }
-            });
+            })?
+            .transform_data(|plan| optimize_subqueries(plan, config));
         }
         LogicalPlan::TableScan(table_scan) => {
             let TableScan {
@@ -271,7 +275,8 @@ fn optimize_projections(
             let new_scan =
                 TableScan::try_new(table_name, source, Some(projection), filters, fetch)?;
 
-            return Ok(Transformed::yes(LogicalPlan::TableScan(new_scan)));
+            return Transformed::yes(LogicalPlan::TableScan(new_scan))
+                .transform_data(|plan| optimize_subqueries(plan, config));
         }
         // Other node types are handled below
         _ => {}
@@ -463,6 +468,9 @@ fn optimize_projections(
         )
     })?;
 
+    let transformed_plan =
+        transformed_plan.transform_data(|plan| optimize_subqueries(plan, config))?;
+
     // If any of the children are transformed, we need to potentially update the plan's schema
     if transformed_plan.transformed {
         transformed_plan.map_data(|plan| plan.recompute_schema())
@@ -473,6 +481,19 @@ fn optimize_projections(
 
 /// Merges consecutive projections.
 ///
+/// Optimizes uncorrelated subquery plans embedded in expressions of the given
+/// plan node (e.g., `Expr::ScalarSubquery`). `map_children` only visits direct
+/// plan inputs, so subqueries must be handled separately.
+fn optimize_subqueries(
+    plan: LogicalPlan,
+    config: &dyn OptimizerConfig,
+) -> Result<Transformed<LogicalPlan>> {
+    plan.map_uncorrelated_subqueries(|subquery_plan| {
+        let indices = RequiredIndices::new_for_all_exprs(&subquery_plan);
+        optimize_projections(subquery_plan, config, indices)
+    })
+}
+
 /// Given a projection `proj`, this function attempts to merge it with a previous
 /// projection if it exists and if merging is beneficial. Merging is considered
 /// beneficial when expressions in the current projection are non-trivial and