expression: support GBK charset for builtin function to_base64() (pin…

…gcap#29054)
tisonkun · Oct 25, 2021 · 836b102 · 836b102
1 parent 0ee0d6f
commit 836b102
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 0 deletions.
diff --git a/cmd/explaintest/r/new_character_set_builtin.result b/cmd/explaintest/r/new_character_set_builtin.result
@@ -32,3 +32,14 @@ select lower(a), lower(b) from t;
 lower(a)	lower(b)
 àáèéêìíòóùúüāēěīńňōūǎǐǒǔǖǘǚǜⅺⅻ	àáèéêìíòóùúüāēěīńňōūǎǐǒǔǖǘǚǜⅪⅫ
 set @@tidb_enable_vectorized_expression = false;
+drop table if exists t;
+create table t (a char(20) charset utf8mb4, b char(20) charset gbk, c binary(20));
+insert into t values ('一二三', '一二三', '一二三');
+select to_base64(a), to_base64(b), to_base64(c) from t;
+to_base64(a)	to_base64(b)	to_base64(c)
+5LiA5LqM5LiJ	0ru2/sj9	5LiA5LqM5LiJAAAAAAAAAAAAAAA=
+set @@tidb_enable_vectorized_expression = true;
+select to_base64(a), to_base64(b), to_base64(c) from t;
+to_base64(a)	to_base64(b)	to_base64(c)
+5LiA5LqM5LiJ	0ru2/sj9	5LiA5LqM5LiJAAAAAAAAAAAAAAA=
+set @@tidb_enable_vectorized_expression = false;
diff --git a/cmd/explaintest/t/new_character_set_builtin.test b/cmd/explaintest/t/new_character_set_builtin.test
@@ -20,3 +20,12 @@ select upper(a), upper(b) from t;
 select lower(a), lower(b) from t;
 set @@tidb_enable_vectorized_expression = false;
 
+-- test for builtin function to_base64()
+drop table if exists t;
+create table t (a char(20) charset utf8mb4, b char(20) charset gbk, c binary(20));
+insert into t values ('一二三', '一二三', '一二三');
+select to_base64(a), to_base64(b), to_base64(c) from t;
+set @@tidb_enable_vectorized_expression = true;
+select to_base64(a), to_base64(b), to_base64(c) from t;
+set @@tidb_enable_vectorized_expression = false;
+
diff --git a/expression/builtin_string.go b/expression/builtin_string.go
@@ -3599,6 +3599,12 @@ func (b *builtinToBase64Sig) evalString(row chunk.Row) (d string, isNull bool, e
 	if isNull || err != nil {
 		return "", isNull, err
 	}
+	argTp := b.args[0].GetType()
+	if !types.IsBinaryStr(argTp) {
+		if encodedStr, err := charset.NewEncoding(argTp.Charset).EncodeString(str); err == nil {
+			str = encodedStr
+		}
+	}
 	needEncodeLen := base64NeededEncodedLength(len(str))
 	if needEncodeLen == -1 {
 		return "", true, nil

diff --git a/expression/builtin_string_test.go b/expression/builtin_string_test.go
@@ -2324,6 +2324,28 @@ func (s *testEvaluatorSuite) TestToBase64(c *C) {
 
 	_, err := funcs[ast.ToBase64].getFunction(s.ctx, []Expression{NewZero()})
 	c.Assert(err, IsNil)
+
+	// Test GBK String
+	tbl := []struct {
+		input  string
+		chs    string
+		result string
+	}{
+		{"abc", "gbk", "YWJj"},
+		{"一二三", "gbk", "0ru2/sj9"},
+		{"一二三", "", "5LiA5LqM5LiJ"},
+		{"一二三!", "gbk", "0ru2/sj9IQ=="},
+		{"一二三!", "", "5LiA5LqM5LiJIQ=="},
+	}
+	for _, t := range tbl {
+		err := s.ctx.GetSessionVars().SetSystemVar(variable.CharacterSetConnection, t.chs)
+		c.Assert(err, IsNil)
+		f, err := newFunctionForTest(s.ctx, ast.ToBase64, s.primitiveValsToConstants([]interface{}{t.input})...)
+		c.Assert(err, IsNil)
+		d, err := f.Eval(chunk.Row{})
+		c.Assert(err, IsNil)
+		c.Assert(d.GetString(), Equals, t.result)
+	}
 }
 
 func (s *testEvaluatorSuite) TestToBase64Sig(c *C) {

diff --git a/expression/builtin_string_vec.go b/expression/builtin_string_vec.go
@@ -2424,13 +2424,22 @@ func (b *builtinToBase64Sig) vecEvalString(input *chunk.Chunk, result *chunk.Col
 		return err
 	}
 
+	argTp := b.args[0].GetType()
+	enc := charset.NewEncoding(argTp.Charset)
+	isBinaryStr := types.IsBinaryStr(argTp)
+
 	result.ReserveString(n)
 	for i := 0; i < n; i++ {
 		if buf.IsNull(i) {
 			result.AppendNull()
 			continue
 		}
 		str := buf.GetString(i)
+		if !isBinaryStr {
+			if encodedStr, err := enc.EncodeString(str); err == nil {
+				str = encodedStr
+			}
+		}
 		needEncodeLen := base64NeededEncodedLength(len(str))
 		if needEncodeLen == -1 {
 			result.AppendNull()