expression: fix panic in upper() and lower() func. (pingcap#32505)

close pingcap#32488
wshwsh12 · Mar 7, 2022 · 885a71b · 885a71b
1 parent a8cc9ba
commit 885a71b
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 6 deletions.
diff --git a/expression/builtin_string_vec.go b/expression/builtin_string_vec.go
@@ -43,12 +43,23 @@ func (b *builtinLowerSig) vectorized() bool {
 }
 
 func (b *builtinLowerUTF8Sig) vecEvalString(input *chunk.Chunk, result *chunk.Column) error {
-	if err := b.args[0].VecEvalString(b.ctx, input, result); err != nil {
+	n := input.NumRows()
+	buf, err := b.bufAllocator.get()
+	if err != nil {
+		return err
+	}
+	defer b.bufAllocator.put(buf)
+	if err := b.args[0].VecEvalString(b.ctx, input, buf); err != nil {
 		return err
 	}
+	result.ReserveString(n)
 	enc := charset.FindEncoding(b.args[0].GetType().Charset)
-	for i := 0; i < input.NumRows(); i++ {
-		result.SetRaw(i, []byte(enc.ToLower(result.GetString(i))))
+	for i := 0; i < n; i++ {
+		if buf.IsNull(i) {
+			result.AppendNull()
+		} else {
+			result.AppendString(enc.ToLower(buf.GetString(i)))
+		}
 	}
 	return nil
 }
@@ -142,12 +153,23 @@ func (b *builtinStringIsNullSig) vectorized() bool {
 }
 
 func (b *builtinUpperUTF8Sig) vecEvalString(input *chunk.Chunk, result *chunk.Column) error {
-	if err := b.args[0].VecEvalString(b.ctx, input, result); err != nil {
+	n := input.NumRows()
+	buf, err := b.bufAllocator.get()
+	if err != nil {
+		return err
+	}
+	defer b.bufAllocator.put(buf)
+	if err := b.args[0].VecEvalString(b.ctx, input, buf); err != nil {
 		return err
 	}
+	result.ReserveString(n)
 	enc := charset.FindEncoding(b.args[0].GetType().Charset)
-	for i := 0; i < input.NumRows(); i++ {
-		result.SetRaw(i, []byte(enc.ToUpper(result.GetString(i))))
+	for i := 0; i < n; i++ {
+		if buf.IsNull(i) {
+			result.AppendNull()
+		} else {
+			result.AppendString(enc.ToUpper(buf.GetString(i)))
+		}
 	}
 	return nil
 }

diff --git a/expression/integration_test.go b/expression/integration_test.go
@@ -6995,3 +6995,19 @@ func TestIssue30264(t *testing.T) {
 	// compare JSON/JSON, return JSON type
 	tk.MustQuery("select greatest(cast(a as JSON), cast('3' as JSON)) from t1;").Check(testkit.Rows("3"))
 }
+
+func TestIssue32488(t *testing.T) {
+	store, clean := testkit.CreateMockStore(t)
+	defer clean()
+
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table t(a varchar(32)) DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;")
+	tk.MustExec("insert into t values('ʞ'), ('İ');")
+	tk.MustExec("set @@tidb_enable_vectorized_expression = false;")
+	tk.MustQuery("select binary upper(a), lower(a) from t order by upper(a);").Check([][]interface{}{{"İ i"}, {"Ʞ ʞ"}})
+	tk.MustQuery("select distinct upper(a), lower(a) from t order by upper(a);").Check([][]interface{}{{"İ i"}, {"Ʞ ʞ"}})
+	tk.MustExec("set @@tidb_enable_vectorized_expression = true;")
+	tk.MustQuery("select binary upper(a), lower(a) from t order by upper(a);").Check([][]interface{}{{"İ i"}, {"Ʞ ʞ"}})
+	tk.MustQuery("select distinct upper(a), lower(a) from t order by upper(a);").Check([][]interface{}{{"İ i"}, {"Ʞ ʞ"}})
+}