From 885a71bd5611211f46726017f561aff29423d98c Mon Sep 17 00:00:00 2001
From: Hangjie Mo <mohangjie1995@gmail.com>
Date: Mon, 21 Feb 2022 17:35:43 +0800
Subject: [PATCH] expression: fix panic in upper() and lower() func. (#32505)

close pingcap/tidb#32488
---
 expression/builtin_string_vec.go | 34 ++++++++++++++++++++++++++------
 expression/integration_test.go   | 16 +++++++++++++++
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/expression/builtin_string_vec.go b/expression/builtin_string_vec.go
index 75cf620d3202b..6ca8364eb7f6b 100644
--- 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
index fbf80802f5f84..de0286f283714 100644
--- 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"}, {"Ʞ ʞ"}})
+}