substrait_explain/extensions/
args.rs

1//! Text-format data structures used by registered advanced extension handlers.
2//!
3//! These types describe the arguments accepted by custom relation types,
4//! enhancements, and optimization hints. Relation extensions can additionally
5//! describe output columns.
6//!
7//! The interface presented to extension handlers is structured rather than
8//! textual: handlers read and write values such as [`ExtensionArgs`], [`Expr`],
9//! and [`proto::Type`]. `substrait-explain` handles the surrounding
10//! parsing/textification. Some values need plan context before they reach a
11//! handler; for example, an expression argument like `add($0, $1)` is parsed
12//! using [`SimpleExtensions`](crate::extensions::SimpleExtensions) to resolve
13//! the text function name to the protobuf function anchor, and formatted by
14//! resolving that anchor back to a text name.
15//!
16//! The extension-facing interface for Substrait objects (e.g. [`proto::Type`])
17//! should map directly to Substrait protobuf concepts. Sometimes that means
18//! storing the protobuf type directly, as named output columns do with
19//! [`proto::Type`]; sometimes it means using a small wrapper, as
20//! expression-compatible arguments do with [`Expr`] around
21//! [`proto::Expression`].
22//!
23//! Untyped scalar literals (e.g. `2`, `2.435`, `'string'`) are kept as
24//! extension scalar values so text rendering can preserve scalar syntax even in
25//! verbose output, while handlers that accept expressions can still widen them
26//! into default Substrait literal expressions.
27
28use std::collections::HashSet;
29use std::fmt;
30
31use indexmap::IndexMap;
32use substrait::proto;
33use substrait::proto::expression::field_reference::ReferenceType;
34use substrait::proto::expression::literal::LiteralType;
35use substrait::proto::expression::{RexType, reference_segment};
36
37use super::ExtensionError;
38use crate::textify::expressions::Reference;
39
40/// Kind of relation addendum in the text format.
41///
42/// Addenda are `+`-prefixed lines attached to relations. They are syntax-level
43/// constructs, distinct from [`crate::extensions::registry::ExtensionType`],
44/// which describes registry namespaces.
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub(crate) enum AddendumKind {
47    Enhancement,
48    Optimization,
49}
50
51impl AddendumKind {
52    pub(crate) fn prefix(self) -> &'static str {
53        match self {
54            AddendumKind::Enhancement => "Enh",
55            AddendumKind::Optimization => "Opt",
56        }
57    }
58}
59
60/// A Substrait expression carried as an extension argument or output column.
61///
62/// Boxed because `proto::Expression` is large (multiple `Vec` fields in
63/// variants like `ScalarFunction`).
64#[derive(Debug, Clone)]
65pub struct Expr(Box<proto::Expression>);
66
67impl Expr {
68    /// Borrow the underlying Substrait expression protobuf.
69    pub fn as_proto(&self) -> &proto::Expression {
70        self.0.as_ref()
71    }
72
73    /// Clone the underlying Substrait expression protobuf.
74    pub fn to_proto(&self) -> proto::Expression {
75        self.as_proto().clone()
76    }
77
78    /// If this expression is a direct field reference (`$N`), return it.
79    pub fn as_direct_reference(&self) -> Option<Reference> {
80        let Some(RexType::Selection(field_ref)) = self.as_proto().rex_type.as_ref() else {
81            return None;
82        };
83        let Some(ReferenceType::DirectReference(segment)) = field_ref.reference_type.as_ref()
84        else {
85            return None;
86        };
87        let Some(reference_segment::ReferenceType::StructField(field)) =
88            segment.reference_type.as_ref()
89        else {
90            return None;
91        };
92        if field.child.is_some() {
93            return None;
94        }
95        Some(Reference(field.field))
96    }
97}
98
99impl From<proto::Expression> for Expr {
100    fn from(expr: proto::Expression) -> Self {
101        Expr(Box::new(expr))
102    }
103}
104
105impl From<proto::expression::Literal> for Expr {
106    fn from(literal: proto::expression::Literal) -> Self {
107        proto::Expression {
108            rex_type: Some(RexType::Literal(literal)),
109        }
110        .into()
111    }
112}
113
114impl From<Reference> for Expr {
115    fn from(reference: Reference) -> Self {
116        proto::Expression::from(reference).into()
117    }
118}
119
120impl From<Expr> for proto::Expression {
121    fn from(expr: Expr) -> Self {
122        *expr.0
123    }
124}
125
126impl From<i64> for Expr {
127    fn from(value: i64) -> Self {
128        proto::expression::Literal {
129            literal_type: Some(LiteralType::I64(value)),
130            nullable: false,
131            type_variation_reference: 0,
132        }
133        .into()
134    }
135}
136
137impl From<f64> for Expr {
138    fn from(value: f64) -> Self {
139        proto::expression::Literal {
140            literal_type: Some(LiteralType::Fp64(value)),
141            nullable: false,
142            type_variation_reference: 0,
143        }
144        .into()
145    }
146}
147
148impl From<bool> for Expr {
149    fn from(value: bool) -> Self {
150        proto::expression::Literal {
151            literal_type: Some(LiteralType::Boolean(value)),
152            nullable: false,
153            type_variation_reference: 0,
154        }
155        .into()
156    }
157}
158
159impl From<String> for Expr {
160    fn from(value: String) -> Self {
161        proto::expression::Literal {
162            literal_type: Some(LiteralType::String(value)),
163            nullable: false,
164            type_variation_reference: 0,
165        }
166        .into()
167    }
168}
169
170impl From<&str> for Expr {
171    fn from(value: &str) -> Self {
172        value.to_string().into()
173    }
174}
175
176/// Represents extension arguments plus optional output columns.
177///
178/// Named arguments are stored in an [`IndexMap`] whose iteration order
179/// determines display order. Extension [`super::Explainable::to_args()`]
180/// implementations should insert named arguments in the order they should
181/// appear in the text format.
182#[derive(Debug, Clone, Default)]
183pub struct ExtensionArgs {
184    /// Positional arguments.
185    pub positional: Vec<ExtensionValue>,
186    /// Named arguments, displayed in the order they were inserted
187    pub named: IndexMap<String, ExtensionValue>,
188    /// Output columns for custom relation types.
189    pub output_columns: Vec<ExtensionColumn>,
190}
191
192/// Helper struct for extracting named arguments with validation.
193///
194/// Tracks which arguments have been consumed. Callers **must** call
195/// [`check_exhausted`](ArgsExtractor::check_exhausted) before dropping to
196/// verify no unexpected arguments remain. In debug builds, dropping without
197/// calling `check_exhausted` will panic (matching the [`RuleIter`](crate::parser::RuleIter) pattern).
198pub struct ArgsExtractor<'a> {
199    args: &'a ExtensionArgs,
200    consumed: HashSet<&'a str>,
201    checked: bool,
202}
203
204impl<'a> ArgsExtractor<'a> {
205    /// Create a new extractor for the given arguments
206    pub fn new(args: &'a ExtensionArgs) -> Self {
207        Self {
208            args,
209            consumed: HashSet::new(),
210            checked: false,
211        }
212    }
213
214    /// Get a named argument value, marking it as consumed if found.
215    pub fn get_named_arg(&mut self, name: &str) -> Option<&'a ExtensionValue> {
216        match self.args.named.get_key_value(name) {
217            Some((k, value)) => {
218                self.consumed.insert(k);
219                Some(value)
220            }
221            None => None,
222        }
223    }
224
225    /// Get a named argument value or return an error
226    /// Marks the argument as consumed if found
227    pub fn expect_named_arg<T>(&mut self, name: &str) -> Result<T, ExtensionError>
228    where
229        T: TryFrom<&'a ExtensionValue>,
230        T::Error: Into<ExtensionError>,
231    {
232        match self.get_named_arg(name) {
233            Some(value) => T::try_from(value).map_err(Into::into),
234            None => Err(ExtensionError::MissingArgument {
235                name: name.to_string(),
236            }),
237        }
238    }
239
240    /// Get a named argument value or default
241    /// Marks the argument as consumed if it exists in the source args
242    pub fn get_named_or<T>(&mut self, name: &str, default: T) -> Result<T, ExtensionError>
243    where
244        T: TryFrom<&'a ExtensionValue>,
245        T::Error: Into<ExtensionError>,
246    {
247        match self.get_named_arg(name) {
248            Some(value) => T::try_from(value).map_err(Into::into),
249            None => Ok(default),
250        }
251    }
252
253    /// Check that all named arguments in the source have been consumed,
254    /// returning an error if not.
255    ///
256    /// Must be called before the extractor is dropped, to validate that all
257    /// args are correctly handled. In debug builds, dropping without calling
258    /// this method will panic.
259    pub fn check_exhausted(&mut self) -> Result<(), ExtensionError> {
260        self.checked = true;
261
262        let mut unknown_args = Vec::new();
263        for name in self.args.named.keys() {
264            if !self.consumed.contains(name.as_str()) {
265                unknown_args.push(name.as_str());
266            }
267        }
268
269        if unknown_args.is_empty() {
270            Ok(())
271        } else {
272            // Sort for stable error messages
273            unknown_args.sort();
274            Err(ExtensionError::InvalidArgument(format!(
275                "Unknown named arguments: {}",
276                unknown_args.join(", ")
277            )))
278        }
279    }
280}
281
282impl Drop for ArgsExtractor<'_> {
283    fn drop(&mut self) {
284        if self.checked || std::thread::panicking() {
285            return;
286        }
287        // If we get here, the caller forgot to call check_exhausted().
288        debug_assert!(
289            false,
290            "ArgsExtractor dropped without calling check_exhausted()"
291        );
292    }
293}
294
295/// A tuple-valued extension argument.
296///
297/// Tuple values preserve positional order and can be iterated by value or by
298/// reference.
299#[derive(Debug, Clone)]
300pub struct TupleValue(Vec<ExtensionValue>);
301
302impl TupleValue {
303    pub fn len(&self) -> usize {
304        self.0.len()
305    }
306
307    pub fn is_empty(&self) -> bool {
308        self.0.is_empty()
309    }
310
311    pub fn iter(&self) -> std::slice::Iter<'_, ExtensionValue> {
312        self.0.iter()
313    }
314}
315
316impl<'a> IntoIterator for &'a TupleValue {
317    type Item = &'a ExtensionValue;
318    type IntoIter = std::slice::Iter<'a, ExtensionValue>;
319
320    fn into_iter(self) -> Self::IntoIter {
321        self.0.iter()
322    }
323}
324
325impl IntoIterator for TupleValue {
326    type Item = ExtensionValue;
327    type IntoIter = std::vec::IntoIter<ExtensionValue>;
328
329    fn into_iter(self) -> Self::IntoIter {
330        self.0.into_iter()
331    }
332}
333
334impl FromIterator<ExtensionValue> for TupleValue {
335    fn from_iter<I: IntoIterator<Item = ExtensionValue>>(iter: I) -> Self {
336        TupleValue(iter.into_iter().collect())
337    }
338}
339
340impl From<Vec<ExtensionValue>> for TupleValue {
341    fn from(items: Vec<ExtensionValue>) -> Self {
342        TupleValue(items)
343    }
344}
345
346/// Represents a value in extension arguments.
347///
348/// These values are the structured form of text-format extension arguments,
349/// fully resolved - i.e. any additional context (such as function anchors etc)
350/// are part of this struct itself.
351#[derive(Debug, Clone)]
352pub enum ExtensionValue {
353    /// Untyped literals. These are not input or output with types (e.g. `2`,
354    /// not `2:i64`), and suitable for protobuf extension fields that are not
355    /// substrait types.
356    String(String),
357    Integer(i64),
358    Float(f64),
359    Boolean(bool),
360
361    /// Substrait expression value, including typed literals and field references.
362    ///
363    /// Use `TryFrom<&ExtensionValue> for Expr` when a handler accepts either an
364    /// expression or a scalar value widened into an expression.
365    Expr(Expr),
366    /// Enum value (e.g. &CORE, &Inner) — the string holds the identifier
367    /// without the `&` prefix
368    Enum(String),
369    /// Tuple of values, e.g. (&HASH, &RANGE) or (42, 'hello')
370    Tuple(TupleValue),
371    // TODO: Consider adding support for types as arguments. May need dedicated
372    // syntax (`:typename`, perhaps?), as type names may not be distinguishable
373    // from identifiers
374}
375
376/// The variant kind of an [`ExtensionValue`], used in diagnostics.
377#[derive(Debug, Clone, Copy, PartialEq, Eq)]
378pub enum ExtensionValueKind {
379    String,
380    Integer,
381    Float,
382    Boolean,
383    Reference,
384    Enum,
385    Tuple,
386    Expression,
387}
388
389impl fmt::Display for ExtensionValueKind {
390    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
391        match self {
392            ExtensionValueKind::String => write!(f, "string"),
393            ExtensionValueKind::Integer => write!(f, "integer"),
394            ExtensionValueKind::Float => write!(f, "float"),
395            ExtensionValueKind::Boolean => write!(f, "boolean"),
396            ExtensionValueKind::Reference => write!(f, "reference"),
397            ExtensionValueKind::Enum => write!(f, "enum"),
398            ExtensionValueKind::Tuple => write!(f, "tuple"),
399            ExtensionValueKind::Expression => write!(f, "expression"),
400        }
401    }
402}
403
404impl ExtensionValue {
405    /// Return the variant kind of this value for structured diagnostics.
406    pub fn kind(&self) -> ExtensionValueKind {
407        match self {
408            ExtensionValue::String(_) => ExtensionValueKind::String,
409            ExtensionValue::Integer(_) => ExtensionValueKind::Integer,
410            ExtensionValue::Float(_) => ExtensionValueKind::Float,
411            ExtensionValue::Boolean(_) => ExtensionValueKind::Boolean,
412            ExtensionValue::Expr(_) => ExtensionValueKind::Expression,
413            ExtensionValue::Enum(_) => ExtensionValueKind::Enum,
414            ExtensionValue::Tuple(_) => ExtensionValueKind::Tuple,
415        }
416    }
417}
418
419impl From<Expr> for ExtensionValue {
420    fn from(expr: Expr) -> Self {
421        ExtensionValue::Expr(expr)
422    }
423}
424
425impl From<proto::Expression> for ExtensionValue {
426    fn from(expr: proto::Expression) -> Self {
427        Expr::from(expr).into()
428    }
429}
430
431impl From<proto::expression::Literal> for ExtensionValue {
432    fn from(literal: proto::expression::Literal) -> Self {
433        Expr::from(literal).into()
434    }
435}
436
437impl From<Reference> for ExtensionValue {
438    fn from(reference: Reference) -> Self {
439        Expr::from(reference).into()
440    }
441}
442
443impl From<i64> for ExtensionValue {
444    fn from(value: i64) -> Self {
445        ExtensionValue::Integer(value)
446    }
447}
448
449impl From<f64> for ExtensionValue {
450    fn from(value: f64) -> Self {
451        ExtensionValue::Float(value)
452    }
453}
454
455impl From<bool> for ExtensionValue {
456    fn from(value: bool) -> Self {
457        ExtensionValue::Boolean(value)
458    }
459}
460
461impl From<String> for ExtensionValue {
462    fn from(value: String) -> Self {
463        ExtensionValue::String(value)
464    }
465}
466
467impl From<&str> for ExtensionValue {
468    fn from(value: &str) -> Self {
469        ExtensionValue::String(value.to_string())
470    }
471}
472
473fn invalid_type(expected: ExtensionValueKind, actual: &ExtensionValue) -> ExtensionError {
474    ExtensionError::InvalidArgumentType {
475        expected,
476        actual: actual.kind(),
477    }
478}
479
480impl<'a> TryFrom<&'a ExtensionValue> for &'a str {
481    type Error = ExtensionError;
482
483    fn try_from(value: &'a ExtensionValue) -> Result<&'a str, Self::Error> {
484        match value {
485            ExtensionValue::String(s) => Ok(s),
486            v => Err(invalid_type(ExtensionValueKind::String, v)),
487        }
488    }
489}
490
491impl TryFrom<ExtensionValue> for String {
492    type Error = ExtensionError;
493
494    fn try_from(value: ExtensionValue) -> Result<String, Self::Error> {
495        <&str>::try_from(&value).map(ToOwned::to_owned)
496    }
497}
498
499/// Helper for extracting the identifier from an [`ExtensionValue::Enum`].
500pub struct EnumValue(pub String);
501
502impl<'a> TryFrom<&'a ExtensionValue> for EnumValue {
503    type Error = ExtensionError;
504
505    fn try_from(value: &'a ExtensionValue) -> Result<EnumValue, Self::Error> {
506        match value {
507            ExtensionValue::Enum(s) => Ok(EnumValue(s.clone())),
508            v => Err(invalid_type(ExtensionValueKind::Enum, v)),
509        }
510    }
511}
512
513impl<'a> TryFrom<&'a ExtensionValue> for &'a TupleValue {
514    type Error = ExtensionError;
515
516    fn try_from(value: &'a ExtensionValue) -> Result<&'a TupleValue, Self::Error> {
517        match value {
518            ExtensionValue::Tuple(tv) => Ok(tv),
519            v => Err(invalid_type(ExtensionValueKind::Tuple, v)),
520        }
521    }
522}
523
524impl TryFrom<&ExtensionValue> for i64 {
525    type Error = ExtensionError;
526
527    fn try_from(value: &ExtensionValue) -> Result<i64, Self::Error> {
528        match value {
529            ExtensionValue::Integer(i) => Ok(*i),
530            v => Err(invalid_type(ExtensionValueKind::Integer, v)),
531        }
532    }
533}
534
535impl TryFrom<&ExtensionValue> for f64 {
536    type Error = ExtensionError;
537
538    fn try_from(value: &ExtensionValue) -> Result<f64, Self::Error> {
539        match value {
540            ExtensionValue::Float(f) => Ok(*f),
541            v => Err(invalid_type(ExtensionValueKind::Float, v)),
542        }
543    }
544}
545
546impl TryFrom<&ExtensionValue> for bool {
547    type Error = ExtensionError;
548
549    fn try_from(value: &ExtensionValue) -> Result<bool, Self::Error> {
550        match value {
551            ExtensionValue::Boolean(b) => Ok(*b),
552            v => Err(invalid_type(ExtensionValueKind::Boolean, v)),
553        }
554    }
555}
556
557impl TryFrom<&ExtensionValue> for Reference {
558    type Error = ExtensionError;
559
560    fn try_from(value: &ExtensionValue) -> Result<Reference, Self::Error> {
561        match value {
562            ExtensionValue::Expr(expr) => expr
563                .as_direct_reference()
564                .ok_or_else(|| invalid_type(ExtensionValueKind::Reference, value)),
565            v => Err(invalid_type(ExtensionValueKind::Reference, v)),
566        }
567    }
568}
569
570impl TryFrom<&ExtensionValue> for Expr {
571    type Error = ExtensionError;
572
573    fn try_from(value: &ExtensionValue) -> Result<Expr, Self::Error> {
574        match value {
575            ExtensionValue::Expr(e) => Ok(e.clone()),
576            // Untyped extension scalars are intentionally expression-compatible:
577            // `arg=2` carries no syntax that distinguishes "configuration
578            // integer" from "i64 literal expression". Scalar-specific
579            // extraction (`i64`, `&str`, `bool`, etc.) still requires the scalar
580            // variants, while expression extraction widens them to default
581            // non-nullable Substrait literal expressions.
582            ExtensionValue::Integer(i) => Ok(Expr::from(*i)),
583            ExtensionValue::Float(f) => Ok(Expr::from(*f)),
584            ExtensionValue::String(s) => Ok(Expr::from(s.as_str())),
585            ExtensionValue::Boolean(b) => Ok(Expr::from(*b)),
586            v => Err(invalid_type(ExtensionValueKind::Expression, v)),
587        }
588    }
589}
590
591/// Represents an output column specification.
592///
593/// These values mirror the text-format output column forms. Named columns keep
594/// the parsed Substrait type protobuf so handlers can convert directly to
595/// relation schemas.
596#[derive(Debug, Clone)]
597pub enum ExtensionColumn {
598    /// Named column with a parsed Substrait type (e.g. `name:i64?`).
599    Named {
600        /// Column name as it appears in the extension relation output.
601        name: String,
602        /// Parsed Substrait type for the column.
603        ///
604        /// This uses the protobuf field name, hence the raw identifier.
605        r#type: proto::Type,
606    },
607    /// Expression-compatible output column, including field references.
608    Expr(Expr),
609}
610
611impl ExtensionArgs {
612    /// Push a positional extension argument.
613    pub fn push<T>(&mut self, value: T)
614    where
615        T: Into<ExtensionValue>,
616    {
617        self.positional.push(value.into());
618    }
619
620    /// Insert a named extension argument, returning any previous value.
621    pub fn insert<K, V>(&mut self, name: K, value: V) -> Option<ExtensionValue>
622    where
623        K: Into<String>,
624        V: Into<ExtensionValue>,
625    {
626        self.named.insert(name.into(), value.into())
627    }
628
629    /// Create an extractor for validating named arguments
630    pub fn extractor(&self) -> ArgsExtractor<'_> {
631        ArgsExtractor::new(self)
632    }
633}
substrait_explain/extensions/args.rs

substrait_explain/extensions/
args.rs
