jsonpath: add jsonpath conditional evaluation

normanchenn · normanchenn · commit 1157d89db49d · 2025-03-19T10:02:58.000-04:00
This commit adds support for evaluating jsonpath conditionals.
Conditionals have the structure of `$ ? (predicate)`, and will filter
the current json objects based on whether they satisfy the predicate
expression. Additionally, this commit introduces `@`, which represents
the current json object so filters can be used to reference the object
being evaluated.

Epic: None
Release note (sql change): Add jsonpath filters, which take on the form
`$ ? (predicate)`, allowing results to be filtered.
diff --git a/pkg/sql/logictest/testdata/logic_test/jsonb_path_query b/pkg/sql/logictest/testdata/logic_test/jsonb_path_query
@@ -108,7 +108,6 @@ SELECT jsonb_path_query(data, 'strict $.aa.aaa.aaaa') FROM a
 query empty
 SELECT jsonb_path_query('{}', '$.a')
 
-
 statement ok
 CREATE TABLE b (j JSONB)
 
@@ -196,7 +195,6 @@ SELECT jsonb_path_query('[1, 2, 3, 4, 5]', '$[1 to 3, 2, 1 to 3]');
 query empty
 SELECT jsonb_path_query('[1, 2, 3, 4, 5]', '$[3 to 1]');
 
-
 query T rowsort
 SELECT jsonb_path_query('[1, 2, 3, 4, 5]', '$[4 to 4]');
 ----
@@ -255,7 +253,6 @@ SELECT jsonb_path_query('{"a": [1, 2, 3, 4, 5]}', 'strict $[3 to 1]');
 query empty
 SELECT jsonb_path_query('{"a": [1, 2, 3]}', '$.a.b');
 
-
 statement error pgcode 2203A jsonpath member accessor can only be applied to an object
 SELECT jsonb_path_query('{"a": [1, 2, 3]}', 'strict $.a.b');
 
