substrait_explain/textify/
plan.rs

1use std::fmt;
2
3use substrait::proto;
4
5use super::Textify;
6use crate::extensions::{ExtensionRegistry, SimpleExtensions};
7use crate::parser::{PLAN_HEADER, VERSION_HEADER};
8use crate::textify::foundation::ErrorAccumulator;
9use crate::textify::{OutputOptions, ScopedContext};
10
11#[derive(Debug, Clone)]
12pub(crate) struct PlanWriter<'a, E: ErrorAccumulator + Default> {
13    options: &'a OutputOptions,
14    version: Option<&'a proto::Version>,
15    extensions: SimpleExtensions,
16    relations: &'a [proto::PlanRel],
17    errors: E,
18    extension_registry: &'a ExtensionRegistry,
19}
20
21impl<'a, E: ErrorAccumulator + Default + Clone> PlanWriter<'a, E> {
22    pub(crate) fn new(
23        options: &'a OutputOptions,
24        plan: &'a proto::Plan,
25        extension_registry: &'a ExtensionRegistry,
26    ) -> (Self, E) {
27        let (extensions, errs) =
28            SimpleExtensions::from_extensions(&plan.extension_urns, &plan.extensions);
29
30        let errors = E::default();
31        for err in errs {
32            errors.push(err.into());
33        }
34
35        let relations = plan.relations.as_slice();
36
37        (
38            Self {
39                options,
40                version: plan.version.as_ref(),
41                extensions,
42                relations,
43                errors: errors.clone(),
44                extension_registry,
45            },
46            errors,
47        )
48    }
49
50    pub(crate) fn scope(&'a self) -> ScopedContext<'a, E> {
51        ScopedContext::new(
52            self.options,
53            &self.errors,
54            &self.extensions,
55            self.extension_registry,
56        )
57    }
58
59    /// Write the `=== Version` section. Emits nothing unless the plan
60    /// carries a version that is not entirely empty
61    pub(crate) fn write_version(&self, w: &mut impl fmt::Write) -> fmt::Result {
62        let Some(version) = self.version else {
63            return Ok(());
64        };
65        if version == &proto::Version::default() {
66            return Ok(());
67        }
68
69        writeln!(
70            w,
71            "{VERSION_HEADER} {}.{}.{}",
72            version.major_number, version.minor_number, version.patch_number
73        )?;
74        if !version.producer.is_empty() {
75            writeln!(
76                w,
77                "{}producer: {}",
78                self.options.indent,
79                version.producer.trim()
80            )?;
81        }
82        if !version.git_hash.is_empty() {
83            writeln!(
84                w,
85                "{}git_hash: {}",
86                self.options.indent,
87                version.git_hash.trim()
88            )?;
89        }
90        Ok(())
91    }
92
93    pub(crate) fn write_extensions(&self, w: &mut impl fmt::Write) -> fmt::Result {
94        self.extensions.write(w, &self.options.indent)
95    }
96
97    pub(crate) fn write_relations(&self, w: &mut impl fmt::Write) -> fmt::Result {
98        // We always write the plan header, even if there are no relations.
99        writeln!(w, "{PLAN_HEADER}")?;
100        let scope = self.scope();
101        for (i, relation) in self.relations.iter().enumerate() {
102            if i > 0 {
103                writeln!(w)?;
104                writeln!(w)?;
105            }
106            relation.textify(&scope, w)?;
107        }
108        Ok(())
109    }
110}
111
112impl<'a, E: ErrorAccumulator + Default> fmt::Display for PlanWriter<'a, E> {
113    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
114        self.write_version(f)?;
115        self.write_extensions(f)?;
116        if !self.extensions.is_empty() {
117            writeln!(f)?;
118        }
119        self.write_relations(f)?;
120        writeln!(f)?;
121        Ok(())
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use std::fmt::Write;
128
129    use pext::simple_extension_declaration::{ExtensionFunction, MappingType};
130    use substrait::proto::expression::{RexType, ScalarFunction};
131    use substrait::proto::function_argument::ArgType;
132    use substrait::proto::read_rel::{NamedTable, ReadType};
133    use substrait::proto::r#type::{I64, Kind, Nullability, Struct};
134    use substrait::proto::{
135        Expression, FunctionArgument, NamedStruct, ReadRel, Type, extensions as pext,
136    };
137
138    use super::*;
139    use crate::parser::expressions::FieldIndex;
140    use crate::textify::ErrorQueue;
141
142    /// Test a fairly basic plan with an extension, read, and project.
143    ///
144    /// This has a manually constructed plan, rather than using the parser; more
145    /// complete testing is in the integration tests.
146    #[test]
147    fn test_plan_writer() {
148        let mut plan = proto::Plan::default();
149
150        // Add extension URN
151        plan.extension_urns.push(pext::SimpleExtensionUrn {
152            extension_urn_anchor: 1,
153            urn: "https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml".to_string(),
154        });
155
156        // Add extension function declaration
157        plan.extensions.push(pext::SimpleExtensionDeclaration {
158            mapping_type: Some(MappingType::ExtensionFunction(ExtensionFunction {
159                extension_urn_reference: 1,
160                function_anchor: 10,
161                name: "add".to_string(),
162            })),
163        });
164
165        // Create read relation
166        let read_rel = ReadRel {
167            read_type: Some(ReadType::NamedTable(NamedTable {
168                names: vec!["table1".to_string()],
169                ..Default::default()
170            })),
171            base_schema: Some(NamedStruct {
172                names: vec!["col1".to_string(), "col2".to_string()],
173                r#struct: Some(Struct {
174                    types: vec![
175                        Type {
176                            kind: Some(Kind::I32(proto::r#type::I32 {
177                                nullability: Nullability::Nullable as i32,
178                                type_variation_reference: 0,
179                            })),
180                        },
181                        Type {
182                            kind: Some(Kind::I32(proto::r#type::I32 {
183                                nullability: Nullability::Nullable as i32,
184                                type_variation_reference: 0,
185                            })),
186                        },
187                    ],
188                    ..Default::default()
189                }),
190            }),
191            ..Default::default()
192        };
193
194        // Create project relation with add function
195        let add_function = ScalarFunction {
196            function_reference: 10,
197            arguments: vec![
198                FunctionArgument {
199                    arg_type: Some(ArgType::Value(Expression {
200                        rex_type: Some(RexType::Selection(Box::new(
201                            FieldIndex(0).to_field_reference(),
202                        ))),
203                    })),
204                },
205                FunctionArgument {
206                    arg_type: Some(ArgType::Value(Expression {
207                        rex_type: Some(RexType::Selection(Box::new(
208                            FieldIndex(1).to_field_reference(),
209                        ))),
210                    })),
211                },
212            ],
213            options: vec![],
214            output_type: Some(Type {
215                kind: Some(Kind::I64(I64 {
216                    nullability: Nullability::Required as i32,
217                    type_variation_reference: 0,
218                })),
219            }),
220            #[allow(deprecated)]
221            args: vec![],
222        };
223
224        let project_rel = proto::ProjectRel {
225            expressions: vec![Expression {
226                rex_type: Some(RexType::ScalarFunction(add_function)),
227            }],
228            input: Some(Box::new(proto::Rel {
229                rel_type: Some(proto::rel::RelType::Read(Box::new(read_rel))),
230            })),
231            common: None,
232            advanced_extension: None,
233        };
234
235        // Add relations to plan
236        plan.relations.push(proto::PlanRel {
237            rel_type: Some(proto::plan_rel::RelType::Rel(proto::Rel {
238                rel_type: Some(proto::rel::RelType::Project(Box::new(project_rel))),
239            })),
240        });
241
242        let options = OutputOptions::default();
243        let extension_registry = ExtensionRegistry::new();
244        let (writer, errors) = PlanWriter::<ErrorQueue>::new(&options, &plan, &extension_registry);
245        let mut output = String::new();
246        write!(output, "{writer}").unwrap();
247
248        // Assert that there are no errors
249        let errors: Vec<_> = errors.into();
250        assert!(errors.is_empty(), "Expected no errors, got: {errors:?}");
251
252        let expected = r#"
253=== Extensions
254URNs:
255  @  1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml
256Functions:
257  # 10 @  1: add
258
259=== Plan
260Project[$0, $1, add($0, $1):i64]
261  Read[table1 => col1:i32?, col2:i32?]
262"#
263        .trim_start();
264
265        assert_eq!(
266            output, expected,
267            "Output:\n---\n{output}\n---\nExpected:\n---\n{expected}\n---"
268        );
269    }
270}
substrait_explain/textify/plan.rs

substrait_explain/textify/
plan.rs
