datafusion_physical_expr/
higher_order_function.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//! Declaration of built-in (higher order) functions.
19//! This module contains built-in functions' enumeration and metadata.
20//!
21//! Generally, a function has:
22//! * a signature
23//! * a return type, that is a function of the incoming argument's types
24//! * the computation, that must accept each valid signature
25//!
26//! * Signature: see `Signature`
27//! * Return type: a function `(arg_types) -> return_type`. E.g. for array_transform, ([[f32]], v -> v*2) -> [f32], ([[f32]], v -> v > 3.0) -> [bool].
28//!
29//! This module also has a set of coercion rules to improve user experience: if an argument i32 is passed
30//! to a function that supports f64, it is coerced to f64.
31
32use std::fmt::{self, Debug, Formatter};
33use std::hash::{Hash, Hasher};
34use std::sync::Arc;
35
36use crate::PhysicalExpr;
37use crate::expressions::{LambdaExpr, Literal};
38
39use arrow::array::{Array, RecordBatch};
40use arrow::datatypes::{DataType, FieldRef, Schema};
41use datafusion_common::config::{ConfigEntry, ConfigOptions};
42use datafusion_common::datatype::FieldExt;
43use datafusion_common::utils::remove_list_null_values;
44use datafusion_common::{
45    Result, ScalarValue, exec_err, internal_datafusion_err, internal_err,
46    plan_datafusion_err, plan_err,
47};
48use datafusion_expr::type_coercion::functions::value_fields_with_higher_order_udf;
49use datafusion_expr::{
50    ColumnarValue, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, HigherOrderUDF,
51    LambdaArgument, LambdaParametersProgress, ValueOrLambda, Volatility, expr_vec_fmt,
52};
53
54/// Per-argument classification cached at construction time.
55///
56/// Walking the wrapped lambda tree and scanning a `Vec<usize>` of lambda
57/// positions used to be done on every `evaluate` call. Both costs collapse
58/// to a single up-front pass by storing the classification (and the resolved
59/// inner [`LambdaExpr`]) here.
60enum ArgSlot {
61    /// A regular value-producing expression at this position.
62    Value,
63    /// A lambda position. Stores the inner [`LambdaExpr`] pre-extracted from
64    /// any wrapper expressions that may have been introduced via
65    /// [`PhysicalExpr::with_new_children`] tree rewrites.
66    Lambda(Arc<LambdaExpr>),
67}
68
69/// Physical expression of a higher order function
70pub struct HigherOrderFunctionExpr {
71    /// A shared instance of the higher-order function
72    fun: Arc<HigherOrderUDF>,
73    /// The name of the higher-order function
74    name: String,
75    /// List of expressions to feed to the function as arguments
76    ///
77    /// For example, for `array_transform([2, 3], v -> v != 2)`, this will be:
78    ///
79    /// ```text
80    /// ListExpression [2,3]
81    /// LambdaExpression
82    ///     parameters: ["v"]
83    ///     body:
84    ///         BinaryExpression (!=)
85    ///             left:
86    ///                 LambdaVariableExpression("v", Field::new("", Int32, false))
87    ///             right:
88    ///                 LiteralExpression(2)
89    /// ```
90    args: Vec<Arc<dyn PhysicalExpr>>,
91    /// Per-arg classification, parallel to `args`. Length always equals
92    /// `args.len()`. Lambda variants carry the resolved inner [`LambdaExpr`]
93    /// so `evaluate` doesn't walk through wrapper nodes.
94    slots: Vec<ArgSlot>,
95    /// The output field associated this expression
96    ///
97    /// For example, for `array_transform([2, 3], v -> v != 2)`, this will be
98    /// `Field::new("", DataType::new_list(DataType::Boolean, true), true)`
99    return_field: FieldRef,
100    /// The config options at execution time
101    config_options: Arc<ConfigOptions>,
102}
103
104impl Debug for HigherOrderFunctionExpr {
105    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
106        let lambda_positions: Vec<_> = self
107            .slots
108            .iter()
109            .enumerate()
110            .filter_map(|(i, slot)| matches!(slot, ArgSlot::Lambda(_)).then_some(i))
111            .collect();
112        f.debug_struct("HigherOrderFunctionExpr")
113            .field("fun", &"<FUNC>")
114            .field("name", &self.name)
115            .field("args", &self.args)
116            .field("lambda_positions", &lambda_positions)
117            .field("return_field", &self.return_field)
118            .finish()
119    }
120}
121
122impl HigherOrderFunctionExpr {
123    /// Create a new Higher Order function
124    ///
125    /// Note that lambda arguments must be present directly in args as [LambdaExpr],
126    /// and not as a wrapped child of any arg
127    pub fn try_new_with_schema(
128        fun: Arc<HigherOrderUDF>,
129        args: Vec<Arc<dyn PhysicalExpr>>,
130        schema: &Schema,
131        config_options: Arc<ConfigOptions>,
132    ) -> Result<Self> {
133        let name = fun.name().to_string();
134        let mut slots = Vec::with_capacity(args.len());
135        let arg_fields = args
136            .iter()
137            .map(|e| match e.downcast_ref::<LambdaExpr>() {
138                Some(lambda) => {
139                    slots.push(ArgSlot::Lambda(Arc::new(lambda.clone())));
140                    Ok(ValueOrLambda::Lambda(lambda.body().return_field(schema)?))
141                }
142                None => {
143                    slots.push(ArgSlot::Value);
144                    Ok(ValueOrLambda::Value(e.return_field(schema)?))
145                }
146            })
147            .collect::<Result<Vec<_>>>()?;
148
149        // verify that input data types is consistent with function's `HigherOrderTypeSignature`
150        value_fields_with_higher_order_udf(&arg_fields, fun.as_ref())?;
151
152        let arguments = args
153            .iter()
154            .map(|e| e.downcast_ref::<Literal>().map(|literal| literal.value()))
155            .collect::<Vec<_>>();
156
157        let ret_args = HigherOrderReturnFieldArgs {
158            arg_fields: &arg_fields,
159            scalar_arguments: &arguments,
160        };
161
162        let return_field = fun.return_field_from_args(ret_args)?;
163
164        Ok(Self {
165            fun,
166            name,
167            args,
168            slots,
169            return_field,
170            config_options,
171        })
172    }
173
174    /// Get the higher order function implementation
175    pub fn fun(&self) -> &HigherOrderUDF {
176        self.fun.as_ref()
177    }
178
179    /// The name for this expression
180    pub fn name(&self) -> &str {
181        &self.name
182    }
183
184    /// Input arguments
185    pub fn args(&self) -> &[Arc<dyn PhysicalExpr>] {
186        &self.args
187    }
188
189    /// Data type produced by this expression
190    pub fn return_type(&self) -> &DataType {
191        self.return_field.data_type()
192    }
193
194    pub fn nullable(&self) -> bool {
195        self.return_field.is_nullable()
196    }
197
198    pub fn config_options(&self) -> &ConfigOptions {
199        &self.config_options
200    }
201
202    /// Resolve every lambda's parameter list. Returns an empty `Vec` when
203    /// there are no lambdas, avoiding the [`datafusion_expr::HigherOrderUDFImpl::lambda_parameters`]
204    /// virtual call entirely.
205    fn resolve_lambda_parameters(
206        &self,
207        fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>],
208    ) -> Result<Vec<Vec<FieldRef>>> {
209        let num_lambdas = self
210            .slots
211            .iter()
212            .filter(|s| matches!(s, ArgSlot::Lambda(_)))
213            .count();
214        if num_lambdas == 0 {
215            return Ok(Vec::new());
216        }
217        match self.fun().lambda_parameters(0, fields)? {
218            LambdaParametersProgress::Partial(_) => plan_err!(
219                "{} lambda_parameters returned a partial result when the return type of all it's lambdas were provided",
220                self.name()
221            ),
222            LambdaParametersProgress::Complete(items) => {
223                // functions can support multiple lambdas where some trailing ones are optional,
224                // but to simplify the implementor, lambda_parameters returns the parameters of all of them,
225                // so we can't do equality check. one example is spark reduce:
226                // https://spark.apache.org/docs/latest/api/sql/index.html#reduce
227                if items.len() < num_lambdas {
228                    return exec_err!(
229                        "{} invocation defined {num_lambdas} but lambda_parameters returned only {}",
230                        self.name(),
231                        items.len()
232                    );
233                }
234                Ok(items)
235            }
236        }
237    }
238}
239
240impl fmt::Display for HigherOrderFunctionExpr {
241    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
242        write!(f, "{}({})", self.name, expr_vec_fmt!(self.args))
243    }
244}
245
246impl PartialEq for HigherOrderFunctionExpr {
247    fn eq(&self, o: &Self) -> bool {
248        if std::ptr::eq(self, o) {
249            // The equality implementation is somewhat expensive, so let's short-circuit when possible.
250            return true;
251        }
252        // `slots` is a deterministic function of `fun` and `args`, so it's
253        // not part of the comparison.
254        let Self {
255            fun,
256            name,
257            args,
258            slots: _,
259            return_field,
260            config_options,
261        } = self;
262        fun.eq(&o.fun)
263            && name.eq(&o.name)
264            && args.eq(&o.args)
265            && return_field.eq(&o.return_field)
266            && (Arc::ptr_eq(config_options, &o.config_options)
267                || sorted_config_entries(config_options)
268                    == sorted_config_entries(&o.config_options))
269    }
270}
271impl Eq for HigherOrderFunctionExpr {}
272impl Hash for HigherOrderFunctionExpr {
273    fn hash<H: Hasher>(&self, state: &mut H) {
274        let Self {
275            fun,
276            name,
277            args,
278            slots: _,
279            return_field,
280            config_options: _, // expensive to hash, and often equal
281        } = self;
282        fun.hash(state);
283        name.hash(state);
284        args.hash(state);
285        return_field.hash(state);
286    }
287}
288
289fn sorted_config_entries(config_options: &ConfigOptions) -> Vec<ConfigEntry> {
290    let mut entries = config_options.entries();
291    entries.sort_by(|l, r| l.key.cmp(&r.key));
292    entries
293}
294
295impl PhysicalExpr for HigherOrderFunctionExpr {
296    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
297        let mut arg_fields = Vec::with_capacity(self.args.len());
298        let mut fields = Vec::with_capacity(self.args.len());
299        for (arg, slot) in self.args.iter().zip(&self.slots) {
300            match slot {
301                ArgSlot::Lambda(lambda) => {
302                    let field = lambda.body().return_field(batch.schema_ref())?;
303                    arg_fields.push(ValueOrLambda::Lambda(Arc::clone(&field)));
304                    fields.push(ValueOrLambda::Lambda(Some(field)));
305                }
306                ArgSlot::Value => {
307                    let field = arg.return_field(batch.schema_ref())?;
308                    arg_fields.push(ValueOrLambda::Value(Arc::clone(&field)));
309                    fields.push(ValueOrLambda::Value(field));
310                }
311            }
312        }
313
314        let mut lambda_parameters = self.resolve_lambda_parameters(&fields)?.into_iter();
315
316        let args = self
317            .args
318            .iter()
319            .zip(&self.slots)
320            .map(|(arg, slot)| match slot {
321                ArgSlot::Lambda(lambda) => {
322                    let lambda_params = lambda_parameters.next().ok_or_else(|| {
323                        internal_datafusion_err!(
324                            "params len should have been checked above"
325                        )
326                    })?;
327
328                    if lambda.params().len() > lambda_params.len() {
329                        return exec_err!(
330                            "lambda defined {} params but higher-order function support only {}",
331                            lambda.params().len(),
332                            lambda_params.len()
333                        );
334                    }
335
336                    let params = std::iter::zip(lambda.params(), lambda_params)
337                        .map(|(name, param)| param.renamed(name.as_str()))
338                        .collect();
339
340                    // lambda.projection may include indexes of nested lambda variables not present on this batch
341                    let projection = lambda
342                        .projection()
343                        .iter()
344                        .copied()
345                        .filter(|i| *i < batch.num_columns())
346                        .collect::<Vec<_>>();
347
348                    Ok(ValueOrLambda::Lambda(LambdaArgument::new(
349                        params,
350                        Arc::clone(lambda.projected_body()),
351                        if projection.is_empty() {
352                            None
353                        } else {
354                            Some(batch.project(&projection)?)
355                        },
356                        lambda.used_param_indices(),
357                    )))
358                }
359                ArgSlot::Value => {
360                    let value = arg.evaluate(batch)?;
361
362                    let value = if self.fun.clear_null_values()
363                        && matches!(
364                            value.data_type(),
365                            DataType::List(_) | DataType::LargeList(_)
366                        )
367                    {
368                        let arr = value.into_array(batch.num_rows())?;
369                        if arr.null_count() == 0 {
370                            ColumnarValue::Array(arr)
371                        } else {
372                            ColumnarValue::Array(remove_list_null_values(&arr)?)
373                        }
374                    } else {
375                        value
376                    };
377
378                    Ok(ValueOrLambda::Value(value))
379                }
380            })
381            .collect::<Result<Vec<_>>>()?;
382
383        let input_empty = args.is_empty();
384        let input_all_scalar = args
385            .iter()
386            .all(|arg| matches!(arg, ValueOrLambda::Value(ColumnarValue::Scalar(_))));
387
388        // evaluate the function
389        let output = self.fun.invoke_with_args(HigherOrderFunctionArgs {
390            args,
391            arg_fields,
392            number_rows: batch.num_rows(),
393            return_field: Arc::clone(&self.return_field),
394            config_options: Arc::clone(&self.config_options),
395        })?;
396
397        if let ColumnarValue::Array(array) = &output
398            && array.len() != batch.num_rows()
399        {
400            // If the arguments are a non-empty slice of scalar values, we can assume that
401            // returning a one-element array is equivalent to returning a scalar.
402            let preserve_scalar = array.len() == 1 && !input_empty && input_all_scalar;
403            return if preserve_scalar {
404                ScalarValue::try_from_array(array, 0).map(ColumnarValue::Scalar)
405            } else {
406                internal_err!(
407                    "higher-order function {} returned a different number of rows than expected. Expected: {}, Got: {}",
408                    self.name,
409                    batch.num_rows(),
410                    array.len()
411                )
412            };
413        }
414        Ok(output)
415    }
416
417    fn return_field(&self, _input_schema: &Schema) -> Result<FieldRef> {
418        Ok(Arc::clone(&self.return_field))
419    }
420
421    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
422        self.args.iter().collect()
423    }
424
425    fn with_new_children(
426        self: Arc<Self>,
427        children: Vec<Arc<dyn PhysicalExpr>>,
428    ) -> Result<Arc<dyn PhysicalExpr>> {
429        if children.len() != self.args.len() {
430            return internal_err!(
431                "HigherOrderFunctionExpr expects exactly {} child, got {}",
432                self.args.len(),
433                children.len()
434            );
435        }
436
437        // Re-derive `slots` for the new children using the original slot kinds
438        // as the source of truth for which positions must (still) be lambdas.
439        let mut new_slots = Vec::with_capacity(children.len());
440        for (i, child) in children.iter().enumerate() {
441            match &self.slots[i] {
442                ArgSlot::Lambda(_) => {
443                    let lambda = wrapped_lambda(child).ok_or_else(|| {
444                        plan_datafusion_err!(
445                            "{} unable to unwrap lambda from {} at position {i}",
446                            &children[i],
447                            self.name()
448                        )
449                    })?;
450                    new_slots.push(ArgSlot::Lambda(Arc::new(lambda.clone())));
451                }
452                ArgSlot::Value => {
453                    if child.is::<LambdaExpr>() {
454                        return plan_err!(
455                            "{} received a lambda via with_new_children at position {i} that wasn't a lambda before",
456                            self.name()
457                        );
458                    }
459                    new_slots.push(ArgSlot::Value);
460                }
461            }
462        }
463
464        Ok(Arc::new(HigherOrderFunctionExpr {
465            name: self.name.clone(),
466            fun: Arc::clone(&self.fun),
467            args: children,
468            slots: new_slots,
469            return_field: Arc::clone(&self.return_field),
470            config_options: Arc::clone(&self.config_options),
471        }))
472    }
473
474    fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result {
475        write!(f, "{}(", self.name)?;
476        for (i, expr) in self.args.iter().enumerate() {
477            if i > 0 {
478                write!(f, ", ")?;
479            }
480            expr.fmt_sql(f)?;
481        }
482        write!(f, ")")
483    }
484
485    fn is_volatile_node(&self) -> bool {
486        self.fun.signature().volatility == Volatility::Volatile
487    }
488}
489
490fn wrapped_lambda(expr: &Arc<dyn PhysicalExpr>) -> Option<&LambdaExpr> {
491    let mut current = expr;
492
493    loop {
494        if let Some(lambda) = current.downcast_ref::<LambdaExpr>() {
495            return Some(lambda);
496        } else if current.is::<HigherOrderFunctionExpr>() {
497            return None;
498        }
499
500        match current.children().as_slice() {
501            [single_child] => current = *single_child,
502            _ => return None,
503        }
504    }
505}
506
507#[cfg(test)]
508mod tests {
509    use std::sync::Arc;
510
511    use super::*;
512    use crate::HigherOrderFunctionExpr;
513    use crate::create_physical_expr;
514    use crate::expressions::Column;
515    use crate::expressions::NoOp;
516    use crate::expressions::lambda;
517    use crate::expressions::not;
518    use arrow::array::RecordBatchOptions;
519    use arrow::array::{ArrayRef, Int32Array};
520    use arrow::datatypes::{DataType, Field, Schema};
521    use datafusion_common::Result;
522    use datafusion_common::assert_contains;
523    use datafusion_expr::execution_props::ExecutionProps;
524    use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
525    use datafusion_expr::{
526        HigherOrderFunctionArgs, HigherOrderSignature, HigherOrderUDF, HigherOrderUDFImpl,
527    };
528    use datafusion_expr_common::columnar_value::ColumnarValue;
529    use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
530    use datafusion_physical_expr_common::physical_expr::is_volatile;
531
532    /// Test helper to create a mock UDF with a specific volatility
533    #[derive(Debug, PartialEq, Eq, Hash)]
534    struct MockHigherOrderUDF {
535        signature: HigherOrderSignature,
536    }
537
538    impl HigherOrderUDFImpl for MockHigherOrderUDF {
539        fn name(&self) -> &str {
540            "mock_function"
541        }
542
543        fn signature(&self) -> &HigherOrderSignature {
544            &self.signature
545        }
546
547        fn lambda_parameters(
548            &self,
549            _step: usize,
550            _fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>],
551        ) -> Result<LambdaParametersProgress> {
552            // Offer two params; single-param lambdas just ignore the second.
553            Ok(LambdaParametersProgress::Complete(vec![vec![
554                Arc::new(Field::new("", DataType::Int32, true)),
555                Arc::new(Field::new("", DataType::Int32, true)),
556            ]]))
557        }
558
559        fn return_field_from_args(
560            &self,
561            args: HigherOrderReturnFieldArgs,
562        ) -> Result<FieldRef> {
563            match &args.arg_fields[0] {
564                ValueOrLambda::Lambda(field) | ValueOrLambda::Value(field) => {
565                    Ok(Arc::clone(field))
566                }
567            }
568        }
569
570        fn invoke_with_args(
571            &self,
572            args: HigherOrderFunctionArgs,
573        ) -> Result<ColumnarValue> {
574            match &args.args[0] {
575                ValueOrLambda::Lambda(lambda) => lambda.evaluate(
576                    &[
577                        // Sentinel for the first param, distinct from the second's value.
578                        &|| {
579                            Ok(Arc::new(Int32Array::from(vec![-1000; args.number_rows]))
580                                as ArrayRef)
581                        },
582                        &|| {
583                            Ok(Arc::new(Int32Array::from_iter_values(
584                                (0..args.number_rows as i32).map(|i| 10 * (i + 1)),
585                            )) as ArrayRef)
586                        },
587                    ],
588                    |arrays| Ok(arrays.to_vec()),
589                ),
590                ValueOrLambda::Value(value) => Ok(value.clone()),
591            }
592        }
593    }
594
595    #[test]
596    fn test_higher_order_function_volatile_node() {
597        // Create a volatile UDF
598        let volatile_udf = Arc::new(HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
599            signature: HigherOrderSignature::variadic_any(Volatility::Volatile),
600        }));
601
602        // Create a non-volatile UDF
603        let stable_udf = Arc::new(HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
604            signature: HigherOrderSignature::variadic_any(Volatility::Stable),
605        }));
606
607        let schema = Schema::new(vec![Field::new("a", DataType::Float32, false)]);
608        let args = vec![Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>];
609        let config_options = Arc::new(ConfigOptions::new());
610
611        // Test volatile function
612        let volatile_expr = HigherOrderFunctionExpr::try_new_with_schema(
613            volatile_udf,
614            args.clone(),
615            &schema,
616            Arc::clone(&config_options),
617        )
618        .unwrap();
619
620        assert!(volatile_expr.is_volatile_node());
621        let volatile_arc: Arc<dyn PhysicalExpr> = Arc::new(volatile_expr);
622        assert!(is_volatile(&volatile_arc));
623
624        // Test non-volatile function
625        let stable_expr = HigherOrderFunctionExpr::try_new_with_schema(
626            stable_udf,
627            args,
628            &schema,
629            config_options,
630        )
631        .unwrap();
632
633        assert!(!stable_expr.is_volatile_node());
634        let stable_arc: Arc<dyn PhysicalExpr> = Arc::new(stable_expr);
635        assert!(!is_volatile(&stable_arc));
636    }
637
638    #[test]
639    fn test_higher_order_function_wrapped_lambda() {
640        let fun = Arc::new(HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
641            signature: HigherOrderSignature::variadic_any(Volatility::Stable),
642        }));
643
644        let expected = ScalarValue::Int32(Some(42));
645
646        let hof = HigherOrderFunctionExpr::try_new_with_schema(
647            fun,
648            vec![lambda(["a"], Arc::new(Literal::new(expected.clone()))).unwrap()],
649            &Schema::empty(),
650            Arc::new(ConfigOptions::new()),
651        )
652        .unwrap();
653
654        let new_children = vec![not(Arc::clone(&hof.args[0])).unwrap()];
655        let wrapped = Arc::new(hof).with_new_children(new_children).unwrap();
656
657        let result = wrapped
658            .evaluate(
659                &RecordBatch::try_new_with_options(
660                    Arc::new(Schema::empty()),
661                    vec![],
662                    &RecordBatchOptions::new().with_row_count(Some(0)),
663                )
664                .unwrap(),
665            )
666            .unwrap();
667
668        let ColumnarValue::Scalar(result) = result else {
669            unreachable!()
670        };
671
672        assert_eq!(result, expected);
673    }
674
675    #[test]
676    fn test_higher_order_function_badly_wrapped_lambda() {
677        let fun = Arc::new(HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
678            signature: HigherOrderSignature::variadic_any(Volatility::Stable),
679        }));
680
681        let hof = HigherOrderFunctionExpr::try_new_with_schema(
682            fun,
683            vec![
684                not(
685                    lambda(["a"], Arc::new(Literal::new(ScalarValue::Int32(Some(42)))))
686                        .unwrap(),
687                )
688                .unwrap(),
689            ],
690            &Schema::empty(),
691            Arc::new(ConfigOptions::new()),
692        )
693        .unwrap();
694
695        let result = hof
696            .evaluate(
697                &RecordBatch::try_new_with_options(
698                    Arc::new(Schema::empty()),
699                    vec![],
700                    &RecordBatchOptions::new().with_row_count(Some(0)),
701                )
702                .unwrap(),
703            )
704            .unwrap_err();
705
706        assert_contains!(
707            result.to_string(),
708            "LambdaExpr::evaluate() should not be called"
709        );
710    }
711
712    #[test]
713    fn test_higher_order_function_unexpected_lambda() {
714        let fun = Arc::new(HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
715            signature: HigherOrderSignature::variadic_any(Volatility::Stable),
716        }));
717
718        let hof = HigherOrderFunctionExpr::try_new_with_schema(
719            fun,
720            vec![Arc::new(NoOp::new())],
721            &Schema::empty(),
722            Arc::new(ConfigOptions::new()),
723        )
724        .unwrap();
725
726        let result = Arc::new(hof)
727            .with_new_children(vec![lambda(["a"], Arc::new(NoOp::new())).unwrap()])
728            .unwrap_err();
729
730        assert_contains!(
731            result.to_string(),
732            "mock_function received a lambda via with_new_children at position 0 that wasn't a lambda before"
733        );
734    }
735
736    /// Exercises the real planner end to end (not hand-picked indices) to
737    /// check the "captures before own-params" layout invariant.
738    #[test]
739    fn test_higher_order_function_two_lambda_params_capture_and_unused_param() {
740        use datafusion_common::DFSchema;
741        use datafusion_expr::expr::{HigherOrderFunction, LambdaVariable};
742        use datafusion_expr::{Expr, col, lambda as logical_lambda};
743
744        let fun = Arc::new(HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
745            signature: HigherOrderSignature::variadic_any(Volatility::Stable),
746        }));
747
748        // Body uses capture "a" and param "v"; param "k" is left unused.
749        let v = Expr::LambdaVariable(LambdaVariable::new(
750            "v".to_string(),
751            Some(Arc::new(Field::new("v", DataType::Int32, true))),
752        ));
753        let body = col("a") + v;
754        let lambda_expr = logical_lambda(["k", "v"], body);
755
756        let schema = DFSchema::from_unqualified_fields(
757            vec![Field::new("a", DataType::Int32, false)].into(),
758            std::collections::HashMap::new(),
759        )
760        .unwrap();
761
762        let physical_expr = create_physical_expr(
763            &Expr::HigherOrderFunction(HigherOrderFunction::new(fun, vec![lambda_expr])),
764            &schema,
765            &ExecutionProps::new(),
766            &PhysicalPlanningContext::default(),
767        )
768        .unwrap();
769
770        let batch = RecordBatch::try_new(
771            Arc::clone(schema.inner()),
772            vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef],
773        )
774        .unwrap();
775
776        let result = physical_expr.evaluate(&batch).unwrap();
777        let ColumnarValue::Array(result) = result else {
778            unreachable!()
779        };
780
781        // a + v; k's sentinel (-1000) must not leak into the result.
782        let expected = Int32Array::from(vec![11, 22, 33]);
783        assert_eq!(result.as_ref(), &expected);
784    }
785}
datafusion_physical_expr/higher_order_function.rs

datafusion_physical_expr/
higher_order_function.rs
