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