fix apply join transpose

risingwavelabs · Oct 22, 2024 · 9a5f4ba · 9a5f4ba
1 parent 5ca0a73
commit 9a5f4ba
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 14 deletions.
diff --git a/src/frontend/src/optimizer/logical_optimization.rs b/src/frontend/src/optimizer/logical_optimization.rs
@@ -200,10 +200,7 @@ static GENERAL_UNNESTING_TRANS_APPLY_WITH_SHARE: LazyLock<OptimizationStage> =
         OptimizationStage::new(
             "General Unnesting(Translate Apply)",
             vec![
-                TranslateApplyRule::create(true),
-                // Separate the project from a join if necessary because `ApplyJoinTransposeRule`
-                // can't handle a join with `output_indices`.
-                ProjectJoinSeparateRule::create(),
+                TranslateApplyRule::create(true)
             ],
             ApplyOrder::TopDown,
         )
@@ -214,10 +211,7 @@ static GENERAL_UNNESTING_TRANS_APPLY_WITHOUT_SHARE: LazyLock<OptimizationStage>
         OptimizationStage::new(
             "General Unnesting(Translate Apply)",
             vec![
-                TranslateApplyRule::create(false),
-                // Separate the project from a join if necessary because `ApplyJoinTransposeRule`
-                // can't handle a join with `output_indices`.
-                ProjectJoinSeparateRule::create(),
+                TranslateApplyRule::create(false)
             ],
             ApplyOrder::TopDown,
         )

diff --git a/src/frontend/src/optimizer/rule/apply_join_transpose_rule.rs b/src/frontend/src/optimizer/rule/apply_join_transpose_rule.rs
@@ -24,7 +24,7 @@ use crate::expr::{
     InputRef,
 };
 use crate::optimizer::plan_node::generic::GenericPlanRef;
-use crate::optimizer::plan_node::{LogicalApply, LogicalFilter, LogicalJoin, PlanTreeNodeBinary};
+use crate::optimizer::plan_node::{LogicalApply, LogicalFilter, LogicalJoin, PlanTreeNode, PlanTreeNodeBinary};
 use crate::optimizer::plan_visitor::{ExprCorrelatedIdFinder, PlanCorrelatedIdFinder};
 use crate::optimizer::rule::apply_offset_rewriter::ApplyCorrelatedIndicesConverter;
 use crate::optimizer::PlanRef;
@@ -122,10 +122,12 @@ impl Rule for ApplyJoinTransposeRule {
             return None;
         }
 
-        assert!(
-            join.output_indices_are_trivial(),
-            "ApplyJoinTransposeRule requires the join containing no output indices, so make sure ProjectJoinSeparateRule is always applied before this rule"
-        );
+        // ApplyJoinTransposeRule requires the join containing no output indices, so make sure ProjectJoinSeparateRule is always applied before this rule.
+        // As this rule will be applied until we reach the fixed point, if the join has output indices, apply ProjectJoinSeparateRule first and return is safety.
+        if !join.output_indices_are_trivial() {
+            let new_apply_right = crate::optimizer::rule::ProjectJoinSeparateRule::create().apply(join.clone().into()).unwrap();
+            return Some(apply.clone_with_inputs(&[apply_left, new_apply_right]));
+        }
 
         let (push_left, push_right) = match join.join_type() {
             // `LeftSemi`, `LeftAnti`, `LeftOuter` can only push to left side if it's right side has

diff --git a/src/frontend/src/optimizer/rule/project_join_separate_rule.rs b/src/frontend/src/optimizer/rule/project_join_separate_rule.rs
@@ -25,7 +25,7 @@ impl ProjectJoinSeparateRule {
 impl Rule for ProjectJoinSeparateRule {
     fn apply(&self, plan: PlanRef) -> Option<PlanRef> {
         let join = plan.as_logical_join()?;
-        if join.is_full_out() {
+        if join.output_indices_are_trivial() {
             None
         } else {
             let (left, right, on, join_type, output_indices) = join.clone().decompose();