planner: fix index merge union case couldn't sink limit into index merge reader (#48590)

pkg/planner/core/casetest/physicalplantest/physical_plan_test.go

@@ -2232,6 +2232,7 @@ func TestIndexMergeSinkLimit(t *testing.T) {
 	tk.MustExec("insert into t2 values(1,2,1),(2,1,1),(3,3,1)")
 	tk.MustExec("create table t(a int, j json, index kj((cast(j as signed array))))")
 	tk.MustExec("insert into t values(1, '[1,2,3]')")
+	tk.MustExec("CREATE TABLE `t3` (\n  `id` int(11) NOT NULL,\n  `aid` bigint(20) DEFAULT NULL,\n  `c1` varchar(255) DEFAULT NULL,\n  `c2` varchar(255) DEFAULT NULL,\n  `d` int(11) DEFAULT NULL,\n  PRIMARY KEY (`id`) /*T![clustered_index] CLUSTERED */,\n  KEY `aid_c1` (`aid`,`c1`),\n  KEY `aid_c2` (`aid`,`c2`)\n)")
 
 	for i, ts := range input {
 		testdata.OnRecord(func() {

pkg/planner/core/casetest/physicalplantest/testdata/plan_suite_in.json

@@ -1049,7 +1049,8 @@
       "explain format = 'brief' select /*+ use_index(t, kj) */ * from t where json_overlaps(j, '[1, 2, 3]') limit 1; -- index merge union case, sink limit above selection above index merge reader, because json_overlaps can't be pushed down",
       "explain format = 'brief' select /*+ use_index(t, kj) */ * from t where (1 member of (j) and a=1 ) limit 1; -- index merge union case, sink limit to table side, because selection exists on table side",
       "explain format = 'brief' select /*+ use_index(t, kj) */ * from t where json_contains(j, '[1, 2, 3]') and a=1 limit 1; -- index merge intersection case, sink limit to table side because selection exists on table side",
-      "explain format = 'brief' select /*+ use_index(t, kj) */ * from t where json_overlaps(j, '[1, 2, 3]') and a=1 limit 1; -- index merge union case, sink limit above selection above index merge reader, because json_overlaps can't be pushed down"
+      "explain format = 'brief' select /*+ use_index(t, kj) */ * from t where json_overlaps(j, '[1, 2, 3]') and a=1 limit 1; -- index merge union case, sink limit above selection above index merge reader, because json_overlaps can't be pushed down",
+      "explain select /*+ USE_INDEX_MERGE(t3, aid_c1, aid_c2) */ * from t3 where (aid = 1 and c1='aaa') or (aid = 1 and c2='bbb') limit 1; -- index merge union and intersection case from issue 48588"
     ]
   }
 ]

pkg/planner/core/casetest/physicalplantest/testdata/plan_suite_out.json

@@ -7949,6 +7949,18 @@
         "Warning": [
           "Warning 1105 Scalar function 'json_overlaps'(signature: Unspecified, return type: bigint(20)) is not supported to push down to storage layer now."
         ]
+      },
+      {
+        "SQL": "explain select /*+ USE_INDEX_MERGE(t3, aid_c1, aid_c2) */ * from t3 where (aid = 1 and c1='aaa') or (aid = 1 and c2='bbb') limit 1; -- index merge union and intersection case from issue 48588",
+        "Plan": [
+          "IndexMerge_20 1.00 root  type: union, limit embedded(offset:0, count:1)",
+          "├─Limit_18(Build) 0.01 cop[tikv]  offset:0, count:1",
+          "│ └─IndexRangeScan_11 0.01 cop[tikv] table:t3, index:aid_c1(aid, c1) range:[1 \"aaa\",1 \"aaa\"], keep order:false, stats:pseudo",
+          "├─Limit_19(Build) 0.01 cop[tikv]  offset:0, count:1",
+          "│ └─IndexRangeScan_12 0.01 cop[tikv] table:t3, index:aid_c2(aid, c2) range:[1 \"bbb\",1 \"bbb\"], keep order:false, stats:pseudo",
+          "└─TableRowIDScan_13(Probe) 1.00 cop[tikv] table:t3 keep order:false, stats:pseudo"
+        ],
+        "Warning": null
       }
     ]
   }

pkg/planner/core/find_best_task.go

@@ -405,7 +405,10 @@ func getTaskPlanCost(t task, op *physicalOptimizeOp) (float64, bool, error) {
 	}
 
 	// use the new cost interface
-	var taskType property.TaskType
+	var (
+		taskType         property.TaskType
+		indexPartialCost float64
+	)
 	switch t.(type) {
 	case *rootTask:
 		taskType = property.RootTaskType
@@ -442,13 +445,28 @@ func getTaskPlanCost(t task, op *physicalOptimizeOp) (float64, bool, error) {
 				taskType = property.MppTaskType
 			}
 		}
+
+		// Detail reason ref about comment in function `convertToIndexMergeScan`
+		// for cop task with {indexPlan=nil, tablePlan=xxx, idxMergePartPlans=[x,x,x], indexPlanFinished=true} we should
+		// plus the partial index plan cost into the final cost. Because t.plan() the below code used only calculate the
+		// cost about table plan.
+		if cop.indexPlanFinished && len(cop.idxMergePartPlans) != 0 {
+			for _, partialScan := range cop.idxMergePartPlans {
+				partialCost, err := getPlanCost(partialScan, taskType, NewDefaultPlanCostOption().WithOptimizeTracer(op))
+				if err != nil {
+					return 0, false, err
+				}
+				indexPartialCost += partialCost
+			}
+		}
 	case *mppTask:
 		taskType = property.MppTaskType
 	default:
 		return 0, false, errors.New("unknown task type")
 	}
 	if t.plan() == nil {
 		// It's a very special case for index merge case.
+		// t.plan() == nil in index merge COP case, it means indexPlanFinished is false in other words.
 		cost := 0.0
 		copTsk := t.(*copTask)
 		for _, partialScan := range copTsk.idxMergePartPlans {
@@ -461,7 +479,7 @@ func getTaskPlanCost(t task, op *physicalOptimizeOp) (float64, bool, error) {
 		return cost, false, nil
 	}
 	cost, err := getPlanCost(t.plan(), taskType, NewDefaultPlanCostOption().WithOptimizeTracer(op))
-	return cost, false, err
+	return cost + indexPartialCost, false, err
 }
 
 type physicalOptimizeOp struct {
@@ -1374,6 +1392,11 @@ func (ds *DataSource) convertToIndexMergeScan(prop *property.PhysicalProperty, c
 	if remainingFilters != nil {
 		cop.rootTaskConds = remainingFilters
 	}
+	// after we lift the limitation of intersection and cop-type task in the code in this
+	// function above, we could set its index plan finished as true once we found its table
+	// plan is pure table scan below.
+	// And this will cause cost underestimation when we estimate the cost of the entire cop
+	// task plan in function `getTaskPlanCost`.
 	if prop.TaskTp == property.RootTaskType {
 		cop.indexPlanFinished = true
 		task = cop.convertToRootTask(ds.SCtx())

pkg/planner/core/task.go

@@ -913,7 +913,7 @@ func (p *PhysicalLimit) attach2Task(tasks ...task) task {
 					sunk = p.sinkIntoIndexMerge(t)
 				}
 			} else {
-				// otherwise, suspend the limit out of index merge reader.
+				// Otherwise, suspend the limit out of index merge reader.
 				t = cop.convertToRootTask(p.SCtx())
 				sunk = p.sinkIntoIndexMerge(t)
 			}