substrait/parse/text/simple_extensions/
registry.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Substrait Extension Registry
4//!
5//! This module provides registries for Substrait extensions:
6//! - **Global Registry**: Immutable, reusable across plans, URI+name based lookup
7//! - **Local Registry**: Per-plan, anchor-based, references Global Registry (TODO)
8//!
9//! Currently only type definitions are supported. Function parsing will be added in a future update.
10//!
11//! This module is only available when the `parse` feature is enabled.
12
13use std::collections::{HashMap, hash_map::Entry};
14
15use super::{ExtensionFile, SimpleExtensions, SimpleExtensionsError, types::CustomType};
16use crate::urn::Urn;
17
18/// Extension Registry that manages Substrait extensions
19///
20/// This registry is immutable and reusable across multiple plans.
21/// It provides URN + name based lookup for extension types. Function parsing will be added in a future update.
22#[derive(Debug)]
23pub struct Registry {
24    /// Pre-validated extension files
25    extensions: HashMap<Urn, SimpleExtensions>,
26}
27
28impl Registry {
29    /// Create a new Global Registry from validated extension files.
30    ///
31    /// Any duplicate URNs will raise an error.
32    pub fn new<I: IntoIterator<Item = ExtensionFile>>(
33        extensions: I,
34    ) -> Result<Self, SimpleExtensionsError> {
35        let mut map = HashMap::new();
36        for ExtensionFile { urn, extension } in extensions {
37            match map.entry(urn.clone()) {
38                Entry::Occupied(_) => return Err(SimpleExtensionsError::DuplicateUrn(urn)),
39                Entry::Vacant(entry) => {
40                    entry.insert(extension);
41                }
42            }
43        }
44        Ok(Self { extensions: map })
45    }
46
47    /// Get an iterator over all extension files in this registry
48    pub fn extensions(&self) -> impl Iterator<Item = (&Urn, &SimpleExtensions)> {
49        self.extensions.iter()
50    }
51
52    /// Create a Global Registry from the built-in core extensions.
53    #[cfg(feature = "extensions")]
54    pub fn from_core_extensions() -> Self {
55        use crate::extensions::EXTENSIONS;
56
57        // Parse the core extensions from the raw extensions format to the parsed format
58        let extensions: HashMap<Urn, SimpleExtensions> = EXTENSIONS
59            .iter()
60            .map(|(orig_urn, simple_extensions)| {
61                let ExtensionFile { urn, extension } = ExtensionFile::create(simple_extensions.clone())
62                    .unwrap_or_else(|err| panic!("Core extensions should be valid, but failed to create extension file for {orig_urn}: {err}"));
63                debug_assert_eq!(orig_urn, &urn);
64                (urn, extension)
65            })
66            .collect();
67
68        Self { extensions }
69    }
70
71    fn get_extension(&self, urn: &Urn) -> Option<&SimpleExtensions> {
72        self.extensions.get(urn)
73    }
74
75    /// Get a type by URN and name
76    pub fn get_type(&self, urn: &Urn, name: &str) -> Option<&CustomType> {
77        self.get_extension(urn)?.get_type(name)
78    }
79
80    /// Get a scalar function by URN and name.
81    ///
82    /// TODO: Add support for retrieving functions by their full signature shorthand
83    /// (e.g., "add:i32_i32").
84    pub fn get_scalar_function(&self, urn: &Urn, name: &str) -> Option<&super::ScalarFunction> {
85        self.get_extension(urn)?.get_scalar_function(name)
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::{ExtensionFile, Registry};
92    use crate::parse::text::simple_extensions::{
93        SimpleExtensionsError, scalar_functions::ScalarFunctionError, types::ExtensionTypeError,
94    };
95    use crate::text::simple_extensions::{SimpleExtensions, SimpleExtensionsTypesItem};
96    use crate::urn::Urn;
97    use std::str::FromStr;
98
99    fn extension_file(urn: &str, type_names: &[&str]) -> ExtensionFile {
100        let types = type_names
101            .iter()
102            .map(|name| SimpleExtensionsTypesItem {
103                name: (*name).to_string(),
104                deprecated: None,
105                description: None,
106                metadata: Default::default(),
107                parameters: None,
108                structure: None,
109                variadic: None,
110            })
111            .collect();
112
113        let raw = SimpleExtensions {
114            scalar_functions: vec![],
115            aggregate_functions: vec![],
116            window_functions: vec![],
117            dependencies: Default::default(),
118            metadata: Default::default(),
119            type_variations: vec![],
120            types,
121            urn: urn.to_string(),
122        };
123
124        ExtensionFile::create(raw).expect("valid extension file")
125    }
126
127    #[test]
128    fn test_registry_iteration() {
129        let urns = vec![
130            "extension:example.com:first",
131            "extension:example.com:second",
132        ];
133        let registry =
134            Registry::new(urns.iter().map(|&urn| extension_file(urn, &["type"]))).unwrap();
135
136        let collected: Vec<&Urn> = registry.extensions().map(|(urn, _)| urn).collect();
137        assert_eq!(collected.len(), 2);
138        for urn in urns {
139            assert!(
140                collected
141                    .iter()
142                    .any(|candidate| candidate.to_string() == urn)
143            );
144        }
145    }
146
147    #[test]
148    fn test_type_lookup() {
149        let urn = Urn::from_str("extension:example.com:test").unwrap();
150        let registry =
151            Registry::new(vec![extension_file(&urn.to_string(), &["test_type"])]).unwrap();
152        let other_urn = Urn::from_str("extension:example.com:other").unwrap();
153
154        let cases = vec![
155            (&urn, "test_type", true),
156            (&urn, "missing", false),
157            (&other_urn, "test_type", false),
158        ];
159
160        for (query_urn, type_name, expected) in cases {
161            assert_eq!(
162                registry.get_type(query_urn, type_name).is_some(),
163                expected,
164                "unexpected lookup result for {query_urn}:{type_name}"
165            );
166        }
167    }
168
169    #[cfg(feature = "extensions")]
170    #[test]
171    fn test_from_core_extensions() {
172        let registry = Registry::from_core_extensions();
173        assert!(registry.extensions().count() > 0);
174
175        // Test that functions_geometry.yaml loaded correctly with its geometry type
176        let urn = Urn::from_str("extension:io.substrait:functions_geometry").unwrap();
177        let core_extension = registry
178            .get_extension(&urn)
179            .expect("Should find functions_geometry extension");
180
181        let geometry_type = core_extension.get_type("geometry");
182        assert!(
183            geometry_type.is_some(),
184            "Should find 'geometry' type in functions_geometry extension"
185        );
186
187        // Also test the registry's get_type method with the actual URN
188        let type_via_registry = registry.get_type(&urn, "geometry");
189        assert!(type_via_registry.is_some());
190
191        // `unsigned_integers` was added to the catalog in Substrait v0.101.0.
192        let unsigned_integers_urn =
193            Urn::from_str("extension:io.substrait:unsigned_integers").unwrap();
194        assert!(
195            registry.get_extension(&unsigned_integers_urn).is_some(),
196            "unsigned_integers should be a core extension"
197        );
198    }
199
200    #[test]
201    fn test_unknown_type_without_prefix_fails() {
202        use crate::text::simple_extensions;
203
204        // Function that references a type without u! prefix - should fail with UnknownTypeName
205        let invalid_extension = SimpleExtensions {
206            scalar_functions: vec![simple_extensions::ScalarFunction {
207                name: "bad_function".to_string(),
208                description: None,
209                metadata: Default::default(),
210                deprecated: None,
211                impls: vec![simple_extensions::ScalarFunctionImplsItem {
212                    args: None,
213                    deprecated: None,
214                    description: None,
215                    options: None,
216                    variadic: None,
217                    session_dependent: None,
218                    deterministic: None,
219                    nullability: None,
220                    return_: simple_extensions::ReturnValue(simple_extensions::Type::String(
221                        "point".to_string(), // Missing u! prefix - this is an error, not NYI
222                    )),
223                    implementation: None,
224                }],
225            }],
226            aggregate_functions: vec![],
227            window_functions: vec![],
228            dependencies: Default::default(),
229            metadata: Default::default(),
230            type_variations: vec![],
231            types: vec![],
232            urn: "extension:example.com:invalid".to_string(),
233        };
234
235        let result = ExtensionFile::create(invalid_extension);
236        assert!(
237            result.is_err(),
238            "Should fail when type is missing u! prefix"
239        );
240
241        match result {
242            Err(SimpleExtensionsError::ScalarFunctionError(ScalarFunctionError::TypeError(
243                ExtensionTypeError::UnknownTypeName { name },
244            ))) => {
245                assert_eq!(name, "point");
246            }
247            other => panic!("Expected UnknownTypeName error, got {:?}", other),
248        }
249    }
250
251    /// Helper to create a minimal extension with a scalar function returning a custom type
252    fn extension_with_custom_type_reference(
253        urn: &str,
254        function_name: &str,
255        return_type: &str,
256        defined_types: Vec<&str>,
257    ) -> SimpleExtensions {
258        use crate::text::simple_extensions;
259
260        SimpleExtensions {
261            scalar_functions: vec![simple_extensions::ScalarFunction {
262                name: function_name.to_string(),
263                description: None,
264                metadata: Default::default(),
265                deprecated: None,
266                impls: vec![simple_extensions::ScalarFunctionImplsItem {
267                    args: None,
268                    deprecated: None,
269                    description: None,
270                    options: None,
271                    variadic: None,
272                    session_dependent: None,
273                    deterministic: None,
274                    nullability: None,
275                    return_: simple_extensions::ReturnValue(simple_extensions::Type::String(
276                        return_type.to_string(),
277                    )),
278                    implementation: None,
279                }],
280            }],
281            aggregate_functions: vec![],
282            window_functions: vec![],
283            dependencies: Default::default(),
284            metadata: Default::default(),
285            type_variations: vec![],
286            types: defined_types
287                .into_iter()
288                .map(|name| SimpleExtensionsTypesItem {
289                    name: name.to_string(),
290                    deprecated: None,
291                    description: None,
292                    metadata: Default::default(),
293                    parameters: None,
294                    structure: None,
295                    variadic: None,
296                })
297                .collect(),
298            urn: urn.to_string(),
299        }
300    }
301
302    #[test]
303    fn test_custom_type_reference_valid() {
304        let extension = extension_with_custom_type_reference(
305            "extension:example.com:valid",
306            "get_point",
307            "u!point",
308            vec!["point"],
309        );
310
311        let result = ExtensionFile::create(extension);
312        assert!(
313            result.is_ok(),
314            "Should succeed when referenced type exists with u! prefix"
315        );
316    }
317
318    #[test]
319    fn test_custom_type_reference_missing() {
320        let extension = extension_with_custom_type_reference(
321            "extension:example.com:invalid",
322            "get_rectangle",
323            "u!rectangle",
324            vec![], // rectangle type not defined
325        );
326
327        let result = ExtensionFile::create(extension);
328        assert!(
329            result.is_err(),
330            "Should fail when referenced type doesn't exist"
331        );
332
333        match result {
334            Err(SimpleExtensionsError::UnresolvedTypeReference { type_name }) => {
335                assert_eq!(type_name, "rectangle");
336            }
337            other => panic!("Expected UnresolvedTypeReference error, got {:?}", other),
338        }
339    }
340
341    #[cfg(feature = "extensions")]
342    #[test]
343    fn test_scalar_function_parses_completely() {
344        use super::super::{
345            argument::ArgumentsItem,
346            scalar_functions::{Impl, NullabilityHandling, Options},
347            types::*,
348        };
349        use crate::parse::Parse;
350        use crate::text::simple_extensions;
351        use std::collections::HashMap;
352
353        let registry = Registry::from_core_extensions();
354        let functions_arithmetic_urn =
355            Urn::from_str("extension:io.substrait:functions_arithmetic").unwrap();
356
357        let add = registry
358            .get_scalar_function(&functions_arithmetic_urn, "add")
359            .expect("add function should exist");
360
361        // Verify function-level metadata
362        assert_eq!(add.name, "add");
363        assert_eq!(add.description, Some("Add two values.".to_string()));
364        assert!(
365            !add.impls.is_empty(),
366            "add should have at least one implementation"
367        );
368
369        // Create the expected first implementation (i8 + i8 -> i8)
370        let mut ctx = super::super::extensions::TypeContext::default();
371        let expected_impl = Impl {
372            args: vec![
373                ArgumentsItem::ValueArgument(
374                    simple_extensions::ValueArg {
375                        name: Some("x".to_string()),
376                        description: None,
377                        value: simple_extensions::Type::String("i8".to_string()),
378                        constant: None,
379                    }
380                    .parse(&mut ctx)
381                    .unwrap(),
382                ),
383                ArgumentsItem::ValueArgument(
384                    simple_extensions::ValueArg {
385                        name: Some("y".to_string()),
386                        description: None,
387                        value: simple_extensions::Type::String("i8".to_string()),
388                        constant: None,
389                    }
390                    .parse(&mut ctx)
391                    .unwrap(),
392                ),
393            ],
394            options: Options({
395                let mut map = HashMap::new();
396                map.insert(
397                    "overflow".to_string(),
398                    vec![
399                        "SILENT".to_string(),
400                        "SATURATE".to_string(),
401                        "ERROR".to_string(),
402                    ],
403                );
404                map
405            }),
406            variadic: None,
407            session_dependent: false,
408            deterministic: true,
409            nullability: NullabilityHandling::Mirror,
410            return_type: ConcreteType {
411                kind: ConcreteTypeKind::Builtin(BasicBuiltinType::I8),
412                nullable: false,
413            },
414            implementation: HashMap::new(),
415        };
416
417        assert_eq!(&add.impls[0], &expected_impl);
418    }
419}
substrait/parse/text/simple_extensions/registry.rs

substrait/parse/text/simple_extensions/
registry.rs
