Skip to content

Commit f6bfa56

Browse files
normanchennnormanchenn
committed
jsonpath: add arithmetic evaluation
This commit adds functionality to evaluate arithmetic expressions within jsonpath queries. Addition, subtraction, multiplication, division and modulo operators are now supported, which use the `apd` library to do the calculations. Epic: None Release note (sql change): Add support for addition, subtraction, multiplication, division and modulo operators within jsonpath queries.
1 parent 6b825e3 commit f6bfa56

File tree

7 files changed

+225
-14
lines changed

7 files changed

+225
-14
lines changed

‎pkg/sql/logictest/testdata/logic_test/jsonb_path_query‎

Lines changed: 60 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -701,13 +701,72 @@ SELECT jsonb_path_query('{"a": [[[[[[{"b": 1}]]]]]]}', '$.a ? (@.b == 1)');
701701
query empty
702702
SELECT jsonb_path_query('{"a": [[[{"b": 1}], [{"b": 2}]]]}', '$.a ? (@.b == 1)');
703703

704+
query T
705+
SELECT jsonb_path_query('{}', '1 + 2');
706+
----
707+
3
708+
709+
query T
710+
SELECT jsonb_path_query('{}', '1 - 2');
711+
----
712+
-1
713+
714+
query T
715+
SELECT jsonb_path_query('{}', '1 * 2');
716+
----
717+
2
718+
719+
query T
720+
SELECT jsonb_path_query('{}', '1 / 2');
721+
----
722+
0.50000000000000000000
723+
724+
query T
725+
SELECT jsonb_path_query('{}', '3 % 2');
726+
----
727+
1
728+
729+
query T
730+
SELECT jsonb_path_query('{"a": 4, "b": 5}', '$.a + $.b');
731+
----
732+
9
733+
734+
query T
735+
SELECT jsonb_path_query('{"a": 4, "b": 5, "c": [9, 8, 7]}', '$.c[0] + $.c[1]');
736+
----
737+
17
738+
739+
query T
740+
SELECT jsonb_path_query('{"a": 4, "b": 5, "c": [9, 8, 7]}', '$.c[$.b - $.a]');
741+
----
742+
8
743+
744+
query T
745+
SELECT jsonb_path_query('{"a": 4, "b": 5, "c": [9, 8, 7]}', '$.c[$.b - $.a] + $var', '{"var": 10}');
746+
----
747+
18
748+
749+
statement error pgcode 22012 pq: division by zero
750+
SELECT jsonb_path_query('{}', '1 / 0');
751+
752+
statement error pgcode 22038 pq: left operand of jsonpath operator / is not a single numeric value
753+
SELECT jsonb_path_query('[1, 2]', '$[*] / 2');
754+
755+
statement error pgcode 22038 pq: right operand of jsonpath operator \+ is not a single numeric value
756+
SELECT jsonb_path_query('[1, 2]', '2 + $[*]');
757+
758+
statement error pgcode 22038 pq: left operand of jsonpath operator / is not a single numeric value
759+
SELECT jsonb_path_query('{"a": "hello"}', '$.a / 2');
760+
761+
statement error pgcode 22038 pq: right operand of jsonpath operator \+ is not a single numeric value
762+
SELECT jsonb_path_query('{"a": null}', '2 + $.a');
763+
704764
# when string literals are supported
705765
# query T rowsort
706766
# SELECT jsonb_path_query('{"data": [{"val": "a", "num": 1}, {"val": "b", "num": 2}, {"val": "a", "num": 3}]}'::jsonb, '$.data ? (@.val == "a")'::jsonpath);
707767
# ----
708768
# {"num": 1, "val": "a"}
709769
# {"num": 3, "val": "a"}
710-
711770
# select jsonb_path_query('[1, 2, 3, 4, 5]', '$[-1]');
712771
# select jsonb_path_query('[1, 2, 3, 4, 5]', 'strict $[-1]');
713772

