Skip to main content

substrait_explain/textify/
addenda.rs

1//! Addendum line textification support.
2//!
3//! Addenda are `+`-prefixed lines emitted between a relation header and its
4//! child relations. This module owns their textifier-side shape and canonical
5//! ordering.
6
7use std::fmt;
8
9use substrait::proto::extensions::AdvancedExtension;
10
11use crate::FormatError;
12use crate::extensions::any::AnyRef;
13use crate::extensions::registry::ExtensionError;
14use crate::extensions::{AddendumKind, ExtensionArgs};
15use crate::textify::foundation::{PlanError, Scope, Textify};
16
17/// All addenda associated with a relation, in canonical text-format order.
18#[derive(Default)]
19pub(super) struct AddendumLines {
20    // Addenda are in textification order: extension table first (if it exists),
21    // then enhancements, then optimizations.
22    lines: Vec<AddendumLine>,
23}
24
25impl AddendumLines {
26    pub(super) fn from_advanced_extension<S: Scope>(
27        ctx: &S,
28        advanced_extension: Option<&AdvancedExtension>,
29    ) -> Self {
30        let mut lines = Self::default();
31        if let Some(advanced_extension) = advanced_extension {
32            lines.extend_from_advanced_extension(ctx, advanced_extension);
33        }
34        lines
35    }
36
37    pub(super) fn extension_table<S: Scope>(
38        ctx: &S,
39        extension_table: Result<(String, ExtensionArgs), ExtensionError>,
40        advanced_extension: Option<&AdvancedExtension>,
41    ) -> Self {
42        let mut lines = Self {
43            lines: vec![AddendumLine::extension_table(extension_table)],
44        };
45        if let Some(advanced_extension) = advanced_extension {
46            lines.extend_from_advanced_extension(ctx, advanced_extension);
47        }
48        lines
49    }
50
51    pub(super) fn none() -> Self {
52        Self::default()
53    }
54
55    /// Add [`AddendumLine`]s to this `AddendumLines` instance for each
56    /// enhancement and optimization present in the [`AdvancedExtension`].
57    fn extend_from_advanced_extension<S: Scope>(
58        &mut self,
59        ctx: &S,
60        advanced_extension: &AdvancedExtension,
61    ) {
62        if let Some(enhancement) = &advanced_extension.enhancement {
63            self.lines
64                .push(AddendumLine::enhancement(ctx, AnyRef::from(enhancement)));
65        }
66        self.lines.extend(
67            advanced_extension
68                .optimization
69                .iter()
70                .map(|optimization| AddendumLine::optimization(ctx, AnyRef::from(optimization))),
71        );
72    }
73}
74
75impl Textify for AddendumLines {
76    fn name() -> &'static str {
77        "AddendumLines"
78    }
79
80    fn textify<S: Scope, W: fmt::Write>(&self, ctx: &S, w: &mut W) -> fmt::Result {
81        for line in &self.lines {
82            writeln!(w)?;
83            line.textify(ctx, w)?;
84        }
85        Ok(())
86    }
87}
88
89enum AddendumLine {
90    Decoded {
91        kind: AddendumKind,
92        name: String,
93        args: ExtensionArgs,
94    },
95    DecodeError {
96        kind: AddendumKind,
97        error: ExtensionError,
98    },
99}
100
101impl AddendumLine {
102    fn extension_table(result: Result<(String, ExtensionArgs), ExtensionError>) -> Self {
103        Self::from_decode_result(AddendumKind::ExtensionTable, result)
104    }
105
106    fn enhancement<S: Scope>(ctx: &S, detail: AnyRef<'_>) -> Self {
107        Self::from_decode_result(
108            AddendumKind::Enhancement,
109            ctx.extension_registry().decode_enhancement(detail),
110        )
111    }
112
113    fn optimization<S: Scope>(ctx: &S, detail: AnyRef<'_>) -> Self {
114        Self::from_decode_result(
115            AddendumKind::Optimization,
116            ctx.extension_registry().decode_optimization(detail),
117        )
118    }
119
120    fn from_decode_result(
121        kind: AddendumKind,
122        result: Result<(String, ExtensionArgs), ExtensionError>,
123    ) -> Self {
124        match result {
125            Ok((name, args)) => Self::Decoded { kind, name, args },
126            Err(error) => Self::DecodeError { kind, error },
127        }
128    }
129}
130
131impl Textify for AddendumLine {
132    fn name() -> &'static str {
133        "AddendumLine"
134    }
135
136    fn textify<S: Scope, W: fmt::Write>(&self, ctx: &S, w: &mut W) -> fmt::Result {
137        match self {
138            AddendumLine::Decoded { kind, name, args } => {
139                let indent = ctx.indent();
140                let prefix = kind.prefix();
141
142                if !args.output_columns.is_empty() {
143                    let (message, description) = match kind {
144                        AddendumKind::Enhancement | AddendumKind::Optimization => (
145                            "addendum",
146                            "output_columns cannot be represented in addendum syntax",
147                        ),
148                        AddendumKind::ExtensionTable => (
149                            "addendum",
150                            "output_columns cannot be represented in extension table addendum syntax",
151                        ),
152                    };
153                    write!(
154                        w,
155                        "{indent}+ {prefix}[{}]",
156                        ctx.failure(FormatError::Format(PlanError::invalid(
157                            message,
158                            Some(name.clone()),
159                            description,
160                        )))
161                    )
162                } else {
163                    write!(w, "{indent}+ {prefix}:{name}[{}]", ctx.display(args))
164                }
165            }
166            AddendumLine::DecodeError { kind, error } => {
167                let indent = ctx.indent();
168                let prefix = kind.prefix();
169                write!(w, "{indent}+ {prefix}[{}]", ctx.failure(error.clone()))
170            }
171        }
172    }
173}