datafusion_expr/
physical_planning_context.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::fmt;
19use std::hash::{Hash, Hasher};
20use std::sync::{Arc, Mutex};
21
22use datafusion_common::{HashMap, Result, ScalarValue, TableReference, internal_err};
23
24/// Context used while converting a logical plan subtree into a physical plan.
25///
26/// Unlike [`ExecutionProps`](crate::execution_props::ExecutionProps), which
27/// applies to the overall planning and execution of a query, this context can
28/// differ between recursively planned subtrees. It currently carries:
29///
30/// * the state needed to create physical expressions for
31///   [`Expr::ScalarSubquery`] nodes that read from a shared
32///   [`ScalarSubqueryResults`] container, and
33/// * the qualifiers assigned to the [`Expr::LambdaVariable`]s that are in scope.
34///
35/// The physical planner builds this context from the set of uncorrelated scalar
36/// subqueries it has scheduled for a subtree. It is then passed explicitly
37/// through `create_physical_expr` so that function can find the slot index for
38/// each [`Subquery`]. While planning the body of a lambda,
39/// `create_physical_expr` extends the context with the lambda's parameters via
40/// [`Self::with_qualified_lambda_variables`].
41///
42/// An empty [`PhysicalPlanningContext`] (the [`Default`]) is what every
43/// non-physical-planner caller passes; if such a caller encounters a scalar
44/// subquery, `create_physical_expr` returns a `not_impl_err`.
45///
46/// [`Expr::ScalarSubquery`]: crate::Expr::ScalarSubquery
47/// [`Expr::LambdaVariable`]: crate::Expr::LambdaVariable
48/// [`Subquery`]: crate::logical_plan::Subquery
49#[derive(Clone, Debug, Default)]
50pub struct PhysicalPlanningContext {
51    /// Behind an `Arc` because the context is cloned for each lambda body that
52    /// is planned, and the indexes are the same for the whole subtree.
53    indexes: Arc<HashMap<crate::logical_plan::Subquery, SubqueryIndex>>,
54    results: ScalarSubqueryResults,
55    /// Maps each lambda variable name in scope to the qualifier generated for
56    /// its lambda during physical planning.
57    lambda_variable_qualifier: HashMap<String, TableReference>,
58}
59
60impl PhysicalPlanningContext {
61    /// Create a [`PhysicalPlanningContext`] from an index map and a shared
62    /// results container. The index map must use the same indices as slots in
63    /// `results`.
64    pub fn new(
65        indexes: HashMap<crate::logical_plan::Subquery, SubqueryIndex>,
66        results: ScalarSubqueryResults,
67    ) -> Self {
68        Self {
69            indexes: Arc::new(indexes),
70            results,
71            lambda_variable_qualifier: HashMap::new(),
72        }
73    }
74
75    /// Returns the slot index assigned to `subquery`, if any.
76    pub fn index_of(
77        &self,
78        subquery: &crate::logical_plan::Subquery,
79    ) -> Option<SubqueryIndex> {
80        self.indexes.get(subquery).copied()
81    }
82
83    /// Returns the shared results container.
84    pub fn results(&self) -> &ScalarSubqueryResults {
85        &self.results
86    }
87
88    /// Adds a mapping for each variable to the given qualifier. Existing
89    /// variables with conflicting names are shadowed.
90    pub fn with_qualified_lambda_variables(
91        mut self,
92        qualifier: &TableReference,
93        variables: &[String],
94    ) -> Self {
95        for var in variables {
96            self.lambda_variable_qualifier
97                .entry_ref(var)
98                .insert(qualifier.clone());
99        }
100
101        self
102    }
103
104    /// Returns the qualifier of the lambda variable `name`, if it is in scope.
105    pub fn lambda_variable_qualifier(&self, name: &str) -> Option<&TableReference> {
106        self.lambda_variable_qualifier.get(name)
107    }
108}
109
110/// Index of a scalar subquery within a [`ScalarSubqueryResults`] container.
111#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
112pub struct SubqueryIndex(usize);
113
114impl SubqueryIndex {
115    /// Creates a new subquery index.
116    pub const fn new(index: usize) -> Self {
117        Self(index)
118    }
119
120    /// Returns the underlying slot index.
121    pub const fn as_usize(self) -> usize {
122        self.0
123    }
124}
125
126/// Shared results container for uncorrelated scalar subqueries.
127///
128/// Each entry corresponds to one scalar subquery, identified by its index.
129/// Each slot is populated at execution time by `ScalarSubqueryExec`, read by
130/// `ScalarSubqueryExpr` instances that share this container, and cleared when
131/// the plan is reset for re-execution.
132#[derive(Clone, Default)]
133pub struct ScalarSubqueryResults {
134    slots: Arc<Vec<Mutex<Option<ScalarValue>>>>,
135}
136
137impl ScalarSubqueryResults {
138    /// Creates a new shared results container with `n` empty slots.
139    pub fn new(n: usize) -> Self {
140        Self {
141            slots: Arc::new((0..n).map(|_| Mutex::new(None)).collect()),
142        }
143    }
144
145    /// Returns the scalar value stored at `index`, if it has been populated.
146    pub fn get(&self, index: SubqueryIndex) -> Option<ScalarValue> {
147        let slot = self.slots.get(index.as_usize())?;
148        slot.lock().unwrap().clone()
149    }
150
151    /// Stores `value` in the slot at `index`.
152    pub fn set(&self, index: SubqueryIndex, value: ScalarValue) -> Result<()> {
153        let Some(slot) = self.slots.get(index.as_usize()) else {
154            return internal_err!(
155                "ScalarSubqueryResults: result index {} is out of bounds",
156                index.as_usize()
157            );
158        };
159
160        let mut slot = slot.lock().unwrap();
161        if slot.is_some() {
162            return internal_err!(
163                "ScalarSubqueryResults: result for index {} was already populated",
164                index.as_usize()
165            );
166        }
167        *slot = Some(value);
168
169        Ok(())
170    }
171
172    /// Clears all populated results so the container can be reused.
173    pub fn clear(&self) {
174        for slot in self.slots.iter() {
175            *slot.lock().unwrap() = None;
176        }
177    }
178
179    /// Returns true if `this` and `other` point to the same shared container.
180    pub fn ptr_eq(this: &Self, other: &Self) -> bool {
181        Arc::ptr_eq(&this.slots, &other.slots)
182    }
183}
184
185impl fmt::Debug for ScalarSubqueryResults {
186    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
187        f.debug_list()
188            .entries(self.slots.iter().map(|slot| slot.lock().unwrap().clone()))
189            .finish()
190    }
191}
192
193impl PartialEq for ScalarSubqueryResults {
194    fn eq(&self, other: &Self) -> bool {
195        Self::ptr_eq(self, other)
196    }
197}
198
199impl Eq for ScalarSubqueryResults {}
200
201impl Hash for ScalarSubqueryResults {
202    fn hash<H: Hasher>(&self, state: &mut H) {
203        Arc::as_ptr(&self.slots).hash(state);
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210
211    #[test]
212    fn scalar_subquery_results_set_and_get() -> Result<()> {
213        let results = ScalarSubqueryResults::new(1);
214        assert_eq!(results.get(SubqueryIndex::new(0)), None);
215
216        results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(42)))?;
217        assert_eq!(
218            results.get(SubqueryIndex::new(0)),
219            Some(ScalarValue::Int32(Some(42)))
220        );
221        assert!(
222            results
223                .set(SubqueryIndex::new(0), ScalarValue::Int32(Some(7)))
224                .is_err()
225        );
226
227        Ok(())
228    }
229
230    #[test]
231    fn lambda_variables_shadow_outer_scope() {
232        let outer = TableReference::bare("lambda_1");
233        let inner = TableReference::bare("lambda_2");
234
235        let ctx = PhysicalPlanningContext::default()
236            .with_qualified_lambda_variables(&outer, &["x".to_string(), "y".to_string()])
237            .with_qualified_lambda_variables(&inner, &["y".to_string()]);
238
239        assert_eq!(ctx.lambda_variable_qualifier("x"), Some(&outer));
240        assert_eq!(ctx.lambda_variable_qualifier("y"), Some(&inner));
241        assert_eq!(ctx.lambda_variable_qualifier("z"), None);
242    }
243
244    #[test]
245    fn scalar_subquery_results_clear() -> Result<()> {
246        let results = ScalarSubqueryResults::new(1);
247        results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(42)))?;
248
249        results.clear();
250
251        assert_eq!(results.get(SubqueryIndex::new(0)), None);
252        results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(7)))?;
253        assert_eq!(
254            results.get(SubqueryIndex::new(0)),
255            Some(ScalarValue::Int32(Some(7)))
256        );
257
258        Ok(())
259    }
260}
datafusion_expr/physical_planning_context.rs

datafusion_expr/
physical_planning_context.rs
