Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/expr/src/logical_plan/display.rs
Line
Count
Source (jump to first uncovered line)
1
// Licensed to the Apache Software Foundation (ASF) under one
2
// or more contributor license agreements.  See the NOTICE file
3
// distributed with this work for additional information
4
// regarding copyright ownership.  The ASF licenses this file
5
// to you under the Apache License, Version 2.0 (the
6
// "License"); you may not use this file except in compliance
7
// with the License.  You may obtain a copy of the License at
8
//
9
//   http://www.apache.org/licenses/LICENSE-2.0
10
//
11
// Unless required by applicable law or agreed to in writing,
12
// software distributed under the License is distributed on an
13
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
// KIND, either express or implied.  See the License for the
15
// specific language governing permissions and limitations
16
// under the License.
17
//! This module provides logic for displaying LogicalPlans in various styles
18
19
use std::collections::HashMap;
20
use std::fmt;
21
22
use crate::{
23
    expr_vec_fmt, Aggregate, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr,
24
    Filter, Join, Limit, LogicalPlan, Partitioning, Prepare, Projection, RecursiveQuery,
25
    Repartition, Sort, Subquery, SubqueryAlias, TableProviderFilterPushDown, TableScan,
26
    Unnest, Values, Window,
27
};
28
29
use crate::dml::CopyTo;
30
use arrow::datatypes::Schema;
31
use datafusion_common::display::GraphvizBuilder;
32
use datafusion_common::tree_node::{TreeNodeRecursion, TreeNodeVisitor};
33
use datafusion_common::{Column, DataFusionError};
34
use serde_json::json;
35
36
/// Formats plans with a single line per node. For example:
37
///
38
/// Projection: id
39
///    Filter: state Eq Utf8(\"CO\")\
40
///       CsvScan: employee.csv projection=Some([0, 3])";
41
pub struct IndentVisitor<'a, 'b> {
42
    f: &'a mut fmt::Formatter<'b>,
43
    /// If true, includes summarized schema information
44
    with_schema: bool,
45
    /// The current indent
46
    indent: usize,
47
}
48
49
impl<'a, 'b> IndentVisitor<'a, 'b> {
50
    /// Create a visitor that will write a formatted LogicalPlan to f. If `with_schema` is
51
    /// true, includes schema information on each line.
52
0
    pub fn new(f: &'a mut fmt::Formatter<'b>, with_schema: bool) -> Self {
53
0
        Self {
54
0
            f,
55
0
            with_schema,
56
0
            indent: 0,
57
0
        }
58
0
    }
59
}
60
61
impl<'n, 'a, 'b> TreeNodeVisitor<'n> for IndentVisitor<'a, 'b> {
62
    type Node = LogicalPlan;
63
64
0
    fn f_down(
65
0
        &mut self,
66
0
        plan: &'n LogicalPlan,
67
0
    ) -> datafusion_common::Result<TreeNodeRecursion> {
68
0
        if self.indent > 0 {
69
0
            writeln!(self.f)?;
70
0
        }
71
0
        write!(self.f, "{:indent$}", "", indent = self.indent * 2)?;
72
0
        write!(self.f, "{}", plan.display())?;
73
0
        if self.with_schema {
74
0
            write!(
75
0
                self.f,
76
0
                " {}",
77
0
                display_schema(&plan.schema().as_ref().to_owned().into())
78
0
            )?;
79
0
        }
80
81
0
        self.indent += 1;
82
0
        Ok(TreeNodeRecursion::Continue)
83
0
    }
84
85
0
    fn f_up(
86
0
        &mut self,
87
0
        _plan: &'n LogicalPlan,
88
0
    ) -> datafusion_common::Result<TreeNodeRecursion> {
89
0
        self.indent -= 1;
90
0
        Ok(TreeNodeRecursion::Continue)
91
0
    }
92
}
93
94
/// Print the schema in a compact representation to `buf`
95
///
96
/// For example: `foo:Utf8` if `foo` can not be null, and
97
/// `foo:Utf8;N` if `foo` is nullable.
98
///
99
/// ```
100
/// use arrow::datatypes::{Field, Schema, DataType};
101
/// # use datafusion_expr::logical_plan::display_schema;
102
/// let schema = Schema::new(vec![
103
///     Field::new("id", DataType::Int32, false),
104
///     Field::new("first_name", DataType::Utf8, true),
105
///  ]);
106
///
107
///  assert_eq!(
108
///      "[id:Int32, first_name:Utf8;N]",
109
///      format!("{}", display_schema(&schema))
110
///  );
111
/// ```
112
0
pub fn display_schema(schema: &Schema) -> impl fmt::Display + '_ {
113
    struct Wrapper<'a>(&'a Schema);
