.

Signed-off-by: AilinKid <314806019@qq.com>
pingcap · Jan 24, 2025 · 874e4a4 · 874e4a4
1 parent bd2869b
commit 874e4a4
Show file tree

Hide file tree

Showing 6 changed files with 55 additions and 4 deletions.
diff --git a/pkg/planner/cascades/BUILD.bazel b/pkg/planner/cascades/BUILD.bazel
@@ -9,9 +9,11 @@ go_library(
         "//pkg/planner/cascades/base",
         "//pkg/planner/cascades/base/cascadesctx",
         "//pkg/planner/cascades/memo",
+        "//pkg/planner/cascades/rule",
         "//pkg/planner/cascades/task",
         "//pkg/planner/core/base",
         "//pkg/util/intest",
+        "@com_github_bits_and_blooms_bitset//:bitset",
     ],
 )
 

diff --git a/pkg/planner/cascades/base/cascadesctx/BUILD.bazel b/pkg/planner/cascades/base/cascadesctx/BUILD.bazel
@@ -8,5 +8,6 @@ go_library(
     deps = [
         "//pkg/planner/cascades/base",
         "//pkg/planner/cascades/memo",
+        "@com_github_bits_and_blooms_bitset//:bitset",
     ],
 )
diff --git a/pkg/planner/cascades/memo/memo.go b/pkg/planner/cascades/memo/memo.go
@@ -147,6 +147,24 @@ func (mm *Memo) CopyIn(target *Group, lp base.LogicalPlan) (*GroupExpression, er
 	return groupExpr, nil
 }
 
+// RemoveOut remove the old invalid GE out of target group, make sure insert first, then delete.
+func (mm *Memo) RemoveOut(target *Group, lp base.LogicalPlan) {
+	intest.Assert(target != nil)
+	intest.Assert(lp != nil)
+	ge := lp.(*GroupExpression)
+	intest.Assert(ge != nil)
+	// delete from group
+	target.Delete(ge)
+	// delete from global
+	mm.hash2GlobalGroupExpr.Remove(ge)
+	// maintain the parentGERef.
+	for _, childG := range ge.Inputs {
+		childG.removeParentGEs(ge)
+	}
+	// mark current ge as abandoned in case of it has been used in pushed task.
+	ge.SetAbandoned()
+}
+
 // GetGroups gets all groups in the memo.
 func (mm *Memo) GetGroups() *list.List {
 	return mm.groups

diff --git a/pkg/planner/cascades/rule/apply/decorrelate_apply/BUILD.bazel b/pkg/planner/cascades/rule/apply/decorrelate_apply/BUILD.bazel
@@ -1,8 +1,11 @@
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "decorrelate_apply",
-    srcs = ["xf_decorrelate_apply.go"],
+    srcs = [
+        "xf_decorrelate_apply_base.go",
+        "xf_decorrelate_simple_apply.go",
+    ],
     importpath = "github.com/pingcap/tidb/pkg/planner/cascades/rule/apply/decorrelate_apply",
     visibility = ["//visibility:public"],
     deps = [
@@ -15,3 +18,25 @@ go_library(
         "//pkg/util/plancodec",
     ],
 )
+
+go_test(
+    name = "decorrelate_apply_test",
+    timeout = "short",
+    srcs = ["xf_decorrelate_apply_test.go"],
+    flaky = True,
+    deps = [
+        ":decorrelate_apply",
+        "//pkg/expression",
+        "//pkg/parser/ast",
+        "//pkg/planner/cascades",
+        "//pkg/planner/cascades/base",
+        "//pkg/planner/cascades/memo",
+        "//pkg/planner/core/base",
+        "//pkg/planner/core/operator/logicalop",
+        "//pkg/types",
+        "//pkg/util/mock",
+        "//pkg/util/plancodec",
+        "@com_github_pingcap_failpoint//:failpoint",
+        "@com_github_stretchr_testify//require",
+    ],
+)
diff --git a/pkg/planner/cascades/rule/ruleset/BUILD.bazel b/pkg/planner/cascades/rule/ruleset/BUILD.bazel
@@ -6,8 +6,11 @@ go_library(
     importpath = "github.com/pingcap/tidb/pkg/planner/cascades/rule/ruleset",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/planner/cascades/memo",
         "//pkg/planner/cascades/pattern",
         "//pkg/planner/cascades/rule",
         "//pkg/planner/cascades/rule/apply/decorrelate_apply",
+        "//pkg/planner/core/operator/logicalop",
+        "@com_github_bits_and_blooms_bitset//:bitset",
     ],
 )
diff --git a/pkg/planner/cascades/task/task_apply_rule.go b/pkg/planner/cascades/task/task_apply_rule.go
@@ -107,8 +107,7 @@ func (a *ApplyRuleTask) Execute() error {
 		if !a.rule.PreCheck(holder) {
 			continue
 		}
-		// todo: remove the src GE, cause some src GE is intermediary.
-		newExprs, _, err := a.rule.XForm(holder)
+		newExprs, remove, err := a.rule.XForm(holder)
 		if err != nil {
 			return err
 		}
@@ -120,6 +119,9 @@ func (a *ApplyRuleTask) Execute() error {
 			// YAMS only care about logical plan now.
 			a.Push(NewOptGroupExpressionTask(a.ctx, newGroupExpr))
 		}
+		if remove {
+			a.ctx.GetMemo().RemoveOut(a.gE.GetGroup(), a.gE)
+		}
 	}
 	a.gE.SetExplored(a.rule.ID())
 	return nil