@@ -613,6 +610,95 @@ SELECT jsonb_path_query('{"a": 5, "b": 10}', '(1.5 > 1.2 && (!($.a == 1) || $.b
 ----
 true
 
+query T rowsort
+SELECT jsonb_path_query('{"a": [1,2,3]}', '$.a ? (1 == 1)');
+----
+1
+2
+3
+
+query empty
+SELECT jsonb_path_query('{"a": [1,2,3]}', '$.a ? (1 != 1)');
+
+query T
+SELECT jsonb_path_query('{"a": [1,2,3]}', 'strict $.a ? (1 == 1)');
+----
+[1, 2, 3]
+
+query empty
+SELECT jsonb_path_query('{"a": [1,2,3]}', 'strict $.a ? (1 != 1)');
+
+query T rowsort
+SELECT jsonb_path_query('{"a": [{"b": 1, "c": "hello"}, {"b": 2, "c": "world"}, {"b": 1, "c": "!"}]}', '$.a[*] ? (@.b == 1)');
+----
+{"b": 1, "c": "hello"}
+{"b": 1, "c": "!"}
+
+query empty
+SELECT jsonb_path_query('{"a": [{"b": 1, "c": "hello"}, {"b": 2, "c": "world"}, {"b": 1, "c": "!"}]}', 'strict $.a ? (@.b == 1)');
+
+query T rowsort
+SELECT jsonb_path_query('{"a": [{"b": 1, "c": "hello"}, {"b": 2, "c": "world"}, {"b": 1, "c": "!"}]}', '$.a ? (@.b == 1)');
+----
+{"b": 1, "c": "hello"}
+{"b": 1, "c": "!"}
+
+query T rowsort
+SELECT jsonb_path_query('{"a": [[{"b": 1, "c": "hello"}, {"b": 2, "c": "world"}, {"b": 1, "c": "!"}], [{"b": 1, "c": "hello"}, {"b": 2, "c": "world"}, {"b": 1, "c": "!"}]]}', '$.a ? (@.b == 1)');
+----
+[{"b": 1, "c": "hello"}, {"b": 2, "c": "world"}, {"b": 1, "c": "!"}]
+[{"b": 1, "c": "hello"}, {"b": 2, "c": "world"}, {"b": 1, "c": "!"}]
+
+query T rowsort
+SELECT jsonb_path_query('{"a": [[{"b": 1, "c": "hello"}, {"b": 2, "c": "world"}, {"b": 1, "c": "!"}], [{"b": 1, "c": "hello"}, {"b": 2, "c": "world"}, {"b": 1, "c": "!"}]]}', '$.a[*] ? (@.b == 1)');
+----
+{"b": 1, "c": "hello"}
+{"b": 1, "c": "!"}
+{"b": 1, "c": "hello"}
+{"b": 1, "c": "!"}
+
+query empty
+SELECT jsonb_path_query('{"a": [[{"b": 1, "c": "hello"}, {"b": 2, "c": "world"}, {"b": 1, "c": "!"}], [{"b": 1, "c": "hello"}, {"b": 2, "c": "world"}, {"b": 1, "c": "!"}]]}', 'strict $.a ? (@.b == 1)');
+
+query empty
+SELECT jsonb_path_query('{"a": [[{"b": 1, "c": "hello"}, {"b": 2, "c": "world"}, {"b": 1, "c": "!"}], [{"b": 1, "c": "hello"}, {"b": 2, "c": "world"}, {"b": 1, "c": "!"}]]}', 'strict $.a[*] ? (@.b == 1)');
+
+query T rowsort
+SELECT jsonb_path_query('{"a": [1,2,3,4,5]}', '$.a ? (@ > 3)');
+----
+4
+5
+
+query T rowsort
+SELECT jsonb_path_query('{"a": [{"b": 1, "c": 10}, {"b": 2, "c": 20}, {"b": 3, "c": 30}]}', '$.a ? (@.c > 15)');
+----
+{"b": 2, "c": 20}
+{"b": 3, "c": 30}
+
+query T rowsort
+SELECT jsonb_path_query('{"a": [{"b": "x", "c": true}, {"b": "y", "c": false}, {"b": "z", "c": true}]}', '$.a ? (@.c == true)');
+----
+{"b": "x", "c": true}
+{"b": "z", "c": true}
+
+query T
+SELECT jsonb_path_query('{"c": {"a": 1, "b":1}}', '$.c ? ($.c.a == @.b)');
+----
+{"a": 1, "b": 1}
+
+query empty
+SELECT jsonb_path_query('{"a": [1,2,3]}', '$.a ? (@ > 10)');
+
+query empty
+SELECT jsonb_path_query('{"a": [{"b": 1, "c": 10}, {"b": 2, "c": 20}]}', '$.a ? (@.c > 100)');
+
+# when string literals are supported
+# query T rowsort
+# SELECT jsonb_path_query('{"data": [{"val": "a", "num": 1}, {"val": "b", "num": 2}, {"val": "a", "num": 3}]}'::jsonb, '$.data ? (@.val == "a")'::jsonpath);
+# ----
+# {"num": 1, "val": "a"}
+# {"num": 3, "val": "a"}
+
 # select jsonb_path_query('[1, 2, 3, 4, 5]', '$[-1]');
 # select jsonb_path_query('[1, 2, 3, 4, 5]', 'strict $[-1]');
 
diff --git a/pkg/sql/scanner/jsonpath_scan.go b/pkg/sql/scanner/jsonpath_scan.go
@@ -95,6 +95,9 @@ func (s *JSONPathScanner) Scan(lval ScanSymType) {
 			return
 		}
 		return
+	case '@':
+		lval.SetID(lexbase.CURRENT)
+		return
 	default:
 		if sqllexbase.IsDigit(ch) {
 			s.scanNumber(lval, ch)
diff --git a/pkg/util/jsonpath/eval/BUILD.bazel b/pkg/util/jsonpath/eval/BUILD.bazel
@@ -5,6 +5,7 @@ go_library(
     srcs = [
         "array.go",
         "eval.go",
+        "filter.go",
         "key.go",
         "operation.go",
         "scalar.go",
diff --git a/pkg/util/jsonpath/eval/eval.go b/pkg/util/jsonpath/eval/eval.go
@@ -71,6 +71,8 @@ func (ctx *jsonpathCtx) eval(jp jsonpath.Path, current []json.JSON) ([]json.JSON
 		return current, nil
 	case jsonpath.Root:
 		return []json.JSON{ctx.root}, nil
+	case jsonpath.Current:
+		return current, nil
 	case jsonpath.Key:
 		return ctx.evalKey(p, current)
 	case jsonpath.Wildcard:
@@ -85,11 +87,21 @@ func (ctx *jsonpathCtx) eval(jp jsonpath.Path, current []json.JSON) ([]json.JSON
 		return []json.JSON{resolved}, nil
 	case jsonpath.Operation:
 		return ctx.evalOperation(p, current)
+	case jsonpath.Filter:
+		return ctx.evalFilter(p, current)
 	default:
 		return nil, errUnimplemented
 	}
 }
 
+func (ctx *jsonpathCtx) unwrap(input json.JSON) []json.JSON {
+	if !ctx.strict && input.Type() == json.ArrayJSONType {
+		array, _ := input.AsArray()
+		return array
+	}
+	return []json.JSON{input}
+}
+
 func (ctx *jsonpathCtx) evalAndUnwrap(path jsonpath.Path, inputs []json.JSON) ([]json.JSON, error) {
 	results, err := ctx.eval(path, inputs)
 	if err != nil {
@@ -100,12 +112,12 @@ func (ctx *jsonpathCtx) evalAndUnwrap(path jsonpath.Path, inputs []json.JSON) ([
 	}
 	var unwrapped []json.JSON
 	for _, result := range results {
-		if result.Type() == json.ArrayJSONType {
-			array, _ := result.AsArray()
-			unwrapped = append(unwrapped, array...)
-		} else {
-			unwrapped = append(unwrapped, result)
-		}
+		unwrapped = append(unwrapped, ctx.unwrap(result)...)
 	}
 	return unwrapped, nil
 }
+
+func (ctx *jsonpathCtx) unwrapAndEval(path jsonpath.Path, input json.JSON) ([]json.JSON, error) {
+	unwrapped := ctx.unwrap(input)
+	return ctx.eval(path, unwrapped)
+}
diff --git a/pkg/util/jsonpath/eval/filter.go b/pkg/util/jsonpath/eval/filter.go
@@ -0,0 +1,40 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package eval
+
+import (
+	"github.com/cockroachdb/cockroach/pkg/util/json"
+	"github.com/cockroachdb/cockroach/pkg/util/jsonpath"
+	"github.com/cockroachdb/errors"
+)
+
+func (ctx *jsonpathCtx) evalFilter(p jsonpath.Filter, current []json.JSON) ([]json.JSON, error) {
+	// TODO(normanchenn): clean up.
+	var unwrapped []json.JSON
+	for _, j := range current {
+		unwrapped = append(unwrapped, ctx.unwrap(j)...)
+	}
+	current = unwrapped
+	var filtered []json.JSON
+	// unwrap before
+	for _, j := range current {
+		results, err := ctx.eval(p.Condition, []json.JSON{j})
+		if err != nil {
+			// Postgres doesn't error when there's a structural error within
+			// filter condition, and will return nothing instead.
+			return []json.JSON{}, nil //nolint:returnerrcheck
+		}
+		if len(results) != 1 || !isBool(results[0]) {
+			return nil, errors.New("filter condition must evaluate to a boolean")
+		}
+
+		condition, _ := results[0].AsBool()
+		if condition {
+			filtered = append(filtered, j)
+		}
+	}
+	return filtered, nil
+}
diff --git a/pkg/util/jsonpath/eval/operation.go b/pkg/util/jsonpath/eval/operation.go
@@ -53,19 +53,19 @@ func convertToBool(j json.JSON) jsonpathBool {
 }
 
 func (ctx *jsonpathCtx) evalOperation(
-	p jsonpath.Operation, current []json.JSON,
+	op jsonpath.Operation, current []json.JSON,
 ) ([]json.JSON, error) {
-	switch p.Type {
+	switch op.Type {
 	case jsonpath.OpLogicalAnd, jsonpath.OpLogicalOr, jsonpath.OpLogicalNot:
-		res, err := ctx.evalLogical(p, current)
+		res, err := ctx.evalLogical(op, current)
 		if err != nil {
 			return nil, err
 		}
 		return convertFromBool(res), nil
 	case jsonpath.OpCompEqual, jsonpath.OpCompNotEqual,
 		jsonpath.OpCompLess, jsonpath.OpCompLessEqual,
 		jsonpath.OpCompGreater, jsonpath.OpCompGreaterEqual:
-		res, err := ctx.evalComparison(p, current)
+		res, err := ctx.evalComparison(op, current)
 		if err != nil {
 			return nil, err
 		}
diff --git a/pkg/util/jsonpath/expr.go b/pkg/util/jsonpath/expr.go
@@ -92,3 +92,19 @@ func (a ArrayList) String() string {
 	sb.WriteString("]")
 	return sb.String()
 }
+
+type Filter struct {
+	Condition Path
+}
+
+var _ Path = Filter{}
+
+func (f Filter) String() string {
+	return fmt.Sprintf("?(%s)", f.Condition)
+}
+
+type Current struct{}
+
+var _ Path = Current{}
+
+func (c Current) String() string { return "@" }
diff --git a/pkg/util/jsonpath/parser/jsonpath.y b/pkg/util/jsonpath/parser/jsonpath.y
@@ -169,6 +169,8 @@ func unaryOp(op jsonpath.OperationType, left jsonpath.Path) jsonpath.Operation {
 %token <str> OR
 %token <str> NOT
 
+%token <str> CURRENT
+
 %type <jsonpath.Jsonpath> jsonpath
 %type <jsonpath.Path> expr_or_predicate
 %type <jsonpath.Path> expr
@@ -251,6 +253,10 @@ path_primary:
   {
     $$.val = jsonpath.Root{}
   }
+| CURRENT
+  {
+    $$.val = jsonpath.Current{}
+  }
 | scalar_value
   {
     $$.val = $1.path()
@@ -267,6 +273,10 @@ accessor_op:
   {
     $$.val = $1.path()
   }
+| '?' '(' predicate ')'
+  {
+    $$.val = jsonpath.Filter{Condition: $3.path()}
+  }
 ;
 
 key:
diff --git a/pkg/util/jsonpath/parser/testdata/jsonpath b/pkg/util/jsonpath/parser/testdata/jsonpath
@@ -342,6 +342,27 @@ parse
 ----
 ((1 == 1) || (1 != 1)) -- normalized!
 
+parse
+$.abc ? ($.a[1] > 2)
+----
+$."abc"?(($."a"[1] > 2)) -- normalized!
+
+# TODO(normanchenn): this should be not allowed
+parse
+@
+----
+@
+
+parse
+$.a[*] ? (@.b > 100)
+----
+$."a"[*]?((@."b" > 100)) -- normalized!
+
+parse
+$.a[*] ? (@.b > 100 || (@.c < 100))
+----
+$."a"[*]?(((@."b" > 100) || (@."c" < 100))) -- normalized!
+
 # postgres allows floats as array indexes
 # parse
 # $.abc[1.0]

Original file line number	Diff line number	Diff line change
`@@ -169,6 +169,8 @@ func unaryOp(op jsonpath.OperationType, left jsonpath.Path) jsonpath.Operation {`
`169`	`169`	`%token <str> OR`
`170`	`170`	`%token <str> NOT`
`171`	`171`
	`172`	`+%token <str> CURRENT`
	`173`	`+`
`172`	`174`	`%type <jsonpath.Jsonpath> jsonpath`
`173`	`175`	`%type <jsonpath.Path> expr_or_predicate`
`174`	`176`	`%type <jsonpath.Path> expr`
`@@ -251,6 +253,10 @@ path_primary:`
`251`	`253`	`{`
`252`	`254`	`$$.val = jsonpath.Root{}`
`253`	`255`	`}`
	`256`	`+\| CURRENT`
	`257`	`+ {`
	`258`	`+ $$.val = jsonpath.Current{}`
	`259`	`+ }`
`254`	`260`	`\| scalar_value`
`255`	`261`	`{`
`256`	`262`	`$$.val = $1.path()`
`@@ -267,6 +273,10 @@ accessor_op:`
`267`	`273`	`{`
`268`	`274`	`$$.val = $1.path()`
`269`	`275`	`}`
	`276`	`+\| '?' '(' predicate ')'`
	`277`	`+ {`
	`278`	`+ $$.val = jsonpath.Filter{Condition: $3.path()}`
	`279`	`+ }`
`270`	`280`	`;`
`271`	`281`
`272`	`282`	`key:`