jsonpath: add support for non-integer array indices

This commit adds support for non-integer array indices, which adds
Postgres compatibility in the JSONPath feature set. Non-integer array
indices are rounded towards zero before being used as array indices
(floor for positive numbers, ceil for negative numbers).

Release note (sql change): Add support for non-integer array indices in
JSONPath queries (ex. `SELECT jsonb_path_query('[1, 2, 3]', '$[2.5]');`).
Indices are rounded towards 0.

Author: normanchenn

pkg/sql/logictest/testdata/logic_test/jsonb_path_query

@@ -264,22 +264,25 @@ SELECT jsonb_path_query('[1, 2, 3, 4, 5]', '$[1]');
 ----
 2
 
-# TODO(normanchenn): this will be changed after floats are supported.
-statement error pq: jsonpath array subscript is not a single numeric value
-SELECT jsonb_path_query('{"a": [10, 9, 8, 7]}', '$.a[0.99]')
+query T
+SELECT jsonb_path_query('{"a": [10, 9, 8, 7]}', '$.a[0.99]');
+----
+10
 
-# TODO(normanchenn): this will be changed after floats are supported.
-statement error pq: jsonpath array subscript is not a single numeric value
-SELECT jsonb_path_query('{"a": [10, 9, 8, 7]}', '$.a[1.01]')
+query T
+SELECT jsonb_path_query('{"a": [10, 9, 8, 7]}', '$.a[1.01]');
+----
+9
 
 query T
 SELECT jsonb_path_query('{"a": [10, 9, 8, 7]}', '$.a[$varInt]', '{"varInt": 1, "varFloat": 2.3, "varString": "a", "varBool": true, "varNull": null}')
 ----
 9
 
-# TODO(normanchenn): this will be changed after floats are supported.
-statement error pgcode 22033 jsonpath array subscript is not a single numeric value
-SELECT jsonb_path_query('{"a": [10, 9, 8, 7]}', '$.a[$varFloat]', '{"varInt": 1, "varFloat": 2.3, "varString": "a", "varBool": true, "varNull": null}')
+query T
+SELECT jsonb_path_query('{"a": [10, 9, 8, 7]}', '$.a[$varFloat]', '{"varInt": 1, "varFloat": 2.3, "varString": "a", "varBool": true, "varNull": null}');
+----
+8
 
 statement error pgcode 22033 jsonpath array subscript is not a single numeric value
 SELECT jsonb_path_query('{"a": [10, 9, 8, 7]}', '$.a[$varString]', '{"varInt": 1, "varFloat": 2.3, "varString": "a", "varBool": true, "varNull": null}')
@@ -1552,3 +1555,48 @@ query T
 SELECT jsonb_path_query('[1,2,3]', '$[last ? (@.type() == "number")]');
 ----
 3
+
+query T
+SELECT jsonb_path_query('[1, 2, 3]', '$[1.1]');
+----
+2
+
+query T
+SELECT jsonb_path_query('[1, 2, 3]', '$[1.5]');
+----
+2
+
+query T
+SELECT jsonb_path_query('[1, 2, 3]', '$[1.99999999]');
+----
+2
+
+query T
+SELECT jsonb_path_query('[1, 2, 3]', '$[2.1]');
+----
+3
+
+query T
+SELECT jsonb_path_query('[1, 2, 3]', '$[2.5]');
+----
+3
+
+query T
+SELECT jsonb_path_query('[1, 2, 3]', '$[2.99999999]');
+----
+3
+
+query T
+SELECT jsonb_path_query('[1, 2, 3]', '$[-0.1]');
+----
+1
+
+query T
+SELECT jsonb_path_query('[1, 2, 3]', '$[-0.5]');
+----
+1
+
+query T
+SELECT jsonb_path_query('[1, 2, 3]', '$[-0.99999999]');
+----
+1

pkg/util/jsonpath/eval/array.go

@@ -8,8 +8,10 @@ package eval
 import (
 	"math"
 
+	"github.com/cockroachdb/apd/v3"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/util/json"
 	"github.com/cockroachdb/cockroach/pkg/util/jsonpath"
 	"github.com/cockroachdb/errors"
@@ -125,11 +127,33 @@ func asInt(j json.JSON) (int, error) {
 	if !ok {
 		return 0, errIndexNotSingleNumValue
 	}
-	i64, err := d.Int64()
+
+	// First, try direct conversion to int64.
+	if i64, err := d.Int64(); err == nil {
+		return validateInt32Range(i64)
+	}
+
+	// If direct conversion fails, truncate the non-integer towards zero.
+	var err error
+	dec := &apd.Decimal{}
+	if d.Sign() == 1 {
+		_, err = tree.ExactCtx.Floor(dec, d)
+	} else {
+		_, err = tree.ExactCtx.Ceil(dec, d)
+	}
+	if err != nil {
+		return 0, errIndexNotSingleNumValue
+	}
+
+	i64, err := dec.Int64()
 	if err != nil {
 		return 0, errIndexNotSingleNumValue
 	}
-	// Postgres returns an error if the index is outside int32 range.
+	return validateInt32Range(i64)
+}
+
+// validateInt32Range checks if the int64 value is within int32 range.
+func validateInt32Range(i64 int64) (int, error) {
 	if i64 < math.MinInt32 || i64 > math.MaxInt32 {
 		return 0, errInvalidSubscript
 	}

pkg/util/jsonpath/parser/testdata/jsonpath

@@ -599,38 +599,44 @@ $.a.type()
 ----
 $."a".type() -- normalized!
 
-# postgres allows floats as array indexes
-# parse
-# $.abc[1.0]
-# ----
+# Postgres outputs $."abc"[1.0]. This might be a limitation of the apd library.
+parse
+$.abc[1.0]
+----
+$."abc"[1] -- normalized!
 
-# parse
-# $.abc[1.1]
-# ----
+parse
+$.abc[1.1]
+----
+$."abc"[1.1] -- normalized!
 
 # Postgres outputs $."abc"[1000000000].
 parse
 $.abc[1e9]
 ----
 $."abc"[1E+9] -- normalized!
 
-# parse
-# $.abc[0.0]
-# ----
+# Postgres outputs $."abc"[0.0].
+parse
+$.abc[0.0]
+----
+$."abc"[0] -- normalized!
 
 # TODO(normanchenn): Related to the TODOs in pkg/util/jsonpath/operation.go.
 parse
 $.abc[-0]
 ----
 $."abc"[(-0)] -- normalized!
 
-# parse
-# $.abc[-1.99]
-# ----
+parse
+$.abc[-1.99]
+----
+$."abc"[(-1.99)] -- normalized!
 
-# parse
-# $[1.999999999999999]
-# ----
+parse
+$[1.999999999999999]
+----
+$[1.999999999999999]
 
 parse
 $[null]