114
115
    impl<'a> fmt::Display for Wrapper<'a> {
116
0
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
117
0
            write!(f, "[")?;
118
0
            for (idx, field) in self.0.fields().iter().enumerate() {
119
0
                if idx > 0 {
120
0
                    write!(f, ", ")?;
121
0
                }
122
0
                let nullable_str = if field.is_nullable() { ";N" } else { "" };
123
0
                write!(
124
0
                    f,
125
0
                    "{}:{:?}{}",
126
0
                    field.name(),
127
0
                    field.data_type(),
128
0
                    nullable_str
129
0
                )?;
130
            }
131
0
            write!(f, "]")
132
0
        }
133
    }
134
0
    Wrapper(schema)
135
0
}
136
137
/// Formats plans for graphical display using the `DOT` language. This
138
/// format can be visualized using software from
139
/// [`graphviz`](https://graphviz.org/)
140
pub struct GraphvizVisitor<'a, 'b> {
141
    f: &'a mut fmt::Formatter<'b>,
142
    graphviz_builder: GraphvizBuilder,
143
    /// If true, includes summarized schema information
144
    with_schema: bool,
145
146
    /// Holds the ids (as generated from `graphviz_builder` of all
147
    /// parent nodes
148
    parent_ids: Vec<usize>,
149
}
150
151
impl<'a, 'b> GraphvizVisitor<'a, 'b> {
152
0
    pub fn new(f: &'a mut fmt::Formatter<'b>) -> Self {
153
0
        Self {
154
0
            f,
155
0
            graphviz_builder: GraphvizBuilder::default(),
156
0
            with_schema: false,
157
0
            parent_ids: Vec::new(),
158
0
        }
159
0
    }
160
161
    /// Sets a flag which controls if the output schema is displayed
162
0
    pub fn set_with_schema(&mut self, with_schema: bool) {
163
0
        self.with_schema = with_schema;
164
0
    }
165
166
0
    pub fn pre_visit_plan(&mut self, label: &str) -> fmt::Result {
167
0
        self.graphviz_builder.start_cluster(self.f, label)
168
0
    }
169
170
0
    pub fn post_visit_plan(&mut self) -> fmt::Result {
171
0
        self.graphviz_builder.end_cluster(self.f)
172
0
    }
173
174
0
    pub fn start_graph(&mut self) -> fmt::Result {
175
0
        self.graphviz_builder.start_graph(self.f)
176
0
    }
177
178
0
    pub fn end_graph(&mut self) -> fmt::Result {
179
0
        self.graphviz_builder.end_graph(self.f)
180
0
    }
181
}
182
183
impl<'n, 'a, 'b> TreeNodeVisitor<'n> for GraphvizVisitor<'a, 'b> {
184
    type Node = LogicalPlan;
185
186
0
    fn f_down(
187
0
        &mut self,
188
0
        plan: &'n LogicalPlan,
189
0
    ) -> datafusion_common::Result<TreeNodeRecursion> {
190
0
        let id = self.graphviz_builder.next_id();
191
192
        // Create a new graph node for `plan` such as
193
        // id [label="foo"]
194
0
        let label = if self.with_schema {
195
0
            format!(
196
0
                r"{}\nSchema: {}",
197
0
                plan.display(),
198
0
                display_schema(&plan.schema().as_ref().to_owned().into())
199
0
            )
200
        } else {
201
0
            format!("{}", plan.display())
202
        };
203
204
0
        self.graphviz_builder
205
0
            .add_node(self.f, id, &label, None)
206
0
            .map_err(|_e| DataFusionError::Internal("Fail to format".to_string()))?;
207
208
        // Create an edge to our parent node, if any
209
        //  parent_id -> id
210
0
        if let Some(parent_id) = self.parent_ids.last() {
211
0
            self.graphviz_builder
212
0
                .add_edge(self.f, *parent_id, id)
213
0
                .map_err(|_e| DataFusionError::Internal("Fail to format".to_string()))?;
214
0
        }
215
216
0
        self.parent_ids.push(id);
217
0
        Ok(TreeNodeRecursion::Continue)
218
0
    }
219
220
0
    fn f_up(
221
0
        &mut self,
222
0
        _plan: &LogicalPlan,
223
0
    ) -> datafusion_common::Result<TreeNodeRecursion> {
224
0
        // always be non-empty as pre_visit always pushes
225
0
        // So it should always be Ok(true)
226
0
        let res = self.parent_ids.pop();
227
0
        res.ok_or(DataFusionError::Internal("Fail to format".to_string()))
228
0
            .map(|_| TreeNodeRecursion::Continue)
229
0
    }
230
}
231
232
/// Formats plans to display as postgresql plan json format.
233
///
234
/// There are already many existing visualizer for this format, for example [dalibo](https://explain.dalibo.com/).
235
/// Unfortunately, there is no formal spec for this format, but it is widely used in the PostgreSQL community.
236
///
237
/// Here is an example of the format:
238
///
239
/// ```json
240
/// [
241
///     {
242
///         "Plan": {
243
///             "Node Type": "Sort",
244
///             "Output": [
245
///                 "question_1.id",
246
///                 "question_1.title",
247
///                 "question_1.text",
248
///                 "question_1.file",
249
///                 "question_1.type",
250
///                 "question_1.source",
251
///                 "question_1.exam_id"
252
///             ],
253
///             "Sort Key": [
254
///                 "question_1.id"
255
///             ],
256
///             "Plans": [
257
///                 {
258
///                     "Node Type": "Seq Scan",
259
///                     "Parent Relationship": "Left",
260
///                     "Relation Name": "question",
261
///                     "Schema": "public",
262
///                     "Alias": "question_1",
263
///                     "Output": [
264
///                        "question_1.id",
265
///                         "question_1.title",
266
///                        "question_1.text",
267
///                         "question_1.file",
268
///                         "question_1.type",
269
///                         "question_1.source",
270
///                         "question_1.exam_id"
271
///                     ],
272
///                     "Filter": "(question_1.exam_id = 1)"
273
///                 }
274
///             ]
275
///         }
276
///     }
277
/// ]
278
/// ```
279
pub struct PgJsonVisitor<'a, 'b> {
280
    f: &'a mut fmt::Formatter<'b>,
281
282
    /// A mapping from plan node id to the plan node json representation.
283
    objects: HashMap<u32, serde_json::Value>,
284
285
    next_id: u32,
286
287
    /// If true, includes summarized schema information
288
    with_schema: bool,
289
290
    /// Holds the ids (as generated from `graphviz_builder` of all
291
    /// parent nodes
292
    parent_ids: Vec<u32>,
293
}
294
295
impl<'a, 'b> PgJsonVisitor<'a, 'b> {
296
0
    pub fn new(f: &'a mut fmt::Formatter<'b>) -> Self {
297
0
        Self {
298
0
            f,
299
0
            objects: HashMap::new(),
300
0
            next_id: 0,
301
0
            with_schema: false,
302
0
            parent_ids: Vec::new(),
303
0
        }
304
0
    }
305
306
    /// Sets a flag which controls if the output schema is displayed
307
0
    pub fn with_schema(&mut self, with_schema: bool) {
308
0
        self.with_schema = with_schema;
309
0
    }
310
311
    /// Converts a logical plan node to a json object.
312
0
    fn to_json_value(node: &LogicalPlan) -> serde_json::Value {
313
0
        match node {
314
            LogicalPlan::EmptyRelation(_) => {
315
0
                json!({
316
0
                    "Node Type": "EmptyRelation",
317
0
                })
318
            }
319
0
            LogicalPlan::RecursiveQuery(RecursiveQuery { is_distinct, .. }) => {
320
0
                json!({
321
0
                    "Node Type": "RecursiveQuery",
322
0
                    "Is Distinct": is_distinct,
323
0
                })
324
            }
325
0
            LogicalPlan::Values(Values { ref values, .. }) => {
326
0
                let str_values = values
327
0
                    .iter()
328
0
                    // limit to only 5 values to avoid horrible display
329
0
                    .take(5)
330
0
                    .map(|row| {
331
0
                        let item = row
332
0
                            .iter()
333
0
                            .map(|expr| expr.to_string())
334
0
                            .collect::<Vec<_>>()
335
0
                            .join(", ");
336
0
                        format!("({item})")
337
0
                    })
338
0
                    .collect::<Vec<_>>()
339
0
                    .join(", ");
340
341
0
                let eclipse = if values.len() > 5 { "..." } else { "" };
342
343
0
                let values_str = format!("{}{}", str_values, eclipse);
344
0
                json!({
345
0
                    "Node Type": "Values",
346
0
                    "Values": values_str
347
0
                })
348
            }
349
            LogicalPlan::TableScan(TableScan {
350
0
                ref source,
351
0
                ref table_name,
352
0
                ref filters,
353
0
                ref fetch,
354
0
                ..
355
0
            }) => {
356
0
                let mut object = json!({
357
0
                    "Node Type": "TableScan",
358
0
                    "Relation Name": table_name.table(),
359
0
                });
360
361
0
                if let Some(s) = table_name.schema() {
362
0
                    object["Schema"] = serde_json::Value::String(s.to_string());
363
0
                }
364
365
0
                if let Some(c) = table_name.catalog() {
366
0
                    object["Catalog"] = serde_json::Value::String(c.to_string());
367
0
                }
368
369
0
                if !filters.is_empty() {
370
0
                    let mut full_filter = vec![];
371
0
                    let mut partial_filter = vec![];
372
0
                    let mut unsupported_filters = vec![];
373
0
                    let filters: Vec<&Expr> = filters.iter().collect();
374
375
0
                    if let Ok(results) = source.supports_filters_pushdown(&filters) {
376
0
                        filters.iter().zip(results.iter()).for_each(
377
0
                            |(x, res)| match res {
378
0
                                TableProviderFilterPushDown::Exact => full_filter.push(x),
379
                                TableProviderFilterPushDown::Inexact => {
380
0
                                    partial_filter.push(x)
381
                                }
382
                                TableProviderFilterPushDown::Unsupported => {
383
0
                                    unsupported_filters.push(x)
384
                                }
385
0
                            },
386
0
                        );
387
0
                    }
388
389
0
                    if !full_filter.is_empty() {
390
0
                        object["Full Filters"] =
391
0
                            serde_json::Value::String(expr_vec_fmt!(full_filter));
392
0
                    };
393
0
                    if !partial_filter.is_empty() {
394
0
                        object["Partial Filters"] =
395
0
                            serde_json::Value::String(expr_vec_fmt!(partial_filter));
396
0
                    }
397
0
                    if !unsupported_filters.is_empty() {
398
0
                        object["Unsupported Filters"] =
399
0
                            serde_json::Value::String(expr_vec_fmt!(unsupported_filters));
400
0
                    }
401
0
                }
402
403
0
                if let Some(f) = fetch {
404
0
                    object["Fetch"] = serde_json::Value::Number((*f).into());
405
0
                }
406
407
0
                object
408
            }
409
0
            LogicalPlan::Projection(Projection { ref expr, .. }) => {
410
0
                json!({
411
0
                    "Node Type": "Projection",
412
0
                    "Expressions": expr.iter().map(|e| e.to_string()).collect::<Vec<_>>()
413
0
                })
414
            }
415
0
            LogicalPlan::Dml(DmlStatement { table_name, op, .. }) => {
416
0
                json!({
417
0
                    "Node Type": "Projection",
418
0
                    "Operation": op.name(),
419
0
                    "Table Name": table_name.table()
420
0
                })
421
            }
422
            LogicalPlan::Copy(CopyTo {
423
                input: _,
424
0
                output_url,
425
0
                file_type,
426
0
                partition_by: _,
427
0
                options,
428
0
            }) => {
429
0
                let op_str = options
430
0
                    .iter()
431
0
                    .map(|(k, v)| format!("{}={}", k, v))
432
0
                    .collect::<Vec<_>>()
433
0
                    .join(", ");
434
0
                json!({
435
0
                    "Node Type": "CopyTo",
436
0
                    "Output URL": output_url,
437
0
                    "File Type": format!("{}", file_type.get_ext()),
438
0
                    "Options": op_str
439
0
                })
440
            }
441
0
            LogicalPlan::Ddl(ddl) => {
442
0
                json!({
443
0
                    "Node Type": "Ddl",
444
0
                    "Operation": format!("{}", ddl.display())
445
0
                })
446
            }
447
            LogicalPlan::Filter(Filter {
448
0
                predicate: ref expr,
449
0
                ..
450
0
            }) => {
451
0
                json!({
452
0
                    "Node Type": "Filter",
453
0
                    "Condition": format!("{}", expr)
454
0
                })
455
            }
456
            LogicalPlan::Window(Window {
457
0
                ref window_expr, ..
458
0
            }) => {
459
0
                json!({
460
0
                    "Node Type": "WindowAggr",
461
0
                    "Expressions": expr_vec_fmt!(window_expr)
462
0
                })
463
            }
464
            LogicalPlan::Aggregate(Aggregate {
465
0
                ref group_expr,
466
0
                ref aggr_expr,
467
0
                ..
468
0
            }) => {
469
0
                json!({
470
0
                    "Node Type": "Aggregate",
471
0
                    "Group By": expr_vec_fmt!(group_expr),
472
0
                    "Aggregates": expr_vec_fmt!(aggr_expr)
473
0
                })
474
            }
475
0
            LogicalPlan::Sort(Sort { expr, fetch, .. }) => {
476
0
                let mut object = json!({
477
0
                    "Node Type": "Sort",
478
0
                    "Sort Key": expr_vec_fmt!(expr),
479
0
                });
480
481
0
                if let Some(fetch) = fetch {
482
0
                    object["Fetch"] = serde_json::Value::Number((*fetch).into());
483
0
                }
484
485
0
                object
486
            }
487
            LogicalPlan::Join(Join {
488
0
                on: ref keys,
489
0
                filter,
490
0
                join_constraint,
491
0
                join_type,
492
0
                ..
493
0
            }) => {
494
0
                let join_expr: Vec<String> =
495
0
                    keys.iter().map(|(l, r)| format!("{l} = {r}")).collect();
496
0
                let filter_expr = filter
497
0
                    .as_ref()
498
0
                    .map(|expr| format!(" Filter: {expr}"))
499
0
                    .unwrap_or_else(|| "".to_string());
500
0
                json!({
501
0
                    "Node Type": format!("{} Join", join_type),
502
0
                    "Join Constraint": format!("{:?}", join_constraint),
503
0
                    "Join Keys": join_expr.join(", "),
504
0
                    "Filter": format!("{}", filter_expr)
505
0
                })
506
            }
507
            LogicalPlan::CrossJoin(_) => {
508
0
                json!({
509
0
                    "Node Type": "Cross Join"
510
0
                })
511
            }
512
            LogicalPlan::Repartition(Repartition {
513
0
                partitioning_scheme,
514
0
                ..
515
0
            }) => match partitioning_scheme {
516
0
                Partitioning::RoundRobinBatch(n) => {
517
0
                    json!({
518
0
                        "Node Type": "Repartition",
519
0
                        "Partitioning Scheme": "RoundRobinBatch",
520
0
                        "Partition Count": n
521
0
                    })
522
                }
523
0
                Partitioning::Hash(expr, n) => {
524
0
                    let hash_expr: Vec<String> =
525
0
                        expr.iter().map(|e| format!("{e}")).collect();
526
0
527
0
                    json!({
528
0
                        "Node Type": "Repartition",
529
0
                        "Partitioning Scheme": "Hash",
530
0
                        "Partition Count": n,
531
0
                        "Partitioning Key": hash_expr
532
0
                    })
533
                }
534
0
                Partitioning::DistributeBy(expr) => {
535
0
                    let dist_by_expr: Vec<String> =
536
0
                        expr.iter().map(|e| format!("{e}")).collect();
537
0
                    json!({
538
0
                        "Node Type": "Repartition",
539
0
                        "Partitioning Scheme": "DistributeBy",
540
0
                        "Partitioning Key": dist_by_expr
541
0
                    })
542
                }
543
            },
544
            LogicalPlan::Limit(Limit {
545
0
                ref skip,
546
0
                ref fetch,
547
0
                ..
548
0
            }) => {
549
0
                let mut object = serde_json::json!(
550
0
                    {
551
0
                        "Node Type": "Limit",
552
0
                        "Skip": skip,
553
0
                    }
554
0
                );
555
0
                if let Some(f) = fetch {
556
0
                    object["Fetch"] = serde_json::Value::Number((*f).into());
557
0
                };
558
0
                object
559
            }
560
            LogicalPlan::Subquery(Subquery { .. }) => {
561
0
                json!({
562
0
                    "Node Type": "Subquery"
563
0
                })
564
            }
565
0
            LogicalPlan::SubqueryAlias(SubqueryAlias { ref alias, .. }) => {
566
0
                json!({
567
0
                    "Node Type": "Subquery",
568
0
                    "Alias": alias.table(),
569
0
                })
570
            }
571
0
            LogicalPlan::Statement(statement) => {
572
0
                json!({
573
0
                    "Node Type": "Statement",
574
0
                    "Statement": format!("{}", statement.display())
575
0
                })
576
            }
577
0
            LogicalPlan::Distinct(distinct) => match distinct {
578
                Distinct::All(_) => {
579
0
                    json!({
580
0
                        "Node Type": "DistinctAll"
581
0
                    })
582
                }
583
                Distinct::On(DistinctOn {
584
0
                    on_expr,
585
0
                    select_expr,
586
0
                    sort_expr,
587
0
                    ..
588
0
                }) => {
589
0
                    let mut object = json!({
590
0
                        "Node Type": "DistinctOn",
591
0
                        "On": expr_vec_fmt!(on_expr),
592
0
                        "Select": expr_vec_fmt!(select_expr),
593
0
                    });
594
0
                    if let Some(sort_expr) = sort_expr {
595
0
                        object["Sort"] =
596
0
                            serde_json::Value::String(expr_vec_fmt!(sort_expr));
597
0
                    }
598
599
0
                    object
600
                }
601
            },
602
            LogicalPlan::Explain { .. } => {
603
0
                json!({
604
0
                    "Node Type": "Explain"
605
0
                })
606
            }
607
            LogicalPlan::Analyze { .. } => {
608
0
                json!({
609
0
                    "Node Type": "Analyze"
610
0
                })
611
            }
612
            LogicalPlan::Union(_) => {
613
0
                json!({
614
0
                    "Node Type": "Union"
615
0
                })
616
            }
617
0
            LogicalPlan::Extension(e) => {
618
0
                json!({
619
0
                    "Node Type": e.node.name(),
620
0
                    "Detail": format!("{:?}", e.node)
621
0
                })
622
            }
623
            LogicalPlan::Prepare(Prepare {
624
0
                name, data_types, ..
625
0
            }) => {
626
0
                json!({
627
0
                    "Node Type": "Prepare",
628
0
                    "Name": name,
629
0
                    "Data Types": format!("{:?}", data_types)
630
0
                })
631
            }
632
            LogicalPlan::DescribeTable(DescribeTable { .. }) => {
633
0
                json!({
634
0
                    "Node Type": "DescribeTable"
635
0
                })
636
            }
637
            LogicalPlan::Unnest(Unnest {
638
0
                input: plan,
639
0
                list_type_columns: list_col_indices,
640
0
                struct_type_columns: struct_col_indices,
641
0
                ..
642
0
            }) => {
643
0
                let input_columns = plan.schema().columns();
644
0
                let list_type_columns = list_col_indices
645
0
                    .iter()
646
0
                    .map(|(i, unnest_info)| {
647
0
                        format!(
648
0
                            "{}|depth={:?}",
649
0
                            &input_columns[*i].to_string(),
650
0
                            unnest_info.depth
651
0
                        )
652
0
                    })
653
0
                    .collect::<Vec<String>>();
654
0
                let struct_type_columns = struct_col_indices
655
0
                    .iter()
656
0
                    .map(|i| &input_columns[*i])
657
0
                    .collect::<Vec<&Column>>();
658
0
                json!({
659
0
                    "Node Type": "Unnest",
660
0
                    "ListColumn": expr_vec_fmt!(list_type_columns),
661
0
                    "StructColumn": expr_vec_fmt!(struct_type_columns),
662
0
                })
663
            }
664
        }
665
0
    }
666
}
667
668
impl<'n, 'a, 'b> TreeNodeVisitor<'n> for PgJsonVisitor<'a, 'b> {
669
    type Node = LogicalPlan;
670
671
0
    fn f_down(
672
0
        &mut self,
673
0
        node: &'n LogicalPlan,
674
0
    ) -> datafusion_common::Result<TreeNodeRecursion> {
675
0
        let id = self.next_id;
676
0
        self.next_id += 1;
677
0
        let mut object = Self::to_json_value(node);
678
0
679
0
        object["Plans"] = serde_json::Value::Array(vec![]);
680
0
681
0
        if self.with_schema {
682
0
            object["Output"] = serde_json::Value::Array(
683
0
                node.schema()
684
0
                    .fields()
685
0
                    .iter()
686
0
                    .map(|f| f.name().to_string())
687
0
                    .map(serde_json::Value::String)
688
0
                    .collect(),
689
0
            );
690
0
        };
691
692
0
        self.objects.insert(id, object);
693
0
        self.parent_ids.push(id);
694
0
        Ok(TreeNodeRecursion::Continue)
695
0
    }
696
697
0
    fn f_up(
698
0
        &mut self,
699
0
        _node: &Self::Node,
700
0
    ) -> datafusion_common::Result<TreeNodeRecursion> {
701
0
        let id = self.parent_ids.pop().unwrap();
702
703
0
        let current_node = self.objects.remove(&id).ok_or_else(|| {
704
0
            DataFusionError::Internal("Missing current node!".to_string())
705
0
        })?;
706
707
0
        if let Some(parent_id) = self.parent_ids.last() {
708
0
            let parent_node = self
709
0
                .objects
710
0
                .get_mut(parent_id)
711
0
                .expect("Missing parent node!");
712
0
            let plans = parent_node
713
0
                .get_mut("Plans")
714
0
                .and_then(|p| p.as_array_mut())
715
0
                .expect("Plans should be an array");
716
0
717
0
            plans.push(current_node);
718
0
        } else {
719
            // This is the root node
720
0
            let plan = serde_json::json!([{"Plan": current_node}]);
721
0
            write!(
722
0
                self.f,
723
0
                "{}",
724
0
                serde_json::to_string_pretty(&plan)
725
0
                    .map_err(|e| DataFusionError::External(Box::new(e)))?
726
0
            )?;
727
        }
728
729
0
        Ok(TreeNodeRecursion::Continue)
730
0
    }
731
}
732
733
#[cfg(test)]
734
mod tests {
735
    use arrow::datatypes::{DataType, Field};
736
737
    use super::*;
738
739
    #[test]
740
    fn test_display_empty_schema() {
741
        let schema = Schema::empty();
742
        assert_eq!("[]", format!("{}", display_schema(&schema)));
743
    }
744
745
    #[test]
746
    fn test_display_schema() {
747
        let schema = Schema::new(vec![
748
            Field::new("id", DataType::Int32, false),
749
            Field::new("first_name", DataType::Utf8, true),
750
        ]);
751
752
        assert_eq!(
753
            "[id:Int32, first_name:Utf8;N]",
754
            format!("{}", display_schema(&schema))
755
        );
756
    }
757
}