## 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
    }