refactor(plan_node_fmt): initial rewrite of fmt_with_name based on pretty_xmlish (#10132)

Author: ice1000

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

@@ -15,9 +15,10 @@
 use std::fmt;
 
 use itertools::Itertools;
+use pretty_xmlish::Pretty;
 use risingwave_common::catalog::{FieldDisplay, Schema};
 
-use super::{GenericPlanNode, GenericPlanRef};
+use super::{DistillUnit, GenericPlanNode, GenericPlanRef};
 use crate::optimizer::property::FunctionalDependencySet;
 use crate::OptimizerContextRef;
 
@@ -41,6 +42,22 @@ impl<PlanRef: GenericPlanRef> Dedup<PlanRef> {
             .map(|i| FieldDisplay(self.input.schema().fields.get(*i).unwrap()))
             .collect_vec()
     }
+
+    fn dedup_cols_pretty<'a>(&self) -> Pretty<'a> {
+        Pretty::Array(
+            self.dedup_cols
+                .iter()
+                .map(|i| FieldDisplay(self.input.schema().fields.get(*i).unwrap()))
+                .map(|fd| Pretty::display(&fd))
+                .collect(),
+        )
+    }
+}
+
+impl<PlanRef: GenericPlanRef> DistillUnit for Dedup<PlanRef> {
+    fn distill_with_name<'a>(&self, name: &'a str) -> Pretty<'a> {
+        Pretty::childless_record(name, vec![("dedup_cols", self.dedup_cols_pretty())])
+    }
 }
 
 impl<PlanRef: GenericPlanRef> GenericPlanNode for Dedup<PlanRef> {

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

@@ -13,9 +13,10 @@
 // limitations under the License.
 use std::fmt;
 
+use pretty_xmlish::Pretty;
 use risingwave_common::catalog::{Schema, TableVersionId};
 
-use super::GenericPlanRef;
+use super::{DistillUnit, GenericPlanRef};
 use crate::catalog::TableId;
 use crate::OptimizerContextRef;
 
@@ -69,3 +70,14 @@ impl<PlanRef> Delete<PlanRef> {
         )
     }
 }
+
+impl<PlanRef> DistillUnit for Delete<PlanRef> {
+    fn distill_with_name<'a>(&self, name: &'a str) -> Pretty<'a> {
+        let mut vec = Vec::with_capacity(if self.returning { 2 } else { 1 });
+        vec.push(("table", Pretty::Text(self.table_name.clone().into())));
+        if self.returning {
+            vec.push(("returning", Pretty::display(&true)));
+        }
+        Pretty::childless_record(name, vec)
+    }
+}

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

@@ -14,9 +14,10 @@
 
 use std::fmt;
 
+use pretty_xmlish::Pretty;
 use risingwave_common::catalog::Schema;
 
-use super::{GenericPlanNode, GenericPlanRef};
+use super::{DistillUnit, GenericPlanNode, GenericPlanRef};
 use crate::optimizer::optimizer_context::OptimizerContextRef;
 use crate::optimizer::property::FunctionalDependencySet;
 
@@ -49,11 +50,13 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for Except<PlanRef> {
 impl<PlanRef: GenericPlanRef> Except<PlanRef> {
     pub fn fmt_with_name(&self, f: &mut fmt::Formatter<'_>, name: &str) -> fmt::Result {
         let mut builder = f.debug_struct(name);
-        self.fmt_fields_with_builder(&mut builder);
+        builder.field("all", &self.all);
         builder.finish()
     }
+}
 
-    pub fn fmt_fields_with_builder(&self, builder: &mut fmt::DebugStruct<'_, '_>) {
-        builder.field("all", &self.all);
+impl<PlanRef> DistillUnit for Except<PlanRef> {
+    fn distill_with_name<'a>(&self, name: &'a str) -> Pretty<'a> {
+        Pretty::childless_record(name, vec![("all", Pretty::debug(&self.all))])
     }
 }

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

@@ -15,11 +15,12 @@
 use std::fmt;
 
 use itertools::Itertools;
+use pretty_xmlish::Pretty;
 use risingwave_common::catalog::{Field, FieldDisplay, Schema};
 use risingwave_common::types::DataType;
 use risingwave_common::util::column_index_mapping::ColIndexMapping;
 
-use super::{GenericPlanNode, GenericPlanRef};
+use super::{DistillUnit, GenericPlanNode, GenericPlanRef};
 use crate::optimizer::optimizer_context::OptimizerContextRef;
 use crate::optimizer::property::FunctionalDependencySet;
 