‎pkg/util/jsonpath/eval/BUILD.bazel‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ go_library(
2020
"//pkg/util/json",
2121
"//pkg/util/jsonpath",
2222
"//pkg/util/jsonpath/parser",
23+
"@com_github_cockroachdb_apd_v3//:apd",
2324
"@com_github_cockroachdb_errors//:errors",
2425
],
2526
)

‎pkg/util/jsonpath/eval/eval.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@ func (ctx *jsonpathCtx) eval(
9090
if err != nil {
9191
return nil, err
9292
}
93-
return convertFromBool(res), nil
93+
return []json.JSON{res}, nil
9494
case jsonpath.Filter:
9595
return ctx.evalFilter(path, jsonValue, unwrap)
9696
default:

‎pkg/util/jsonpath/eval/operation.go‎

Lines changed: 67 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,10 @@
66
package eval
77

88
import (
9+
"github.com/cockroachdb/apd/v3"
10+
"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
11+
"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
12+
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
913
"github.com/cockroachdb/cockroach/pkg/util/json"
1014
"github.com/cockroachdb/cockroach/pkg/util/jsonpath"
1115
"github.com/cockroachdb/errors"
@@ -28,14 +32,14 @@ func isBool(j json.JSON) bool {
2832
}
2933
}
3034

31-
func convertFromBool(b jsonpathBool) []json.JSON {
35+
func convertFromBool(b jsonpathBool) json.JSON {
3236
switch b {
3337
case jsonpathBoolTrue:
34-
return []json.JSON{json.TrueJSONValue}
38+
return json.TrueJSONValue
3539
case jsonpathBoolFalse:
36-
return []json.JSON{json.FalseJSONValue}
40+
return json.FalseJSONValue
3741
case jsonpathBoolUnknown:
38-
return []json.JSON{json.NullJSONValue}
42+
return json.NullJSONValue
3943
default:
4044
panic(errors.AssertionFailedf("unhandled jsonpath boolean type"))
4145
}
@@ -54,22 +58,25 @@ func convertToBool(j json.JSON) jsonpathBool {
5458

5559
func (ctx *jsonpathCtx) evalOperation(
5660
op jsonpath.Operation, jsonValue json.JSON,
57-
) (jsonpathBool, error) {
61+
) (json.JSON, error) {
5862
switch op.Type {
5963
case jsonpath.OpLogicalAnd, jsonpath.OpLogicalOr, jsonpath.OpLogicalNot:
6064
res, err := ctx.evalLogical(op, jsonValue)
6165
if err != nil {
62-
return jsonpathBoolUnknown, err
66+
return convertFromBool(jsonpathBoolUnknown), err
6367
}
64-
return res, nil
68+
return convertFromBool(res), nil
6569
case jsonpath.OpCompEqual, jsonpath.OpCompNotEqual,
6670
jsonpath.OpCompLess, jsonpath.OpCompLessEqual,
6771
jsonpath.OpCompGreater, jsonpath.OpCompGreaterEqual:
6872
res, err := ctx.evalComparison(op, jsonValue, true /* unwrapRight */)
6973
if err != nil {
70-
return jsonpathBoolUnknown, err
74+
return convertFromBool(jsonpathBoolUnknown), err
7175
}
72-
return res, nil
76+
return convertFromBool(res), nil
77+
case jsonpath.OpAdd, jsonpath.OpSub, jsonpath.OpMult,
78+
jsonpath.OpDiv, jsonpath.OpMod:
79+
return ctx.evalArithmetic(op, jsonValue)
7380
default:
7481
panic(errors.AssertionFailedf("unhandled operation type"))
7582
}
@@ -234,3 +241,54 @@ func execComparison(l, r json.JSON, op jsonpath.OperationType) (jsonpathBool, er
234241
}
235242
return jsonpathBoolFalse, nil
236243
}
244+
245+
func (ctx *jsonpathCtx) evalArithmetic(
246+
op jsonpath.Operation, jsonValue json.JSON,
247+
) (json.JSON, error) {
248+
left, err := ctx.evalAndUnwrapResult(op.Left, jsonValue, true /* unwrap */)
249+
if err != nil {
250+
return nil, err
251+
}
252+
right, err := ctx.evalAndUnwrapResult(op.Right, jsonValue, true /* unwrap */)
253+
if err != nil {
254+
return nil, err
255+
}
256+
257+
if len(left) != 1 || left[0].Type() != json.NumberJSONType {
258+
return nil, pgerror.Newf(pgcode.SingletonSQLJSONItemRequired,
259+
"left operand of jsonpath operator %s is not a single numeric value",
260+
jsonpath.OperationTypeStrings[op.Type])
261+
}
262+
if len(right) != 1 || right[0].Type() != json.NumberJSONType {
263+
return nil, pgerror.Newf(pgcode.SingletonSQLJSONItemRequired,
264+
"right operand of jsonpath operator %s is not a single numeric value",
265+
jsonpath.OperationTypeStrings[op.Type])
266+
}
267+
268+
leftNum, _ := left[0].AsDecimal()
269+
rightNum, _ := right[0].AsDecimal()
270+
var res apd.Decimal
271+
var cond apd.Condition
272+
switch op.Type {
273+
case jsonpath.OpAdd:
274+
_, err = tree.DecimalCtx.Add(&res, leftNum, rightNum)
275+
case jsonpath.OpSub:
276+
_, err = tree.DecimalCtx.Sub(&res, leftNum, rightNum)
277+
case jsonpath.OpMult:
278+
_, err = tree.DecimalCtx.Mul(&res, leftNum, rightNum)
279+
case jsonpath.OpDiv:
280+
cond, err = tree.DecimalCtx.Quo(&res, leftNum, rightNum)
281+
// Division by zero or 0 / 0.
282+
if cond.DivisionByZero() || cond.DivisionUndefined() {
283+
return nil, tree.ErrDivByZero
284+
}
285+
case jsonpath.OpMod:
286+
_, err = tree.DecimalCtx.Rem(&res, leftNum, rightNum)
287+
default:
288+
panic(errors.AssertionFailedf("unhandled jsonpath arithmetic type"))
289+
}
290+
if err != nil {
291+
return nil, err
292+
}
293+
return json.FromDecimal(res), nil
294+
}

‎pkg/util/jsonpath/operation.go‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,14 @@ const (
1919
OpLogicalAnd
2020
OpLogicalOr
2121
OpLogicalNot
22+
OpAdd
23+
OpSub
24+
OpMult
25+
OpDiv
26+
OpMod
2227
)
2328

24-
var operationTypeStrings = map[OperationType]string{
29+
var OperationTypeStrings = map[OperationType]string{
2530
OpCompEqual: "==",
2631
OpCompNotEqual: "!=",
2732
OpCompLess: "<",
@@ -31,6 +36,11 @@ var operationTypeStrings = map[OperationType]string{
3136
OpLogicalAnd: "&&",
3237
OpLogicalOr: "||",
3338
OpLogicalNot: "!",
39+
OpAdd: "+",
40+
OpSub: "-",
41+
OpMult: "*",
42+
OpDiv: "/",
43+
OpMod: "%",
3444
}
3545

3646
type Operation struct {
@@ -46,7 +56,7 @@ func (o Operation) String() string {
4656
// 1 == 1 && 1 != 1, postgres will output (1 == 1 && 1 != 1), but we output
4757
// ((1 == 1) && (1 != 1)).
4858
if o.Type == OpLogicalNot {
49-
return fmt.Sprintf("%s(%s)", operationTypeStrings[o.Type], o.Left)
59+
return fmt.Sprintf("%s(%s)", OperationTypeStrings[o.Type], o.Left)
5060
}
51-
return fmt.Sprintf("(%s %s %s)", o.Left, operationTypeStrings[o.Type], o.Right)
61+
return fmt.Sprintf("(%s %s %s)", o.Left, OperationTypeStrings[o.Type], o.Right)
5262
}

‎pkg/util/jsonpath/parser/jsonpath.y‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,9 @@ func unaryOp(op jsonpath.OperationType, left jsonpath.Path) jsonpath.Operation {
194194
%left AND
195195
%right NOT
196196

197+
%left '+' '-'
198+
%left '*' '/' '%'
199+
197200
%%
198201

199202
jsonpath:
@@ -235,6 +238,31 @@ expr:
235238
{
236239
$$.val = jsonpath.Paths($1.pathArr())
237240
}
241+
| '(' expr ')'
242+
{
243+
$$.val = $2.path()
244+
}
245+
| expr '+' expr
246+
{
247+
$$.val = binaryOp(jsonpath.OpAdd, $1.path(), $3.path())
248+
}
249+
| expr '-' expr
250+
{
251+
$$.val = binaryOp(jsonpath.OpSub, $1.path(), $3.path())
252+
}
253+
| expr '*' expr
254+
{
255+
$$.val = binaryOp(jsonpath.OpMult, $1.path(), $3.path())
256+
}
257+
| expr '/' expr
258+
{
259+
$$.val = binaryOp(jsonpath.OpDiv, $1.path(), $3.path())
260+
}
261+
| expr '%' expr
262+
{
263+
$$.val = binaryOp(jsonpath.OpMod, $1.path(), $3.path())
264+
}
265+
// TODO(normanchenn): add unary + and -.
238266
;
239267

240268
accessor_expr:

‎pkg/util/jsonpath/parser/testdata/jsonpath‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -363,6 +363,61 @@ $.a[*] ? (@.b > 100 || (@.c < 100))
363363
----
364364
$."a"[*]?(((@."b" > 100) || (@."c" < 100))) -- normalized!
365365

366+
parse
367+
1 + 1
368+
----
369+
(1 + 1) -- normalized!
370+
371+
parse
372+
1 + 1 * 2
373+
----
374+
(1 + (1 * 2)) -- normalized!
375+
376+
parse
377+
1 + 2 - 3 * 4 / 5 % 6
378+
----
379+
((1 + 2) - (((3 * 4) / 5) % 6)) -- normalized!
380+
381+
parse
382+
(1 + 2) * (3 - 4) / 5
383+
----
384+
(((1 + 2) * (3 - 4)) / 5) -- normalized!
385+
386+
parse
387+
1 * 2 + 3 * 4
388+
----
389+
((1 * 2) + (3 * 4)) -- normalized!
390+
391+
parse
392+
1 + 2 * (3 - 4) / (5 + 6) - 7 % 8
393+
----
394+
((1 + ((2 * (3 - 4)) / (5 + 6))) - (7 % 8)) -- normalized!
395+
396+
parse
397+
1 * (2 + 3) - 4 / (5 - 6) % 7
398+
----
399+
((1 * (2 + 3)) - ((4 / (5 - 6)) % 7)) -- normalized!
400+
401+
parse
402+
((1 + 2) * 3) - (4 % 5) * 6
403+
----
404+
(((1 + 2) * 3) - ((4 % 5) * 6)) -- normalized!
405+
406+
parse
407+
1 + 2 - 3 + 4 - 5
408+
----
409+
((((1 + 2) - 3) + 4) - 5) -- normalized!
410+
411+
parse
412+
$.c[$.b - $.a]
413+
----
414+
$."c"[($."b" - $."a")] -- normalized!
415+
416+
parse
417+
$.c[$.b - $.a to $.d - $.b]
418+
----
419+
$."c"[($."b" - $."a") to ($."d" - $."b")] -- normalized!
420+
366421
# postgres allows floats as array indexes
367422
# parse
368423
# $.abc[1.0]

0 commit comments

Comments
 (0)