Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

expression: add builtin function json_length #7739

Merged
merged 6 commits into from
Sep 27, 2018
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 61 additions & 1 deletion expression/builtin_json.go
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@ var (
_ builtinFunc = &builtinJSONRemoveSig{}
_ builtinFunc = &builtinJSONMergeSig{}
_ builtinFunc = &builtinJSONContainsSig{}
_ builtinFunc = &builtinJSONLengthSig{}
)

type jsonTypeFunctionClass struct {
Expand Down Expand Up @@ -772,6 +773,65 @@ type jsonLengthFunctionClass struct {
baseFunctionClass
}

type builtinJSONLengthSig struct {
baseBuiltinFunc
}

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

func (c *jsonLengthFunctionClass) getFunction(ctx sessionctx.Context, args []Expression) (builtinFunc, error) {
return nil, errFunctionNotExists.GenWithStackByArgs("FUNCTION", "JSON_LENGTH")
if err := c.verifyArgs(args); err != nil {
return nil, errors.Trace(err)
}

argTps := make([]types.EvalType, 0, len(args))
argTps = append(argTps, types.ETJson)
if len(args) == 2 {
argTps = append(argTps, types.ETString)
}

bf := newBaseBuiltinFuncWithTp(ctx, args, types.ETInt, argTps...)
sig := &builtinJSONLengthSig{bf}
sig.setPbCode(tipb.ScalarFuncSig_JsonLengthSig)
return sig, nil
}

func (b *builtinJSONLengthSig) evalInt(row chunk.Row) (res int64, isNull bool, err error) {
obj, isNull, err := b.args[0].EvalJSON(b.ctx, row)
if isNull || err != nil {
return res, isNull, errors.Trace(err)
}

if obj.TypeCode != json.TypeCodeObject && obj.TypeCode != json.TypeCodeArray {
return 1, false, nil
}

if len(b.args) == 2 {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how about:

if len(b.args) == 1 {
	return int64(obj.GetElemCount()), false, nil
}
// handle the case that len(b.args) == 2
...

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

actually, I'm ready to extract the common path visit but in a other pr(some other pr are conflicted with that)

path, isNull, err := b.args[1].EvalString(b.ctx, row)
if isNull || err != nil {
return res, isNull, errors.Trace(err)
}

pathExpr, err := json.ParseJSONPathExpr(path)
if err != nil {
return res, true, errors.Trace(err)
}
if pathExpr.ContainsAnyAsterisk() {
return res, true, json.ErrInvalidJSONPathWildcard
}

var exists bool
obj, exists = obj.Extract([]json.PathExpression{pathExpr})
if !exists {
return res, true, nil
}
if obj.TypeCode != json.TypeCodeObject && obj.TypeCode != json.TypeCodeArray {
return 1, false, nil
}
}
return int64(obj.GetElemCount()), false, nil
}
71 changes: 71 additions & 0 deletions expression/builtin_json_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -424,3 +424,74 @@ func (s *testEvaluatorSuite) TestJSONContainsPath(c *C) {
}
}
}

