substrait/parse/proto/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Parsing of [proto](crate::proto) types.
4
5mod version;
6use std::collections::{HashMap, hash_map::Entry};
7
8pub use version::{Version, VersionError};
9
10mod plan_version;
11pub use plan_version::{PlanVersion, PlanVersionError};
12
13use crate::{
14    parse::{Anchor, context::ContextError, proto::extensions::SimpleExtensionUrn},
15    urn::Urn,
16};
17
18pub mod extensions;
19
20/// Tracks the known anchors for simple extensions, as needed for parsing a
21/// Substrait protobuf.
22///
23/// At the moment, this only tracks simple extension URNs by anchor, as further
24/// parsing of simple extensions, or the protobuf plan itself, is not supported
25/// yet.
26#[derive(Default)]
27pub struct ExtensionAnchors {
28    simple_extensions: HashMap<Anchor<SimpleExtensionUrn>, Urn>,
29}
30
31impl super::Context for ExtensionAnchors {}
32
33impl ExtensionAnchors {
34    /// Register a [SimpleExtensionUrn] with this registry, rejecting duplicate
35    /// anchors.
36    pub fn add_simple_extension_urn(
37        &mut self,
38        simple_extension_urn: &SimpleExtensionUrn,
39    ) -> Result<(), ContextError> {
40        let anchor = simple_extension_urn.anchor();
41        match self.simple_extensions.entry(anchor) {
42            Entry::Occupied(_) => Err(ContextError::DuplicateSimpleExtension(anchor)),
43            Entry::Vacant(entry) => {
44                entry.insert(simple_extension_urn.urn().clone());
45                Ok(())
46            }
47        }
48    }
49
50    /// Look up the [Urn] for a previously registered anchor. Returns an error
51    /// if the anchor has not been seen.
52    pub fn get_simple_extension_urn(&self, anchor: &Anchor<SimpleExtensionUrn>) -> Option<&Urn> {
53        self.simple_extensions.get(anchor)
54    }
55}