datafusion_physical_expr/expressions/
lambda.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
18//! Physical lambda expression: [`LambdaExpr`]
19
20use std::hash::Hash;
21use std::sync::Arc;
22
23use crate::{
24    ScalarFunctionExpr,
25    expressions::{Column, LambdaVariable},
26    physical_expr::PhysicalExpr,
27};
28use arrow::{
29    datatypes::{DataType, Schema},
30    record_batch::RecordBatch,
31};
32use datafusion_common::{
33    HashMap, plan_err,
34    tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeVisitor},
35};
36use datafusion_common::{HashSet, Result, internal_err};
37use datafusion_expr::ColumnarValue;
38
39/// Represents a lambda with the given parameters names and body
40#[derive(Debug, Eq, Clone)]
41pub struct LambdaExpr {
42    params: Vec<String>,
43    body: Arc<dyn PhysicalExpr>,
44    projected_body: Arc<dyn PhysicalExpr>,
45    projection: Vec<usize>,
46    used_param_indices: Vec<usize>,
47}
48
49// Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 [https://github.com/apache/datafusion/issues/13196]
50impl PartialEq for LambdaExpr {
51    fn eq(&self, other: &Self) -> bool {
52        self.params.eq(&other.params) && self.body.eq(&other.body)
53    }
54}
55
56impl Hash for LambdaExpr {
57    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
58        self.params.hash(state);
59        self.body.hash(state);
60    }
61}
62
63impl LambdaExpr {
64    /// Create a new lambda expression with the given parameters and body.
65    pub fn try_new(params: Vec<String>, body: Arc<dyn PhysicalExpr>) -> Result<Self> {
66        if !all_unique(&params) {
67            return plan_err!(
68                "lambda params must be unique, got ({})",
69                params.join(", ")
70            );
71        }
72
73        check_async_udf(&body)?;
74
75        Ok(Self::new(params, body))
76    }
77
78    fn new(params: Vec<String>, body: Arc<dyn PhysicalExpr>) -> Self {
79        let own_params: HashSet<String> = params.iter().cloned().collect();
80
81        let mut visitor = CollectUsedVisitor {
82            own_params: &own_params,
83            used_indices: HashSet::new(),
84            used_param_names: HashSet::new(),
85            shadow_stack: Vec::new(),
86        };
87        body.visit(&mut visitor).expect("visitor is infallible");
88        let CollectUsedVisitor {
89            used_indices,
90            used_param_names,
91            ..
92        } = visitor;
93
94        let mut projection = used_indices.into_iter().collect::<Vec<_>>();
95
96        projection.sort();
97
98        let column_index_map = projection
99            .iter()
100            .copied()
101            .enumerate()
102            .map(|(new_idx, original)| (original, new_idx))
103            .collect::<HashMap<_, _>>();
104
105        let projected_body = Arc::clone(&body)
106            .transform_down(|e| {
107                if let Some(column) = e.downcast_ref::<Column>() {
108                    let original = column.index();
109                    let projected = *column_index_map.get(&original).unwrap();
110                    if projected != original {
111                        return Ok(Transformed::yes(Arc::new(Column::new(
112                            column.name(),
113                            projected,
114                        ))));
115                    }
116                } else if let Some(lambda_variable) = e.downcast_ref::<LambdaVariable>() {
117                    let original = lambda_variable.index();
118                    let projected = *column_index_map.get(&original).unwrap();
119                    if projected != original {
120                        return Ok(Transformed::yes(Arc::new(LambdaVariable::new(
121                            projected,
122                            Arc::clone(lambda_variable.field()),
123                        ))));
124                    }
125                }
126                Ok(Transformed::no(e))
127            })
128            .expect("closure should be infallible")
129            .data;
130
131        let used_param_indices = params
132            .iter()
133            .enumerate()
134            .filter(|(_, name)| used_param_names.contains(*name))
135            .map(|(i, _)| i)
136            .collect();
137
138        Self {
139            params,
140            body,
141            projected_body,
142            projection,
143            used_param_indices,
144        }
145    }
146
147    /// Get the lambda's params names
148    pub fn params(&self) -> &[String] {
149        &self.params
150    }
151
152    /// Get the lambda's body
153    pub fn body(&self) -> &Arc<dyn PhysicalExpr> {
154        &self.body
155    }
156
157    #[cfg(feature = "proto")]
158    /// Reconstruct a [`LambdaExpr`] from a proto node.
159    pub fn try_from_proto(
160        node: &datafusion_proto_models::protobuf::PhysicalExprNode,
161        ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
162    ) -> Result<Arc<dyn PhysicalExpr>> {
163        use datafusion_physical_expr_common::expect_expr_variant;
164        use datafusion_proto_models::protobuf;
165
166        let lambda = expect_expr_variant!(
167            node,
168            protobuf::physical_expr_node::ExprType::Lambda,
169            "LambdaExpr",
170        );
171
172        Ok(Arc::new(LambdaExpr::try_new(
173            lambda.params.clone(),
174            ctx.decode_required_expression(lambda.body.as_deref(), "LambdaExpr", "body")?,
175        )?))
176    }
177
178    pub(crate) fn projection(&self) -> &[usize] {
179        &self.projection
180    }
181
182    pub(crate) fn projected_body(&self) -> &Arc<dyn PhysicalExpr> {
183        &self.projected_body
184    }
185
186    /// Indices into [`params`](Self::params) of the parameters the body
187    /// actually references, in declaration order. See `CollectUsedVisitor`
188    /// in this module.
189    ///
190    /// Relies on the planner appending each lambda's own params after
191    /// captures, matching the `captures ++ used_params` layout
192    /// `LambdaArgument::new` builds.
193    pub fn used_param_indices(&self) -> &[usize] {
194        &self.used_param_indices
195    }
196}
197
198/// Walks the body of a [`LambdaExpr`] and collects, on a single pass:
199///
200/// * `used_indices` — every `Column` / `LambdaVariable` index referenced
201///   anywhere in the tree (including inside nested lambdas). This drives
202///   the `projection` used to slice the outer batch.
203/// * `used_param_names` — the subset of *this* lambda's `own_params` that
204///   the body actually references.
205///
206///   A nested lambda can declare its own parameter with the same name as
207///   one of `own_params` — a distinct variable that happens to reuse the
208///   name (variable shadowing). E.g. in
209///   `(k, v) -> func(col, (k, v2) -> k + v2 + v)`, the inner `k` is not
210///   `own_params`' `k`; only `v` should flow up as used, not `k`.
211///
212///   `shadow_stack` holds one frame per nested `LambdaExpr` currently being
213///   visited, each frame being that lambda's own parameter names. A
214///   `LambdaVariable` only counts toward `used_param_names` if its name
215///   isn't in any active frame (i.e. not shadowed).
216///
217/// The stack is maintained via `TreeNodeVisitor`'s `f_down` / `f_up`:
218/// push a frame when entering a nested [`LambdaExpr`], pop it when leaving.
219struct CollectUsedVisitor<'a> {
220    own_params: &'a HashSet<String>,
221    used_indices: HashSet<usize>,
222    used_param_names: HashSet<String>,
223    shadow_stack: Vec<HashSet<String>>,
224}
225
226impl TreeNodeVisitor<'_> for CollectUsedVisitor<'_> {
227    type Node = Arc<dyn PhysicalExpr>;
228
229    fn f_down(&mut self, node: &Self::Node) -> Result<TreeNodeRecursion> {
230        if let Some(col) = node.downcast_ref::<Column>() {
231            self.used_indices.insert(col.index());
232        } else if let Some(var) = node.downcast_ref::<LambdaVariable>() {
233            self.used_indices.insert(var.index());
234
235            let name = var.name();
236            let shadowed = self.shadow_stack.iter().any(|frame| frame.contains(name));
237            if !shadowed && self.own_params.contains(name) {
238                self.used_param_names.insert(name.to_string());
239            }
240        } else if let Some(nested) = node.downcast_ref::<LambdaExpr>() {
241            self.shadow_stack
242                .push(nested.params.iter().cloned().collect());
243        }
244
245        Ok(TreeNodeRecursion::Continue)
246    }
247
248    fn f_up(&mut self, node: &Self::Node) -> Result<TreeNodeRecursion> {
249        if node.downcast_ref::<LambdaExpr>().is_some() {
250            self.shadow_stack.pop();
251        }
252        Ok(TreeNodeRecursion::Continue)
253    }
254}
255
256impl std::fmt::Display for LambdaExpr {
257    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
258        write!(f, "({}) -> {}", self.params.join(", "), self.body)
259    }
260}
261
262impl PhysicalExpr for LambdaExpr {
263    fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
264        Ok(DataType::Null)
265    }
266
267    fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
268        Ok(true)
269    }
270
271    fn evaluate(&self, _batch: &RecordBatch) -> Result<ColumnarValue> {
272        internal_err!("LambdaExpr::evaluate() should not be called")
273    }
274
275    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
276        vec![&self.body]
277    }
278
279    fn with_new_children(
280        self: Arc<Self>,
281        children: Vec<Arc<dyn PhysicalExpr>>,
282    ) -> Result<Arc<dyn PhysicalExpr>> {
283        let [body] = children.as_slice() else {
284            return internal_err!(
285                "LambdaExpr expects exactly 1 child, got {}",
286                children.len()
287            );
288        };
289
290        check_async_udf(body)?;
291
292        Ok(Arc::new(Self::new(self.params.clone(), Arc::clone(body))))
293    }
294
295    fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
296        write!(f, "({}) -> {}", self.params.join(", "), self.body)
297    }
298
299    #[cfg(feature = "proto")]
300    fn try_to_proto(
301        &self,
302        ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
303    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
304        use datafusion_proto_models::protobuf;
305
306        Ok(Some(protobuf::PhysicalExprNode {
307            expr_id: None,
308            expr_type: Some(protobuf::physical_expr_node::ExprType::Lambda(Box::new(
309                protobuf::PhysicalLambdaExprNode {
310                    params: self.params().to_vec(),
311                    body: Some(Box::new(ctx.encode_child(self.body())?)),
312                },
313            ))),
314        }))
315    }
316}
317
318/// Create a lambda expression.
319pub fn lambda(
320    params: impl IntoIterator<Item = impl Into<String>>,
321    body: Arc<dyn PhysicalExpr>,
322) -> Result<Arc<dyn PhysicalExpr>> {
323    Ok(Arc::new(LambdaExpr::try_new(
324        params.into_iter().map(Into::into).collect(),
325        body,
326    )?))
327}
328
329fn all_unique(params: &[String]) -> bool {
330    match params.len() {
331        0 | 1 => true,
332        2 => params[0] != params[1],
333        _ => {
334            let mut set = HashSet::with_capacity(params.len());
335
336            params.iter().all(|p| set.insert(p.as_str()))
337        }
338    }
339}
340
341fn check_async_udf(body: &Arc<dyn PhysicalExpr>) -> Result<()> {
342    if body.exists(|expr| {
343        Ok(expr
344            .downcast_ref::<ScalarFunctionExpr>()
345            .is_some_and(|udf| udf.fun().as_async().is_some()))
346    })? {
347        return plan_err!(
348            "Async functions in lambdas aren't supported, see https://github.com/apache/datafusion/issues/22091"
349        );
350    }
351
352    Ok(())
353}
354
355#[cfg(test)]
356mod tests {
357    use crate::expressions::{Column, LambdaVariable, NoOp, lambda::lambda};
358    use arrow::{
359        array::RecordBatch,
360        datatypes::{DataType, Field, Schema},
361    };
362    use std::sync::Arc;
363
364    use super::LambdaExpr;
365
366    #[test]
367    fn test_lambda_evaluate() {
368        let lambda = lambda(["a"], Arc::new(NoOp::new())).unwrap();
369        let batch = RecordBatch::new_empty(Arc::new(Schema::empty()));
370        assert!(lambda.evaluate(&batch).is_err());
371    }
372
373    #[test]
374    fn test_lambda_duplicate_name() {
375        assert!(lambda(["a", "a"], Arc::new(NoOp::new())).is_err());
376    }
377
378    /// A two-parameter lambda whose body only references the second
379    /// parameter (`v`) must report only `v` as used. The higher-order
380    /// function uses this set to push only `v` into the merged batch, so
381    /// the body's compressed `LambdaVariable` index for `v` lines up with
382    /// the batch layout.
383    #[test]
384    fn test_used_params_collects_only_referenced_param() {
385        let v_field = Arc::new(Field::new("v", DataType::Int32, true));
386        let body = Arc::new(LambdaVariable::new(1, Arc::clone(&v_field)));
387
388        let lambda =
389            LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], body).unwrap();
390
391        assert_eq!(lambda.projection(), &[1]);
392        assert_eq!(lambda.used_param_indices(), &[1]);
393    }
394
395    /// A body that references neither declared parameter reports no used params.
396    #[test]
397    fn test_used_params_all_unused() {
398        let body = Arc::new(NoOp::new());
399
400        let lambda =
401            LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], body).unwrap();
402
403        assert!(lambda.projection().is_empty());
404        assert!(lambda.used_param_indices().is_empty());
405    }
406
407    /// A three-parameter lambda that skips the middle parameter reports only the ends as used.
408    #[test]
409    fn test_used_params_three_params_middle_unused() {
410        let a_field = Arc::new(Field::new("a", DataType::Int32, true));
411        let c_field = Arc::new(Field::new("c", DataType::Int32, true));
412        let body = Arc::new(crate::expressions::BinaryExpr::new(
413            Arc::new(LambdaVariable::new(0, Arc::clone(&a_field))),
414            datafusion_expr::Operator::Plus,
415            Arc::new(LambdaVariable::new(2, Arc::clone(&c_field))),
416        ));
417
418        let lambda = LambdaExpr::try_new(
419            vec!["a".to_string(), "b".to_string(), "c".to_string()],
420            body,
421        )
422        .unwrap();
423
424        assert_eq!(lambda.used_param_indices(), &[0, 2]);
425    }
426
427    /// Referencing params out of declaration order still reports both as used.
428    #[test]
429    fn test_used_params_both_used_in_reverse_reference_order() {
430        let k_field = Arc::new(Field::new("k", DataType::Int32, true));
431        let v_field = Arc::new(Field::new("v", DataType::Int32, true));
432        let body = Arc::new(crate::expressions::BinaryExpr::new(
433            Arc::new(LambdaVariable::new(1, Arc::clone(&v_field))),
434            datafusion_expr::Operator::Plus,
435            Arc::new(LambdaVariable::new(0, Arc::clone(&k_field))),
436        ));
437
438        let lambda =
439            LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], body).unwrap();
440
441        assert_eq!(lambda.projection(), &[0, 1]);
442        assert_eq!(lambda.used_param_indices(), &[0, 1]);
443    }
444
445    /// Inside a nested lambda that re-declares one of the outer parameter
446    /// names, only the non-shadowed outer references should be reported as
447    /// used by the outer lambda. In
448    /// `(k, v) -> func(col, (k, v2) -> k + v2 + v)` the inner `k` shadows
449    /// the outer `k`, so the outer lambda must only see `v` as used.
450    #[test]
451    fn test_used_params_handles_shadowing_inside_nested_lambda() {
452        let outer_k_field = Arc::new(Field::new("k", DataType::Int32, true));
453        let outer_v_field = Arc::new(Field::new("v", DataType::Int32, true));
454        let inner_v2_field = Arc::new(Field::new("v2", DataType::Int32, true));
455
456        // Inner lambda body references "k" (inner's), "v2" (inner's), and
457        // "v" (outer's). Build it directly with the dense compressed
458        // indices the inner LambdaExpr::new would produce: sorted referenced
459        // indices, so the names alone matter here — what matters for
460        // shadow tracking is the names, not the indices.
461        let inner_body: Arc<dyn crate::PhysicalExpr> =
462            Arc::new(crate::expressions::BinaryExpr::new(
463                Arc::new(crate::expressions::BinaryExpr::new(
464                    Arc::new(LambdaVariable::new(1, Arc::clone(&outer_k_field))),
465                    datafusion_expr::Operator::Plus,
466                    Arc::new(LambdaVariable::new(2, Arc::clone(&inner_v2_field))),
467                )),
468                datafusion_expr::Operator::Plus,
469                Arc::new(LambdaVariable::new(0, Arc::clone(&outer_v_field))),
470            ));
471        let inner_lambda = Arc::new(
472            LambdaExpr::try_new(vec!["k".to_string(), "v2".to_string()], inner_body)
473                .unwrap(),
474        );
475
476        // Outer body wraps the inner lambda in a binary op next to a
477        // regular column reference so the walk has something non-trivial
478        // to descend through. The outer body references the inner lambda
479        // via `inner_lambda`.
480        let outer_body: Arc<dyn crate::PhysicalExpr> =
481            Arc::new(crate::expressions::BinaryExpr::new(
482                Arc::new(Column::new("col", 0)),
483                datafusion_expr::Operator::Plus,
484                inner_lambda,
485            ));
486
487        let outer_lambda =
488            LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], outer_body)
489                .unwrap();
490
491        assert_eq!(
492            outer_lambda.used_param_indices(),
493            &[1],
494            "only outer's `v` (index 1) should be reported as used; `k` (index 0) is \
495             shadowed inside the nested lambda"
496        );
497    }
498}
datafusion_physical_expr/expressions/lambda.rs

datafusion_physical_expr/expressions/
lambda.rs
