improve lambda_parameters infinite loop guard

Author: gstvg

datafusion/expr/src/higher_order_function.rs

@@ -75,15 +75,22 @@ pub struct HigherOrderSignature {
     pub volatility: Volatility,
     /// Whether [HigherOrderUDF::coerce_values_for_lambdas] should be called
     pub coerce_values_for_lambdas: bool,
+    /// The max number of times to call [HigherOrderUDF::lambda_parameters] before raising an error.
+    /// Used to guard against implementations that causes an infinite loop by endlessly returning
+    /// [LambdaParametersProgress::Partial]. Defaults to 256
+    pub lambda_parameters_max_iterations: usize,
 }
 
+const LAMBDA_PARAMETERS_MAX_ITERATIONS: usize = 256;
+
 impl HigherOrderSignature {
     /// Creates a new `HigherOrderSignature` from a given type signature and volatility.
     pub fn new(type_signature: HigherOrderTypeSignature, volatility: Volatility) -> Self {
         HigherOrderSignature {
             type_signature,
             volatility,
             coerce_values_for_lambdas: false,
+            lambda_parameters_max_iterations: LAMBDA_PARAMETERS_MAX_ITERATIONS,
         }
     }
 
@@ -93,6 +100,7 @@ impl HigherOrderSignature {
             type_signature: HigherOrderTypeSignature::UserDefined,
             volatility,
             coerce_values_for_lambdas: false,
+            lambda_parameters_max_iterations: LAMBDA_PARAMETERS_MAX_ITERATIONS,
         }
     }
 
@@ -102,6 +110,7 @@ impl HigherOrderSignature {
             type_signature: HigherOrderTypeSignature::VariadicAny,
             volatility,
             coerce_values_for_lambdas: false,
+            lambda_parameters_max_iterations: LAMBDA_PARAMETERS_MAX_ITERATIONS,
         }
     }
 
@@ -111,6 +120,7 @@ impl HigherOrderSignature {
             type_signature: HigherOrderTypeSignature::Any(arg_count),
             volatility,
             coerce_values_for_lambdas: false,
+            lambda_parameters_max_iterations: LAMBDA_PARAMETERS_MAX_ITERATIONS,
         }
     }
 

datafusion/sql/src/expr/function.rs

@@ -513,11 +513,13 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                     }
                 }
 
+                let limit = fm.signature().lambda_parameters_max_iterations;
+
                 step += 1;
 
-                if step > 256 {
+                if step > limit {
                     return plan_err!(
-                        "{} lambda_parameters called 256 times without completion",
+                        "{} lambda_parameters called {limit} times without completion",
                         fm.name()
                     );
                 }