planner: add MppVersion for mpp task; support data compression in Exc…

…hange Operator; (pingcap#40132)

ref pingcap/tiflash#6620, close pingcap#40494

distsql/distsql.go

@@ -42,7 +42,7 @@ func DispatchMPPTasks(ctx context.Context, sctx sessionctx.Context, tasks []*kv.
  ctx = WithSQLKvExecCounterInterceptor(ctx, sctx.GetSessionVars().StmtCtx)
  _, allowTiFlashFallback := sctx.GetSessionVars().AllowFallbackToTiKV[kv.TiFlash]
  ctx = SetTiFlashMaxThreadsInContext(ctx, sctx)
- resp := sctx.GetMPPClient().DispatchMPPTasks(ctx, sctx.GetSessionVars().KVVars, tasks, allowTiFlashFallback, startTs, mppQueryID)
+ resp := sctx.GetMPPClient().DispatchMPPTasks(ctx, sctx.GetSessionVars().KVVars, tasks, allowTiFlashFallback, startTs, mppQueryID, sctx.GetSessionVars().ChooseMppVersion())
  if resp == nil {
  return nil, errors.New("client returns nil response")
  }

executor/builder.go

@@ -3401,6 +3401,7 @@ func (b *executorBuilder) buildMPPGather(v *plannercore.PhysicalTableReader) Exe
  b.err = err
  return nil
  }
+
  gather := &MPPGather{
  baseExecutor: newBaseExecutor(b.ctx, v.Schema(), v.ID()),
  is: b.is,

executor/mpp_gather.go

@@ -93,10 +93,14 @@ func (e *MPPGather) appendMPPDispatchReq(pf *plannercore.Fragment) error {
  if err != nil {
  return errors.Trace(err)
  }
+
  logutil.BgLogger().Info("Dispatch mpp task", zap.Uint64("timestamp", mppTask.StartTs),
  zap.Int64("ID", mppTask.ID), zap.Uint64("QueryTs", mppTask.MppQueryID.QueryTs), zap.Uint64("LocalQueryId", mppTask.MppQueryID.LocalQueryID),
  zap.Uint64("ServerID", mppTask.MppQueryID.ServerID), zap.String("address", mppTask.Meta.GetAddress()),
- zap.String("plan", plannercore.ToString(pf.ExchangeSender)))
+ zap.String("plan", plannercore.ToString(pf.ExchangeSender)),
+ zap.Int64("mpp-version", mppTask.MppVersion.ToInt64()),
+ zap.String("exchange-compression-mode", pf.ExchangeSender.CompressionMode.Name()),
+ )
  req := &kv.MPPDispatchRequest{
  Data: pbData,
  Meta: mppTask.Meta,

executor/set_test.go

@@ -2071,3 +2071,59 @@ func TestSetChunkReuseVariable(t *testing.T) {
  // error value
  tk.MustGetErrCode("set @@tidb_enable_reuse_chunk=s;", errno.ErrWrongValueForVar)
 }
+
+func TestSetMppVersionVariable(t *testing.T) {
+ store := testkit.CreateMockStore(t)
+ tk := testkit.NewTestKit(t, store)
+ tk.MustQuery("select @@session.mpp_version").Check(testkit.Rows("UNSPECIFIED"))
+ tk.MustExec("SET SESSION mpp_version = -1")
+ tk.MustQuery("select @@session.mpp_version").Check(testkit.Rows("-1"))
+ tk.MustExec("SET SESSION mpp_version = 0")
+ tk.MustQuery("select @@session.mpp_version").Check(testkit.Rows("0"))
+ tk.MustExec("SET SESSION mpp_version = 1")
+ tk.MustQuery("select @@session.mpp_version").Check(testkit.Rows("1"))
+ tk.MustExec("SET SESSION mpp_version = unspecified")
+ tk.MustQuery("select @@session.mpp_version").Check(testkit.Rows("unspecified"))
+ {
+ tk.MustGetErrMsg("SET SESSION mpp_version = 2", "incorrect value: 2. mpp_version options: -1 (unspecified), 0, 1")
+ }
+ {
+ tk.MustExec("SET GLOBAL mpp_version = 1")
+ tk.MustQuery("select @@global.mpp_version").Check(testkit.Rows("1"))
+ tk.MustExec("SET GLOBAL mpp_version = -1")
+ tk.MustQuery("select @@global.mpp_version").Check(testkit.Rows("-1"))
+ }
+}
+
+func TestSetMppExchangeCompressionModeVariable(t *testing.T) {
+ store := testkit.CreateMockStore(t)
+ tk := testkit.NewTestKit(t, store)
+ tk.MustGetErrMsg(
+ "SET SESSION mpp_exchange_compression_mode = 123",
+ "incorrect value: `123`. mpp_exchange_compression_mode options: NONE, FAST, HIGH_COMPRESSION, UNSPECIFIED")
+ tk.MustQuery("select @@session.mpp_exchange_compression_mode").Check(testkit.Rows("UNSPECIFIED"))
+
+ tk.MustExec("SET SESSION mpp_exchange_compression_mode = none")
+ tk.MustQuery("select @@session.mpp_exchange_compression_mode").Check(testkit.Rows("none"))
+ tk.MustExec("SET SESSION mpp_exchange_compression_mode = fast")
+ tk.MustQuery("select @@session.mpp_exchange_compression_mode").Check(testkit.Rows("fast"))
+ tk.MustExec("SET SESSION mpp_exchange_compression_mode = HIGH_COMPRESSION")
+ tk.MustQuery("select @@session.mpp_exchange_compression_mode").Check(testkit.Rows("HIGH_COMPRESSION"))
+
+ {
+ tk.MustExec("SET GLOBAL mpp_exchange_compression_mode = none")
+ tk.MustQuery("select @@global.mpp_exchange_compression_mode").Check(testkit.Rows("none"))
+ }
+ {
+ tk.MustExec("SET mpp_version = 0")
+ tk.MustExec("SET mpp_exchange_compression_mode = unspecified")
+ require.Equal(t, len(tk.Session().GetSessionVars().StmtCtx.GetWarnings()), 0)
+ }
+ {
+ tk.MustExec("SET mpp_version = 0")
+ tk.MustExec("SET mpp_exchange_compression_mode = HIGH_COMPRESSION")
+ warnings := tk.Session().GetSessionVars().StmtCtx.GetWarnings()
+ require.Equal(t, len(warnings), 1)
+ require.Equal(t, warnings[0].Err.Error(), "mpp exchange compression won't work under current mpp version 0")
+ }
+}

executor/tiflashtest/BUILD.bazel

@@ -14,6 +14,7 @@ go_test(
  "//config",
  "//domain",
  "//executor",
+ "//kv",
  "//meta/autoid",
  "//parser/terror",
  "//planner/core",

executor/tiflashtest/tiflash_test.go

@@ -30,6 +30,7 @@ import (
  "github.com/pingcap/tidb/config"
  "github.com/pingcap/tidb/domain"
  "github.com/pingcap/tidb/executor"
+ "github.com/pingcap/tidb/kv"
  "github.com/pingcap/tidb/parser/terror"
  plannercore "github.com/pingcap/tidb/planner/core"
  "github.com/pingcap/tidb/store/mockstore"
@@ -629,10 +630,41 @@ func TestDispatchTaskRetry(t *testing.T) {
  require.NoError(t, err)
  tk.MustExec("set @@session.tidb_enforce_mpp=ON")
  require.Nil(t, failpoint.Enable("github.com/pingcap/tidb/store/mockstore/unistore/mppDispatchTimeout", "3*return(true)"))
- tk.MustQuery("select count(*) from t").Check(testkit.Rows("4"))
+ tk.MustQuery("select count(*) from t group by b").Check(testkit.Rows("4"))
  require.Nil(t, failpoint.Disable("github.com/pingcap/tidb/store/mockstore/unistore/mppDispatchTimeout"))
 }
 
+func TestMppVersionError(t *testing.T) {
+ store := testkit.CreateMockStore(t, withMockTiFlash(2))
+ tk := testkit.NewTestKit(t, store)
+ tk.MustExec("use test")
+ tk.MustExec("drop table if exists t")
+ tk.MustExec("create table t(a int not null primary key, b int not null)")
+ tk.MustExec("alter table t set tiflash replica 1")
+ tk.MustExec("insert into t values(1,0),(2,0),(3,0),(4,0)")
+ tb := external.GetTableByName(t, tk, "test", "t")
+ err := domain.GetDomain(tk.Session()).DDL().UpdateTableReplicaInfo(tk.Session(), tb.Meta().ID, true)
+ require.NoError(t, err)
+ tk.MustExec("set @@session.tidb_enforce_mpp=ON")
+ {
+ item := fmt.Sprintf("return(%d)", kv.GetNewestMppVersion()+1)
+ require.Nil(t, failpoint.Enable("github.com/pingcap/tidb/store/mockstore/unistore/MppVersionError", item))
+ }
+ {
+ err := tk.QueryToErr("select count(*) from t group by b")
+ require.Error(t, err)
+ }
+ require.Nil(t, failpoint.Disable("github.com/pingcap/tidb/store/mockstore/unistore/MppVersionError"))
+ {
+ item := fmt.Sprintf("return(%d)", kv.GetNewestMppVersion())
+ require.Nil(t, failpoint.Enable("github.com/pingcap/tidb/store/mockstore/unistore/MppVersionError", item))
+ }
+ {
+ tk.MustQuery("select count(*) from t group by b").Check(testkit.Rows("4"))
+ }
+ require.Nil(t, failpoint.Disable("github.com/pingcap/tidb/store/mockstore/unistore/MppVersionError"))
+}
+
 func TestCancelMppTasks(t *testing.T) {
  var hang = "github.com/pingcap/tidb/store/mockstore/unistore/mppRecvHang"
  store := testkit.CreateMockStore(t, withMockTiFlash(2))
@@ -1325,15 +1357,15 @@ func TestDisaggregatedTiFlashQuery(t *testing.T) {
  require.NoError(t, err)
  tk.MustQuery("explain select * from t1 where c1 < 2").Check(testkit.Rows(
  "PartitionUnion_10 9970.00 root ",
- "├─TableReader_15 3323.33 root data:ExchangeSender_14",
+ "├─TableReader_15 3323.33 root MppVersion: 1, data:ExchangeSender_14",
  "│ └─ExchangeSender_14 3323.33 mpp[tiflash] ExchangeType: PassThrough",
  "│ └─Selection_13 3323.33 mpp[tiflash] lt(test.t1.c1, 2)",
  "│ └─TableFullScan_12 10000.00 mpp[tiflash] table:t1, partition:p0 keep order:false, stats:pseudo",
- "├─TableReader_19 3323.33 root data:ExchangeSender_18",
+ "├─TableReader_19 3323.33 root MppVersion: 1, data:ExchangeSender_18",
  "│ └─ExchangeSender_18 3323.33 mpp[tiflash] ExchangeType: PassThrough",
  "│ └─Selection_17 3323.33 mpp[tiflash] lt(test.t1.c1, 2)",
  "│ └─TableFullScan_16 10000.00 mpp[tiflash] table:t1, partition:p1 keep order:false, stats:pseudo",
- "└─TableReader_23 3323.33 root data:ExchangeSender_22",
+ "└─TableReader_23 3323.33 root MppVersion: 1, data:ExchangeSender_22",
  " └─ExchangeSender_22 3323.33 mpp[tiflash] ExchangeType: PassThrough",
  " └─Selection_21 3323.33 mpp[tiflash] lt(test.t1.c1, 2)",
  " └─TableFullScan_20 10000.00 mpp[tiflash] table:t1, partition:p2 keep order:false, stats:pseudo"))

kv/BUILD.bazel

@@ -84,6 +84,7 @@ go_test(
  "@com_github_pingcap_kvproto//pkg/coprocessor",
  "@com_github_pingcap_kvproto//pkg/deadlock",
  "@com_github_pingcap_kvproto//pkg/kvrpcpb",
+ "@com_github_pingcap_tipb//go-tipb",
  "@com_github_stretchr_testify//assert",
  "@com_github_stretchr_testify//require",
  "@com_github_tikv_client_go_v2//oracle",

kv/mpp.go

@@ -16,11 +16,65 @@ package kv
 
 import (
  "context"
+ "strconv"
+ "strings"
  "time"
 
  "github.com/pingcap/kvproto/pkg/mpp"
+ "github.com/pingcap/tipb/go-tipb"
 )
 
+// MppVersion indicates the mpp-version used to build mpp plan
+type MppVersion int64
+
+const (
+ // MppVersionV0 supports TiFlash version [~, ~]
+ MppVersionV0 MppVersion = iota
+
+ // MppVersionV1 supports TiFlash version [v6.6.x, ~]
+ MppVersionV1
+
+ // MppVersionV2
+ // MppVersionV3
+
+ mppVersionMax
+
+ newestMppVersion MppVersion = mppVersionMax - 1
+
+ // MppVersionUnspecified means the illegal or unspecified version, it only used in TiDB.
+ MppVersionUnspecified MppVersion = -1
+
+ // MppVersionUnspecifiedName denotes name of UNSPECIFIED mpp version
+ MppVersionUnspecifiedName string = "UNSPECIFIED"
+)
+
+// ToInt64 transforms MppVersion to int64
+func (v MppVersion) ToInt64() int64 {
+ return int64(v)
+}
+
+// ToMppVersion transforms string to MppVersion
+func ToMppVersion(name string) (MppVersion, bool) {
+ name = strings.ToUpper(name)
+ if name == MppVersionUnspecifiedName {
+ return MppVersionUnspecified, true
+ }
+ v, err := strconv.ParseInt(name, 10, 64)
+ if err != nil {
+ return MppVersionUnspecified, false
+ }
+ version := MppVersion(v)
+ if version >= MppVersionUnspecified && version <= newestMppVersion {
+ return version, true
+ }
+ return MppVersionUnspecified, false
+}
+
+// GetNewestMppVersion returns the mpp-version can be used in mpp plan
+func GetNewestMppVersion() MppVersion {
+ return newestMppVersion
+}
+
 // MPPTaskMeta means the meta info such as location of a mpp task.
 type MPPTaskMeta interface {
  // GetAddress indicates which node this task should execute on.
@@ -40,7 +94,8 @@ type MPPTask struct {
  ID int64 // mppTaskID
  StartTs uint64
  MppQueryID MPPQueryID
- TableID int64 // physical table id
+ TableID int64 // physical table id
+ MppVersion MppVersion // mpp version
 
  PartitionTableIDs []int64
  IsDisaggregatedTiFlashStaticPrune bool
@@ -54,6 +109,7 @@ func (t *MPPTask) ToPB() *mpp.TaskMeta {
  LocalQueryId: t.MppQueryID.LocalQueryID,
  ServerId: t.MppQueryID.ServerID,
  TaskId: t.ID,
+ MppVersion: t.MppVersion.ToInt64(),
  }
  if t.ID != -1 {
  meta.Address = t.Meta.GetAddress()
@@ -94,9 +150,8 @@ type MPPClient interface {
  // ConstructMPPTasks schedules task for a plan fragment.
  // TODO:: This interface will be refined after we support more executors.
  ConstructMPPTasks(context.Context, *MPPBuildTasksRequest, time.Duration) ([]MPPTaskMeta, error)
-
  // DispatchMPPTasks dispatches ALL mpp requests at once, and returns an iterator that transfers the data.
- DispatchMPPTasks(ctx context.Context, vars interface{}, reqs []*MPPDispatchRequest, needTriggerFallback bool, startTs uint64, mppQueryID MPPQueryID) Response
+ DispatchMPPTasks(ctx context.Context, vars interface{}, reqs []*MPPDispatchRequest, needTriggerFallback bool, startTs uint64, mppQueryID MPPQueryID, mppVersion MppVersion) Response
 }
 
 // MPPBuildTasksRequest request the stores allocation for a mpp plan fragment.
@@ -107,3 +162,56 @@ type MPPBuildTasksRequest struct {
 
  PartitionIDAndRanges []PartitionIDAndRanges
 }
+
+// ExchangeCompressionMode means the compress method used in exchange operator
+type ExchangeCompressionMode int
+
+const (
+ // ExchangeCompressionModeNONE indicates no compression
+ ExchangeCompressionModeNONE ExchangeCompressionMode = iota
+ // ExchangeCompressionModeFast indicates fast compression/decompression speed, compression ratio is lower than HC mode
+ ExchangeCompressionModeFast
+ // ExchangeCompressionModeHC indicates high compression (HC) ratio mode
+ ExchangeCompressionModeHC
+ // ExchangeCompressionModeUnspecified indicates unspecified compress method, let TiDB choose one
+ ExchangeCompressionModeUnspecified
+
+ // RecommendedExchangeCompressionMode indicates recommended compression mode
+ RecommendedExchangeCompressionMode ExchangeCompressionMode = ExchangeCompressionModeFast
+
+ exchangeCompressionModeUnspecifiedName string = "UNSPECIFIED"
+)
+
+// Name returns the name of ExchangeCompressionMode
+func (t ExchangeCompressionMode) Name() string {
+ if t == ExchangeCompressionModeUnspecified {
+ return exchangeCompressionModeUnspecifiedName
+ }
+ return t.ToTipbCompressionMode().String()
+}
+
+// ToExchangeCompressionMode returns the ExchangeCompressionMode from name
+func ToExchangeCompressionMode(name string) (ExchangeCompressionMode, bool) {
+ name = strings.ToUpper(name)
+ if name == exchangeCompressionModeUnspecifiedName {
+ return ExchangeCompressionModeUnspecified, true
+ }
+ value, ok := tipb.CompressionMode_value[name]
+ if ok {
+ return ExchangeCompressionMode(value), true
+ }
+ return ExchangeCompressionModeNONE, false
+}
+
+// ToTipbCompressionMode returns tipb.CompressionMode from kv.ExchangeCompressionMode
+func (t ExchangeCompressionMode) ToTipbCompressionMode() tipb.CompressionMode {
+ switch t {
+ case ExchangeCompressionModeNONE:
+ return tipb.CompressionMode_NONE
+ case ExchangeCompressionModeFast:
+ return tipb.CompressionMode_FAST
+ case ExchangeCompressionModeHC:
+ return tipb.CompressionMode_HIGH_COMPRESSION
+ }
+ return tipb.CompressionMode_NONE
+}

kv/version_test.go

@@ -17,6 +17,7 @@ package kv
 import (
  "testing"
 
+ "github.com/pingcap/tipb/go-tipb"
  "github.com/stretchr/testify/assert"
 )
 
@@ -30,3 +31,57 @@ func TestVersion(t *testing.T) {
  assert.True(t, eq == 0)
  assert.True(t, MinVersion.Cmp(MaxVersion) < 0)
 }
+
+func TestMppVersion(t *testing.T) {
+ assert.Equal(t, int64(1), GetNewestMppVersion().ToInt64())
+ {
+ v, ok := ToMppVersion("unspecified")
+ assert.True(t, ok)
+ assert.Equal(t, v, MppVersionUnspecified)
+ }
+ {
+ v, ok := ToMppVersion("-1")
+ assert.True(t, ok)
+ assert.Equal(t, v, MppVersionUnspecified)
+ }
+ {
+ v, ok := ToMppVersion("0")
+ assert.True(t, ok)
+ assert.Equal(t, v, MppVersionV0)
+ }
+ {
+ v, ok := ToMppVersion("1")
+ assert.True(t, ok)
+ assert.Equal(t, v, MppVersionV1)
+ }
+}
+
+func TestExchangeCompressionMode(t *testing.T) {
+ assert.Equal(t, "UNSPECIFIED", ExchangeCompressionModeUnspecified.Name())
+ {
+ a, ok := ToExchangeCompressionMode("UNSPECIFIED")
+ assert.Equal(t, a, ExchangeCompressionModeUnspecified)
+ assert.True(t, ok)
+ }
+ assert.Equal(t, "NONE", ExchangeCompressionModeNONE.Name())
+ {
+ a, ok := ToExchangeCompressionMode("NONE")
+ assert.Equal(t, a, ExchangeCompressionModeNONE)
+ assert.True(t, ok)
+ }
+ assert.Equal(t, "FAST", ExchangeCompressionModeFast.Name())
+ {
+ a, ok := ToExchangeCompressionMode("FAST")
+ assert.Equal(t, a, ExchangeCompressionModeFast)
+ assert.True(t, ok)
+ }
+ assert.Equal(t, "HIGH_COMPRESSION", ExchangeCompressionModeHC.Name())
+ {
+ a, ok := ToExchangeCompressionMode("HIGH_COMPRESSION")
+ assert.Equal(t, a, ExchangeCompressionModeHC)
+ assert.True(t, ok)
+ }
+ // default `FAST`
+ assert.Equal(t, ExchangeCompressionModeFast, RecommendedExchangeCompressionMode)
+ assert.Equal(t, tipb.CompressionMode_FAST, RecommendedExchangeCompressionMode.ToTipbCompressionMode())
+}