expression: fix the behavior when in func meet collation (#15609)

pingcap · Mar 26, 2020 · db789d7 · db789d7
1 parent 01db67b
commit db789d7
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 4 deletions.
diff --git a/expression/builtin_other.go b/expression/builtin_other.go
@@ -23,6 +23,7 @@ import (
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/types/json"
 	"github.com/pingcap/tidb/util/chunk"
+	"github.com/pingcap/tidb/util/collate"
 	"github.com/pingcap/tidb/util/set"
 	"github.com/pingcap/tidb/util/stringutil"
 	"github.com/pingcap/tipb/go-tipb"
@@ -242,6 +243,7 @@ type builtinInStringSig struct {
 func (b *builtinInStringSig) buildHashMapForConstArgs(ctx sessionctx.Context) error {
 	b.nonConstArgs = []Expression{b.args[0]}
 	b.hashSet = set.NewStringSet()
+	collator := collate.GetCollator(b.collation)
 	for i := 1; i < len(b.args); i++ {
 		if b.args[i].ConstItem(b.ctx.GetSessionVars().StmtCtx) {
 			val, isNull, err := b.args[i].EvalString(ctx, chunk.Row{})
@@ -252,7 +254,7 @@ func (b *builtinInStringSig) buildHashMapForConstArgs(ctx sessionctx.Context) er
 				b.hasNull = true
 				continue
 			}
-			b.hashSet.Insert(val)
+			b.hashSet.Insert(string(collator.Key(val)))
 		} else {
 			b.nonConstArgs = append(b.nonConstArgs, b.args[i])
 		}
@@ -280,9 +282,10 @@ func (b *builtinInStringSig) evalInt(row chunk.Row) (int64, bool, error) {
 	}
 
 	args := b.args
+	collator := collate.GetCollator(b.collation)
 	if len(b.hashSet) != 0 {
 		args = b.nonConstArgs
-		if b.hashSet.Exist(arg0) {
+		if b.hashSet.Exist(string(collator.Key(arg0))) {
 			return 1, false, nil
 		}
 	}
@@ -297,7 +300,7 @@ func (b *builtinInStringSig) evalInt(row chunk.Row) (int64, bool, error) {
 			hasNull = true
 			continue
 		}
-		if arg0 == evaledArg {
+		if types.CompareString(arg0, evaledArg, b.collation) == 0 {
 			return 1, false, nil
 		}
 	}

diff --git a/expression/builtin_other_test.go b/expression/builtin_other_test.go
@@ -24,6 +24,7 @@ import (
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/types/json"
 	"github.com/pingcap/tidb/util/chunk"
+	"github.com/pingcap/tidb/util/collate"
 	"github.com/pingcap/tidb/util/hack"
 )
 
@@ -124,6 +125,22 @@ func (s *testEvaluatorSuite) TestInFunc(c *C) {
 		c.Assert(err, IsNil)
 		c.Assert(d.GetValue(), Equals, tc.res, Commentf("%v", types.MakeDatums(tc.args)))
 	}
+	collate.SetNewCollationEnabledForTest(true)
+	strD1 := types.NewCollationStringDatum("a", "utf8_general_ci", 0)
+	strD2 := types.NewCollationStringDatum("Á", "utf8_general_ci", 0)
+	fn, err := fc.getFunction(s.ctx, s.datumsToConstants([]types.Datum{strD1, strD2}))
+	c.Assert(err, IsNil)
+	d, isNull, err := fn.evalInt(chunk.Row{})
+	c.Assert(isNull, IsFalse)
+	c.Assert(err, IsNil)
+	c.Assert(d, Equals, int64(1), Commentf("%v, %v", strD1, strD2))
+	chk1 := chunk.NewChunkWithCapacity(nil, 1)
+	chk1.SetNumVirtualRows(1)
+	chk2 := chunk.NewChunkWithCapacity([]*types.FieldType{types.NewFieldType(mysql.TypeTiny)}, 1)
+	err = fn.vecEvalInt(chk1, chk2.Column(0))
+	c.Assert(err, IsNil)
+	c.Assert(chk2.Column(0).GetInt64(0), Equals, int64(1))
+	collate.SetNewCollationEnabledForTest(false)
 }
 
 func (s *testEvaluatorSuite) TestRowFunc(c *C) {

diff --git a/expression/builtin_other_vec_generated.go b/expression/builtin_other_vec_generated.go
diff --git a/expression/generator/other_vec.go b/expression/generator/other_vec.go
@@ -51,6 +51,7 @@ const builtinOtherImports = `import (
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/types/json"
 	"github.com/pingcap/tidb/util/chunk"
+	"github.com/pingcap/tidb/util/collate"
 )
 `
 
@@ -116,6 +117,7 @@ var builtinInTmpl = template.Must(template.New("builtinInTmpl").Parse(`
 {{ range . }}
 {{ $InputInt := (eq .Input.TypeName "Int") }}
 {{ $InputJSON := (eq .Input.TypeName "JSON")}}
+{{ $InputString := (eq .Input.TypeName "String") }}
 {{ $InputFixed := ( .Input.Fixed ) }}
 {{ $UseHashKey := ( or (eq .Input.TypeName "Decimal") (eq .Input.TypeName "JSON") )}}
 func (b *{{.SigName}}) vecEvalInt(input *chunk.Chunk, result *chunk.Column) error {
@@ -144,6 +146,9 @@ func (b *{{.SigName}}) vecEvalInt(input *chunk.Chunk, result *chunk.Column) erro
 	args := b.args
 	{{- if not $InputJSON}}
 	if len(b.hashSet) != 0 {
+		{{- if $InputString }}
+			collator := collate.GetCollator(b.collation)
+		{{- end }}
 		args = b.nonConstArgs
 		for i := 0; i < n; i++ {
 			if buf0.IsNull(i) {
@@ -178,6 +183,11 @@ func (b *{{.SigName}}) vecEvalInt(input *chunk.Chunk, result *chunk.Column) erro
 						r64s[i] = 1
 						result.SetNull(i, false)
 					}
+				{{- else if $InputString }}
+					if _, ok := b.hashSet[string(collator.Key(arg0))]; ok {
+						r64s[i] = 1
+						result.SetNull(i, false)
+					}
 				{{- else }}
 					if _, ok := b.hashSet[arg0]; ok {
 						r64s[i] = 1