@@ -100,8 +101,30 @@ impl<PlanRef: GenericPlanRef> GenericPlanNode for Expand<PlanRef> {
     }
 }
 
+impl<PlanRef: GenericPlanRef> DistillUnit for Expand<PlanRef> {
+    fn distill_with_name<'a>(&self, name: &'a str) -> Pretty<'a> {
+        Pretty::childless_record(name, vec![("column_subsets", self.column_subsets_pretty())])
+    }
+}
+
 impl<PlanRef: GenericPlanRef> Expand<PlanRef> {
-    pub fn column_subsets_display(&self) -> Vec<Vec<FieldDisplay<'_>>> {
+    fn column_subsets_pretty<'a>(&self) -> Pretty<'a> {
+        Pretty::Array(
+            self.column_subsets
+                .iter()
+                .map(|subset| {
+                    subset
+                        .iter()
+                        .map(|&i| FieldDisplay(self.input.schema().fields.get(i).unwrap()))
+                        .map(|d| Pretty::display(&d))
+                        .collect()
+                })
+                .map(Pretty::Array)
+                .collect(),
+        )
+    }
+
+    fn column_subsets_display(&self) -> Vec<Vec<FieldDisplay<'_>>> {
         self.column_subsets
             .iter()
             .map(|subset| {

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

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use pretty_xmlish::Pretty;
 use risingwave_common::catalog::Schema;
 
 use super::{stream, EqJoinPredicate};
@@ -57,6 +58,10 @@ pub use update::*;
 mod delete;
 pub use delete::*;
 
+pub trait DistillUnit {
+    fn distill_with_name<'a>(&self, name: &'a str) -> Pretty<'a>;
+}
+
 pub trait GenericPlanRef {
     fn schema(&self) -> &Schema;
     fn logical_pk(&self) -> &[usize];

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

@@ -15,14 +15,15 @@
 use std::fmt;
 
 use itertools::Itertools;
+use pretty_xmlish::Pretty;
 use risingwave_common::catalog::{Field, Schema};
 use risingwave_common::types::DataType;
 use risingwave_common::util::column_index_mapping::ColIndexMapping;
 use risingwave_common::util::sort_util::{ColumnOrder, ColumnOrderDisplay};
 use risingwave_expr::function::window::{Frame, WindowFuncKind};
 use risingwave_pb::expr::PbWindowFunction;
 
-use super::{GenericPlanNode, GenericPlanRef};
+use super::{DistillUnit, GenericPlanNode, GenericPlanRef};
 use crate::expr::{InputRef, InputRefDisplay};
 use crate::optimizer::property::FunctionalDependencySet;
 use crate::utils::ColIndexMappingRewriteExt;
@@ -178,6 +179,20 @@ impl<PlanRef: GenericPlanRef> OverWindow<PlanRef> {
     }
 }
 
+impl<PlanRef: GenericPlanRef> DistillUnit for OverWindow<PlanRef> {
+    fn distill_with_name<'a>(&self, name: &'a str) -> Pretty<'a> {
+        let f = |func| {
+            Pretty::debug(&PlanWindowFunctionDisplay {
+                window_function: func,
+                input_schema: self.input.schema(),
+            })
+        };
+        let wf = Pretty::Array(self.window_functions.iter().map(f).collect());
+        let vec = vec![("window_functions", wf)];
+        Pretty::childless_record(name, vec)
+    }
+}
+
 impl<PlanRef: GenericPlanRef> GenericPlanNode for OverWindow<PlanRef> {
     fn functional_dependency(&self) -> FunctionalDependencySet {
         let mapping =

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

@@ -14,10 +14,11 @@
 
 use std::fmt;
 
+use pretty_xmlish::Pretty;
 use risingwave_common::catalog::{Field, Schema};
 use risingwave_common::types::DataType;
 
-use super::{GenericPlanNode, GenericPlanRef};
+use super::{DistillUnit, GenericPlanNode, GenericPlanRef};
 use crate::expr::{Expr, ExprDisplay, ExprImpl, ExprRewriter};
 use crate::optimizer::optimizer_context::OptimizerContextRef;
 use crate::optimizer::plan_node::batch::BatchPlanRef;
@@ -58,6 +59,13 @@ impl<PlanRef> ProjectSet<PlanRef> {
     }
 }
 
+impl<PlanRef> DistillUnit for ProjectSet<PlanRef> {
+    fn distill_with_name<'a>(&self, name: &'a str) -> Pretty<'a> {
+        let fields = vec![("select_list", Pretty::debug(&self.select_list))];
+        Pretty::childless_record(name, fields)
+    }
+}
+
 impl<PlanRef: GenericPlanRef> GenericPlanNode for ProjectSet<PlanRef> {
     fn schema(&self) -> Schema {
         let input_schema = self.input.schema();

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

@@ -15,11 +15,12 @@
 use std::collections::HashSet;
 use std::fmt;
 
+use pretty_xmlish::Pretty;
 use risingwave_common::catalog::Schema;
 use risingwave_common::util::sort_util::OrderType;
 
 use super::super::utils::TableCatalogBuilder;
-use super::{stream, GenericPlanNode, GenericPlanRef};
+use super::{stream, DistillUnit, GenericPlanNode, GenericPlanRef};
 use crate::optimizer::optimizer_context::OptimizerContextRef;
 use crate::optimizer::property::{FunctionalDependencySet, Order, OrderDisplay};
 use crate::TableCatalog;
@@ -146,6 +147,27 @@ impl<PlanRef: GenericPlanRef> TopN<PlanRef> {
     }
 }
 
+impl<PlanRef: GenericPlanRef> DistillUnit for TopN<PlanRef> {
+    fn distill_with_name<'a>(&self, name: &'a str) -> Pretty<'a> {
+        let mut vec = Vec::with_capacity(5);
+        let input_schema = self.input.schema();
+        let order_d = Pretty::display(&OrderDisplay {
+            order: &self.order,
+            input_schema,
+        });
+        vec.push(("order", order_d));
+        vec.push(("limit", Pretty::debug(&self.limit_attr.limit())));
+        vec.push(("offset", Pretty::debug(&self.offset)));
+        if self.limit_attr.with_ties() {
+            vec.push(("with_ties", Pretty::debug(&true)));
+        }
+        if !self.group_key.is_empty() {
+            vec.push(("group_key", Pretty::debug(&self.group_key)));
+        }
+        Pretty::childless_record(name, vec)
+    }
+}
+
 impl<PlanRef: GenericPlanRef> GenericPlanNode for TopN<PlanRef> {
     fn schema(&self) -> Schema {
         self.input.schema().clone()

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

@@ -16,8 +16,10 @@ use std::fmt;
 use std::hash::Hash;
 
 use educe::Educe;
+use pretty_xmlish::Pretty;
 use risingwave_common::catalog::TableVersionId;
 
+use super::DistillUnit;
 use crate::catalog::TableId;
 use crate::expr::{ExprImpl, ExprRewriter};
 
@@ -76,3 +78,15 @@ impl<PlanRef: Eq + Hash> Update<PlanRef> {
             .collect();
     }
 }
+
+impl<PlanRef: Eq + Hash> DistillUnit for Update<PlanRef> {
+    fn distill_with_name<'a>(&self, name: &'a str) -> Pretty<'a> {
+        let mut vec = Vec::with_capacity(if self.returning { 3 } else { 2 });
+        vec.push(("table", Pretty::Text(self.table_name.clone().into())));
+        vec.push(("exprs", Pretty::debug(&self.exprs)));
+        if self.returning {
+            vec.push(("returning", Pretty::display(&true)));
+        }
+        Pretty::childless_record(name, vec)
+    }
+}

src/frontend/src/optimizer/plan_node/logical_except.rs

@@ -50,10 +50,6 @@ impl LogicalExcept {
         self.core.fmt_with_name(f, name)
     }
 
-    pub fn fmt_fields_with_builder(&self, builder: &mut fmt::DebugStruct<'_, '_>) {
-        self.core.fmt_fields_with_builder(builder)
-    }
-
     pub fn all(&self) -> bool {
         self.core.all
     }

src/frontend/src/optimizer/plan_node/logical_project.rs

@@ -88,11 +88,6 @@ impl LogicalProject {
         self.core.fmt_with_name(f, name, self.base.schema())
     }
 
-    pub fn fmt_fields_with_builder(&self, builder: &mut fmt::DebugStruct<'_, '_>) {
-        self.core
-            .fmt_fields_with_builder(builder, self.base.schema())
-    }
-
     pub fn is_identity(&self) -> bool {
         self.core.is_identity()
     }