*: support to execute CTE on MPP side

DEPS.bzl

@@ -3450,8 +3450,8 @@ def go_deps():
         name = "com_github_pingcap_tipb",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/pingcap/tipb",
-        sum = "h1:ltplM2dLXcIAwlleA5v4gke6m6ZeHpvUA3qYX9dCC18=",
-        version = "v0.0.0-20230427024529-aed92caf20b9",
+        sum = "h1:ZVehx2Mand1frpxzJud7FUOonbLZeXXQpEsNdjnEAJA=",
+        version = "v0.0.0-20230516140330-b3e432c40cb3",
     )
     go_repository(
         name = "com_github_pkg_browser",

cmd/explaintest/r/explain_cte.result

@@ -191,7 +191,7 @@ CTE_0	100.00	root		Non-Recursive CTE
         └─IndexFullScan_31	10000.00	cop[tikv]	table:t2, index:c1(c1)	keep order:false, stats:pseudo
 explain with recursive cte1(c1) as (select c1 from t1 union select c1 from t2 limit 0 offset 0) select * from cte1;
 id	estRows	task	access object	operator info
-CTEFullScan_18	0.00	root	CTE:cte1	data:CTE_0
+CTEFullScan_19	0.00	root	CTE:cte1	data:CTE_0
 CTE_0	0.00	root		Non-Recursive CTE
 └─TableDual_16(Seed Part)	0.00	root		rows:0
 CREATE TABLE `customer` (

distsql/select_result.go

@@ -558,12 +558,11 @@ func (r *selectResult) updateCopRuntimeStats(ctx context.Context, copStats *copr
 	}
 	if hasExecutor {
 		var recorededPlanIDs = make(map[int]int)
-		for i, detail := range r.selectResp.GetExecutionSummaries() {
+		for _, detail := range r.selectResp.GetExecutionSummaries() {
 			if detail != nil && detail.TimeProcessedNs != nil &&
 				detail.NumProducedRows != nil && detail.NumIterations != nil {
-				planID := r.copPlanIDs[i]
 				recorededPlanIDs[r.ctx.GetSessionVars().StmtCtx.RuntimeStatsColl.
-					RecordOneCopTask(planID, r.storeType.Name(), callee, detail)] = 0
+					RecordOneCopTask(-1, r.storeType.Name(), callee, detail)] = 0
 			}
 		}
 		num := uint64(0)

executor/BUILD.bazel

@@ -229,6 +229,7 @@ go_library(
         "@com_github_pingcap_kvproto//pkg/encryptionpb",
         "@com_github_pingcap_kvproto//pkg/kvrpcpb",
         "@com_github_pingcap_kvproto//pkg/metapb",
+        "@com_github_pingcap_kvproto//pkg/mpp",
         "@com_github_pingcap_kvproto//pkg/resource_manager",
         "@com_github_pingcap_kvproto//pkg/tikvpb",
         "@com_github_pingcap_log//:log",

executor/mpp_gather.go

@@ -16,10 +16,12 @@ package executor
 
 import (
 	"context"
+	"fmt"
 	"time"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/kvproto/pkg/mpp"
 	"github.com/pingcap/tidb/distsql"
 	"github.com/pingcap/tidb/infoschema"
 	"github.com/pingcap/tidb/kv"
@@ -99,6 +101,10 @@ func (e *MPPGather) appendMPPDispatchReq(pf *plannercore.Fragment) error {
 		if err != nil {
 			return errors.Trace(err)
 		}
+		err = e.fixTaskForCTEStorageAndReader(dagReq.RootExecutor, mppTask.Meta)
+		if err != nil {
+			return err
+		}
 		pbData, err := dagReq.Marshal()
 		if err != nil {
 			return errors.Trace(err)
@@ -127,6 +133,88 @@ func (e *MPPGather) appendMPPDispatchReq(pf *plannercore.Fragment) error {
 	return nil
 }
 
+// fixTaskForCTEStorageAndReader fixes the upstream/downstream tasks for the producers and consumers.
+// After we split the fragments. A CTE producer in the fragment will holds all the task address of the consumers.
+// For example, the producer has two task on node_1 and node_2. As we know that each consumer also has two task on the same nodes(node_1 and node_2)
+// We need to prune address of node_2 for producer's task on node_1 since we just want the producer task on the node_1 only send to the consumer tasks on the node_1.
+// And the same for the task on the node_2.
+// And the same for the consumer task. We need to prune the unnecessary task address of its producer tasks(i.e. the downstream tasks).
+func (e *MPPGather) fixTaskForCTEStorageAndReader(exec *tipb.Executor, meta kv.MPPTaskMeta) error {
+	children := make([]*tipb.Executor, 0, 2)
+	switch exec.Tp {
+	case tipb.ExecType_TypeTableScan, tipb.ExecType_TypePartitionTableScan, tipb.ExecType_TypeIndexScan:
+	case tipb.ExecType_TypeSelection:
+		children = append(children, exec.Selection.Child)
+	case tipb.ExecType_TypeAggregation, tipb.ExecType_TypeStreamAgg:
+		children = append(children, exec.Aggregation.Child)
+	case tipb.ExecType_TypeTopN:
+		children = append(children, exec.TopN.Child)
+	case tipb.ExecType_TypeLimit:
+		children = append(children, exec.Limit.Child)
+	case tipb.ExecType_TypeExchangeSender:
+		children = append(children, exec.ExchangeSender.Child)
+		if len(exec.ExchangeSender.UpstreamCteTaskMeta) == 0 {
+			break
+		}
+		actualUpStreamTasks := make([][]byte, 0, len(exec.ExchangeSender.UpstreamCteTaskMeta))
+		actualTIDs := make([]int64, 0, len(exec.ExchangeSender.UpstreamCteTaskMeta))
+		for _, tasksFromOneConsumer := range exec.ExchangeSender.UpstreamCteTaskMeta {
+			for _, taskBytes := range tasksFromOneConsumer.EncodedTasks {
+				taskMeta := &mpp.TaskMeta{}
+				err := taskMeta.Unmarshal(taskBytes)
+				if err != nil {
+					return err
+				}
+				if taskMeta.Address != meta.GetAddress() {
+					continue
+				}
+				actualUpStreamTasks = append(actualUpStreamTasks, taskBytes)
+				actualTIDs = append(actualTIDs, taskMeta.TaskId)
+			}
+		}
+		logutil.BgLogger().Warn("refine tunnel for cte producer task", zap.String("the final tunnel", fmt.Sprintf("up stream consumer tasks: %v", actualTIDs)))
+		exec.ExchangeSender.EncodedTaskMeta = actualUpStreamTasks
+	case tipb.ExecType_TypeExchangeReceiver:
+		if len(exec.ExchangeReceiver.OriginalCtePrdocuerTaskMeta) == 0 {
+			break
+		}
+		exec.ExchangeReceiver.EncodedTaskMeta = [][]byte{}
+		actualTIDs := make([]int64, 0, 4)
+		for _, taskBytes := range exec.ExchangeReceiver.OriginalCtePrdocuerTaskMeta {
+			taskMeta := &mpp.TaskMeta{}
+			err := taskMeta.Unmarshal(taskBytes)
+			if err != nil {
+				return err
+			}
+			if taskMeta.Address != meta.GetAddress() {
+				continue
+			}
+			exec.ExchangeReceiver.EncodedTaskMeta = append(exec.ExchangeReceiver.EncodedTaskMeta, taskBytes)
+			actualTIDs = append(actualTIDs, taskMeta.TaskId)
+		}
+		logutil.BgLogger().Warn("refine tunnel for cte consumer task", zap.String("the final tunnel", fmt.Sprintf("down stream producer task: %v", actualTIDs)))
+	case tipb.ExecType_TypeJoin:
+		children = append(children, exec.Join.Children...)
+	case tipb.ExecType_TypeProjection:
+		children = append(children, exec.Projection.Child)
+	case tipb.ExecType_TypeWindow:
+		children = append(children, exec.Window.Child)
+	case tipb.ExecType_TypeSort:
+		children = append(children, exec.Sort.Child)
+	case tipb.ExecType_TypeExpand:
+		children = append(children, exec.Expand.Child)
+	default:
+		return errors.Errorf("unknown new tipb protocol %d", exec.Tp)
+	}
+	for _, child := range children {
+		err := e.fixTaskForCTEStorageAndReader(child, meta)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func collectPlanIDS(plan plannercore.PhysicalPlan, ids []int) []int {
 	ids = append(ids, plan.ID())
 	for _, child := range plan.Children() {

go.mod

@@ -78,7 +78,7 @@ require (
 	github.com/pingcap/log v1.1.1-0.20230317032135-a0d097d16e22
 	github.com/pingcap/sysutil v1.0.1-0.20230407040306-fb007c5aff21
 	github.com/pingcap/tidb/parser v0.0.0-20211011031125-9b13dc409c5e
-	github.com/pingcap/tipb v0.0.0-20230427024529-aed92caf20b9
+	github.com/pingcap/tipb v0.0.0-20230516140330-b3e432c40cb3
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.15.1
 	github.com/prometheus/client_model v0.4.0

go.sum

@@ -798,8 +798,8 @@ github.com/pingcap/log v1.1.1-0.20230317032135-a0d097d16e22 h1:2SOzvGvE8beiC1Y4g
 github.com/pingcap/log v1.1.1-0.20230317032135-a0d097d16e22/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pingcap/sysutil v1.0.1-0.20230407040306-fb007c5aff21 h1:QV6jqlfOkh8hqvEAgwBZa+4bSgO0EeKC7s5c6Luam2I=
 github.com/pingcap/sysutil v1.0.1-0.20230407040306-fb007c5aff21/go.mod h1:QYnjfA95ZaMefyl1NO8oPtKeb8pYUdnDVhQgf+qdpjM=
-github.com/pingcap/tipb v0.0.0-20230427024529-aed92caf20b9 h1:ltplM2dLXcIAwlleA5v4gke6m6ZeHpvUA3qYX9dCC18=
-github.com/pingcap/tipb v0.0.0-20230427024529-aed92caf20b9/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs=
+github.com/pingcap/tipb v0.0.0-20230516140330-b3e432c40cb3 h1:ZVehx2Mand1frpxzJud7FUOonbLZeXXQpEsNdjnEAJA=
+github.com/pingcap/tipb v0.0.0-20230516140330-b3e432c40cb3/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs=
 github.com/pkg/browser v0.0.0-20180916011732-0a3d74bf9ce4 h1:49lOXmGaUpV9Fz3gd7TFZY106KVlPVa5jcYD1gaQf98=
 github.com/pkg/browser v0.0.0-20180916011732-0a3d74bf9ce4/go.mod h1:4OwLy04Bl9Ef3GJJCoec+30X3LQs/0/m4HFRt/2LUSA=
 github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=

planner/core/BUILD.bazel

@@ -63,6 +63,7 @@ go_library(
         "rule_partition_processor.go",
         "rule_predicate_push_down.go",
         "rule_predicate_simplification.go",
+        "rule_push_down_sequence.go",
         "rule_result_reorder.go",
         "rule_semi_join_rewrite.go",
         "rule_topn_push_down.go",

planner/core/access_object.go

@@ -230,6 +230,9 @@ func (o OtherAccessObject) SetIntoPB(pb *tipb.ExplainOperator) {
 	if pb == nil {
 		return
 	}
+	if o == "" {
+		return
+	}
 	pb.AccessObjects = []*tipb.AccessObject{
 		{
 			AccessObject: &tipb.AccessObject_OtherObject{OtherObject: string(o)},
@@ -396,7 +399,10 @@ func (p *PhysicalTableReader) accessObject(sctx sessionctx.Context) AccessObject
 		return DynamicPartitionAccessObjects(nil)
 	}
 	if len(p.PartitionInfos) == 0 {
-		ts := p.TablePlans[0].(*PhysicalTableScan)
+		ts, ok := p.TablePlans[0].(*PhysicalTableScan)
+		if !ok {
+			return OtherAccessObject("")
+		}
 		asName := ""
 		if ts.TableAsName != nil && len(ts.TableAsName.O) > 0 {
 			asName = ts.TableAsName.O

planner/core/casetest/enforce_mpp_test.go

@@ -598,3 +598,51 @@ func TestMPPNullAwareSemiJoinPushDown(t *testing.T) {
 		require.Equal(t, output[i].Warn, testdata.ConvertSQLWarnToStrings(tk.Session().GetSessionVars().StmtCtx.GetWarnings()))
 	}
 }
+
+func TestMPPSharedCTEScan(t *testing.T) {
+	store := testkit.CreateMockStore(t, internal.WithMockTiFlash(2))
+	tk := testkit.NewTestKit(t, store)
+
+	// test table
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("drop table if exists s")
+	tk.MustExec("create table t(a int, b int, c int)")
+	tk.MustExec("create table s(a int, b int, c int)")
+	tk.MustExec("alter table t set tiflash replica 1")
+	tk.MustExec("alter table s set tiflash replica 1")
+
+	tb := external.GetTableByName(t, tk, "test", "t")
+	err := domain.GetDomain(tk.Session()).DDL().UpdateTableReplicaInfo(tk.Session(), tb.Meta().ID, true)
+	require.NoError(t, err)
+
+	tb = external.GetTableByName(t, tk, "test", "s")
+	err = domain.GetDomain(tk.Session()).DDL().UpdateTableReplicaInfo(tk.Session(), tb.Meta().ID, true)
+	require.NoError(t, err)
+
+	var input []string
+	var output []struct {
+		SQL  string
+		Plan []string
+		Warn []string
+	}
+
+	tk.MustExec("set @@tidb_enforce_mpp='on'")
+	tk.MustExec("set @@tidb_opt_enable_mpp_shared_cte_execution='on'")
+
+	enforceMPPSuiteData := GetEnforceMPPSuiteData()
+	enforceMPPSuiteData.LoadTestCases(t, &input, &output)
+	for i, tt := range input {
+		testdata.OnRecord(func() {
+			output[i].SQL = tt
+		})
+		testdata.OnRecord(func() {
+			output[i].SQL = tt
+			output[i].Plan = testdata.ConvertRowsToStrings(tk.MustQuery(tt).Rows())
+			output[i].Warn = testdata.ConvertSQLWarnToStrings(tk.Session().GetSessionVars().StmtCtx.GetWarnings())
+		})
+		res := tk.MustQuery(tt)
+		res.Check(testkit.Rows(output[i].Plan...))
+		require.Equal(t, output[i].Warn, testdata.ConvertSQLWarnToStrings(tk.Session().GetSessionVars().StmtCtx.GetWarnings()))
+	}
+}

planner/core/casetest/testdata/binary_plan_suite_out.json

@@ -297,6 +297,7 @@
         "BinaryPlan": {
           "main": {
             "name": "CTEFullScan_17",
+            "cost": 0.8982000000000001,
             "est_rows": 1.8,
             "act_rows": 5,
             "task_type": 1,
@@ -363,6 +364,7 @@
                   "operator_info": "cast(plus(Column#3, 1), bigint(1) BINARY)->Column#5"
                 }
               ],
+              "cost": 0.8982000000000001,
               "est_rows": 1.8,
               "act_rows": 5,
               "task_type": 1,

planner/core/casetest/testdata/enforce_mpp_suite_in.json

@@ -172,5 +172,31 @@
       "EXPLAIN select *, (t.a, t.b) not in (select s.a, s.b from s) from t; -- 7. left anti semi join, two join key",
       "EXPLAIN select *, (t.a, t.b) not in (select s.a, s.b from s where t.c < s.c) from t; -- 8. left anti semi join, two join key + other condition"
     ]
+  },
+  {
+    "name": "TestMPPSharedCTEScan",
+    "cases": [
+      // The most simple case.
+      "explain format = 'brief' with c1 as (select * from t) select * from c1, c1 c2 where c1.a=c2.b ",
+      "explain format = 'brief' with c1 as (select /*+ read_from_storage(tikv[t]) */ * from t) select * from c1, c1 c2 where c1.a=c2.b ",
+      "explain format = 'brief' with c1 as (select * from t) select c1.* from c1, c1 c2 where c1.b=c2.c",
+      // Can work when there's global limit/topn
+      "explain format = 'brief' with c1 as (select * from t) select * from c1, c1 c2 where c1.a=c2.b limit 10",
+      "explain format = 'brief' with c1 as (select * from t) select * from c1, c1 c2 where c1.a=c2.b order by c1.a limit 10",
+      // The c2 references c1, c1 can mpp, and then c2 can mpp, so the main query can mpp.
+      "explain format = 'brief' with c1 as (select * from t), c2 as (select c1.* from c1, c1 c2 where c1.b=c2.c) select * from c2 c1, c2 where c1.a=c2.b",
+      // The same SQL, c1 forces to read tikv. So c2 cannot MPP, then the whole SQL.
+      "explain format = 'brief' with c1 as (select /*+ read_from_storage(tikv[t]) */ * from t), c2 as (select c1.* from c1, c1 c2 where c1.b=c2.c) select * from c2 c1, c2 where c1.a=c2.b",
+      // The two WITH satement can all be MPP.
+      "explain format = 'brief' with c1 as (select * from t), c2 as (select c1.* from c1, c1 c2 where c1.b=c2.c) select * from c2 c1, c2, (with c3 as (select * from t) select c3.* from c3, c3 c4 where c3.c=c4.b) c3 where c1.a=c2.b and c1.a=c3.a",
+      // The outer one will fail to use MPP. But the inner WITH statement can. But we haven't implemented the least common ancestor to detect the best position of the Sequence. So the whole SQL cannot MPP.
+      "explain format = 'brief' with c1 as (select /*+ read_from_storage(tikv[t]) */ * from t), c2 as (select c1.* from c1, c1 c2 where c1.b=c2.c) select * from c2 c1, c2, (with c3 as (select * from t) select c3.* from c3, c3 c4 where c3.c=c4.b) c3 where c1.a=c2.b and c1.a=c3.a",
+      // The inner one will fail. So the whole SQL cannot MPP.
+      "explain format = 'brief' with c1 as (select * from t), c2 as (select c1.* from c1, c1 c2 where c1.b=c2.c) select * from c2 c1, c2, (with c3 as (select /*+ read_from_storage(tikv[t]) */ * from t) select c3.* from c3, c3 c4 where c3.c=c4.b) c3 where c1.a=c2.b and c1.a=c3.a",
+      // A little change that the inner WITH statement references the outer's c1.
+      "explain format = 'brief' with c1 as (select * from t), c2 as (select c1.* from c1, c1 c2 where c1.b=c2.c) select * from c2 c1, c2, (with c3 as (select * from c1) select c3.* from c3, c3 c4 where c3.c=c4.b) c3 where c1.a=c2.b and c1.a=c3.a",
+      // The outer one will fail to use MPP. Since the inner one is references the outer one, the whole SQL cannot MPP.
+      "explain format = 'brief' with c1 as (select /*+ read_from_storage(tikv[t]) */ * from t), c2 as (select c1.* from c1, c1 c2 where c1.b=c2.c) select * from c2 c1, c2, (with c3 as (select * from c1) select c3.* from c3, c3 c4 where c3.c=c4.b) c3 where c1.a=c2.b and c1.a=c3.a"
+    ]
   }
 ]