Adds support for subquery representation in plans

Signed-off-by: coldWater <[email protected]>

Author: forsaken628

src/query/sql/src/planner/format/display.rs

@@ -157,6 +157,40 @@ where
             tree.children.extend(stats);
         }
 
+        let subquerys = op.get_subquery(vec![]);
+        if !subquerys.is_empty() {
+            let subquerys = subquerys
+                .into_iter()
+                .map(|subquery| {
+                    let children = vec![
+                        subquery
+                            .compare_op
+                            .map(|compare| FormatTreeNode::new(format!("compare_op: {compare:?}"))),
+                        Some(FormatTreeNode::new(format!(
+                            "output_column: {}",
+                            self.id_humanizer
+                                .humanize_column_id(subquery.output_column.index)
+                        ))),
+                        subquery.projection_index.map(|id| {
+                            FormatTreeNode::new(format!(
+                                "projection_index: {}",
+                                self.id_humanizer.humanize_column_id(id)
+                            ))
+                        }),
+                        Some(subquery.subquery.to_format_tree(self.id_humanizer)?),
+                    ];
+                    Ok(FormatTreeNode::with_children(
+                        format!("Subquery ({:?})", subquery.typ),
+                        children.into_iter().filter_map(|x| x).collect(),
+                    ))
+                })
+                .collect::<Result<Vec<_>>>()?;
+            tree.children.push(FormatTreeNode::with_children(
+                "subquerys".to_string(),
+                subquerys,
+            ));
+        }
+
         let children = s_expr
             .children()
             .map(|s_expr| self.humanize_s_expr(s_expr))

src/query/sql/src/planner/optimizer/ir/expr/s_expr.rs

@@ -175,10 +175,7 @@ impl SExpr {
     /// Check if contain subquery
     #[recursive::recursive]
     pub(crate) fn has_subquery(&self) -> bool {
-        if !self.plan.has_subquery() {
-            return self.children.iter().any(|child| child.has_subquery());
-        }
-        true
+        self.plan.has_subquery() || self.children.iter().any(|child| child.has_subquery())
     }
 
     //

src/query/sql/src/planner/plans/operator.rs

@@ -20,6 +20,7 @@ use databend_common_exception::Result;
 use educe::Educe;
 
 use super::MutationSource;
+use super::SubqueryExpr;
 use crate::optimizer::ir::PhysicalProperty;
 use crate::optimizer::ir::RelExpr;
 use crate::optimizer::ir::RelationalProperty;
@@ -209,6 +210,76 @@ impl RelOperator {
             RelOperator::MutationSource(_) => false,
         }
     }
+
+    pub fn get_subquery(&self, mut result: Vec<SubqueryExpr>) -> Vec<SubqueryExpr> {
+        match self {
+            RelOperator::Scan(_)
+            | RelOperator::Limit(_)
+            | RelOperator::Exchange(_)
+            | RelOperator::UnionAll(_)
+            | RelOperator::Sort(_)
+            | RelOperator::DummyTableScan(_)
+            | RelOperator::ConstantTableScan(_)
+            | RelOperator::ExpressionScan(_)
+            | RelOperator::CacheScan(_)
+            | RelOperator::AsyncFunction(_)
+            | RelOperator::RecursiveCteScan(_)
+            | RelOperator::Mutation(_)
+            | RelOperator::CompactBlock(_)
+            | RelOperator::MutationSource(_) => (),
+            RelOperator::Join(op) => {
+                for condition in &op.equi_conditions {
+                    result = condition.left.get_subquery(result);
+                    result = condition.right.get_subquery(result);
+                }
+                for expr in &op.non_equi_conditions {
+                    result = expr.get_subquery(result);
+                }
+            }
+            RelOperator::EvalScalar(op) => {
+                for item in &op.items {
+                    result = item.scalar.get_subquery(result);
+                }
+            }
+            RelOperator::Filter(op) => {
+                for pred in &op.predicates {
+                    result = pred.get_subquery(result);
+                }
+            }
+            RelOperator::Aggregate(op) => {
+                for item in &op.group_items {
+                    result = item.scalar.get_subquery(result);
+                }
+                for func in &op.aggregate_functions {
+                    result = func.scalar.get_subquery(result);
+                }
+            }
+            RelOperator::Window(op) => {
+                for order in &op.order_by {
+                    result = order.order_by_item.scalar.get_subquery(result);
+                }
+                for expr in &op.partition_by {
+                    result = expr.scalar.get_subquery(result);
+                }
+                if let WindowFuncType::Aggregate(agg) = &op.function {
+                    for expr in agg.exprs() {
+                        result = expr.get_subquery(result);
+                    }
+                }
+            }
+            RelOperator::ProjectSet(op) => {
+                for srf in &op.srfs {
+                    result = srf.scalar.get_subquery(result);
+                }
+            }
+            RelOperator::Udf(op) => {
+                for item in &op.items {
+                    result = item.scalar.get_subquery(result);
+                }
+            }
+        }
+        result
+    }
 }
 
 impl Operator for RelOperator {

src/query/sql/src/planner/plans/scalar_expr.rs

@@ -420,6 +420,23 @@ impl ScalarExpr {
         has_subquery.visit(self).unwrap();
         has_subquery.has_subquery
     }
+
+    pub fn get_subquery(&self, result: Vec<SubqueryExpr>) -> Vec<SubqueryExpr> {
+        struct GetSubquery {
+            subquerys: Vec<SubqueryExpr>,
+        }
+
+        impl<'a> Visitor<'a> for GetSubquery {
+            fn visit_subquery(&mut self, subquery: &'a SubqueryExpr) -> Result<()> {
+                self.subquerys.push(subquery.clone());
+                Ok(())
+            }
+        }
+
+        let mut get_subquery = GetSubquery { subquerys: result };
+        get_subquery.visit(self).unwrap();
+        get_subquery.subquerys
+    }
 }
 
 impl From<BoundColumnRef> for ScalarExpr {

tests/sqllogictests/suites/mode/standalone/explain/subquery.test

@@ -296,6 +296,31 @@ HashJoin
     ├── push downs: [filters: [], limit: NONE]
     └── estimated rows: 1.00
 
+query T
+explain raw select * from numbers(1) as t where exists (select number as a from numbers(1) where number = t.number)
+----
+EvalScalar
+├── scalars: [t.number (#0) AS (#0)]
+└── Filter
+    ├── filters: [SUBQUERY]
+    ├── subquerys
+    │   └── Subquery (Exists)
+    │       ├── output_column: system.numbers.number (#1)
+    │       └── EvalScalar
+    │           ├── scalars: [numbers.number (#1) AS (#1)]
+    │           └── Filter
+    │               ├── filters: [eq(numbers.number (#1), t.number (#0))]
+    │               └── Scan
+    │                   ├── table: system.numbers (#1)
+    │                   ├── filters: []
+    │                   ├── order by: []
+    │                   └── limit: NONE
+    └── Scan
+        ├── table: system.numbers (#0)
+        ├── filters: []
+        ├── order by: []
+        └── limit: NONE
+
 query T
 explain select t.number from numbers(1) as t where exists (select * from numbers(1) where number = t.number and number = 0 and t.number < 10)
 ----