From c35fa2eba04532c42d624aec19b6463583fa6ec1 Mon Sep 17 00:00:00 2001 From: ti-srebot <66930949+ti-srebot@users.noreply.github.com> Date: Tue, 5 Jul 2022 12:33:01 +0800 Subject: [PATCH] planner, expression: pushdown AggFuncMode to coprocessor (#31392) (#32276) close pingcap/tidb#30923 --- expression/aggregation/agg_to_pb.go | 43 ++++++- expression/aggregation/agg_to_pb_test.go | 14 +- go.mod | 14 +- go.sum | 27 ++-- planner/cascades/transformation_rules.go | 4 +- planner/core/plan_test.go | 155 +++++++++++++++++++++++ planner/core/task.go | 40 +++--- 7 files changed, 253 insertions(+), 44 deletions(-) diff --git a/expression/aggregation/agg_to_pb.go b/expression/aggregation/agg_to_pb.go index e07c598bcf005..6f7edaefaffa3 100644 --- a/expression/aggregation/agg_to_pb.go +++ b/expression/aggregation/agg_to_pb.go @@ -80,7 +80,46 @@ func AggFuncToPBExpr(sc *stmtctx.StatementContext, client kv.Client, aggFunc *Ag } children = append(children, pbArg) } - return &tipb.Expr{Tp: tp, Children: children, FieldType: expression.ToPBFieldType(aggFunc.RetTp), HasDistinct: aggFunc.HasDistinct} + return &tipb.Expr{Tp: tp, Children: children, FieldType: expression.ToPBFieldType(aggFunc.RetTp), HasDistinct: aggFunc.HasDistinct, AggFuncMode: AggFunctionModeToPB(aggFunc.Mode)} +} + +// AggFunctionModeToPB converts aggregate function mode to PB. +func AggFunctionModeToPB(mode AggFunctionMode) (pbMode *tipb.AggFunctionMode) { + pbMode = new(tipb.AggFunctionMode) + switch mode { + case CompleteMode: + *pbMode = tipb.AggFunctionMode_CompleteMode + case FinalMode: + *pbMode = tipb.AggFunctionMode_FinalMode + case Partial1Mode: + *pbMode = tipb.AggFunctionMode_Partial1Mode + case Partial2Mode: + *pbMode = tipb.AggFunctionMode_Partial2Mode + case DedupMode: + *pbMode = tipb.AggFunctionMode_DedupMode + } + return pbMode +} + +// PBAggFuncModeToAggFuncMode converts pb to aggregate function mode. +func PBAggFuncModeToAggFuncMode(pbMode *tipb.AggFunctionMode) (mode AggFunctionMode) { + // Default mode of the aggregate function is PartialMode. + mode = Partial1Mode + if pbMode != nil { + switch *pbMode { + case tipb.AggFunctionMode_CompleteMode: + mode = CompleteMode + case tipb.AggFunctionMode_FinalMode: + mode = FinalMode + case tipb.AggFunctionMode_Partial1Mode: + mode = Partial1Mode + case tipb.AggFunctionMode_Partial2Mode: + mode = Partial2Mode + case tipb.AggFunctionMode_DedupMode: + mode = DedupMode + } + } + return mode } // PBExprToAggFuncDesc converts pb to aggregate function. @@ -125,7 +164,7 @@ func PBExprToAggFuncDesc(ctx sessionctx.Context, aggFunc *tipb.Expr, fieldTps [] base.WrapCastForAggArgs(ctx) return &AggFuncDesc{ baseFuncDesc: base, - Mode: Partial1Mode, + Mode: PBAggFuncModeToAggFuncMode(aggFunc.AggFuncMode), HasDistinct: false, }, nil } diff --git a/expression/aggregation/agg_to_pb_test.go b/expression/aggregation/agg_to_pb_test.go index 5030959eed08b..5ee987f1e76f1 100644 --- a/expression/aggregation/agg_to_pb_test.go +++ b/expression/aggregation/agg_to_pb_test.go @@ -77,13 +77,13 @@ func (s *testEvaluatorSuite) TestAggFunc2Pb(c *C) { } jsons := []string{ - `{"tp":3002,"children":[{"tp":201,"val":"gAAAAAAAAAE=","sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":false}],"sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":%v}`, - `{"tp":3001,"children":[{"tp":201,"val":"gAAAAAAAAAE=","sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":false}],"sig":0,"field_type":{"tp":8,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":%v}`, - `{"tp":3003,"children":[{"tp":201,"val":"gAAAAAAAAAE=","sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":false}],"sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":%v}`, - "null", - `{"tp":3005,"children":[{"tp":201,"val":"gAAAAAAAAAE=","sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":false}],"sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":%v}`, - `{"tp":3004,"children":[{"tp":201,"val":"gAAAAAAAAAE=","sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":false}],"sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":%v}`, - `{"tp":3006,"children":[{"tp":201,"val":"gAAAAAAAAAE=","sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":false}],"sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":%v}`, + `{"tp":3002,"children":[{"tp":201,"val":"gAAAAAAAAAE=","sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":false}],"sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":%v,"aggFuncMode":0}`, + `{"tp":3001,"children":[{"tp":201,"val":"gAAAAAAAAAE=","sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":false}],"sig":0,"field_type":{"tp":8,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":%v,"aggFuncMode":0}`, + `{"tp":3003,"children":[{"tp":201,"val":"gAAAAAAAAAE=","sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":false}],"sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":%v,"aggFuncMode":0}`, + `{"tp":3007,"val":"AAAAAAAABAA=","children":[{"tp":201,"val":"gAAAAAAAAAE=","sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":false}],"sig":0,"field_type":{"tp":15,"flag":0,"flen":1,"decimal":-1,"collate":-46,"charset":"utf8mb4"},"has_distinct":%v,"aggFuncMode":0}`, + `{"tp":3005,"children":[{"tp":201,"val":"gAAAAAAAAAE=","sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":false}],"sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":%v,"aggFuncMode":0}`, + `{"tp":3004,"children":[{"tp":201,"val":"gAAAAAAAAAE=","sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":false}],"sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":%v,"aggFuncMode":0}`, + `{"tp":3006,"children":[{"tp":201,"val":"gAAAAAAAAAE=","sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":false}],"sig":0,"field_type":{"tp":5,"flag":0,"flen":-1,"decimal":-1,"collate":63,"charset":""},"has_distinct":%v,"aggFuncMode":0}`, } for i, funcName := range funcNames { for _, hasDistinct := range []bool{true, false} { diff --git a/go.mod b/go.mod index 797fa36d69993..118e2fb27a7df 100644 --- a/go.mod +++ b/go.mod @@ -16,7 +16,7 @@ require ( github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2 github.com/fatih/color v1.10.0 // indirect github.com/fsouza/fake-gcs-server v1.17.0 // indirect - github.com/go-sql-driver/mysql v1.5.0 + github.com/go-sql-driver/mysql v1.6.0 github.com/go-yaml/yaml v2.1.0+incompatible github.com/gogo/protobuf v1.3.2 github.com/golang/protobuf v1.3.4 @@ -39,16 +39,16 @@ require ( github.com/pingcap/badger v1.5.1-0.20200908111422-2e78ee155d19 github.com/pingcap/br v5.1.0-alpha.0.20210604015827-db22c6d284a0+incompatible github.com/pingcap/check v0.0.0-20200212061837-5e12011dc712 - github.com/pingcap/errors v0.11.5-0.20201126102027-b0a155152ca3 + github.com/pingcap/errors v0.11.5-0.20210425183316-da1aaba5fb63 github.com/pingcap/failpoint v0.0.0-20210316064728-7acb0f0a3dfd github.com/pingcap/fn v0.0.0-20200306044125-d5540d389059 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 github.com/pingcap/kvproto v0.0.0-20210722113201-768c114017e8 - github.com/pingcap/log v0.0.0-20210317133921-96f4fcab92a4 + github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7 github.com/pingcap/parser v0.0.0-20210618053735-57843e8185c4 github.com/pingcap/sysutil v0.0.0-20210315073920-cc0985d983a3 github.com/pingcap/tidb-tools v4.0.9-0.20201127090955-2707c97b3853+incompatible - github.com/pingcap/tipb v0.0.0-20211025074540-e1c7362eeeb4 + github.com/pingcap/tipb v0.0.0-20220704075215-f654f5a3fece github.com/prometheus/client_golang v1.5.1 github.com/prometheus/client_model v0.2.0 github.com/prometheus/common v0.9.1 @@ -68,8 +68,8 @@ require ( go.etcd.io/etcd v0.5.0-alpha.5.0.20200824191128-ae9734ed278b go.uber.org/atomic v1.8.0 go.uber.org/automaxprocs v1.2.0 - go.uber.org/goleak v0.10.0 - go.uber.org/zap v1.17.0 + go.uber.org/goleak v1.1.10 + go.uber.org/zap v1.18.1 golang.org/x/net v0.0.0-20210316092652-d523dce5a7f4 golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9 golang.org/x/sys v0.0.0-20210324051608-47abb6519492 @@ -77,7 +77,7 @@ require ( golang.org/x/tools v0.1.0 google.golang.org/grpc v1.27.1 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect - modernc.org/mathutil v1.2.2 // indirect + modernc.org/mathutil v1.4.1 // indirect sourcegraph.com/sourcegraph/appdash v0.0.0-20190731080439-ebfcffb1b5c0 sourcegraph.com/sourcegraph/appdash-data v0.0.0-20151005221446-73f23eafcf67 ) diff --git a/go.sum b/go.sum index 7ff55ac2bdbdc..d9e62be9f00b8 100644 --- a/go.sum +++ b/go.sum @@ -60,6 +60,8 @@ github.com/armon/consul-api v0.0.0-20180202201655-eb2c6b5be1b6/go.mod h1:grANhF5 github.com/aws/aws-sdk-go v1.30.19/go.mod h1:5zCpMtNQVjRREroY7sYe8lOMRSxkhG6MZveU8YkpAk0= github.com/aws/aws-sdk-go v1.35.3 h1:r0puXncSaAfRt7Btml2swUo74Kao+vKhO3VLjwDjK54= github.com/aws/aws-sdk-go v1.35.3/go.mod h1:H7NKnBqNVzoTJpGfLrQkkD+ytBA93eiDYi/+8rV9s48= +github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8= +github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA= github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q= github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8= github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= @@ -189,8 +191,9 @@ github.com/go-playground/overalls v0.0.0-20180201144345-22ec1a223b7c/go.mod h1:U github.com/go-playground/universal-translator v0.16.0/go.mod h1:1AnU7NaIRDWWzGEKwgtJRd2xk99HeFyHw3yid4rvQIY= github.com/go-sql-driver/mysql v1.3.0/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG5ZlKdlhCg5w= github.com/go-sql-driver/mysql v1.4.1/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG5ZlKdlhCg5w= -github.com/go-sql-driver/mysql v1.5.0 h1:ozyZYNQW3x3HtqT1jira07DN2PArx2v7/mN66gGcHOs= github.com/go-sql-driver/mysql v1.5.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg= +github.com/go-sql-driver/mysql v1.6.0 h1:BCTh4TKNUYmOmMUcQ3IipzF5prigylS7XXjEkfCHuOE= +github.com/go-sql-driver/mysql v1.6.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg= github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY= github.com/go-yaml/yaml v2.1.0+incompatible h1:RYi2hDdss1u4YE7GwixGzWwVo47T8UQwnTLB6vQiq+o= github.com/go-yaml/yaml v2.1.0+incompatible/go.mod h1:w2MrLa16VYP0jy6N7M5kHaCkaLENm+P+Tv+MfurjSw0= @@ -421,8 +424,9 @@ github.com/pingcap/errors v0.11.4/go.mod h1:Oi8TUi2kEtXXLMJk9l1cGmz20kV3TaQ0usTw github.com/pingcap/errors v0.11.5-0.20190809092503-95897b64e011/go.mod h1:Oi8TUi2kEtXXLMJk9l1cGmz20kV3TaQ0usTwv5KuLY8= github.com/pingcap/errors v0.11.5-0.20200917111840-a15ef68f753d/go.mod h1:g4vx//d6VakjJ0mk7iLBlKA8LFavV/sAVINT/1PFxeQ= github.com/pingcap/errors v0.11.5-0.20201029093017-5a7df2af2ac7/go.mod h1:G7x87le1poQzLB/TqvTJI2ILrSgobnq4Ut7luOwvfvI= -github.com/pingcap/errors v0.11.5-0.20201126102027-b0a155152ca3 h1:LllgC9eGfqzkfubMgjKIDyZYaa609nNWAyNZtpy2B3M= github.com/pingcap/errors v0.11.5-0.20201126102027-b0a155152ca3/go.mod h1:G7x87le1poQzLB/TqvTJI2ILrSgobnq4Ut7luOwvfvI= +github.com/pingcap/errors v0.11.5-0.20210425183316-da1aaba5fb63 h1:+FZIDR/D97YOPik4N4lPDaUcLDF/EQPogxtlHB2ZZRM= +github.com/pingcap/errors v0.11.5-0.20210425183316-da1aaba5fb63/go.mod h1:X2r9ueLEUZgtx2cIogM0v4Zj5uvvzhuuiu7Pn8HzMPg= github.com/pingcap/failpoint v0.0.0-20191029060244-12f4ac2fd11d/go.mod h1:DNS3Qg7bEDhU6EXNHF+XSv/PGznQaMJ5FWvctpm6pQI= github.com/pingcap/failpoint v0.0.0-20200702092429-9f69995143ce/go.mod h1:w4PEZ5y16LeofeeGwdgZB4ddv9bLyDuIX+ljstgKZyk= github.com/pingcap/failpoint v0.0.0-20210316064728-7acb0f0a3dfd h1:I8IeI8MNiZVKnwuXhcIIzz6pREcOSbq18Q31KYIzFVM= @@ -440,8 +444,9 @@ github.com/pingcap/kvproto v0.0.0-20210722113201-768c114017e8/go.mod h1:IOdRDPLy github.com/pingcap/log v0.0.0-20191012051959-b742a5d432e9/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8= github.com/pingcap/log v0.0.0-20200511115504-543df19646ad/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8= github.com/pingcap/log v0.0.0-20201112100606-8f1e84a3abc8/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8= -github.com/pingcap/log v0.0.0-20210317133921-96f4fcab92a4 h1:ERrF0fTuIOnwfGbt71Ji3DKbOEaP189tjym50u8gpC8= github.com/pingcap/log v0.0.0-20210317133921-96f4fcab92a4/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8= +github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7 h1:k2BbABz9+TNpYRwsCCFS8pEEnFVOdbgEjL/kTlLuzZQ= +github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM= github.com/pingcap/parser v0.0.0-20210525032559-c37778aff307/go.mod h1:xZC8I7bug4GJ5KtHhgAikjTfU4kBv1Sbo3Pf1MZ6lVw= github.com/pingcap/parser v0.0.0-20210618053735-57843e8185c4 h1:NASsbyMTNW8pbYfoO/YTykO6MQJiNRa094lwCPU6R2Q= github.com/pingcap/parser v0.0.0-20210618053735-57843e8185c4/go.mod h1:xZC8I7bug4GJ5KtHhgAikjTfU4kBv1Sbo3Pf1MZ6lVw= @@ -451,8 +456,8 @@ github.com/pingcap/sysutil v0.0.0-20210315073920-cc0985d983a3/go.mod h1:tckvA041 github.com/pingcap/tidb-dashboard v0.0.0-20210312062513-eef5d6404638/go.mod h1:OzFN8H0EDMMqeulPhPMw2i2JaiZWOKFQ7zdRPhENNgo= github.com/pingcap/tidb-tools v4.0.9-0.20201127090955-2707c97b3853+incompatible h1:ceznmu/lLseGHP/jKyOa/3u/5H3wtLLLqkH2V3ssSjg= github.com/pingcap/tidb-tools v4.0.9-0.20201127090955-2707c97b3853+incompatible/go.mod h1:XGdcy9+yqlDSEMTpOXnwf3hiTeqrV6MN/u1se9N8yIM= -github.com/pingcap/tipb v0.0.0-20211025074540-e1c7362eeeb4 h1:9Ef4j3DLmUidURfob0tf94v+sqvozqdCTr7e5hi19qU= -github.com/pingcap/tipb v0.0.0-20211025074540-e1c7362eeeb4/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs= +github.com/pingcap/tipb v0.0.0-20220704075215-f654f5a3fece h1:Jn9/JTBcD7Y5LEKEjY7gKCifSIauJ/2Ihl99HYQ6VKs= +github.com/pingcap/tipb v0.0.0-20220704075215-f654f5a3fece/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs= github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= @@ -626,8 +631,9 @@ go.uber.org/automaxprocs v1.2.0 h1:+RUihKM+nmYUoB9w0D0Ov5TJ2PpFO2FgenTxMJiZBZA= go.uber.org/automaxprocs v1.2.0/go.mod h1:YfO3fm683kQpzETxlTGZhGIVmXAhaw3gxeBADbpZtnU= go.uber.org/dig v1.8.0/go.mod h1:X34SnWGr8Fyla9zQNO2GSO2D+TIuqB14OS8JhYocIyw= go.uber.org/fx v1.10.0/go.mod h1:vLRicqpG/qQEzno4SYU86iCwfT95EZza+Eba0ItuxqY= -go.uber.org/goleak v0.10.0 h1:G3eWbSNIskeRqtsN/1uI5B+eP73y3JUuBsv9AZjehb4= go.uber.org/goleak v0.10.0/go.mod h1:VCZuO8V8mFPlL0F5J5GK1rtHV3DrFcQ1R8ryq7FK0aI= +go.uber.org/goleak v1.1.10 h1:z+mqJhf6ss6BSfSM671tgKyZBFPTTJM+HLxnhPC3wu0= +go.uber.org/goleak v1.1.10/go.mod h1:8a7PlsEVH3e/a/GLqe5IIrQx6GzcnRmZEufDUTk4A7A= go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0= go.uber.org/multierr v1.3.0/go.mod h1:VgVr7evmIr6uPjLBxg28wmKNXyqE9akIJ5XnfpiKl+4= go.uber.org/multierr v1.4.0/go.mod h1:VgVr7evmIr6uPjLBxg28wmKNXyqE9akIJ5XnfpiKl+4= @@ -642,8 +648,8 @@ go.uber.org/zap v1.10.0/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q= go.uber.org/zap v1.12.0/go.mod h1:zwrFLgMcdUuIBviXEYEH1YKNaOBnKXsx2IPda5bBwHM= go.uber.org/zap v1.15.0/go.mod h1:Mb2vm2krFEG5DV0W9qcHBYFtp/Wku1cvYaqPsS/WYfc= go.uber.org/zap v1.16.0/go.mod h1:MA8QOfq0BHJwdXa996Y4dYkAqRKB8/1K1QMMZVaNZjQ= -go.uber.org/zap v1.17.0 h1:MTjgFu6ZLKvY6Pvaqk97GlxNBuMpV4Hy/3P6tRGlI2U= -go.uber.org/zap v1.17.0/go.mod h1:MXVU+bhUf/A7Xi2HNOnopQOrmycQ5Ih87HtOu4q5SSo= +go.uber.org/zap v1.18.1 h1:CSUJ2mjFszzEWt4CdKISEuChVIXGBn3lAPwkRGyVrc4= +go.uber.org/zap v1.18.1/go.mod h1:xg/QME4nWcxGxrpdeYfq7UvYrLh66cuVKdrbD1XF/NI= golang.org/x/crypto v0.0.0-20180723164146-c126467f60eb/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4= golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4= golang.org/x/crypto v0.0.0-20181203042331-505ab145d0a9/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4= @@ -818,6 +824,7 @@ golang.org/x/tools v0.0.0-20191029041327-9cc4af7d6b2c/go.mod h1:b+2E5dAYhXwXZwtn golang.org/x/tools v0.0.0-20191029190741-b9c20aec41a5/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20191030062658-86caa796c7ab/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20191107010934-f79515f33823/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= +golang.org/x/tools v0.0.0-20191108193012-7d206e10da11/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20191113191852-77e3bb0ad9e7/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20191114200427-caa0b0f7d508/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20191115202509-3a792d9c32b2/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= @@ -941,8 +948,8 @@ honnef.co/go/tools v0.0.1-2020.1.3/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9 honnef.co/go/tools v0.2.0 h1:ws8AfbgTX3oIczLPNPCu5166oBg9ST2vNs0rcht+mDE= honnef.co/go/tools v0.2.0/go.mod h1:lPVVZ2BS5TfnjLyizF7o7hv7j9/L+8cZY2hLyjP9cGY= k8s.io/klog v1.0.0/go.mod h1:4Bi6QPql/J/LkTDqv7R/cd3hPo4k2DG6Ptcz060Ez5I= -modernc.org/mathutil v1.2.2 h1:+yFk8hBprV+4c0U9GjFtL+dV3N8hOJ8JCituQcMShFY= -modernc.org/mathutil v1.2.2/go.mod h1:mZW8CKdRPY1v87qxC/wUdX5O1qDzXMP5TH3wjfpga6E= +modernc.org/mathutil v1.4.1 h1:ij3fYGe8zBF4Vu+g0oT7mB06r8sqGWKuJu1yXeR4by8= +modernc.org/mathutil v1.4.1/go.mod h1:mZW8CKdRPY1v87qxC/wUdX5O1qDzXMP5TH3wjfpga6E= rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8= rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0= rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA= diff --git a/planner/cascades/transformation_rules.go b/planner/cascades/transformation_rules.go index 9961509299a52..9ddfb0077c8d1 100644 --- a/planner/cascades/transformation_rules.go +++ b/planner/cascades/transformation_rules.go @@ -438,7 +438,7 @@ func (r *PushAggDownGather) OnTransform(old *memo.ExprIter) (newExprs []*memo.Gr gbyItems := make([]expression.Expression, len(agg.GroupByItems)) copy(gbyItems, agg.GroupByItems) - partialPref, finalPref, funcMap := plannercore.BuildFinalModeAggregation(agg.SCtx(), + partialPref, finalPref, firstRowFuncMap := plannercore.BuildFinalModeAggregation(agg.SCtx(), &plannercore.AggInfo{ AggFuncs: aggFuncs, GroupByItems: gbyItems, @@ -446,7 +446,7 @@ func (r *PushAggDownGather) OnTransform(old *memo.ExprIter) (newExprs []*memo.Gr }, true, false) // Remove unnecessary FirstRow. partialPref.AggFuncs = - plannercore.RemoveUnnecessaryFirstRow(agg.SCtx(), finalPref.AggFuncs, finalPref.GroupByItems, partialPref.AggFuncs, partialPref.GroupByItems, partialPref.Schema, funcMap) + plannercore.RemoveUnnecessaryFirstRow(agg.SCtx(), finalPref.GroupByItems, partialPref.AggFuncs, partialPref.GroupByItems, partialPref.Schema, firstRowFuncMap) partialAgg := plannercore.LogicalAggregation{ AggFuncs: partialPref.AggFuncs, diff --git a/planner/core/plan_test.go b/planner/core/plan_test.go index 2f395c823542d..61b4023529b26 100644 --- a/planner/core/plan_test.go +++ b/planner/core/plan_test.go @@ -17,18 +17,27 @@ import ( "bytes" "fmt" "strings" + "testing" . "github.com/pingcap/check" + "github.com/pingcap/parser/ast" "github.com/pingcap/parser/model" + "github.com/pingcap/parser/mysql" "github.com/pingcap/tidb/domain" + "github.com/pingcap/tidb/expression" + "github.com/pingcap/tidb/expression/aggregation" "github.com/pingcap/tidb/kv" "github.com/pingcap/tidb/planner/core" + "github.com/pingcap/tidb/planner/util" + "github.com/pingcap/tidb/sessionctx" "github.com/pingcap/tidb/sessionctx/variable" + "github.com/pingcap/tidb/types" "github.com/pingcap/tidb/util/israce" "github.com/pingcap/tidb/util/plancodec" "github.com/pingcap/tidb/util/testkit" "github.com/pingcap/tidb/util/testleak" "github.com/pingcap/tidb/util/testutil" + "github.com/stretchr/testify/require" ) var _ = Suite(&testPlanNormalize{}) @@ -589,3 +598,149 @@ func (s *testPlanNormalize) BenchmarkEncodePlan(c *C) { core.EncodePlan(p) } } + +func TestBuildFinalModeAggregation(t *testing.T) { + aggSchemaBuilder := func(sctx sessionctx.Context, aggFuncs []*aggregation.AggFuncDesc) *expression.Schema { + schema := expression.NewSchema(make([]*expression.Column, 0, len(aggFuncs))...) + for _, agg := range aggFuncs { + newCol := &expression.Column{ + UniqueID: sctx.GetSessionVars().AllocPlanColumnID(), + RetType: agg.RetTp, + } + schema.Append(newCol) + } + return schema + } + isFinalAggMode := func(mode aggregation.AggFunctionMode) bool { + return mode == aggregation.FinalMode || mode == aggregation.CompleteMode + } + checkResult := func(sctx sessionctx.Context, aggFuncs []*aggregation.AggFuncDesc, groubyItems []expression.Expression) { + for partialIsCop := 0; partialIsCop < 2; partialIsCop++ { + for isMPPTask := 0; isMPPTask < 2; isMPPTask++ { + partial, final, _ := core.BuildFinalModeAggregation(sctx, &core.AggInfo{ + AggFuncs: aggFuncs, + GroupByItems: groubyItems, + Schema: aggSchemaBuilder(sctx, aggFuncs), + }, partialIsCop == 0, isMPPTask == 0) + if partial != nil { + for _, aggFunc := range partial.AggFuncs { + if partialIsCop == 0 { + require.True(t, !isFinalAggMode(aggFunc.Mode)) + } else { + require.True(t, isFinalAggMode(aggFunc.Mode)) + } + } + } + if final != nil { + for _, aggFunc := range final.AggFuncs { + require.True(t, isFinalAggMode(aggFunc.Mode)) + } + } + } + } + } + + ctx := core.MockContext() + + aggCol := &expression.Column{ + Index: 0, + RetType: types.NewFieldType(mysql.TypeLonglong), + } + gbyCol := &expression.Column{ + Index: 1, + RetType: types.NewFieldType(mysql.TypeLonglong), + } + orderCol := &expression.Column{ + Index: 2, + RetType: types.NewFieldType(mysql.TypeLonglong), + } + + emptyGroupByItems := make([]expression.Expression, 0, 1) + groupByItems := make([]expression.Expression, 0, 1) + groupByItems = append(groupByItems, gbyCol) + + orderByItems := make([]*util.ByItems, 0, 1) + orderByItems = append(orderByItems, &util.ByItems{ + Expr: orderCol, + Desc: true, + }) + + aggFuncs := make([]*aggregation.AggFuncDesc, 0, 5) + desc, err := aggregation.NewAggFuncDesc(ctx, ast.AggFuncMax, []expression.Expression{aggCol}, false) + require.NoError(t, err) + aggFuncs = append(aggFuncs, desc) + desc, err = aggregation.NewAggFuncDesc(ctx, ast.AggFuncFirstRow, []expression.Expression{aggCol}, false) + require.NoError(t, err) + aggFuncs = append(aggFuncs, desc) + desc, err = aggregation.NewAggFuncDesc(ctx, ast.AggFuncCount, []expression.Expression{aggCol}, false) + require.NoError(t, err) + aggFuncs = append(aggFuncs, desc) + desc, err = aggregation.NewAggFuncDesc(ctx, ast.AggFuncSum, []expression.Expression{aggCol}, false) + require.NoError(t, err) + aggFuncs = append(aggFuncs, desc) + desc, err = aggregation.NewAggFuncDesc(ctx, ast.AggFuncAvg, []expression.Expression{aggCol}, false) + require.NoError(t, err) + aggFuncs = append(aggFuncs, desc) + + aggFuncsWithDistinct := make([]*aggregation.AggFuncDesc, 0, 2) + desc, err = aggregation.NewAggFuncDesc(ctx, ast.AggFuncAvg, []expression.Expression{aggCol}, true) + require.NoError(t, err) + aggFuncsWithDistinct = append(aggFuncsWithDistinct, desc) + desc, err = aggregation.NewAggFuncDesc(ctx, ast.AggFuncCount, []expression.Expression{aggCol}, true) + require.NoError(t, err) + aggFuncsWithDistinct = append(aggFuncsWithDistinct, desc) + + groupConcatAggFuncs := make([]*aggregation.AggFuncDesc, 0, 4) + groupConcatWithoutDistinctWithoutOrderBy, err := aggregation.NewAggFuncDesc(ctx, ast.AggFuncGroupConcat, []expression.Expression{aggCol, aggCol}, false) + require.NoError(t, err) + groupConcatAggFuncs = append(groupConcatAggFuncs, groupConcatWithoutDistinctWithoutOrderBy) + groupConcatWithoutDistinctWithOrderBy, err := aggregation.NewAggFuncDesc(ctx, ast.AggFuncGroupConcat, []expression.Expression{aggCol, aggCol}, false) + require.NoError(t, err) + groupConcatWithoutDistinctWithOrderBy.OrderByItems = orderByItems + groupConcatAggFuncs = append(groupConcatAggFuncs, groupConcatWithoutDistinctWithOrderBy) + groupConcatWithDistinctWithoutOrderBy, err := aggregation.NewAggFuncDesc(ctx, ast.AggFuncGroupConcat, []expression.Expression{aggCol, aggCol}, true) + require.NoError(t, err) + groupConcatAggFuncs = append(groupConcatAggFuncs, groupConcatWithDistinctWithoutOrderBy) + groupConcatWithDistinctWithOrderBy, err := aggregation.NewAggFuncDesc(ctx, ast.AggFuncGroupConcat, []expression.Expression{aggCol, aggCol}, true) + require.NoError(t, err) + groupConcatWithDistinctWithOrderBy.OrderByItems = orderByItems + groupConcatAggFuncs = append(groupConcatAggFuncs, groupConcatWithDistinctWithOrderBy) + + // case 1 agg without distinct + checkResult(ctx, aggFuncs, emptyGroupByItems) + checkResult(ctx, aggFuncs, groupByItems) + + // case 2 agg with distinct + checkResult(ctx, aggFuncsWithDistinct, emptyGroupByItems) + checkResult(ctx, aggFuncsWithDistinct, groupByItems) + + // case 3 mixed with distinct and without distinct + mixedAggFuncs := make([]*aggregation.AggFuncDesc, 0, 10) + mixedAggFuncs = append(mixedAggFuncs, aggFuncs...) + mixedAggFuncs = append(mixedAggFuncs, aggFuncsWithDistinct...) + checkResult(ctx, mixedAggFuncs, emptyGroupByItems) + checkResult(ctx, mixedAggFuncs, groupByItems) + + // case 4 group concat + for _, groupConcatAggFunc := range groupConcatAggFuncs { + checkResult(ctx, []*aggregation.AggFuncDesc{groupConcatAggFunc}, emptyGroupByItems) + checkResult(ctx, []*aggregation.AggFuncDesc{groupConcatAggFunc}, groupByItems) + } + checkResult(ctx, groupConcatAggFuncs, emptyGroupByItems) + checkResult(ctx, groupConcatAggFuncs, groupByItems) + + // case 5 mixed group concat and other agg funcs + for _, groupConcatAggFunc := range groupConcatAggFuncs { + funcs := make([]*aggregation.AggFuncDesc, 0, 10) + funcs = append(funcs, groupConcatAggFunc) + funcs = append(funcs, aggFuncs...) + checkResult(ctx, funcs, emptyGroupByItems) + checkResult(ctx, funcs, groupByItems) + funcs = append(funcs, aggFuncsWithDistinct...) + checkResult(ctx, funcs, emptyGroupByItems) + checkResult(ctx, funcs, groupByItems) + } + mixedAggFuncs = append(mixedAggFuncs, groupConcatAggFuncs...) + checkResult(ctx, mixedAggFuncs, emptyGroupByItems) + checkResult(ctx, mixedAggFuncs, groupByItems) +} diff --git a/planner/core/task.go b/planner/core/task.go index 882c0a588b276..75f1b7ea052a2 100644 --- a/planner/core/task.go +++ b/planner/core/task.go @@ -1488,11 +1488,13 @@ type AggInfo struct { // BuildFinalModeAggregation splits either LogicalAggregation or PhysicalAggregation to finalAgg and partial1Agg, // returns the information of partial and final agg. -// partialIsCop means whether partial agg is a cop task. +// partialIsCop means whether partial agg is a cop task. When partialIsCop is false, +// we do not set the AggMode for partialAgg cause it may be split further when +// building the aggregate executor(e.g. buildHashAgg will split the AggDesc further for parallel executing). +// firstRowFuncMap is a map between partial first_row to final first_row, will be used in RemoveUnnecessaryFirstRow func BuildFinalModeAggregation( - sctx sessionctx.Context, original *AggInfo, partialIsCop bool, isMPPTask bool) (partial, final *AggInfo, funcMap map[*aggregation.AggFuncDesc]*aggregation.AggFuncDesc) { - - funcMap = make(map[*aggregation.AggFuncDesc]*aggregation.AggFuncDesc, len(original.AggFuncs)) + sctx sessionctx.Context, original *AggInfo, partialIsCop bool, isMPPTask bool) (partial, final *AggInfo, firstRowFuncMap map[*aggregation.AggFuncDesc]*aggregation.AggFuncDesc) { + firstRowFuncMap = make(map[*aggregation.AggFuncDesc]*aggregation.AggFuncDesc, len(original.AggFuncs)) partial = &AggInfo{ AggFuncs: make([]*aggregation.AggFuncDesc, 0, len(original.AggFuncs)), GroupByItems: original.GroupByItems, @@ -1630,16 +1632,19 @@ func BuildFinalModeAggregation( sumAgg.RetTp = partial.Schema.Columns[partialCursor-1].GetType() partial.AggFuncs = append(partial.AggFuncs, cntAgg, sumAgg) } else if aggFunc.Name == ast.AggFuncApproxCountDistinct { - approxCountDistinctAgg := *aggFunc - approxCountDistinctAgg.Name = ast.AggFuncApproxCountDistinct - approxCountDistinctAgg.RetTp = partial.Schema.Columns[partialCursor-1].GetType() - partial.AggFuncs = append(partial.AggFuncs, &approxCountDistinctAgg) + newAggFunc := aggFunc.Clone() + newAggFunc.Name = aggFunc.Name + newAggFunc.RetTp = partial.Schema.Columns[partialCursor-1].GetType() + partial.AggFuncs = append(partial.AggFuncs, newAggFunc) } else { - partial.AggFuncs = append(partial.AggFuncs, aggFunc) + partialFuncDesc := aggFunc.Clone() + partial.AggFuncs = append(partial.AggFuncs, partialFuncDesc) + if aggFunc.Name == ast.AggFuncFirstRow { + firstRowFuncMap[partialFuncDesc] = finalAggFunc + } } finalAggFunc.Mode = aggregation.FinalMode - funcMap[aggFunc] = finalAggFunc } finalAggFunc.Args = args @@ -1647,6 +1652,11 @@ func BuildFinalModeAggregation( final.AggFuncs[i] = finalAggFunc } partial.Schema.Append(partialGbySchema.Columns...) + if partialIsCop { + for _, f := range partial.AggFuncs { + f.Mode = aggregation.Partial1Mode + } + } return } @@ -1720,7 +1730,7 @@ func (p *basePhysicalAgg) newPartialAggregate(copTaskType kv.StoreType, isMPPTas if !CheckAggCanPushCop(p.ctx, p.AggFuncs, p.GroupByItems, copTaskType) { return nil, p.self } - partialPref, finalPref, funcMap := BuildFinalModeAggregation(p.ctx, &AggInfo{ + partialPref, finalPref, firstRowFuncMap := BuildFinalModeAggregation(p.ctx, &AggInfo{ AggFuncs: p.AggFuncs, GroupByItems: p.GroupByItems, Schema: p.Schema().Clone(), @@ -1730,8 +1740,7 @@ func (p *basePhysicalAgg) newPartialAggregate(copTaskType kv.StoreType, isMPPTas } // Remove unnecessary FirstRow. partialPref.AggFuncs = RemoveUnnecessaryFirstRow(p.ctx, - finalPref.AggFuncs, finalPref.GroupByItems, - partialPref.AggFuncs, partialPref.GroupByItems, partialPref.Schema, funcMap) + finalPref.GroupByItems, partialPref.AggFuncs, partialPref.GroupByItems, partialPref.Schema, firstRowFuncMap) if copTaskType == kv.TiDB { // For partial agg of TiDB cop task, since TiDB coprocessor reuse the TiDB executor, // and TiDB aggregation executor won't output the group by value, @@ -1788,12 +1797,11 @@ func genFirstRowAggForGroupBy(ctx sessionctx.Context, groupByItems []expression. // Can optimize the schema to [count(b), a] , and change the index to get value. func RemoveUnnecessaryFirstRow( sctx sessionctx.Context, - finalAggFuncs []*aggregation.AggFuncDesc, finalGbyItems []expression.Expression, partialAggFuncs []*aggregation.AggFuncDesc, partialGbyItems []expression.Expression, partialSchema *expression.Schema, - funcMap map[*aggregation.AggFuncDesc]*aggregation.AggFuncDesc) []*aggregation.AggFuncDesc { + firstRowFuncMap map[*aggregation.AggFuncDesc]*aggregation.AggFuncDesc) []*aggregation.AggFuncDesc { partialCursor := 0 newAggFuncs := make([]*aggregation.AggFuncDesc, 0, len(partialAggFuncs)) @@ -1813,7 +1821,7 @@ func RemoveUnnecessaryFirstRow( } if gbyExpr.Equal(sctx, aggFunc.Args[0]) { canOptimize = true - funcMap[aggFunc].Args[0] = finalGbyItems[j] + firstRowFuncMap[aggFunc].Args[0] = finalGbyItems[j] break } }