database/tidb
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

expression: make count distinct multi column aware of new collation (#27111)

Browse Source
This commit is contained in:
xufei
2021-08-13 19:57:58 +08:00
committed by GitHub
parent cd5499e994
commit 08995d915d
2 changed files with 28 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
executor/aggfuncs/func_count_distinct.go
View File

@ -319,6 +319,10 @@ func (e *countOriginalWithDistinct) UpdatePartialResult(sctx sessionctx.Context,
p := (*partialResult4CountWithDistinct)(pr)
encodedBytes := make([]byte, 0)
collators := make([]collate.Collator, 0, len(e.args))
for _, arg := range e.args {
collators = append(collators, collate.GetCollator(arg.GetType().Collate))
}
// Decimal struct is the biggest type we will use.
buf := make([]byte, types.MyDecimalStructSize)
@ -328,7 +332,7 @@ func (e *countOriginalWithDistinct) UpdatePartialResult(sctx sessionctx.Context,
encodedBytes = encodedBytes[:0]
for i := 0; i < len(e.args) && !hasNull; i++ {
encodedBytes, isNull, err = evalAndEncode(sctx, e.args[i], row, buf, encodedBytes)
encodedBytes, isNull, err = evalAndEncode(sctx, e.args[i], collators[i], row, buf, encodedBytes)
if err != nil {
return memDelta, err
}
@ -349,7 +353,7 @@ func (e *countOriginalWithDistinct) UpdatePartialResult(sctx sessionctx.Context,
// evalAndEncode eval one row with an expression and encode value to bytes.
func evalAndEncode(
sctx sessionctx.Context, arg expression.Expression,
sctx sessionctx.Context, arg expression.Expression, collator collate.Collator,
row chunk.Row, buf, encodedBytes []byte,
) (_ []byte, isNull bool, err error) {
switch tp := arg.GetType().EvalType(); tp {
@ -401,7 +405,7 @@ func evalAndEncode(
if err != nil || isNull {
break
}
encodedBytes = codec.EncodeCompactBytes(encodedBytes, hack.Slice(val))
encodedBytes = codec.EncodeCompactBytes(encodedBytes, collator.Key(val))
default:
return nil, false, errors.Errorf("unsupported column type for encode %d", tp)
}
@ -784,6 +788,10 @@ func (e *approxCountDistinctOriginal) UpdatePartialResult(sctx sessionctx.Contex
encodedBytes := make([]byte, 0)
// Decimal struct is the biggest type we will use.
buf := make([]byte, types.MyDecimalStructSize)
collators := make([]collate.Collator, 0, len(e.args))
for _, arg := range e.args {
collators = append(collators, collate.GetCollator(arg.GetType().Collate))
}
for _, row := range rowsInGroup {
var err error
@ -791,7 +799,7 @@ func (e *approxCountDistinctOriginal) UpdatePartialResult(sctx sessionctx.Contex
encodedBytes = encodedBytes[:0]
for i := 0; i < len(e.args) && !hasNull; i++ {
encodedBytes, isNull, err = evalAndEncode(sctx, e.args[i], row, buf, encodedBytes)
encodedBytes, isNull, err = evalAndEncode(sctx, e.args[i], collators[i], row, buf, encodedBytes)
if err != nil {
return memDelta, err
}

16
expression/integration_test.go
View File

@ -7026,6 +7026,22 @@ func (s *testIntegrationSerialSuite) TestIssue16668(c *C) {
tk.MustQuery("select count(distinct(b)) from tx").Check(testkit.Rows("4"))
}
func (s *testIntegrationSerialSuite) TestIssue27091(c *C) {
collate.SetNewCollationEnabledForTest(true)
defer collate.SetNewCollationEnabledForTest(false)
tk := testkit.NewTestKit(c, s.store)
tk.MustExec("use test")
tk.MustExec("drop table if exists tx")
tk.MustExec("CREATE TABLE `tx` ( `a` int(11) NOT NULL,`b` varchar(5) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT NULL, `c` varchar(5) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin DEFAULT NULL)")
tk.MustExec("insert into tx values (1, 'a', 'a'), (2, 'A ', 'a '), (3, 'A', 'A'), (4, 'a ', 'A ')")
tk.MustQuery("select count(distinct b) from tx").Check(testkit.Rows("1"))
tk.MustQuery("select count(distinct c) from tx").Check(testkit.Rows("2"))
tk.MustQuery("select count(distinct b, c) from tx where a < 3").Check(testkit.Rows("1"))
tk.MustQuery("select approx_count_distinct(b) from tx").Check(testkit.Rows("1"))
tk.MustQuery("select approx_count_distinct(c) from tx").Check(testkit.Rows("2"))
tk.MustQuery("select approx_count_distinct(b, c) from tx where a < 3").Check(testkit.Rows("1"))
}
func (s *testIntegrationSerialSuite) TestCollateStringFunction(c *C) {
collate.SetNewCollationEnabledForTest(true)
defer collate.SetNewCollationEnabledForTest(false)