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::{
15    AddendumKind, Expr, ExtensionArgs, ExtensionColumn, ExtensionValue, TupleValue,
16};
17use crate::textify::foundation::{PlanError, Scope, Textify};
18use crate::textify::types::{Name, escaped};
19
20impl Textify for TupleValue {
21    fn name() -> &'static str {
22        "TupleValue"
23    }
24
25    fn textify<S: Scope, W: fmt::Write>(&self, ctx: &S, w: &mut W) -> fmt::Result {
26        write!(w, "(")?;
27        if self.len() == 1 {
28            self.iter().next().unwrap().textify(ctx, w)?;
29            write!(w, ",")?;
30        } else {
31            write!(w, "{}", ctx.separated(self, ", "))?;
32        }
33        write!(w, ")")
34    }
35}
36
37impl Textify for ExtensionValue {
38    fn name() -> &'static str {
39        "ExtensionValue"
40    }
41
42    fn textify<S: Scope, W: fmt::Write>(&self, ctx: &S, w: &mut W) -> fmt::Result {
43        match self {
44            ExtensionValue::String(s) => write!(w, "'{}'", escaped(s)),
45            ExtensionValue::Integer(i) => write!(w, "{i}"),
46            ExtensionValue::Float(f) => write!(w, "{f}"),
47            ExtensionValue::Boolean(b) => write!(w, "{b}"),
48            ExtensionValue::Expr(expr) => expr.textify(ctx, w),
49            ExtensionValue::Enum(e) => write!(w, "&{e}"),
50            ExtensionValue::Tuple(tv) => tv.textify(ctx, w),
51        }
52    }
53}
54
55impl Textify for Expr {
56    fn name() -> &'static str {
57        "Expr"
58    }
59
60    fn textify<S: Scope, W: fmt::Write>(&self, ctx: &S, w: &mut W) -> fmt::Result {
61        write!(w, "{}", ctx.display(self.as_proto()))
62    }
63}
64
65impl Textify for ExtensionColumn {
66    fn name() -> &'static str {
67        "ExtensionColumn"
68    }
69
70    fn textify<S: Scope, W: fmt::Write>(&self, ctx: &S, w: &mut W) -> fmt::Result {
71        match self {
72            ExtensionColumn::Named { name, r#type: ty } => {
73                write!(w, "{}:{}", Name(name), ctx.display(ty))
74            }
75            ExtensionColumn::Expr(expr) => expr.textify(ctx, w),
76        }
77    }
78}
79
80impl Textify for ExtensionArgs {
81    fn name() -> &'static str {
82        "ExtensionArgs"
83    }
84
85    fn textify<S: Scope, W: fmt::Write>(&self, ctx: &S, w: &mut W) -> fmt::Result {
86        let mut has_args = false;
87
88        // Add positional arguments
89        for (i, value) in self.positional.iter().enumerate() {
90            if i > 0 || has_args {
91                write!(w, ", ")?;
92            }
93            value.textify(ctx, w)?;
94            has_args = true;
95        }
96
97        // Add named arguments in display order (IndexMap preserves insertion order)
98        for (name, value) in &self.named {
99            if has_args {
100                write!(w, ", ")?;
101            }
102            write!(w, "{name}=")?;
103            value.textify(ctx, w)?;
104            has_args = true;
105        }
106
107        if !has_args {
108            write!(w, "_")?;
109        }
110
111        // Add output columns if present
112        if !self.output_columns.is_empty() {
113            write!(w, " => {}", ctx.separated(self.output_columns.iter(), ", "))?;
114        }
115
116        Ok(())
117    }
118}
119
120/// Textify a single enhancement or optimization line.
121///
122/// Successful lines include the registered extension name, e.g.
123/// `{indent}+ Enh:Name[args]`; decode failures fall back to a failure token so
124/// the surrounding relation can still be rendered.
125fn format_adv_ext_line<S: Scope, W: fmt::Write>(
126    ctx: &S,
127    w: &mut W,
128    kind: AddendumKind,
129    detail: AnyRef<'_>,
130) -> fmt::Result {
131    let indent = ctx.indent();
132    let registry = ctx.extension_registry();
133    let prefix = kind.prefix();
134    let decode_result = match kind {
135        AddendumKind::Enhancement => registry.decode_enhancement(detail),
136        AddendumKind::Optimization => registry.decode_optimization(detail),
137    };
138    match decode_result {
139        Ok((name, args)) => {
140            if !args.output_columns.is_empty() {
141                write!(
142                    w,
143                    "{indent}+ {prefix}[{}]",
144                    ctx.failure(FormatError::Format(PlanError::invalid(
145                        "adv_extension",
146                        Some(name),
147                        "output_columns cannot be represented in adv_extension syntax",
148                    )))
149                )
150            } else {
151                write!(w, "{indent}+ {prefix}:{name}[{}]", ctx.display(&args))
152            }
153        }
154        Err(error) => {
155            write!(w, "{indent}+ {prefix}[{}]", ctx.failure(error))
156        }
157    }
158}
159
160impl Textify for AdvancedExtension {
161    fn name() -> &'static str {
162        "AdvancedExtension"
163    }
164
165    /// Writes the enhancement line first, if present, followed by optimization
166    /// lines in protobuf order.
167    fn textify<S: Scope, W: fmt::Write>(&self, ctx: &S, w: &mut W) -> fmt::Result {
168        if let Some(enhancement) = &self.enhancement {
169            writeln!(w)?;
170            format_adv_ext_line(ctx, w, AddendumKind::Enhancement, AnyRef::from(enhancement))?;
171        }
172        for optimization in &self.optimization {
173            writeln!(w)?;
174            format_adv_ext_line(
175                ctx,
176                w,
177                AddendumKind::Optimization,
178                AnyRef::from(optimization),
179            )?;
180        }
181        Ok(())
182    }
183}
substrait_explain/textify/extensions.rs

substrait_explain/textify/
extensions.rs
