executor: preserve ParentIdx for correlated index_lookup_pushdown under Apply

pkg/executor/builder.go

@@ -4565,6 +4565,7 @@ func buildNoRangeIndexLookUpReader(b *executorBuilder, v *physicalop.PhysicalInd
 		idxCols:                    is.IdxCols,
 		colLens:                    is.IdxColLens,
 		idxPlans:                   v.IndexPlans,
+		idxPlanUnNatureOrders:      v.IndexPlansUnNatureOrders,
 		tblPlans:                   v.TablePlans,
 		PushedLimit:                v.PushedLimit,
 		idxNetDataSize:             v.GetAvgTableRowSize(),

pkg/executor/distsql.go

@@ -518,13 +518,14 @@ type IndexLookUpExecutor struct {
 
 	indexPaging bool
 
-	corColInIdxSide bool
-	corColInTblSide bool
-	corColInAccess  bool
-	idxPlans        []base.PhysicalPlan
-	tblPlans        []base.PhysicalPlan
-	idxCols         []*expression.Column
-	colLens         []int
+	corColInIdxSide       bool
+	corColInTblSide       bool
+	corColInAccess        bool
+	idxPlans              []base.PhysicalPlan
+	idxPlanUnNatureOrders map[int]int
+	tblPlans              []base.PhysicalPlan
+	idxCols               []*expression.Column
+	colLens               []int
 	// PushedLimit is used to skip the preceding and tailing handles when Limit is sunk into IndexLookUpReader.
 	PushedLimit *physicalop.PushedDownLimit
 
@@ -702,7 +703,11 @@ func (e *IndexLookUpExecutor) open(_ context.Context) error {
 
 	var err error
 	if e.corColInIdxSide {
-		e.dagPB.Executors, err = builder.ConstructListBasedDistExec(e.buildPBCtx, e.idxPlans)
+		if e.indexLookUpPushDown {
+			e.dagPB.Executors, err = builder.ConstructListBasedDistExecForUnNatureOrderPlans(e.buildPBCtx, e.idxPlans, e.idxPlanUnNatureOrders)
+		} else {
+			e.dagPB.Executors, err = builder.ConstructListBasedDistExec(e.buildPBCtx, e.idxPlans)
+		}
 		if err != nil {
 			return err
 		}

pkg/executor/internal/builder/builder_utils.go

@@ -44,6 +44,22 @@ func ConstructListBasedDistExec(pctx *planctx.BuildPBContext, plans []plannercor
 	return executors, nil
 }
 
+// ConstructListBasedDistExecForUnNatureOrderPlans constructs list based executors and
+// sets ParentIdx for those executors whose parent is not the next one in the list.
+func ConstructListBasedDistExecForUnNatureOrderPlans(
+	pctx *planctx.BuildPBContext, plans []plannercore.PhysicalPlan, unNatureOrders map[int]int,
+) ([]*tipb.Executor, error) {
+	executors, err := ConstructListBasedDistExec(pctx, plans)
+	if err != nil {
+		return nil, err
+	}
+	for i, j := range unNatureOrders {
+		parentIdx := uint32(j)
+		executors[i].ParentIdx = &parentIdx
+	}
+	return executors, nil
+}
+
 // ConstructDAGReq constructs DAGRequest for physical plans
 func ConstructDAGReq(ctx sessionctx.Context, plans []plannercore.PhysicalPlan, storeType kv.StoreType) (dagReq *tipb.DAGRequest, err error) {
 	dagReq = &tipb.DAGRequest{}

pkg/store/mockstore/unistore/cophandler/cop_handler.go

@@ -139,7 +139,7 @@ func ExecutorListsToTree(exec []*tipb.Executor) *tipb.Executor {
 		case tipb.ExecType_TypeIndexLookUp:
 			parent.IndexLookup.Children = append(parent.IndexLookup.Children, child)
 		default:
-			panic("unsupported dag executor type")
+			panic("unsupported dag parent executor type: " + parent.Tp.String())
 		}
 	}
 

tests/integrationtest/r/executor/index_lookup_pushdown.result

@@ -1,4 +1,4 @@
-drop table if exists t1, t2, t3, t4, t5, t6, t7, tmp1, tmp2;
+drop table if exists t1, t2, t3, t4, t5, t6, t7, t8, t9, tmp1, tmp2;
 set @@tidb_scatter_region='table';
 create table t1(id int primary key, a varchar(32), b int, c int, index i(a, b));
 create table t2(a varchar(32), b int, c int, d int, e int, primary key(a, b) CLUSTERED, index i(c), unique index u(e)) CHARSET=utf8mb4 COLLATE=utf8mb4_general_ci;
@@ -565,4 +565,43 @@ id	a	b
 1	2	3
 Level	Code	Message
 Warning	1815	hint INDEX_LOOKUP_PUSHDOWN is inapplicable, only leader read is supported
+set @@tidb_replica_read='leader';
+create table t8 (a int, b int, key idx_a(a));
+create table t9 (a int, b int, key idx_a(a));
+insert into t8 values (20,1),(24,5),(23,9);
+insert into t9 values (1,1),(3,2),(6,3),(10,4),(12,5),(15,6),(25,7);
+explain format='plan_tree' select *
+from t8
+where t8.a < (
+select /*+ INDEX_LOOKUP_PUSHDOWN(t9, idx_a) */ sum(t9.b)
+from t9 use index(idx_a)
+where t9.a > t8.b
+)
+order by t8.a, t8.b;
+id	task	access object	operator info
+Sort	root		executor__index_lookup_pushdown.t8.a, executor__index_lookup_pushdown.t8.b
+└─Projection	root		executor__index_lookup_pushdown.t8.a, executor__index_lookup_pushdown.t8.b
+  └─Apply	root		CARTESIAN inner join, other cond:lt(cast(executor__index_lookup_pushdown.t8.a, decimal(10,0) BINARY), Column)
+    ├─TableReader(Build)	root		data:TableFullScan
+    │ └─TableFullScan	cop[tikv]	table:t8	keep order:false, stats:pseudo
+    └─HashAgg(Probe)	root		funcs:sum(Column)->Column
+      └─IndexLookUp	root		
+        ├─HashAgg(Build)	cop[tikv]		funcs:sum(executor__index_lookup_pushdown.t9.b)->Column
+        │ └─LocalIndexLookUp	cop[tikv]		index handle offsets:[1]
+        │   ├─Selection(Build)	cop[tikv]		gt(executor__index_lookup_pushdown.t9.a, executor__index_lookup_pushdown.t8.b)
+        │   │ └─IndexRangeScan	cop[tikv]	table:t9, index:idx_a(a)	range: decided by [gt(executor__index_lookup_pushdown.t9.a, executor__index_lookup_pushdown.t8.b)], keep order:false, stats:pseudo
+        │   └─TableRowIDScan(Probe)	cop[tikv]	table:t9	keep order:false, stats:pseudo
+        └─HashAgg(Probe)	cop[tikv]		funcs:sum(executor__index_lookup_pushdown.t9.b)->Column
+          └─TableRowIDScan	cop[tikv]	table:t9	keep order:false, stats:pseudo
+select *
+from t8
+where t8.a < (
+select /*+ INDEX_LOOKUP_PUSHDOWN(t9, idx_a) */ sum(t9.b)
+from t9 use index(idx_a)
+where t9.a > t8.b
+)
+order by t8.a, t8.b;
+a	b
+20	1
+24	5
 set @@tidb_scatter_region=default;

tests/integrationtest/t/executor/index_lookup_pushdown.test

@@ -1,4 +1,4 @@
-drop table if exists t1, t2, t3, t4, t5, t6, t7, tmp1, tmp2;
+drop table if exists t1, t2, t3, t4, t5, t6, t7, t8, t9, tmp1, tmp2;
 
 # wait regions split after create table to make the test stable
 set @@tidb_scatter_region='table';
@@ -201,5 +201,28 @@ explain select /*+ index_lookup_pushdown(t3, a) */ * from t3 use index(a);
 select /*+ index_lookup_pushdown(t3, a) */ * from t3 use index(a);
 --disable_warnings
 
+## issue #67546, correlated subquery under Apply should work with index_lookup_pushdown,
+set @@tidb_replica_read='leader';
+create table t8 (a int, b int, key idx_a(a));
+create table t9 (a int, b int, key idx_a(a));
+insert into t8 values (20,1),(24,5),(23,9);
+insert into t9 values (1,1),(3,2),(6,3),(10,4),(12,5),(15,6),(25,7);
+explain format='plan_tree' select *
+from t8
+where t8.a < (
+    select /*+ INDEX_LOOKUP_PUSHDOWN(t9, idx_a) */ sum(t9.b)
+    from t9 use index(idx_a)
+    where t9.a > t8.b
+)
+order by t8.a, t8.b;
+select *
+from t8
+where t8.a < (
+    select /*+ INDEX_LOOKUP_PUSHDOWN(t9, idx_a) */ sum(t9.b)
+    from t9 use index(idx_a)
+    where t9.a > t8.b
+)
+order by t8.a, t8.b;
+
 # restore tidb_scatter_region
 set @@tidb_scatter_region=default;

tests/realtikvtest/pushdowntest/index_lookup_pushdown_test.go

@@ -204,6 +204,39 @@ func TestRealTiKVIndexLookUpPushDown(t *testing.T) {
 	v.RunSelectWithCheck(fmt.Sprintf("a > %d and b < %d", randIndexVal(), r.Int63()), 0, r.Intn(50)+1)
 }
 
+// TestCorrelatedIndexLookUpPushDownPreservesParentIdx tests the indexlookup push down on correlated subquery,
+// and make sure the parent index is preserved in the plan after push down.
+// See issue: https://github.com/pingcap/tidb/issues/67546
+func TestCorrelatedIndexLookUpPushDownPreservesParentIdx(t *testing.T) {
+	store := realtikvtest.CreateMockStoreAndSetup(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table touter (a int, b int, key idx_a(a))")
+	tk.MustExec("create table tinner (a int, b int, key idx_a(a))")
+	tk.MustExec("insert into touter values (20,1),(24,5),(23,9)")
+	tk.MustExec("insert into tinner values (1,1),(3,2),(6,3),(10,4),(12,5),(15,6),(25,7)")
+
+	sql := `select *
+from touter
+where touter.a < (
+	select /*+ INDEX_LOOKUP_PUSHDOWN(tinner, idx_a) */ sum(tinner.b)
+	from tinner use index(idx_a)
+	where tinner.a > touter.b
+)
+order by touter.a, touter.b`
+
+	explainRows := tk.MustQuery("explain format='plan_tree' " + sql).Rows()
+	foundLocalIndexLookUp := false
+	for _, row := range explainRows {
+		if strings.Contains(row[0].(string), "LocalIndexLookUp") {
+			foundLocalIndexLookUp = true
+			break
+		}
+	}
+	require.True(t, foundLocalIndexLookUp)
+	tk.MustQuery(sql).Check(testkit.Rows("20 1", "24 5"))
+}
+
 func TestRealTiKVCommonHandleIndexLookUpPushDown(t *testing.T) {
 	store := realtikvtest.CreateMockStoreAndSetup(t)
 	tk := testkit.NewTestKit(t, store)