expression: introduce new function for JSON array (#60728)

ref #60649

Author: joechenrh

pkg/expression/bench_test.go

@@ -299,6 +299,22 @@ func (g *charInt64Gener) gen() any {
 	return int64(nanosecond)
 }
 
+type jsonArrayGener struct {
+	rand *defaultRandGen
+}
+
+func newJSONArrayGener() *jsonArrayGener {
+	return &jsonArrayGener{newDefaultRandGen()}
+}
+
+func (g *jsonArrayGener) gen() any {
+	v := make([]any, 4)
+	for i := range len(v) {
+		v[i] = int64(g.rand.Int())
+	}
+	return types.CreateBinaryJSON(v)
+}
+
 // selectStringGener select one string randomly from the candidates array
 type selectStringGener struct {
 	candidates []string
@@ -1459,7 +1475,10 @@ func genVecBuiltinFuncBenchCase(ctx BuildContext, funcName string, testCase vecE
 	}
 
 	var err error
-	if funcName == ast.Cast {
+	if funcName == ast.JSONSumCrc32 {
+		fc := &jsonSumCRC32FunctionClass{baseFunctionClass{ast.JSONSumCrc32, 1, 1}, fts[0]}
+		baseFunc, err = fc.getFunction(ctx, cols)
+	} else if funcName == ast.Cast {
 		var fc functionClass
 		tp := eType2FieldType(testCase.retEvalType)
 		switch testCase.retEvalType {

pkg/expression/builtin_json.go

@@ -18,6 +18,8 @@ import (
 	"bytes"
 	"context"
 	goJSON "encoding/json"
+	"fmt"
+	"hash/crc32"
 	"strconv"
 	"strings"
 
@@ -35,6 +37,7 @@ import (
 
 var (
 	_ functionClass = &jsonTypeFunctionClass{}
+	_ functionClass = &jsonSumCRC32FunctionClass{}
 	_ functionClass = &jsonExtractFunctionClass{}
 	_ functionClass = &jsonUnquoteFunctionClass{}
 	_ functionClass = &jsonQuoteFunctionClass{}
@@ -64,6 +67,7 @@ var (
 	_ functionClass = &jsonLengthFunctionClass{}
 
 	_ builtinFunc = &builtinJSONTypeSig{}
+	_ builtinFunc = &builtinJSONSumCRC32Sig{}
 	_ builtinFunc = &builtinJSONQuoteSig{}
 	_ builtinFunc = &builtinJSONUnquoteSig{}
 	_ builtinFunc = &builtinJSONArraySig{}
@@ -180,6 +184,112 @@ func (b *builtinJSONTypeSig) evalString(ctx EvalContext, row chunk.Row) (val str
 	return j.Type(), false, nil
 }
 
+type jsonSumCRC32FunctionClass struct {
+	baseFunctionClass
+
+	tp *types.FieldType
+}
+
+func (c *jsonSumCRC32FunctionClass) verifyArgs(ctx EvalContext, args []Expression) error {
+	if err := c.baseFunctionClass.verifyArgs(args); err != nil {
+		return err
+	}
+
+	if args[0].GetType(ctx).EvalType() != types.ETJson {
+		return ErrInvalidTypeForJSON.GenWithStackByArgs(1, "JSON_SUM_CRC32")
+	}
+
+	return nil
+}
+
+func (c *jsonSumCRC32FunctionClass) getFunction(ctx BuildContext, args []Expression) (sig builtinFunc, err error) {
+	if err := c.verifyArgs(ctx.GetEvalCtx(), args); err != nil {
+		return nil, err
+	}
+	arrayType := c.tp.ArrayType()
+	switch arrayType.GetType() {
+	case mysql.TypeYear, mysql.TypeJSON, mysql.TypeFloat, mysql.TypeNewDecimal:
+		return nil, ErrNotSupportedYet.GenWithStackByArgs(fmt.Sprintf("calculating json_sum_crc32 on array of %s", arrayType.String()))
+	}
+	if arrayType.EvalType() == types.ETString && arrayType.GetCharset() != charset.CharsetUTF8MB4 && arrayType.GetCharset() != charset.CharsetBin {
+		return nil, ErrNotSupportedYet.GenWithStackByArgs("unsupported charset")
+	}
+	if arrayType.EvalType() == types.ETString && arrayType.GetFlen() == types.UnspecifiedLength {
+		return nil, ErrNotSupportedYet.GenWithStackByArgs("calculating json_sum_crc32 on array of char/binary BLOBs with unspecified length")
+	}
+
+	bf, err := newBaseBuiltinFunc(ctx, c.funcName, args, c.tp)
+	if err != nil {
+		return nil, err
+	}
+	sig = &builtinJSONSumCRC32Sig{bf}
+	return sig, nil
+}
+
+type builtinJSONSumCRC32Sig struct {
+	baseBuiltinFunc
+}
+
+func (b *builtinJSONSumCRC32Sig) Clone() builtinFunc {
+	newSig := &builtinJSONSumCRC32Sig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	return newSig
+}
+
+func (b *builtinJSONSumCRC32Sig) evalInt(ctx EvalContext, row chunk.Row) (res int64, isNull bool, err error) {
+	val, isNull, err := b.args[0].EvalJSON(ctx, row)
+	if isNull || err != nil {
+		return res, isNull, err
+	}
+
+	if val.TypeCode != types.JSONTypeCodeArray {
+		return 0, false, ErrInvalidTypeForJSON.GenWithStackByArgs(1, "JSON_SUM_CRC32")
+	}
+
+	ft := b.tp.ArrayType()
+	f := convertJSON2Tp(ft.EvalType())
+	if f == nil {
+		return 0, false, ErrNotSupportedYet.GenWithStackByArgs(fmt.Sprintf("calculating sum of %s", ft.String()))
+	}
+
+	var sum int64
+	for i := range val.GetElemCount() {
+		item, err := f(fakeSctx, val.ArrayGetElem(i), ft)
+		if err != nil {
+			if ErrInvalidJSONForFuncIndex.Equal(err) {
+				err = errors.Errorf("Invalid JSON value for CAST to type %s", ft.CompactStr())
+			}
+			return 0, false, err
+		}
+		sum += int64(crc32.ChecksumIEEE(fmt.Appendf(nil, "%v", item)))
+	}
+
+	return sum, false, err
+}
+
+// BuildJSONSumCrc32FunctionWithCheck builds a JSON_SUM_CRC32 ScalarFunction from the Expression and return error if any.
+// The logic is almost the same as build CAST function, except that the return type is fixed to bigint.
+func BuildJSONSumCrc32FunctionWithCheck(ctx BuildContext, expr Expression, tp *types.FieldType) (res Expression, err error) {
+	argType := expr.GetType(ctx.GetEvalCtx())
+	// If source argument's nullable, then target type should be nullable
+	if !mysql.HasNotNullFlag(argType.GetFlag()) {
+		tp.DelFlag(mysql.NotNullFlag)
+	}
+	expr = TryPushCastIntoControlFunctionForHybridType(ctx, expr, tp)
+
+	if tp.EvalType() != types.ETJson || !tp.IsArray() {
+		return nil, errors.Errorf("json_sum_crc32 can only built on JSON array, got type %s", tp.EvalType())
+	}
+
+	fc := &jsonSumCRC32FunctionClass{baseFunctionClass{ast.JSONSumCrc32, 1, 1}, tp}
+	f, err := fc.getFunction(ctx, []Expression{expr})
+	return &ScalarFunction{
+		FuncName: ast.NewCIStr(ast.JSONSumCrc32),
+		RetType:  types.NewFieldType(mysql.TypeLonglong),
+		Function: f,
+	}, err
+}
+
 type jsonExtractFunctionClass struct {
 	baseFunctionClass
 }

pkg/expression/builtin_json_test.go

@@ -16,10 +16,12 @@ package expression
 
 import (
 	"fmt"
+	"hash/crc32"
 	"testing"
 
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/parser/ast"
+	"github.com/pingcap/tidb/pkg/parser/charset"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/parser/terror"
 	"github.com/pingcap/tidb/pkg/testkit/testutil"
@@ -122,6 +124,112 @@ func TestJSONUnquote(t *testing.T) {
 	}
 }
 
+func TestJSONSumCrc32(t *testing.T) {
+	checksumFunc := func(vals []any) int64 {
+		var sum int64
+		for _, v := range vals {
+			sum += int64(crc32.ChecksumIEEE(fmt.Appendf(nil, "%v", v)))
+		}
+		return sum
+	}
+
+	buildJSON := func(vals []any, path string) (Expression, Expression) {
+		v := make(map[string]any)
+		v[path] = vals
+		return datumsToConstants(types.MakeDatums(types.CreateBinaryJSON(v)))[0],
+			datumsToConstants(types.MakeDatums(fmt.Sprintf(`$.%s`, path)))[0]
+	}
+
+	buildType := func(tp byte, unsigned bool) *types.FieldType {
+		b := types.NewFieldTypeBuilder().SetType(tp).SetCharset(charset.CharsetBin).SetCollate(charset.CollationBin).SetArray(true)
+		if unsigned {
+			b.AddFlag(mysql.UnsignedFlag)
+		}
+		if types.IsString(tp) {
+			b.SetFlen(16)
+		}
+		return b.BuildP()
+	}
+
+	ctx := createContext(t)
+	tbl := []struct {
+		input     []any
+		path      string
+		evalError bool
+		tp        *types.FieldType
+	}{
+		{
+			[]any{int64(-1), int64(2), int64(3)},
+			"",
+			false,
+			buildType(mysql.TypeLonglong, false),
+		},
+		{
+			[]any{int64(1), int64(2), int64(3)},
+			"",
+			false,
+			buildType(mysql.TypeLonglong, true),
+		},
+		{
+			[]any{int64(-1), int64(1)},
+			"",
+			true,
+			buildType(mysql.TypeLonglong, true),
+		},
+		{
+			[]any{"a", "b", "c"},
+			"",
+			false,
+			buildType(mysql.TypeVarString, false),
+		},
+		{
+			[]any{float64(1.1), int64(1), float64(3.3)},
+			"",
+			false,
+			buildType(mysql.TypeDouble, false),
+		},
+		{
+			[]any{float64(1.1), int64(1), float64(3.3)},
+			"",
+			true,
+			buildType(mysql.TypeLonglong, false),
+		},
+		{
+			[]any{1.1, 2.2, 3.3},
+			"arr",
+			false,
+			buildType(mysql.TypeDouble, false),
+		},
+		{
+			[]any{1.1, "1.1", 3.3},
+			"arr",
+			true,
+			buildType(mysql.TypeDouble, false),
+		},
+	}
+
+	var err error
+	for _, tt := range tbl {
+		inputExpr := datumsToConstants(types.MakeDatums(types.CreateBinaryJSON(tt.input)))[0]
+		if tt.path != "" {
+			json, extract := buildJSON(tt.input, tt.path)
+			inputExpr, err = NewFunction(ctx, ast.JSONExtract, types.NewFieldType(mysql.TypeJSON), json, extract)
+			require.NoError(t, err)
+		}
+
+		f, err := BuildJSONSumCrc32FunctionWithCheck(ctx, inputExpr, tt.tp)
+		require.NoError(t, err, tt.input)
+
+		val, _, err := f.EvalInt(ctx, chunk.Row{})
+		if tt.evalError {
+			require.Error(t, err)
+		} else {
+			require.NoError(t, err, tt.input)
+			require.EqualValues(t, checksumFunc(tt.input), val)
+		}
+	}
+}
+
 func TestJSONExtract(t *testing.T) {
 	ctx := createContext(t)
 	fc := funcs[ast.JSONExtract]

pkg/expression/builtin_json_vec.go

@@ -17,8 +17,11 @@ package expression
 import (
 	"bytes"
 	goJSON "encoding/json"
+	"fmt"
+	"hash/crc32"
 	"strings"
 
+	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util/chunk"
@@ -914,6 +917,59 @@ func (b *builtinJSONTypeSig) vecEvalString(ctx EvalContext, input *chunk.Chunk,
 	return nil
 }
 
+func (b *builtinJSONSumCRC32Sig) vectorized() bool {
+	return true
+}
+
+func (b *builtinJSONSumCRC32Sig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, result *chunk.Column) error {
+	ft := b.tp.ArrayType()
+	f := convertJSON2Tp(ft.EvalType())
+	if f == nil {
+		return ErrNotSupportedYet.GenWithStackByArgs(fmt.Sprintf("calculating sum of %s", ft.String()))
+	}
+
+	nr := input.NumRows()
+	jsonBuf, err := b.bufAllocator.get()
+	if err != nil {
+		return err
+	}
+
+	defer b.bufAllocator.put(jsonBuf)
+	if err = b.args[0].VecEvalJSON(ctx, input, jsonBuf); err != nil {
+		return err
+	}
+
+	result.ResizeInt64(nr, false)
+	result.MergeNulls(jsonBuf)
+	i64s := result.Int64s()
+
+	for i := range nr {
+		if result.IsNull(i) {
+			continue
+		}
+
+		jsonItem := jsonBuf.GetJSON(i)
+		if jsonItem.TypeCode != types.JSONTypeCodeArray {
+			return ErrInvalidTypeForJSON.GenWithStackByArgs(1, "JSON_SUM_CRC32")
+		}
+
+		var sum int64
+		for j := range jsonItem.GetElemCount() {
+			item, err := f(fakeSctx, jsonItem.ArrayGetElem(j), ft)
+			if err != nil {
+				if ErrInvalidJSONForFuncIndex.Equal(err) {
+					err = errors.Errorf("Invalid JSON value for CAST to type %s", ft.CompactStr())
+				}
+				return err
+			}
+			sum += int64(crc32.ChecksumIEEE(fmt.Appendf(nil, "%v", item)))
+		}
+		i64s[i] = sum
+	}
+
+	return nil
+}
+
 func (b *builtinJSONExtractSig) vectorized() bool {
 	return true
 }

pkg/expression/builtin_json_vec_test.go

@@ -18,6 +18,8 @@ import (
 	"testing"
 
 	"github.com/pingcap/tidb/pkg/parser/ast"
+	"github.com/pingcap/tidb/pkg/parser/charset"
+	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/types"
 )
 
@@ -138,6 +140,13 @@ var vecBuiltinJSONCases = map[string][]vecExprBenchCase{
 	ast.JSONQuote: {
 		{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString}},
 	},
+	ast.JSONSumCrc32: {
+		{
+			retEvalType:        types.ETInt,
+			childrenTypes:      []types.EvalType{types.ETJson},
+			childrenFieldTypes: []*types.FieldType{types.NewFieldTypeBuilder().SetType(mysql.TypeLonglong).SetCharset(charset.CharsetBin).SetCollate(charset.CollationBin).SetArray(true).BuildP()},
+			geners:             []dataGenerator{newJSONArrayGener()}},
+	},
 }
 
 func TestVectorizedBuiltinJSONFunc(t *testing.T) {