substrait/parse/text/simple_extensions/
scalar_functions.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Scalar function definitions parsed from extension YAML.
4
5use std::collections::HashMap;
6
7use crate::text::simple_extensions::{
8    NullabilityHandling as RawNullabilityHandling, Options as RawOptions,
9    ScalarFunction as RawScalarFunction, ScalarFunctionImplsItem as RawImpl, Type as RawType,
10    VariadicBehavior as RawVariadicBehavior, VariadicBehaviorParameterConsistency,
11};
12
13use super::argument::{ArgumentsItem, ArgumentsItemError};
14use super::extensions::TypeContext;
15use super::type_ast::{TypeExpr, TypeParseError};
16use super::types::{ExtensionTypeError, Type};
17use crate::parse::Parse;
18use thiserror::Error;
19
20/// Errors that can occur when parsing scalar functions
21#[derive(Debug, Error)]
22pub enum ScalarFunctionError {
23    /// Scalar function has no implementations
24    #[error("Scalar function '{name}' must have at least one implementation")]
25    NoImplementations {
26        /// The function name
27        name: String,
28    },
29    /// Invalid variadic behavior
30    #[error("Variadic behavior {field} must be a non-negative integer, got {value}")]
31    InvalidVariadicBehavior {
32        /// The field that was invalid (min or max)
33        field: String,
34        /// The invalid value
35        value: f64,
36    },
37    /// Variadic min is greater than max
38    #[error("Variadic min ({min}) must be less than or equal to max ({max})")]
39    VariadicMinGreaterThanMax {
40        /// The minimum value
41        min: u32,
42        /// The maximum value
43        max: u32,
44    },
45    /// Error parsing function argument
46    #[error("Argument error: {0}")]
47    ArgumentError(#[from] ArgumentsItemError),
48    /// Error parsing type in function signature
49    #[error("Type error: {0}")]
50    TypeError(#[from] ExtensionTypeError),
51    /// Error parsing type expression
52    #[error("Type parse error: {0}")]
53    TypeParseError(#[from] TypeParseError),
54    /// Feature not yet implemented
55    #[error("Not yet implemented: {0}")]
56    NotYetImplemented(String),
57}
58
59/// A validated scalar function definition with one or more implementations
60#[derive(Clone, Debug, PartialEq)]
61pub struct ScalarFunction {
62    /// Function name
63    pub name: String,
64    /// Human-readable description
65    pub description: Option<String>,
66    /// Function implementations (overloads)
67    pub impls: Vec<Impl>,
68}
69
70impl ScalarFunction {
71    /// Parse a scalar function from raw YAML with the provided type context.
72    pub(super) fn from_raw(
73        raw: RawScalarFunction,
74        ctx: &mut TypeContext,
75    ) -> Result<Self, ScalarFunctionError> {
76        if raw.impls.is_empty() {
77            return Err(ScalarFunctionError::NoImplementations { name: raw.name });
78        }
79
80        let impls = raw
81            .impls
82            .into_iter()
83            .map(|impl_| Impl::from_raw(impl_, ctx))
84            .collect::<Result<Vec<_>, _>>()?;
85
86        Ok(ScalarFunction {
87            name: raw.name,
88            description: raw.description,
89            impls,
90        })
91    }
92}
93
94/// A single function implementation (overload) with parsed signature types.
95#[derive(Clone, Debug, PartialEq)]
96pub struct Impl {
97    /// Function arguments with types and optional names/descriptions
98    pub args: Vec<ArgumentsItem>,
99    /// Configurable function options (e.g., overflow behavior, rounding modes)
100    pub options: Options,
101    /// Variadic argument behavior.
102    ///
103    /// `None` indicates the function is not variadic.
104    pub variadic: Option<VariadicBehavior>,
105    /// Whether the function output depends on session state (e.g., timezone, locale).
106    ///
107    /// Defaults to `false` per the Substrait spec.
108    pub session_dependent: bool,
109    /// Whether the function is deterministic (same inputs always produce same output).
110    ///
111    /// Defaults to `true` per the Substrait spec.
112    pub deterministic: bool,
113    /// How the function handles null inputs and produces nullable outputs.
114    ///
115    /// Defaults to [`NullabilityHandling::Mirror`] per the Substrait spec.
116    pub nullability: NullabilityHandling,
117    /// The return type, retaining variables and parameters within nested types.
118    ///
119    /// Variables are not bound or resolved during parsing. Type derivation
120    /// expressions remain unsupported.
121    pub return_type: Type,
122    /// Language-specific implementation code (e.g., SQL, C++, Python)
123    ///
124    /// Maps language identifiers to implementation source code snippets.
125    pub implementation: HashMap<String, String>,
126}
127
128impl Impl {
129    /// Parse an implementation from raw YAML with the provided type context.
130    pub(super) fn from_raw(
131        raw: RawImpl,
132        ctx: &mut TypeContext,
133    ) -> Result<Self, ScalarFunctionError> {
134        // Parse and validate the return type
135        let return_type = match raw.return_.0 {
136            RawType::String(s) => {
137                // Multiline strings indicate type derivation expressions
138                // See: https://github.com/substrait-io/substrait-rs/issues/449
139                if s.contains('\n') {
140                    return Err(ScalarFunctionError::NotYetImplemented(
141                        "Type derivation expressions - issue #449".to_string(),
142                    ));
143                }
144                let type_expr = match TypeExpr::parse(&s) {
145                    Err(TypeParseError::UnsupportedNumericExpression(_)) => {
146                        return Err(ScalarFunctionError::NotYetImplemented(
147                            "Numeric expressions in type parameters - issue #449".to_string(),
148                        ));
149                    }
150                    result => result?,
151                };
152                type_expr.visit_references(&mut |name| ctx.linked(name));
153                match Type::try_from(type_expr) {
154                    Ok(expression) => expression,
155                    Err(ExtensionTypeError::InvalidAnyTypeVariable { .. })
156                    | Err(ExtensionTypeError::ExpectedConcreteType { .. })
157                    | Err(ExtensionTypeError::InvalidParameter(_))
158                    | Err(ExtensionTypeError::InvalidParameterKind { .. }) => {
159                        // General numeric expressions are not yet supported.
160                        // See: https://github.com/substrait-io/substrait-rs/issues/452
161                        return Err(ScalarFunctionError::NotYetImplemented(
162                            "Type variables in function signatures - issue #452".to_string(),
163                        ));
164                    }
165                    Err(ExtensionTypeError::UnknownTypeName { name }) => {
166                        return Err(ScalarFunctionError::TypeError(
167                            ExtensionTypeError::UnknownTypeName { name },
168                        ));
169                    }
170                    Err(e) => return Err(ScalarFunctionError::TypeError(e)),
171                }
172            }
173            RawType::Object(_) => {
174                // Struct return types (YAML syntactic sugar) are not yet supported
175                // See: https://github.com/substrait-io/substrait-rs/issues/450
176                return Err(ScalarFunctionError::NotYetImplemented(
177                    "Struct return types - issue #450".to_string(),
178                ));
179            }
180        };
181
182        let variadic = raw.variadic.map(|v| v.try_into()).transpose()?;
183
184        let args = match raw.args {
185            Some(a) => {
186                a.0.into_iter()
187                    .map(|raw_arg| raw_arg.parse(ctx))
188                    .collect::<Result<Vec<_>, _>>()?
189            }
190            None => Vec::new(),
191        };
192
193        Ok(Impl {
194            args,
195            options: raw.options.as_ref().map(Options::from).unwrap_or_default(),
196            variadic,
197            session_dependent: raw.session_dependent.map(|b| b.0).unwrap_or(false),
198            deterministic: raw.deterministic.map(|b| b.0).unwrap_or(true),
199            nullability: raw
200                .nullability
201                .map(Into::into)
202                .unwrap_or(NullabilityHandling::Mirror),
203            return_type,
204            implementation: raw
205                .implementation
206                .map(|i| i.0.into_iter().collect())
207                .unwrap_or_default(),
208        })
209    }
210}
211
212/// Validated variadic behavior with min/max constraints
213#[derive(Clone, Debug, PartialEq)]
214pub struct VariadicBehavior {
215    /// Minimum number of arguments
216    pub min: u32,
217    /// Maximum number of arguments (unlimited when None)
218    pub max: Option<u32>,
219    /// Whether all variadic parameters must have the same type.
220    ///
221    /// `None` when the parameter is not specified in the YAML. We cannot assume a default
222    /// because the Substrait spec does not define default behavior for missing values
223    /// (see <https://github.com/substrait-io/substrait/issues/928>).
224    ///
225    /// TODO: Once the spec defines default behavior, apply it and change this to a non-Option
226    /// type (see issue #454).
227    pub parameter_consistency: Option<ParameterConsistency>,
228}
229
230/// Specifies whether variadic parameters must have consistent types.
231///
232/// When a function's last argument is variadic with a type parameter (e.g., `fn(A, B, C...)`),
233/// this controls type binding for the variadic arguments.
234///
235/// See: <https://github.com/substrait-io/substrait/issues/928>
236#[derive(Clone, Debug, PartialEq)]
237pub enum ParameterConsistency {
238    /// All variadic arguments must have the same concrete type.
239    ///
240    /// For example, if `C` binds to `i32`, all variadic arguments must be `i32`.
241    Consistent,
242    /// Each variadic argument can have a different type.
243    ///
244    /// Each instance of the variadic parameter can bind to different types
245    /// within the constraints of the type parameter.
246    Inconsistent,
247}
248
249impl From<VariadicBehaviorParameterConsistency> for ParameterConsistency {
250    fn from(raw: VariadicBehaviorParameterConsistency) -> Self {
251        match raw {
252            VariadicBehaviorParameterConsistency::Consistent => ParameterConsistency::Consistent,
253            VariadicBehaviorParameterConsistency::Inconsistent => {
254                ParameterConsistency::Inconsistent
255            }
256        }
257    }
258}
259
260impl TryFrom<RawVariadicBehavior> for VariadicBehavior {
261    type Error = ScalarFunctionError;
262
263    fn try_from(raw: RawVariadicBehavior) -> Result<Self, Self::Error> {
264        fn parse_bound(value: f64, field: &str) -> Result<u32, ScalarFunctionError> {
265            if value < 0.0 || value.fract() != 0.0 {
266                return Err(ScalarFunctionError::InvalidVariadicBehavior {
267                    field: field.to_string(),
268                    value,
269                });
270            }
271            Ok(value as u32)
272        }
273
274        let min = raw
275            .min
276            .map(|v| parse_bound(v, "min"))
277            .transpose()?
278            .unwrap_or(0);
279        let max = raw.max.map(|v| parse_bound(v, "max")).transpose()?;
280
281        if let Some(max_val) = max
282            && min > max_val
283        {
284            return Err(ScalarFunctionError::VariadicMinGreaterThanMax { min, max: max_val });
285        }
286
287        Ok(VariadicBehavior {
288            min,
289            max,
290            parameter_consistency: raw.parameter_consistency.map(Into::into),
291        })
292    }
293}
294
295/// How a function handles null inputs and produces nullable outputs
296#[derive(Clone, Debug, PartialEq)]
297pub enum NullabilityHandling {
298    /// Nullability of output mirrors the nullability of input(s)
299    Mirror,
300    /// Function explicitly declares the nullability of its output
301    DeclaredOutput,
302    /// Function handles nulls in a custom way per implementation
303    Discrete,
304}
305
306impl From<RawNullabilityHandling> for NullabilityHandling {
307    fn from(raw: RawNullabilityHandling) -> Self {
308        match raw {
309            RawNullabilityHandling::Mirror => NullabilityHandling::Mirror,
310            RawNullabilityHandling::DeclaredOutput => NullabilityHandling::DeclaredOutput,
311            RawNullabilityHandling::Discrete => NullabilityHandling::Discrete,
312        }
313    }
314}
315
316/// Validated function options
317#[derive(Clone, Debug, Default, PartialEq)]
318pub struct Options(pub HashMap<String, Vec<String>>);
319
320impl From<&RawOptions> for Options {
321    fn from(raw: &RawOptions) -> Self {
322        Options(
323            raw.0
324                .iter()
325                .map(|(k, v)| (k.clone(), v.values.clone()))
326                .collect(),
327        )
328    }
329}
330
331#[cfg(test)]
332mod tests {
333    use super::*;
334
335    #[test]
336    fn test_variadic_invalid_values() {
337        let invalid_cases = vec![
338            (Some(-1.0), None, "negative min"),
339            (None, Some(-2.5), "negative max"),
340            (Some(7.2), None, "non-integer min"),
341            (None, Some(3.5), "non-integer max"),
342            (Some(5.0), Some(3.0), "min greater than max"),
343        ];
344
345        for (min, max, description) in invalid_cases {
346            let raw = RawVariadicBehavior {
347                min,
348                max,
349                parameter_consistency: None,
350            };
351            assert!(
352                VariadicBehavior::try_from(raw).is_err(),
353                "expected error for {}",
354                description
355            );
356        }
357    }
358
359    #[test]
360    fn test_variadic_valid() {
361        let raw = RawVariadicBehavior {
362            min: Some(1.0),
363            max: Some(5.0),
364            parameter_consistency: None,
365        };
366        let result = VariadicBehavior::try_from(raw).unwrap();
367        assert_eq!(result.min, 1);
368        assert_eq!(result.max, Some(5));
369    }
370
371    #[test]
372    fn test_variadic_none_values() {
373        let raw = RawVariadicBehavior {
374            min: None,
375            max: None,
376            parameter_consistency: None,
377        };
378        let result = VariadicBehavior::try_from(raw).unwrap();
379        assert_eq!(result.min, 0);
380        assert_eq!(result.max, None);
381    }
382
383    #[test]
384    fn test_no_implementations_error() {
385        use crate::text::simple_extensions::ScalarFunction as RawScalarFunction;
386
387        let raw = RawScalarFunction {
388            deprecated: None,
389            name: "empty_function".to_string(),
390            description: None,
391            metadata: Default::default(),
392            impls: vec![],
393        };
394
395        let mut ctx = super::super::extensions::TypeContext::default();
396        let result = ScalarFunction::from_raw(raw, &mut ctx);
397
398        assert!(matches!(
399            result,
400            Err(ScalarFunctionError::NoImplementations { name })
401            if name == "empty_function"
402        ));
403    }
404
405    #[test]
406    fn test_scalar_function_with_single_impl() {
407        use crate::text::simple_extensions::{
408            ReturnValue, ScalarFunction as RawScalarFunction, ScalarFunctionImplsItem,
409            Type as RawType,
410        };
411
412        let raw = RawScalarFunction {
413            deprecated: None,
414            name: "add".to_string(),
415            description: Some("Addition function".to_string()),
416            metadata: Default::default(),
417            impls: vec![ScalarFunctionImplsItem {
418                args: None,
419                deprecated: None,
420                description: None,
421                options: None,
422                variadic: None,
423                session_dependent: None,
424                deterministic: None,
425                nullability: None,
426                return_: ReturnValue(RawType::String("i32".to_string())),
427                implementation: None,
428            }],
429        };
430
431        let mut ctx = super::super::extensions::TypeContext::default();
432        let result = ScalarFunction::from_raw(raw, &mut ctx).unwrap();
433
434        assert_eq!(result.name, "add");
435        assert_eq!(result.description, Some("Addition function".to_string()));
436        assert_eq!(result.impls.len(), 1);
437
438        use super::super::types::{Type, TypeKind};
439        assert_eq!(result.impls[0].return_type, Type::new(TypeKind::I32, false));
440    }
441
442    #[test]
443    fn test_options_conversion() {
444        use crate::text::simple_extensions::{Options as RawOptions, OptionsValue};
445        use indexmap::IndexMap;
446
447        let mut raw_map = IndexMap::new();
448        raw_map.insert(
449            "overflow".to_string(),
450            OptionsValue {
451                values: vec!["SILENT".to_string(), "ERROR".to_string()],
452                description: None,
453            },
454        );
455
456        let raw = RawOptions(raw_map);
457        let options = Options::from(&raw);
458
459        assert_eq!(options.0.len(), 1);
460        assert_eq!(
461            options.0.get("overflow").unwrap(),
462            &vec!["SILENT".to_string(), "ERROR".to_string()]
463        );
464    }
465}
substrait/parse/text/simple_extensions/scalar_functions.rs

substrait/parse/text/simple_extensions/
scalar_functions.rs
