refactor(plan_node_fmt): 4 more impls for Distill (#10296)

src/frontend/src/optimizer/plan_node/batch_seq_scan.rs

@@ -183,19 +183,8 @@ impl Distill for BatchSeqScan {
     fn distill<'a>(&self) -> Pretty<'a> {
         let verbose = self.base.ctx.is_explain_verbose();
         let mut vec = Vec::with_capacity(4);
-        vec.push(("table", Pretty::display(&self.logical.table_name)));
-        vec.push((
-            "columns",
-            Pretty::Array(
-                match verbose {
-                    true => self.logical.column_names_with_table_prefix(),
-                    false => self.logical.column_names(),
-                }
-                .into_iter()
-                .map(Pretty::from)
-                .collect(),
-            ),
-        ));
+        vec.push(("table", Pretty::from(self.logical.table_name.clone())));
+        vec.push(("columns", self.logical.columns_pretty(verbose)));
 
         if !self.scan_ranges.is_empty() {
             let range_strs = self.scan_ranges_as_strs(verbose);

src/frontend/src/optimizer/plan_node/generic/scan.rs

@@ -17,6 +17,7 @@ use std::rc::Rc;
 
 use educe::Educe;
 use fixedbitset::FixedBitSet;
+use pretty_xmlish::Pretty;
 use risingwave_common::catalog::{ColumnDesc, Field, Schema, TableDesc};
 use risingwave_common::util::column_index_mapping::ColIndexMapping;
 use risingwave_common::util::sort_util::ColumnOrder;
@@ -264,6 +265,28 @@ impl Scan {
             ctx,
         }
     }
+
+    pub(crate) fn columns_pretty<'a>(&self, verbose: bool) -> Pretty<'a> {
+        Pretty::Array(
+            match verbose {
+                true => self.column_names_with_table_prefix(),
+                false => self.column_names(),
+            }
+            .into_iter()
+            .map(Pretty::from)
+            .collect(),
+        )
+    }
+
+    pub(crate) fn fields_pretty_schema(&self) -> Schema {
+        let fields = self
+            .table_desc
+            .columns
+            .iter()
+            .map(|col| Field::from_with_table_name_prefix(col, &self.table_name))
+            .collect();
+        Schema { fields }
+    }
 }
 
 impl GenericPlanNode for Scan {

src/frontend/src/optimizer/plan_node/logical_now.rs

@@ -15,12 +15,13 @@
 use std::fmt;
 
 use itertools::Itertools;
+use pretty_xmlish::Pretty;
 use risingwave_common::bail;
 use risingwave_common::catalog::{Field, Schema};
 use risingwave_common::error::Result;
 use risingwave_common::types::DataType;
 
-use super::utils::IndicesDisplay;
+use super::utils::{Distill, IndicesDisplay};
 use super::{
     ColPrunable, ColumnPruningContext, ExprRewritable, LogicalFilter, PlanBase, PlanRef,
     PredicatePushdown, RewriteStreamContext, StreamNow, ToBatch, ToStream, ToStreamContext,
@@ -47,6 +48,21 @@ impl LogicalNow {
     }
 }
 
+impl Distill for LogicalNow {
+    fn distill<'a>(&self) -> Pretty<'a> {
+        let vec = if self.base.ctx.is_explain_verbose() {
+            let disp = Pretty::debug(&IndicesDisplay {
+                indices: &(0..self.schema().fields.len()).collect_vec(),
+                input_schema: self.schema(),
+            });
+            vec![("output", disp)]
+        } else {
+            vec![]
+        };
+
+        Pretty::childless_record("LogicalNow", vec)
+    }
+}
 impl fmt::Display for LogicalNow {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let verbose = self.base.ctx.is_explain_verbose();

src/frontend/src/optimizer/plan_node/logical_project_set.rs

@@ -12,13 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::fmt;
-
 use fixedbitset::FixedBitSet;
 use itertools::Itertools;
 use risingwave_common::error::Result;
 use risingwave_common::types::DataType;
 
+use super::utils::impl_distill_by_unit;
 use super::{
     gen_filter_and_pushdown, generic, BatchProjectSet, ColPrunable, ExprRewritable, LogicalProject,
     PlanBase, PlanRef, PlanTreeNodeUnary, PredicatePushdown, StreamProjectSet, ToBatch, ToStream,
@@ -187,13 +186,6 @@ impl LogicalProjectSet {
     pub fn select_list(&self) -> &Vec<ExprImpl> {
         &self.core.select_list
     }
-
-    pub(super) fn fmt_with_name(&self, f: &mut fmt::Formatter<'_>, name: &str) -> fmt::Result {
-        let _verbose = self.base.ctx.is_explain_verbose();
-        // TODO: add verbose display like Project
-
-        self.core.fmt_with_name(f, name)
-    }
 }
 
 impl PlanTreeNodeUnary for LogicalProjectSet {
@@ -225,12 +217,8 @@ impl PlanTreeNodeUnary for LogicalProjectSet {
 }
 
 impl_plan_tree_node_for_unary! {LogicalProjectSet}
-
-impl fmt::Display for LogicalProjectSet {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        self.fmt_with_name(f, "LogicalProjectSet")
-    }
-}
+impl_distill_by_unit!(LogicalProjectSet, core, "LogicalProjectSet");
+// TODO: add verbose display like Project
 
 impl ColPrunable for LogicalProjectSet {
     fn prune_col(&self, required_cols: &[usize], ctx: &mut ColumnPruningContext) -> PlanRef {

src/frontend/src/optimizer/plan_node/logical_scan.rs

@@ -18,11 +18,13 @@ use std::rc::Rc;
 
 use fixedbitset::FixedBitSet;
 use itertools::Itertools;
-use risingwave_common::catalog::{ColumnDesc, Field, Schema, TableDesc};
+use pretty_xmlish::Pretty;
+use risingwave_common::catalog::{ColumnDesc, TableDesc};
 use risingwave_common::error::{ErrorCode, Result, RwError};
 use risingwave_common::util::sort_util::ColumnOrder;
 
 use super::generic::{GenericPlanNode, GenericPlanRef};
+use super::utils::Distill;
 use super::{
     generic, BatchFilter, BatchProject, ColPrunable, ExprRewritable, PlanBase, PlanRef,
     PredicatePushdown, StreamTableScan, ToBatch, ToStream,
@@ -287,6 +289,52 @@ impl LogicalScan {
 
 impl_plan_tree_node_for_leaf! {LogicalScan}
 
+impl Distill for LogicalScan {
+    fn distill<'a>(&self) -> Pretty<'a> {
+        let verbose = self.base.ctx.is_explain_verbose();
+        let mut vec = Vec::with_capacity(5);
+        vec.push(("table", Pretty::from(self.table_name().to_owned())));
+        let key_is_columns =
+            self.predicate().always_true() || self.output_col_idx() == self.required_col_idx();
+        let key = if key_is_columns {
+            "columns"
+        } else {
+            "output_columns"
+        };
+        vec.push((key, self.core.columns_pretty(verbose)));
+        if !key_is_columns {
+            vec.push((
+                "required_columns",
+                Pretty::Array(
+                    self.required_col_idx()
+                        .iter()
+                        .map(|i| {
+                            let col_name = &self.table_desc().columns[*i].name;
+                            Pretty::from(if verbose {
+                                format!("{}.{}", self.table_name(), col_name)
+                            } else {
+                                col_name.to_string()
+                            })
+                        })
+                        .collect(),
+                ),
+            ));
+        }
+
+        if !self.predicate().always_true() {
+            let input_schema = self.core.fields_pretty_schema();
+            vec.push((
+                "predicate",
+                Pretty::display(&ConditionDisplay {
+                    condition: self.predicate(),
+                    input_schema: &input_schema,
+                }),
+            ))
+        }
+
+        Pretty::childless_record("LogicalScan", vec)
+    }
+}
 impl fmt::Display for LogicalScan {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let verbose = self.base.ctx.is_explain_verbose();
@@ -314,26 +362,17 @@ impl fmt::Display for LogicalScan {
                     ", output_columns: [{}], required_columns: [{}]",
                     output_col_names,
                     self.required_col_idx().iter().format_with(", ", |i, f| {
+                        let col_name = &self.table_desc().columns[*i].name;
                         if verbose {
-                            f(&format_args!(
-                                "{}.{}",
-                                self.table_name(),
-                                self.table_desc().columns[*i].name
-                            ))
+                            f(&format_args!("{}.{}", self.table_name(), col_name))
                         } else {
-                            f(&format_args!("{}", self.table_desc().columns[*i].name))
+                            f(&format_args!("{}", col_name))
                         }
                     })
                 )?;
             }
 
-            let fields = self
-                .table_desc()
-                .columns
-                .iter()
-                .map(|col| Field::from_with_table_name_prefix(col, self.table_name()))
-                .collect_vec();
-            let input_schema = Schema { fields };
+            let input_schema = self.core.fields_pretty_schema();
             write!(
                 f,
                 ", predicate: {} }}",

src/frontend/src/optimizer/plan_node/stream_now.rs

@@ -16,14 +16,15 @@ use std::fmt;
 
 use fixedbitset::FixedBitSet;
 use itertools::Itertools;
+use pretty_xmlish::Pretty;
 use risingwave_common::catalog::{Field, Schema};
 use risingwave_common::types::DataType;
 use risingwave_pb::stream_plan::stream_node::NodeBody;
 use risingwave_pb::stream_plan::NowNode;
 
 use super::generic::GenericPlanRef;
 use super::stream::StreamPlanRef;
-use super::utils::{formatter_debug_plan_node, IndicesDisplay, TableCatalogBuilder};
+use super::utils::{formatter_debug_plan_node, Distill, IndicesDisplay, TableCatalogBuilder};
 use super::{ExprRewritable, LogicalNow, PlanBase, StreamNode};
 use crate::optimizer::property::{Distribution, FunctionalDependencySet};
 use crate::stream_fragmenter::BuildFragmentGraphState;
@@ -58,6 +59,21 @@ impl StreamNow {
     }
 }
 
+impl Distill for StreamNow {
+    fn distill<'a>(&self) -> Pretty<'a> {
+        let vec = if self.base.ctx.is_explain_verbose() {
+            let disp = Pretty::debug(&IndicesDisplay {
+                indices: &(0..self.schema().fields.len()).collect_vec(),
+                input_schema: self.schema(),
+            });
+            vec![("output", disp)]
+        } else {
+            vec![]
+        };
+
+        Pretty::childless_record("StreamNow", vec)
+    }
+}
 impl fmt::Display for StreamNow {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let verbose = self.base.ctx.is_explain_verbose();