Skip to content

Commit

Permalink
cherry pick #27111 to release-5.0
Browse files Browse the repository at this point in the history
Signed-off-by: ti-srebot <ti-srebot@pingcap.com>
  • Loading branch information
windtalker authored and ti-srebot committed Aug 23, 2021
1 parent dbc8017 commit 23cfb61
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 4 deletions.
16 changes: 12 additions & 4 deletions executor/aggfuncs/func_count_distinct.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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
}
Expand All @@ -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 {
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -784,14 +788,18 @@ 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
var hasNull, isNull bool
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
}
Expand Down
16 changes: 16 additions & 0 deletions expression/integration_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6814,6 +6814,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)
Expand Down

0 comments on commit 23cfb61

Please sign in to comment.