substrait/parse/text/simple_extensions/
extensions.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Validated simple extensions: [`SimpleExtensions`].
4//!
5//! Both type definitions and scalar function definitions are supported.
6//! Aggregate functions (see #447) and window functions (see #446) are not yet supported.
7
8use indexmap::IndexMap;
9use std::collections::{HashMap, HashSet};
10use std::str::FromStr;
11
12use super::{
13    ExtensionFile, SimpleExtensionsError, scalar_functions::ScalarFunction, types::CustomType,
14};
15use crate::{
16    parse::{Context, Parse},
17    text::simple_extensions::SimpleExtensions as RawExtensions,
18    urn::Urn,
19};
20
21/// The contents (types and functions) in an [`ExtensionFile`](super::file::ExtensionFile).
22///
23/// This structure stores and provides access to the individual objects defined
24/// in an [`ExtensionFile`](super::file::ExtensionFile); [`SimpleExtensions`]
25/// represents the contents of an extensions file.
26#[derive(Clone, Debug, Default)]
27pub struct SimpleExtensions {
28    /// Types defined in this extension file
29    types: HashMap<String, CustomType>,
30    /// Scalar functions defined in this extension file
31    ///
32    /// TODO: Add support for window functions (issue #446) and aggregate functions (issue #447)
33    scalar_functions: HashMap<String, ScalarFunction>,
34}
35
36impl SimpleExtensions {
37    /// Add a type to the context. Name must be unique.
38    pub fn add_type(&mut self, custom_type: &CustomType) -> Result<(), SimpleExtensionsError> {
39        if self.types.contains_key(&custom_type.name) {
40            return Err(SimpleExtensionsError::DuplicateTypeName {
41                name: custom_type.name.clone(),
42            });
43        }
44
45        self.types
46            .insert(custom_type.name.clone(), custom_type.clone());
47        Ok(())
48    }
49
50    /// Get a type by name from the context
51    pub fn get_type(&self, name: &str) -> Option<&CustomType> {
52        self.types.get(name)
53    }
54
55    /// Get an iterator over all types in the context
56    pub fn types(&self) -> impl Iterator<Item = &CustomType> {
57        self.types.values()
58    }
59
60    /// Consume the parsed extension and return its types.
61    pub(crate) fn into_types(self) -> HashMap<String, CustomType> {
62        self.types
63    }
64
65    /// Add a scalar function to the context, merging with existing functions of the same name.
66    ///
67    /// When duplicate function names are encountered, implementations are merged (unioned).
68    /// The existing description is kept if present, otherwise the new description is used.
69    ///
70    /// See: https://github.com/substrait-io/substrait/issues/931
71    pub(super) fn add_scalar_function(&mut self, scalar_function: ScalarFunction) {
72        use std::collections::hash_map::Entry;
73        match self.scalar_functions.entry(scalar_function.name.clone()) {
74            Entry::Vacant(e) => {
75                e.insert(scalar_function);
76            }
77            Entry::Occupied(mut e) => {
78                Self::merge_scalar_function(e.get_mut(), scalar_function);
79            }
80        }
81    }
82
83    /// Merge a new scalar function into an existing one.
84    ///
85    /// Unions the implementations. Keeps the existing description if present,
86    /// otherwise uses the new description.
87    // TODO: Reject conflicting implementations instead of blindly merging.
88    fn merge_scalar_function(existing: &mut ScalarFunction, new: ScalarFunction) {
89        existing.impls.extend(new.impls);
90        existing.description = existing.description.take().or(new.description);
91    }
92
93    /// Get a scalar function by name
94    pub fn get_scalar_function(&self, name: &str) -> Option<&ScalarFunction> {
95        self.scalar_functions.get(name)
96    }
97
98    /// Get an iterator over all scalar functions
99    pub fn scalar_functions(&self) -> impl Iterator<Item = &ScalarFunction> {
100        self.scalar_functions.values()
101    }
102}
103
104/// resolved or unresolved.
105#[derive(Debug, Default)]
106pub(crate) struct TypeContext {
107    /// Types that have been seen so far, now resolved.
108    known: HashSet<String>,
109    /// Types that have been linked to, not yet resolved.
110    linked: HashSet<String>,
111}
112
113impl TypeContext {
114    /// Mark a type as found
115    pub fn found(&mut self, name: &str) {
116        self.linked.remove(name);
117        self.known.insert(name.to_string());
118    }
119
120    /// Mark a type as linked to - some other type or function references it,
121    /// but we haven't seen it.
122    pub fn linked(&mut self, name: &str) {
123        if !self.known.contains(name) {
124            self.linked.insert(name.to_string());
125        }
126    }
127}
128
129impl Context for TypeContext {}
130
131// Implement parsing for the raw text representation to produce an `ExtensionFile`.
132impl Parse<TypeContext> for RawExtensions {
133    // A local type (rather than a `(Urn, SimpleExtensions)` tuple) so the
134    // `Parsed: Into<Self>` round-trip bound can be satisfied without an
135    // orphan-rule violation on the foreign `RawExtensions`.
136    type Parsed = ExtensionFile;
137    type Error = super::SimpleExtensionsError;
138
139    fn parse(self, ctx: &mut TypeContext) -> Result<Self::Parsed, Self::Error> {
140        let RawExtensions {
141            urn,
142            types,
143            scalar_functions,
144            ..
145        } = self;
146        let urn = Urn::from_str(&urn)?;
147        let mut extension = SimpleExtensions::default();
148
149        for type_item in types {
150            let custom_type = Parse::parse(type_item, ctx)?;
151            extension.add_type(&custom_type)?;
152        }
153
154        for scalar_fn in scalar_functions {
155            match ScalarFunction::from_raw(scalar_fn, ctx) {
156                Ok(parsed_fn) => {
157                    extension.add_scalar_function(parsed_fn);
158                }
159                Err(super::scalar_functions::ScalarFunctionError::NotYetImplemented(_)) => {
160                    // Skip functions with unimplemented features (e.g., type derivations)
161                    // These will be supported in a future update
162                    continue;
163                }
164                Err(e) => return Err(e.into()),
165            }
166        }
167
168        if let Some(missing) = ctx.linked.iter().next() {
169            // TODO: Track originating type(s) to improve this error message.
170            return Err(super::SimpleExtensionsError::UnresolvedTypeReference {
171                type_name: missing.clone(),
172            });
173        }
174
175        Ok(ExtensionFile { urn, extension })
176    }
177}
178
179/// Build the raw text representation ([`RawExtensions`]) from a canonical
180/// [`Urn`] and parsed [`SimpleExtensions`].
181///
182/// This is a free function rather than a `From` impl because [`RawExtensions`]
183/// is defined in the `substrait-extensions` crate, so the orphan rule forbids
184/// implementing `From` for it here.
185pub(super) fn to_raw_extensions(urn: Urn, extension: SimpleExtensions) -> RawExtensions {
186    let types = extension
187        .into_types()
188        .into_values()
189        .map(Into::into)
190        .collect();
191
192    RawExtensions {
193        urn: urn.to_string(),
194        aggregate_functions: vec![],
195        dependencies: IndexMap::new(),
196        metadata: Default::default(),
197        scalar_functions: vec![],
198        type_variations: vec![],
199        types,
200        window_functions: vec![],
201    }
202}
substrait/parse/text/simple_extensions/extensions.rs

substrait/parse/text/simple_extensions/
extensions.rs
