Skip to content

Commit

Permalink
expression/json: add builtin function JSON_SEARCH (#8704)
Browse files Browse the repository at this point in the history
  • Loading branch information
@pingyu @alivxxx
pingyu authored and alivxxx committed Apr 18, 2019
1 parent 76e0a91 commit a594287
Show file tree
Hide file tree
Showing 8 changed files with 699 additions and 4 deletions.
117 changes: 116 additions & 1 deletion expression/builtin_json.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import (
"github.com/pingcap/tidb/types"
"github.com/pingcap/tidb/types/json"
"github.com/pingcap/tidb/util/chunk"
"github.com/pingcap/tidb/util/stringutil"
"github.com/pingcap/tipb/go-tipb"
)

Expand Down Expand Up @@ -66,6 +67,7 @@ var (
_ builtinFunc = &builtinJSONMergeSig{}
_ builtinFunc = &builtinJSONContainsSig{}
_ builtinFunc = &builtinJSONDepthSig{}
_ builtinFunc = &builtinJSONSearchSig{}
_ builtinFunc = &builtinJSONKeysSig{}
_ builtinFunc = &builtinJSONKeys2ArgsSig{}
_ builtinFunc = &builtinJSONLengthSig{}
Expand Down Expand Up @@ -876,8 +878,121 @@ type jsonSearchFunctionClass struct {
baseFunctionClass
}

type builtinJSONSearchSig struct {
baseBuiltinFunc
}

func (b *builtinJSONSearchSig) Clone() builtinFunc {
newSig := &builtinJSONSearchSig{}
newSig.cloneFrom(&b.baseBuiltinFunc)
return newSig
}

func (c *jsonSearchFunctionClass) getFunction(ctx sessionctx.Context, args []Expression) (builtinFunc, error) {
return nil, errFunctionNotExists.GenWithStackByArgs("FUNCTION", "JSON_SEARCH")
if err := c.verifyArgs(args); err != nil {
return nil, err
}
// json_doc, one_or_all, search_str[, escape_char[, path] ...])
argTps := make([]types.EvalType, 0, len(args))
argTps = append(argTps, types.ETJson)
for range args[1:] {
argTps = append(argTps, types.ETString)
}
bf := newBaseBuiltinFuncWithTp(ctx, args, types.ETJson, argTps...)
sig := &builtinJSONSearchSig{bf}
sig.setPbCode(tipb.ScalarFuncSig_JsonSearchSig)
return sig, nil
}

func (b *builtinJSONSearchSig) evalJSON(row chunk.Row) (res json.BinaryJSON, isNull bool, err error) {
// json_doc
var obj json.BinaryJSON
obj, isNull, err = b.args[0].EvalJSON(b.ctx, row)
if isNull || err != nil {
return res, isNull, err
}

// one_or_all
var containType string
containType, isNull, err = b.args[1].EvalString(b.ctx, row)
if isNull || err != nil {
return res, isNull, err
}
if containType != json.ContainsPathAll && containType != json.ContainsPathOne {
return res, true, errors.AddStack(json.ErrInvalidJSONContainsPathType)
}

// search_str & escape_char
var searchStr string
searchStr, isNull, err = b.args[2].EvalString(b.ctx, row)
if isNull || err != nil {
return res, isNull, err
}
escape := byte('\\')
if len(b.args) >= 4 {
var escapeStr string
escapeStr, isNull, err = b.args[3].EvalString(b.ctx, row)
if err != nil {
return res, isNull, err
}
if isNull || len(escapeStr) == 0 {
escape = byte('\\')
} else if len(escapeStr) == 1 {
escape = byte(escapeStr[0])
} else {
return res, true, errIncorrectArgs.GenWithStackByArgs("ESCAPE")
}
}
patChars, patTypes := stringutil.CompilePattern(searchStr, escape)

// result
result := make([]interface{}, 0)

// walk json_doc
walkFn := func(fullpath json.PathExpression, bj json.BinaryJSON) (stop bool, err error) {
if bj.TypeCode == json.TypeCodeString && stringutil.DoMatch(string(bj.GetString()), patChars, patTypes) {
result = append(result, fullpath.String())
if containType == json.ContainsPathOne {
return true, nil
}
}
return false, nil
}
if len(b.args) >= 5 { // path...
pathExprs := make([]json.PathExpression, 0, len(b.args)-4)
for i := 4; i < len(b.args); i++ {
var s string
s, isNull, err = b.args[i].EvalString(b.ctx, row)
if isNull || err != nil {
return res, isNull, err
}
var pathExpr json.PathExpression
pathExpr, err = json.ParseJSONPathExpr(s)
if err != nil {
return res, true, err
}
pathExprs = append(pathExprs, pathExpr)
}
err = obj.Walk(walkFn, pathExprs...)
if err != nil {
return res, true, err
}
} else {
err = obj.Walk(walkFn)
if err != nil {
return res, true, err
}
}

// return
switch len(result) {
case 0:
return res, true, nil
case 1:
return json.CreateBinary(result[0]), false, nil
default:
return json.CreateBinary(result), false, nil
}
}

type jsonStorageSizeFunctionClass struct {
Expand Down
74 changes: 74 additions & 0 deletions expression/builtin_json_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -773,3 +773,77 @@ func (s *testEvaluatorSuite) TestJSONArrayAppend(c *C) {
c.Assert(json.CompareBinary(j1, d.GetMysqlJSON()), Equals, 0, comment)
}
}

func (s *testEvaluatorSuite) TestJSONSearch(c *C) {
defer testleak.AfterTest(c)()
fc := funcs[ast.JSONSearch]
jsonString := `["abc", [{"k": "10"}, "def"], {"x":"abc"}, {"y":"bcd"}]`
jsonString2 := `["abc", [{"k": "10"}, "def"], {"x":"ab%d"}, {"y":"abcd"}]`
tbl := []struct {
input []interface{}
expected interface{}
success bool
}{
// simple case
{[]interface{}{jsonString, `one`, `abc`}, `"$[0]"`, true},
{[]interface{}{jsonString, `all`, `abc`}, `["$[0]", "$[2].x"]`, true},
{[]interface{}{jsonString, `all`, `ghi`}, nil, true},
{[]interface{}{jsonString, `all`, `10`}, `"$[1][0].k"`, true},
{[]interface{}{jsonString, `all`, `10`, nil, `$`}, `"$[1][0].k"`, true},
{[]interface{}{jsonString, `all`, `10`, nil, `$[*]`}, `"$[1][0].k"`, true},
{[]interface{}{jsonString, `all`, `10`, nil, `$**.k`}, `"$[1][0].k"`, true},
{[]interface{}{jsonString, `all`, `10`, nil, `$[*][0].k`}, `"$[1][0].k"`, true},
{[]interface{}{jsonString, `all`, `10`, nil, `$[1]`}, `"$[1][0].k"`, true},
{[]interface{}{jsonString, `all`, `10`, nil, `$[1][0]`}, `"$[1][0].k"`, true},
{[]interface{}{jsonString, `all`, `abc`, nil, `$[2]`}, `"$[2].x"`, true},
{[]interface{}{jsonString, `all`, `abc`, nil, `$[2]`, `$[0]`}, `["$[2].x", "$[0]"]`, true},
{[]interface{}{jsonString, `all`, `abc`, nil, `$[2]`, `$[2]`}, `"$[2].x"`, true},

// search pattern
{[]interface{}{jsonString, `all`, `%a%`}, `["$[0]", "$[2].x"]`, true},
{[]interface{}{jsonString, `all`, `%b%`}, `["$[0]", "$[2].x", "$[3].y"]`, true},
{[]interface{}{jsonString, `all`, `%b%`, nil, `$[0]`}, `"$[0]"`, true},
{[]interface{}{jsonString, `all`, `%b%`, nil, `$[2]`}, `"$[2].x"`, true},
{[]interface{}{jsonString, `all`, `%b%`, nil, `$[1]`}, nil, true},
{[]interface{}{jsonString, `all`, `%b%`, ``, `$[1]`}, nil, true},
{[]interface{}{jsonString, `all`, `%b%`, nil, `$[3]`}, `"$[3].y"`, true},
{[]interface{}{jsonString2, `all`, `ab_d`}, `["$[2].x", "$[3].y"]`, true},

// escape char
{[]interface{}{jsonString2, `all`, `ab%d`}, `["$[2].x", "$[3].y"]`, true},
{[]interface{}{jsonString2, `all`, `ab\%d`}, `"$[2].x"`, true},
{[]interface{}{jsonString2, `all`, `ab|%d`, `|`}, `"$[2].x"`, true},

// error handle
{[]interface{}{nil, `all`, `abc`}, nil, true}, // NULL json
{[]interface{}{`a`, `all`, `abc`}, nil, false}, // non json
{[]interface{}{jsonString, `wrong`, `abc`}, nil, false}, // wrong one_or_all
{[]interface{}{jsonString, `all`, nil}, nil, true}, // NULL search_str
{[]interface{}{jsonString, `all`, `abc`, `??`}, nil, false}, // wrong escape_char
{[]interface{}{jsonString, `all`, `abc`, nil, nil}, nil, true}, // NULL path
{[]interface{}{jsonString, `all`, `abc`, nil, `$xx`}, nil, false}, // wrong path
}
for _, t := range tbl {
args := types.MakeDatums(t.input...)
f, err := fc.getFunction(s.ctx, s.datumsToConstants(args))
c.Assert(err, IsNil)
d, err := evalBuiltinFunc(f, chunk.Row{})
if t.success {
c.Assert(err, IsNil)
switch x := t.expected.(type) {
case string:
var j1, j2 json.BinaryJSON
j1, err = json.ParseBinaryFromString(x)
c.Assert(err, IsNil)
j2 = d.GetMysqlJSON()
cmp := json.CompareBinary(j1, j2)
//fmt.Println(j1, j2)
c.Assert(cmp, Equals, 0)
case nil:
c.Assert(d.IsNull(), IsTrue)
}
} else {
c.Assert(err, NotNil)
}
}
}
2 changes: 2 additions & 0 deletions expression/distsql_builtin.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -444,6 +444,8 @@ func getSignatureByPB(ctx sessionctx.Context, sigCode tipb.ScalarFuncSig, tp *ti
f = &builtinJSONLengthSig{base}
case tipb.ScalarFuncSig_JsonDepthSig:
f = &builtinJSONDepthSig{base}
case tipb.ScalarFuncSig_JsonSearchSig:
f = &builtinJSONSearchSig{base}

case tipb.ScalarFuncSig_InInt:
f = &builtinInIntSig{base}
Expand Down
31 changes: 28 additions & 3 deletions expression/distsql_builtin_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ import (
"time"

. "github.com/pingcap/check"
"github.com/pingcap/parser/charset"
"github.com/pingcap/parser/mysql"
"github.com/pingcap/tidb/sessionctx/stmtctx"
"github.com/pingcap/tidb/types"
Expand Down Expand Up @@ -108,6 +109,17 @@ func (s *testEvalSuite) TestEval(c *C) {
),
types.NewIntDatum(3),
},
{
scalarFunctionExpr(tipb.ScalarFuncSig_JsonSearchSig,
toPBFieldType(newJSONFieldType()),
jsonDatumExpr(c, `["abc", [{"k": "10"}, "def"], {"x":"abc"}, {"y":"bcd"}]`),
datumExpr(c, types.NewBytesDatum([]byte(`all`))),
datumExpr(c, types.NewBytesDatum([]byte(`10`))),
datumExpr(c, types.NewBytesDatum([]byte(`\`))),
datumExpr(c, types.NewBytesDatum([]byte(`$**.k`))),
),
newJSONDatum(c, `"$[1][0].k"`),
},
}
sc := new(stmtctx.StatementContext)
for _, tt := range tests {
Expand Down Expand Up @@ -178,12 +190,15 @@ func datumExpr(c *C, d types.Datum) *tipb.Expr {
return expr
}

func jsonDatumExpr(c *C, s string) *tipb.Expr {
var d types.Datum
func newJSONDatum(c *C, s string) (d types.Datum) {
j, err := json.ParseBinaryFromString(s)
c.Assert(err, IsNil)
d.SetMysqlJSON(j)
return datumExpr(c, d)
return d
}

func jsonDatumExpr(c *C, s string) *tipb.Expr {
return datumExpr(c, newJSONDatum(c, s))
}

func columnExpr(columnID int64) *tipb.Expr {
Expand Down Expand Up @@ -214,6 +229,16 @@ func newIntFieldType() *types.FieldType {
}
}

func newJSONFieldType() *types.FieldType {
return &types.FieldType{
Tp: mysql.TypeJSON,
Flen: types.UnspecifiedLength,
Decimal: 0,
Charset: charset.CharsetBin,
Collate: charset.CollationBin,
}
}

func scalarFunctionExpr(sigCode tipb.ScalarFuncSig, retType *tipb.FieldType, args ...*tipb.Expr) *tipb.Expr {
return &tipb.Expr{
Tp: tipb.ExprType_ScalarFunc,
Expand Down
Loading

0 comments on commit a594287

Please sign in to comment.