json: add json_storage_free function (#38544)

close #38517

DEPS.bzl

@@ -2839,8 +2839,8 @@ def go_deps():
         name = "com_github_pingcap_tipb",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/pingcap/tipb",
-        sum = "h1:kWYridgsn8xSKYJ2EkXp7uj5HwJnG5snpY3XP8oYmPU=",
-        version = "v0.0.0-20220824081009-0714a57aff1d",
+        sum = "h1:Yoo8j5xQGxjlsC3yt0ndsiAz0WZXED9rzsKmEN0U0DY=",
+        version = "v0.0.0-20221020071514-cd933387bcb5",
     )
     go_repository(
         name = "com_github_pkg_browser",

executor/fktest/BUILD.bazel

@@ -10,9 +10,12 @@ go_test(
     flaky = True,
     deps = [
         "//config",
+        "//executor",
         "//meta/autoid",
+        "//parser/model",
         "//planner/core",
         "//testkit",
+        "//types",
         "@com_github_stretchr_testify//require",
         "@com_github_tikv_client_go_v2//tikv",
         "@org_uber_go_goleak//:goleak",

executor/showtest/show_test.go

@@ -1478,7 +1478,7 @@ func TestShowBuiltin(t *testing.T) {
 	res := tk.MustQuery("show builtins;")
 	require.NotNil(t, res)
 	rows := res.Rows()
-	const builtinFuncNum = 282
+	const builtinFuncNum = 283
 	require.Equal(t, builtinFuncNum, len(rows))
 	require.Equal(t, rows[0][0].(string), "abs")
 	require.Equal(t, rows[builtinFuncNum-1][0].(string), "yearweek")

expression/builtin.go

@@ -883,6 +883,7 @@ var funcs = map[string]functionClass{
 	ast.JSONPretty:        &jsonPrettyFunctionClass{baseFunctionClass{ast.JSONPretty, 1, 1}},
 	ast.JSONQuote:         &jsonQuoteFunctionClass{baseFunctionClass{ast.JSONQuote, 1, 1}},
 	ast.JSONSearch:        &jsonSearchFunctionClass{baseFunctionClass{ast.JSONSearch, 3, -1}},
+	ast.JSONStorageFree:   &jsonStorageFreeFunctionClass{baseFunctionClass{ast.JSONStorageFree, 1, 1}},
 	ast.JSONStorageSize:   &jsonStorageSizeFunctionClass{baseFunctionClass{ast.JSONStorageSize, 1, 1}},
 	ast.JSONDepth:         &jsonDepthFunctionClass{baseFunctionClass{ast.JSONDepth, 1, 1}},
 	ast.JSONKeys:          &jsonKeysFunctionClass{baseFunctionClass{ast.JSONKeys, 1, 2}},

expression/builtin_json.go

@@ -1392,6 +1392,43 @@ func (b *builtinJSONSearchSig) evalJSON(row chunk.Row) (res types.BinaryJSON, is
 	return obj.Search(containType, searchStr, escape, nil)
 }
 
+type jsonStorageFreeFunctionClass struct {
+	baseFunctionClass
+}
+
+type builtinJSONStorageFreeSig struct {
+	baseBuiltinFunc
+}
+
+func (b *builtinJSONStorageFreeSig) Clone() builtinFunc {
+	newSig := &builtinJSONStorageFreeSig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	return newSig
+}
+
+func (c *jsonStorageFreeFunctionClass) getFunction(ctx sessionctx.Context, args []Expression) (builtinFunc, error) {
+	if err := c.verifyArgs(args); err != nil {
+		return nil, err
+	}
+
+	bf, err := newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETInt, types.ETJson)
+	if err != nil {
+		return nil, err
+	}
+	sig := &builtinJSONStorageFreeSig{bf}
+	sig.setPbCode(tipb.ScalarFuncSig_JsonStorageFreeSig)
+	return sig, nil
+}
+
+func (b *builtinJSONStorageFreeSig) evalInt(row chunk.Row) (res int64, isNull bool, err error) {
+	_, isNull, err = b.args[0].EvalJSON(b.ctx, row)
+	if isNull || err != nil {
+		return res, isNull, err
+	}
+
+	return 0, false, nil
+}
+
 type jsonStorageSizeFunctionClass struct {
 	baseFunctionClass
 }
@@ -1459,6 +1496,10 @@ func (c *jsonDepthFunctionClass) getFunction(ctx sessionctx.Context, args []Expr
 }
 
 func (b *builtinJSONDepthSig) evalInt(row chunk.Row) (res int64, isNull bool, err error) {
+	// as TiDB doesn't support partial update json value, so only check the
+	// json format and whether it's NULL. For NULL return NULL, for invalid json, return
+	// an error, otherwise return 0
+
 	obj, isNull, err := b.args[0].EvalJSON(b.ctx, row)
 	if isNull || err != nil {
 		return res, isNull, err

expression/builtin_json_test.go

@@ -985,6 +985,49 @@ func TestJSONValid(t *testing.T) {
 	}
 }
 
+func TestJSONStorageFree(t *testing.T) {
+	ctx := createContext(t)
+	fc := funcs[ast.JSONStorageFree]
+	tbl := []struct {
+		input    []interface{}
+		expected interface{}
+		success  bool
+	}{
+		// Tests scalar arguments
+		{[]interface{}{`null`}, 0, true},
+		{[]interface{}{`true`}, 0, true},
+		{[]interface{}{`1`}, 0, true},
+		{[]interface{}{`"1"`}, 0, true},
+		// Tests nil arguments
+		{[]interface{}{nil}, nil, true},
+		// Tests valid json documents
+		{[]interface{}{`{}`}, 0, true},
+		{[]interface{}{`{"a":1}`}, 0, true},
+		{[]interface{}{`[{"a":{"a":1},"b":2}]`}, 0, true},
+		{[]interface{}{`{"a": 1000, "b": "wxyz", "c": "[1, 3, 5, 7]"}`}, 0, true},
+		// Tests invalid json documents
+		{[]interface{}{`[{"a":1]`}, 0, false},
+		{[]interface{}{`[{a":1]`}, 0, false},
+	}
+	for _, tt := range tbl {
+		args := types.MakeDatums(tt.input...)
+		f, err := fc.getFunction(ctx, datumsToConstants(args))
+		require.NoError(t, err)
+		d, err := evalBuiltinFunc(f, chunk.Row{})
+		if tt.success {
+			require.NoError(t, err)
+
+			if tt.expected == nil {
+				require.True(t, d.IsNull())
+			} else {
+				require.Equal(t, int64(tt.expected.(int)), d.GetInt64())
+			}
+		} else {
+			require.Error(t, err)
+		}
+	}
+}
+
 func TestJSONStorageSize(t *testing.T) {
 	ctx := createContext(t)
 	fc := funcs[ast.JSONStorageSize]

expression/builtin_json_vec.go

@@ -103,6 +103,33 @@ func vecJSONModify(ctx sessionctx.Context, args []Expression, bufAllocator colum
 	return nil
 }
 
+func (b *builtinJSONStorageFreeSig) vectorized() bool {
+	return true
+}
+
+func (b *builtinJSONStorageFreeSig) vecEvalInt(input *chunk.Chunk, result *chunk.Column) error {
+	n := input.NumRows()
+	buf, err := b.bufAllocator.get()
+	if err != nil {
+		return err
+	}
+	defer b.bufAllocator.put(buf)
+	if err := b.args[0].VecEvalJSON(b.ctx, input, buf); err != nil {
+		return err
+	}
+	result.ResizeInt64(n, false)
+	result.MergeNulls(buf)
+	int64s := result.Int64s()
+	for i := 0; i < n; i++ {
+		if result.IsNull(i) {
+			continue
+		}
+
+		int64s[i] = 0
+	}
+	return nil
+}
+
 func (b *builtinJSONStorageSizeSig) vectorized() bool {
 	return true
 }

expression/integration_test.go

@@ -7755,3 +7755,15 @@ func TestFix38127(t *testing.T) {
 	tk.MustQuery("select from_unixtime(dou, '%Y-%m-%d') from t").Check(testkit.Rows("<nil>"))
 	tk.MustQuery("select from_unixtime(varc, '%Y-%m-%d') from t").Check(testkit.Rows("<nil>"))
 }
+
+func TestJSONStorageFree(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustQuery("select json_storage_free(NULL)").Check(testkit.Rows("<nil>"))
+	tk.MustQuery("select json_storage_free('{}')").Check(testkit.Rows("0"))
+	tk.MustQuery("select json_storage_free('1')").Check(testkit.Rows("0"))
+	tk.MustQuery(`select json_storage_free('{"a": "b"}')`).Check(testkit.Rows("0"))
+	err := tk.ExecToErr(`select json_storage_free('{"c":["a","b"]`)
+	require.Error(t, err, "[json:3140]Invalid JSON text: The document root must not be followed by other values.")
+}

go.mod

@@ -71,7 +71,7 @@ require (
 	github.com/pingcap/log v1.1.0
 	github.com/pingcap/sysutil v0.0.0-20220114020952-ea68d2dbf5b4
 	github.com/pingcap/tidb/parser v0.0.0-20211011031125-9b13dc409c5e
-	github.com/pingcap/tipb v0.0.0-20220824081009-0714a57aff1d
+	github.com/pingcap/tipb v0.0.0-20221020071514-cd933387bcb5
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.13.0
 	github.com/prometheus/client_model v0.2.0

go.sum

@@ -766,8 +766,8 @@ github.com/pingcap/log v1.1.0 h1:ELiPxACz7vdo1qAvvaWJg1NrYFoY6gqAh/+Uo6aXdD8=
 github.com/pingcap/log v1.1.0/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pingcap/sysutil v0.0.0-20220114020952-ea68d2dbf5b4 h1:HYbcxtnkN3s5tqrZ/z3eJS4j3Db8wMphEm1q10lY/TM=
 github.com/pingcap/sysutil v0.0.0-20220114020952-ea68d2dbf5b4/go.mod h1:sDCsM39cGiv2vwunZkaFA917vVkqDTGSPbbV7z4Oops=
-github.com/pingcap/tipb v0.0.0-20220824081009-0714a57aff1d h1:kWYridgsn8xSKYJ2EkXp7uj5HwJnG5snpY3XP8oYmPU=
-github.com/pingcap/tipb v0.0.0-20220824081009-0714a57aff1d/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs=
+github.com/pingcap/tipb v0.0.0-20221020071514-cd933387bcb5 h1:Yoo8j5xQGxjlsC3yt0ndsiAz0WZXED9rzsKmEN0U0DY=
+github.com/pingcap/tipb v0.0.0-20221020071514-cd933387bcb5/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs=
 github.com/pkg/browser v0.0.0-20180916011732-0a3d74bf9ce4 h1:49lOXmGaUpV9Fz3gd7TFZY106KVlPVa5jcYD1gaQf98=
 github.com/pkg/browser v0.0.0-20180916011732-0a3d74bf9ce4/go.mod h1:4OwLy04Bl9Ef3GJJCoec+30X3LQs/0/m4HFRt/2LUSA=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=

parser/ast/functions.go

@@ -341,6 +341,7 @@ const (
 	JSONPretty        = "json_pretty"
 	JSONQuote         = "json_quote"
 	JSONSearch        = "json_search"
+	JSONStorageFree   = "json_storage_free"
 	JSONStorageSize   = "json_storage_size"
 	JSONDepth         = "json_depth"
 	JSONKeys          = "json_keys"

tidb-server/BUILD.bazel

@@ -37,7 +37,6 @@ go_library(
         "//util/deadlockhistory",
         "//util/disk",
         "//util/domainutil",
-        "//util/gctuner",
         "//util/kvcache",
         "//util/logutil",
         "//util/memory",