func (s *testEvaluatorSuite) TestJSONLength(c *C) {
defer testleak.AfterTest(c)()
fc := funcs[ast.JSONLength]
tbl := []struct {
input []interface{}
expected interface{}
success bool
}{
// Tests scalar arguments
{[]interface{}{`null`}, 1, true},
{[]interface{}{`true`}, 1, true},
{[]interface{}{`false`}, 1, true},
{[]interface{}{`1`}, 1, true},
{[]interface{}{`-1`}, 1, true},
{[]interface{}{`1.1`}, 1, true},
{[]interface{}{`"1"`}, 1, true},
{[]interface{}{`"1"`, "$.a"}, 1, true},
{[]interface{}{`null`, "$.a"}, 1, true},
// Tests nil arguments
{[]interface{}{nil}, nil, true},
{[]interface{}{nil, "a"}, nil, true},
{[]interface{}{`{"a": 1}`, nil}, nil, true},
{[]interface{}{nil, nil}, nil, true},
// Tests with path expression
{[]interface{}{`[1,2,[1,[5,[3]]]]`, "$[2]"}, 2, true},
{[]interface{}{`[{"a":1}]`, "$"}, 1, true},
{[]interface{}{`[{"a":1,"b":2}]`, "$[0].a"}, 1, true},
{[]interface{}{`{"a":{"a":1},"b":2}`, "$"}, 2, true},
{[]interface{}{`{"a":{"a":1},"b":2}`, "$.a"}, 1, true},
{[]interface{}{`{"a":{"a":1},"b":2}`, "$.a.a"}, 1, true},
{[]interface{}{`{"a": [1, 2, {"aa": "xx"}]}`, "$.a[2].aa"}, 1, true},
// Tests without path expression
{[]interface{}{`{}`}, 0, true},
{[]interface{}{`{"a":1}`}, 1, true},
{[]interface{}{`{"a":[1]}`}, 1, true},
{[]interface{}{`{"b":2, "c":3}`}, 2, true},
{[]interface{}{`[1]`}, 1, true},
{[]interface{}{`[1,2]`}, 2, true},
{[]interface{}{`[1,2,[1,3]]`}, 3, true},
{[]interface{}{`[1,2,[1,[5,[3]]]]`}, 3, true},
{[]interface{}{`[1,2,[1,[5,{"a":[2,3]}]]]`}, 3, true},
{[]interface{}{`[{"a":1}]`}, 1, true},
{[]interface{}{`[{"a":1,"b":2}]`}, 1, true},
{[]interface{}{`[{"a":{"a":1},"b":2}]`}, 1, true},
// Tests path expression contains any asterisk
{[]interface{}{`{"a": [1, 2, {"aa": "xx"}]}`, "$.*"}, nil, false},
{[]interface{}{`{"a": [1, 2, {"aa": "xx"}]}`, "$[*]"}, nil, false},
{[]interface{}{`{"a": [1, 2, {"aa": "xx"}]}`, "$**.a"}, nil, false},
// Tests path expression does not identify a section of the target document
{[]interface{}{`{"a": [1, 2, {"aa": "xx"}]}`, "$.c"}, nil, true},
{[]interface{}{`{"a": [1, 2, {"aa": "xx"}]}`, "$.a[3]"}, nil, true},
{[]interface{}{`{"a": [1, 2, {"aa": "xx"}]}`, "$.a[2].b"}, nil, true},
}
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)
if t.expected == nil {
c.Assert(d.IsNull(), IsTrue)
} else {
c.Assert(d.GetInt64(), Equals, int64(t.expected.(int)))
}
} else {
c.Assert(err, NotNil)
}
}
}
2 changes: 2 additions & 0 deletions expression/distsql_builtin.go
Original file line number Diff line number Diff line change
Expand Up @@ -434,6 +434,8 @@ func getSignatureByPB(ctx sessionctx.Context, sigCode tipb.ScalarFuncSig, tp *ti
f = &builtinJSONContainsSig{base}
case tipb.ScalarFuncSig_LikeSig:
f = &builtinLikeSig{base}
case tipb.ScalarFuncSig_JsonLengthSig:
f = &builtinJSONLengthSig{base}

case tipb.ScalarFuncSig_InInt:
f = &builtinInIntSig{base}
Expand Down
10 changes: 10 additions & 0 deletions expression/integration_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3316,6 +3316,16 @@ func (s *testIntegrationSuite) TestFuncJSON(c *C) {
json_contains_path('{"a": 1, "b": 2, "c": {"d": 4}}', 'all', '$[*]')
`)
r.Check(testkit.Rows("1 0 1 0"))

r = tk.MustQuery(`select
json_length('1'),
json_length('{}'),
json_length('[]'),
json_length('{"a": 1}'),
json_length('{"a": 1, "b": 2}'),
json_length('[1, 2, 3]')
`)
r.Check(testkit.Rows("1 0 0 1 2 3"))
}

func (s *testIntegrationSuite) TestColumnInfoModified(c *C) {
Expand Down
5 changes: 3 additions & 2 deletions types/json/binary.go
Original file line number Diff line number Diff line change
Expand Up @@ -169,7 +169,8 @@ func (bj BinaryJSON) GetString() []byte {
return bj.Value[lenLen : lenLen+int(strLen)]
}

func (bj BinaryJSON) getElemCount() int {
// GetElemCount gets the count of Object or Array.
func (bj BinaryJSON) GetElemCount() int {
return int(endian.Uint32(bj.Value))
}

Expand All @@ -184,7 +185,7 @@ func (bj BinaryJSON) objectGetKey(i int) []byte {
}

func (bj BinaryJSON) objectGetVal(i int) BinaryJSON {
elemCount := bj.getElemCount()
elemCount := bj.GetElemCount()
return bj.valEntryGet(headerSize + elemCount*keyEntrySize + i*valEntrySize)
}

Expand Down
34 changes: 17 additions & 17 deletions types/json/binary_functions.go
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,7 @@ func (bj BinaryJSON) extractTo(buf []BinaryJSON, pathExpr PathExpression) []Bina
}
return buf
}
elemCount := bj.getElemCount()
elemCount := bj.GetElemCount()
if currentLeg.arrayIndex == arrayIndexAsterisk {
for i := 0; i < elemCount; i++ {
buf = bj.arrayGetElem(i).extractTo(buf, subPathExpr)
Expand All @@ -184,7 +184,7 @@ func (bj BinaryJSON) extractTo(buf []BinaryJSON, pathExpr PathExpression) []Bina
buf = bj.arrayGetElem(currentLeg.arrayIndex).extractTo(buf, subPathExpr)
}
} else if currentLeg.typ == pathLegKey && bj.TypeCode == TypeCodeObject {
elemCount := bj.getElemCount()
elemCount := bj.GetElemCount()
if currentLeg.dotKey == "*" {
for i := 0; i < elemCount; i++ {
buf = bj.objectGetVal(i).extractTo(buf, subPathExpr)
Expand All @@ -198,12 +198,12 @@ func (bj BinaryJSON) extractTo(buf []BinaryJSON, pathExpr PathExpression) []Bina
} else if currentLeg.typ == pathLegDoubleAsterisk {
buf = bj.extractTo(buf, subPathExpr)
if bj.TypeCode == TypeCodeArray {
elemCount := bj.getElemCount()
elemCount := bj.GetElemCount()
for i := 0; i < elemCount; i++ {
buf = bj.arrayGetElem(i).extractTo(buf, pathExpr)
}
} else if bj.TypeCode == TypeCodeObject {
elemCount := bj.getElemCount()
elemCount := bj.GetElemCount()
for i := 0; i < elemCount; i++ {
buf = bj.objectGetVal(i).extractTo(buf, pathExpr)
}
Expand All @@ -213,7 +213,7 @@ func (bj BinaryJSON) extractTo(buf []BinaryJSON, pathExpr PathExpression) []Bina
}

func (bj BinaryJSON) objectSearchKey(key []byte) (BinaryJSON, bool) {
elemCount := bj.getElemCount()
elemCount := bj.GetElemCount()
idx := sort.Search(elemCount, func(i int) bool {
return bytes.Compare(bj.objectGetKey(i), key) >= 0
})
Expand Down Expand Up @@ -371,7 +371,7 @@ func (bm *binaryModifier) doInsert(path PathExpression, newBj BinaryJSON) {
bm.modifyValue = buildBinaryArray([]BinaryJSON{parentBj, newBj})
return
}
elemCount := parentBj.getElemCount()
elemCount := parentBj.GetElemCount()
elems := make([]BinaryJSON, 0, elemCount+1)
for i := 0; i < elemCount; i++ {
elems = append(elems, parentBj.arrayGetElem(i))
Expand All @@ -384,7 +384,7 @@ func (bm *binaryModifier) doInsert(path PathExpression, newBj BinaryJSON) {
return
}
bm.modifyPtr = &parentBj.Value[0]
elemCount := parentBj.getElemCount()
elemCount := parentBj.GetElemCount()
insertKey := hack.Slice(lastLeg.dotKey)
insertIdx := sort.Search(elemCount, func(i int) bool {
return bytes.Compare(parentBj.objectGetKey(i), insertKey) >= 0
Expand Down Expand Up @@ -429,7 +429,7 @@ func (bm *binaryModifier) doRemove(path PathExpression) {
return
}
bm.modifyPtr = &parentBj.Value[0]
elemCount := parentBj.getElemCount()
elemCount := parentBj.GetElemCount()
elems := make([]BinaryJSON, 0, elemCount-1)
for i := 0; i < elemCount; i++ {
if i != lastLeg.arrayIndex {
Expand All @@ -443,7 +443,7 @@ func (bm *binaryModifier) doRemove(path PathExpression) {
return
}
bm.modifyPtr = &parentBj.Value[0]
elemCount := parentBj.getElemCount()
elemCount := parentBj.GetElemCount()
removeKey := hack.Slice(lastLeg.dotKey)
keys := make([][]byte, 0, elemCount+1)
elems := make([]BinaryJSON, 0, elemCount+1)
Expand Down Expand Up @@ -477,7 +477,7 @@ func (bm *binaryModifier) rebuildTo(buf []byte) ([]byte, TypeCode) {
return append(buf, bj.Value...), bj.TypeCode
}
docOff := len(buf)
elemCount := bj.getElemCount()
elemCount := bj.GetElemCount()
var valEntryStart int
if bj.TypeCode == TypeCodeArray {
copySize := headerSize + elemCount*valEntrySize
Expand Down Expand Up @@ -550,8 +550,8 @@ func CompareBinary(left, right BinaryJSON) int {
case TypeCodeString:
cmp = bytes.Compare(left.GetString(), right.GetString())
case TypeCodeArray:
leftCount := left.getElemCount()
rightCount := right.getElemCount()
leftCount := left.GetElemCount()
rightCount := right.GetElemCount()
for i := 0; i < leftCount && i < rightCount; i++ {
elem1 := left.arrayGetElem(i)
elem2 := right.arrayGetElem(i)
Expand Down Expand Up @@ -627,7 +627,7 @@ func mergeBinaryArray(elems []BinaryJSON) BinaryJSON {
if elem.TypeCode != TypeCodeArray {
buf = append(buf, elem)
} else {
childCount := elem.getElemCount()
childCount := elem.GetElemCount()
for j := 0; j < childCount; j++ {
buf = append(buf, elem.arrayGetElem(j))
}
Expand All @@ -640,7 +640,7 @@ func mergeBinaryObject(objects []BinaryJSON) BinaryJSON {
keyValMap := make(map[string]BinaryJSON)
keys := make([][]byte, 0, len(keyValMap))
for _, obj := range objects {
elemCount := obj.getElemCount()
elemCount := obj.GetElemCount()
for i := 0; i < elemCount; i++ {
key := obj.objectGetKey(i)
val := obj.objectGetVal(i)
Expand Down Expand Up @@ -699,7 +699,7 @@ func ContainsBinary(obj, target BinaryJSON) bool {
switch obj.TypeCode {
case TypeCodeObject:
if target.TypeCode == TypeCodeObject {
len := target.getElemCount()
len := target.GetElemCount()
for i := 0; i < len; i++ {
key := target.objectGetKey(i)
val := target.objectGetVal(i)
Expand All @@ -712,15 +712,15 @@ func ContainsBinary(obj, target BinaryJSON) bool {
return false
case TypeCodeArray:
if target.TypeCode == TypeCodeArray {
len := target.getElemCount()
len := target.GetElemCount()
for i := 0; i < len; i++ {
if !ContainsBinary(obj, target.arrayGetElem(i)) {
return false
}
}
return true
}
len := obj.getElemCount()
len := obj.GetElemCount()
for i := 0; i < len; i++ {
if ContainsBinary(obj.arrayGetElem(i), target) {
return true
Expand Down