substrait_explain/textify/
extensions.rs

1//! Extension textification support
2//!
3//! This module provides [`Textify`] implementations for extension-related
4//! types, including [`ExtensionValue`], [`ExtensionColumn`], [`ExtensionArgs`],
5//! and the various extension relation types ([`substrait::proto::ExtensionLeafRel`],
6//! [`substrait::proto::ExtensionSingleRel`], [`substrait::proto::ExtensionMultiRel`]).
7
8use std::fmt;
9
10use substrait::proto::extensions::AdvancedExtension;
11
12use crate::FormatError;
13use crate::extensions::any::AnyRef;
14use crate::extensions::registry::ExtensionType;
15use crate::extensions::{ExtensionArgs, ExtensionColumn, ExtensionValue, TupleValue};
16use crate::textify::foundation::{PlanError, Scope, Textify};
17use crate::textify::types::escaped;
18
19impl Textify for TupleValue {
20    fn name() -> &'static str {
21        "TupleValue"
22    }
23
24    fn textify<S: Scope, W: fmt::Write>(&self, ctx: &S, w: &mut W) -> fmt::Result {
25        write!(w, "(")?;
26        if self.len() == 1 {
27            self.iter().next().unwrap().textify(ctx, w)?;
28            write!(w, ",")?;
29        } else {
30            write!(w, "{}", ctx.separated(self, ", "))?;
31        }
32        write!(w, ")")
33    }
34}
35
36impl Textify for ExtensionValue {
37    fn name() -> &'static str {
38        "ExtensionValue"
39    }
40
41    fn textify<S: Scope, W: fmt::Write>(&self, ctx: &S, w: &mut W) -> fmt::Result {
42        match self {
43            ExtensionValue::String(s) => write!(w, "'{}'", escaped(s)),
44            ExtensionValue::Integer(i) => write!(w, "{i}"),
45            ExtensionValue::Float(f) => write!(w, "{f}"),
46            ExtensionValue::Boolean(b) => write!(w, "{b}"),
47            ExtensionValue::Reference(r) => write!(w, "${r}"),
48            ExtensionValue::Enum(e) => write!(w, "&{e}"),
49            ExtensionValue::Tuple(tv) => tv.textify(ctx, w),
50            ExtensionValue::Expression(e) => write!(w, "{e}"),
51        }
52    }
53}
54
55impl Textify for ExtensionColumn {
56    fn name() -> &'static str {
57        "ExtensionColumn"
58    }
59
60    fn textify<S: Scope, W: fmt::Write>(&self, _ctx: &S, w: &mut W) -> fmt::Result {
61        match self {
62            ExtensionColumn::Named { name, type_spec } => write!(w, "{name}:{type_spec}"),
63            ExtensionColumn::Reference(r) => write!(w, "${r}"),
64            ExtensionColumn::Expression(e) => write!(w, "{e}"),
65        }
66    }
67}
68
69impl Textify for ExtensionArgs {
70    fn name() -> &'static str {
71        "ExtensionArgs"
72    }
73
74    fn textify<S: Scope, W: fmt::Write>(&self, ctx: &S, w: &mut W) -> fmt::Result {
75        let mut has_args = false;
76
77        // Add positional arguments
78        for (i, value) in self.positional.iter().enumerate() {
79            if i > 0 || has_args {
80                write!(w, ", ")?;
81            }
82            value.textify(ctx, w)?;
83            has_args = true;
84        }
85
86        // Add named arguments in display order (IndexMap preserves insertion order)
87        for (name, value) in &self.named {
88            if has_args {
89                write!(w, ", ")?;
90            }
91            write!(w, "{name}=")?;
92            value.textify(ctx, w)?;
93            has_args = true;
94        }
95
96        if !has_args {
97            write!(w, "_")?;
98        }
99
100        // Add output columns if present
101        if !self.output_columns.is_empty() {
102            write!(w, " => {}", ctx.separated(self.output_columns.iter(), ", "))?;
103        }
104
105        Ok(())
106    }
107}
108
109/// Textify a single enhancement or optimization line.
110///
111/// Emits one of:
112/// - `{indent}+ Enh:Name[args]`
113/// - `{indent}+ Opt:Name[args]`
114fn format_adv_ext_line<S: Scope, W: fmt::Write>(
115    ctx: &S,
116    w: &mut W,
117    ext_type: ExtensionType,
118    detail: AnyRef<'_>,
119) -> fmt::Result {
120    let indent = ctx.indent();
121    let registry = ctx.extension_registry();
122    let (prefix, decode_result) = match ext_type {
123        ExtensionType::Enhancement => ("Enh", registry.decode_enhancement(detail)),
124        ExtensionType::Optimization => ("Opt", registry.decode_optimization(detail)),
125        ExtensionType::Relation => unreachable!("Relation extensions don't use adv_ext lines"),
126    };
127    match decode_result {
128        Ok((name, args)) => {
129            if !args.output_columns.is_empty() {
130                write!(
131                    w,
132                    "{indent}+ {prefix}[{}]",
133                    ctx.failure(FormatError::Format(PlanError::invalid(
134                        "adv_extension",
135                        Some(name),
136                        "output_columns cannot be represented in adv_extension syntax",
137                    )))
138                )
139            } else {
140                write!(w, "{indent}+ {prefix}:{name}[{}]", ctx.display(&args))
141            }
142        }
143        Err(error) => {
144            write!(w, "{indent}+ {prefix}[{}]", ctx.failure(error))
145        }
146    }
147}
148
149impl Textify for AdvancedExtension {
150    fn name() -> &'static str {
151        "AdvancedExtension"
152    }
153
154    /// Textify all enhancement and optimization lines for an [`AdvancedExtension`].
155    ///
156    /// Writes one `+ Enh:` line (if an enhancement is present) followed by zero
157    /// or more `+ Opt:` lines, each preceded by a newline.
158    fn textify<S: Scope, W: fmt::Write>(&self, ctx: &S, w: &mut W) -> fmt::Result {
159        if let Some(enhancement) = &self.enhancement {
160            writeln!(w)?;
161            format_adv_ext_line(
162                ctx,
163                w,
164                ExtensionType::Enhancement,
165                AnyRef::from(enhancement),
166            )?;
167        }
168        for optimization in &self.optimization {
169            writeln!(w)?;
170            format_adv_ext_line(
171                ctx,
172                w,
173                ExtensionType::Optimization,
174                AnyRef::from(optimization),
175            )?;
176        }
177        Ok(())
178    }
179}
substrait_explain/textify/extensions.rs

substrait_explain/textify/
extensions.rs
