util,executor: use MutableString as key for DecimalSet (#9913) (#13041)

executor/aggfuncs/func_avg.go

@@ -20,6 +20,7 @@ import (
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/chunk"
+	"github.com/pingcap/tidb/util/hack"
 	"github.com/pingcap/tidb/util/set"
 )
 
@@ -153,7 +154,7 @@ func (e *avgPartial4Decimal) MergePartialResult(sctx sessionctx.Context, src Par
 
 type partialResult4AvgDistinctDecimal struct {
 	partialResult4AvgDecimal
-	valSet set.DecimalSet
+	valSet set.StringSet
 }
 
 type avgOriginal4DistinctDecimal struct {
@@ -162,7 +163,7 @@ type avgOriginal4DistinctDecimal struct {
 
 func (e *avgOriginal4DistinctDecimal) AllocPartialResult() PartialResult {
 	p := &partialResult4AvgDistinctDecimal{
-		valSet: set.NewDecimalSet(),
+		valSet: set.NewStringSet(),
 	}
 	return PartialResult(p)
 }
@@ -171,7 +172,7 @@ func (e *avgOriginal4DistinctDecimal) ResetPartialResult(pr PartialResult) {
 	p := (*partialResult4AvgDistinctDecimal)(pr)
 	p.sum = *types.NewDecFromInt(0)
 	p.count = int64(0)
-	p.valSet = set.NewDecimalSet()
+	p.valSet = set.NewStringSet()
 }
 
 func (e *avgOriginal4DistinctDecimal) UpdatePartialResult(sctx sessionctx.Context, rowsInGroup []chunk.Row, pr PartialResult) error {
@@ -181,18 +182,25 @@ func (e *avgOriginal4DistinctDecimal) UpdatePartialResult(sctx sessionctx.Contex
 		if err != nil {
 			return errors.Trace(err)
 		}
-		if isNull || p.valSet.Exist(input) {
+		if isNull {
 			continue
 		}
-
+		hash, err := input.ToHashKey()
+		if err != nil {
+			return err
+		}
+		decStr := string(hack.String(hash))
+		if p.valSet.Exist(decStr) {
+			continue
+		}
+		p.valSet.Insert(decStr)
 		newSum := new(types.MyDecimal)
 		err = types.DecimalAdd(&p.sum, input, newSum)
 		if err != nil {
 			return errors.Trace(err)
 		}
 		p.sum = *newSum
 		p.count++
-		p.valSet.Insert(input)
 	}
 	return nil
 }

executor/aggfuncs/func_sum.go

@@ -18,6 +18,7 @@ import (
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/chunk"
+	"github.com/pingcap/tidb/util/hack"
 	"github.com/pingcap/tidb/util/set"
 )
 
@@ -38,7 +39,7 @@ type partialResult4SumDistinctFloat64 struct {
 
 type partialResult4SumDistinctDecimal struct {
 	partialResult4SumDecimal
-	valSet set.DecimalSet
+	valSet set.StringSet
 }
 
 type baseSumAggFunc struct {
@@ -222,14 +223,14 @@ type sum4DistinctDecimal struct {
 func (e *sum4DistinctDecimal) AllocPartialResult() PartialResult {
 	p := new(partialResult4SumDistinctDecimal)
 	p.isNull = true
-	p.valSet = set.NewDecimalSet()
+	p.valSet = set.NewStringSet()
 	return PartialResult(p)
 }
 
 func (e *sum4DistinctDecimal) ResetPartialResult(pr PartialResult) {
 	p := (*partialResult4SumDistinctDecimal)(pr)
 	p.isNull = true
-	p.valSet = set.NewDecimalSet()
+	p.valSet = set.NewStringSet()
 }
 
 func (e *sum4DistinctDecimal) UpdatePartialResult(sctx sessionctx.Context, rowsInGroup []chunk.Row, pr PartialResult) error {
@@ -239,10 +240,18 @@ func (e *sum4DistinctDecimal) UpdatePartialResult(sctx sessionctx.Context, rowsI
 		if err != nil {
 			return errors.Trace(err)
 		}
-		if isNull || p.valSet.Exist(input) {
+		if isNull {
 			continue
 		}
-		p.valSet.Insert(input)
+		hash, err := input.ToHashKey()
+		if err != nil {
+			return err
+		}
+		decStr := string(hack.String(hash))
+		if p.valSet.Exist(decStr) {
+			continue
+		}
+		p.valSet.Insert(decStr)
 		if p.isNull {
 			p.val = *input
 			p.isNull = false

executor/executor_test.go

@@ -1145,7 +1145,7 @@ func (s *testSuite) TestUnion(c *C) {
 	tk.MustExec("drop table if exists t")
 	tk.MustExec("create table t(a int, b decimal(6, 3))")
 	tk.MustExec("insert into t values(1, 1.000)")
-	tk.MustQuery("select count(distinct a) from (select a from t union select b from t) tmp;").Check(testkit.Rows("1"))
+	tk.MustQuery("select count(distinct a), sum(distinct a), avg(distinct a) from (select a from t union all select b from t) tmp;").Check(testkit.Rows("1 1.000 1.0000000"))
 }
 
 func (s *testSuite) TestNeighbouringProj(c *C) {