fix: expand on conditional commutative as well (#6484)

* fix: expand on conditional commutative as well

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
Signed-off-by: discord9 <discord9@163.com>

* update sqlness result

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
Signed-off-by: discord9 <discord9@163.com>

* add logging to figure test failure

Signed-off-by: discord9 <discord9@163.com>

* revert

Signed-off-by: discord9 <discord9@163.com>

* feat: stream drop record metrics

Signed-off-by: discord9 <discord9@163.com>

* Revert "feat: stream drop record metrics"

This reverts commit 6a16946a5b8ea37557bbb1b600847d24274d6500.

Signed-off-by: discord9 <discord9@163.com>

* feat: stream drop record metrics

Signed-off-by: discord9 <discord9@163.com>

refactor: move logging to drop too

Signed-off-by: discord9 <discord9@163.com>

fix: drop input stream before collect metrics

Signed-off-by: discord9 <discord9@163.com>

* fix: expand differently

Signed-off-by: discord9 <discord9@163.com>

* test: update sqlness

Signed-off-by: discord9 <discord9@163.com>

* chore: more dbg

Signed-off-by: discord9 <discord9@163.com>

* Revert "feat: stream drop record metrics"

This reverts commit 3eda4a2257928d95cf9c1328ae44fae84cfbb017.

Signed-off-by: discord9 <discord9@163.com>

* test: sqlness redacted

Signed-off-by: discord9 <discord9@163.com>

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
Signed-off-by: discord9 <discord9@163.com>
Co-authored-by: discord9 <discord9@163.com>
Signed-off-by: evenyag <realevenyag@gmail.com>

src/query/src/dist_plan/analyzer.rs

@@ -155,7 +155,23 @@ struct PlanRewriter {
     /// Partition columns of the table in current pass
     partition_cols: Option<Vec<String>>,
     column_requirements: HashSet<Column>,
+    /// Whether to expand on next call
+    /// This is used to handle the case where a plan is transformed, but need to be expanded from it's
+    /// parent node. For example a Aggregate plan is split into two parts in frontend and datanode, and need
+    /// to be expanded from the parent node of the Aggregate plan.
     expand_on_next_call: bool,
+    /// Expanding on next partial/conditional/transformed commutative plan
+    /// This is used to handle the case where a plan is transformed, but still
+    /// need to push down as many node as possible before next partial/conditional/transformed commutative
+    /// plan. I.e.
+    /// ```
+    /// Limit:
+    ///     Sort:
+    /// ```
+    /// where `Limit` is partial commutative, and `Sort` is conditional commutative.
+    /// In this case, we need to expand the `Limit` plan,
+    /// so that we can push down the `Sort` plan as much as possible.
+    expand_on_next_part_cond_trans_commutative: bool,
     new_child_plan: Option<LogicalPlan>,
 }
 
@@ -177,15 +193,38 @@ impl PlanRewriter {
         {
             return true;
         }
+
         if self.expand_on_next_call {
             self.expand_on_next_call = false;
             return true;
         }
+
+        if self.expand_on_next_part_cond_trans_commutative {
+            let comm = Categorizer::check_plan(plan, self.partition_cols.clone());
+            match comm {
+                Commutativity::PartialCommutative => {
+                    // a small difference is that for partial commutative, we still need to
+                    // expand on next call(so `Limit` can be pushed down)
+                    self.expand_on_next_part_cond_trans_commutative = false;
+                    self.expand_on_next_call = true;
+                }
+                Commutativity::ConditionalCommutative(_)
+                | Commutativity::TransformedCommutative { .. } => {
+                    // for conditional commutative and transformed commutative, we can
+                    // expand now
+                    self.expand_on_next_part_cond_trans_commutative = false;
+                    return true;
+                }
+                _ => (),
+            }
+        }
+
         match Categorizer::check_plan(plan, self.partition_cols.clone()) {
             Commutativity::Commutative => {}
             Commutativity::PartialCommutative => {
                 if let Some(plan) = partial_commutative_transformer(plan) {
                     self.update_column_requirements(&plan);
+                    self.expand_on_next_part_cond_trans_commutative = true;
                     self.stage.push(plan)
                 }
             }
@@ -194,6 +233,7 @@ impl PlanRewriter {
                     && let Some(plan) = transformer(plan)
                 {
                     self.update_column_requirements(&plan);
+                    self.expand_on_next_part_cond_trans_commutative = true;
                     self.stage.push(plan)
                 }
             }
@@ -202,7 +242,7 @@ impl PlanRewriter {
                     && let Some(transformer_actions) = transformer(plan)
                 {
                     debug!(
-                        "PlanRewriter: transformed plan: {:#?}\n from {plan}",
+                        "PlanRewriter: transformed plan: {:?}\n from {plan}",
                         transformer_actions.extra_parent_plans
                     );
                     if let Some(last_stage) = transformer_actions.extra_parent_plans.last() {
@@ -226,6 +266,10 @@ impl PlanRewriter {
     }
 
     fn update_column_requirements(&mut self, plan: &LogicalPlan) {
+        debug!(
+            "PlanRewriter: update column requirements for plan: {plan}\n withcolumn_requirements: {:?}",
+            self.column_requirements
+        );
         let mut container = HashSet::new();
         for expr in plan.expressions() {
             // this method won't fail
@@ -235,6 +279,10 @@ impl PlanRewriter {
         for col in container {
             self.column_requirements.insert(col);
         }
+        debug!(
+            "PlanRewriter: updated column requirements: {:?}",
+            self.column_requirements
+        );
     }
 
     fn is_expanded(&self) -> bool {