sql: support JSONB encoding and decoding for forward indexes

pkg/sql/logictest/testdata/logic_test/hash_join

@@ -221,5 +221,5 @@ SELECT * FROM t44797_2 WHERE EXISTS (SELECT * FROM t44797_2 AS l, t44797_3 AS r
 statement ok
 CREATE TABLE table57696(col_table TIME NOT NULL)
 
-statement error unable to encode JSON as a table key\nHINT:.*\n.*35706.*
+statement ok
 WITH cte (col_cte) AS ( SELECT * FROM ( VALUES ( ( 'false':::JSONB, '1970-01-05 16:57:40.000665+00:00':::TIMESTAMPTZ ) ) ) EXCEPT ALL SELECT * FROM ( VALUES ( ( ' [ [[true], [], {}, "b", {}], {"a": []}, {"c": 2.05750813403415} ] ':::JSONB, '1970-01-10 05:23:26.000428+00:00':::TIMESTAMPTZ ) ) ) ) SELECT * FROM cte, table57696

pkg/sql/rowenc/keyside/BUILD.bazel

@@ -8,6 +8,7 @@ go_library(
         "decode.go",
         "doc.go",
         "encode.go",
+        "json.go",
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/sql/rowenc/keyside",
     visibility = ["//visibility:public"],
@@ -19,8 +20,8 @@ go_library(
         "//pkg/util/bitarray",
         "//pkg/util/duration",
         "//pkg/util/encoding",
-        "//pkg/util/errorutil/unimplemented",
         "//pkg/util/ipaddr",
+        "//pkg/util/json",
         "//pkg/util/timetz",
         "//pkg/util/timeutil/pgdate",
         "//pkg/util/uuid",

pkg/sql/rowenc/keyside/decode.go

@@ -22,6 +22,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/duration"
 	"github.com/cockroachdb/cockroach/pkg/util/encoding"
 	"github.com/cockroachdb/cockroach/pkg/util/ipaddr"
+	"github.com/cockroachdb/cockroach/pkg/util/json"
 	"github.com/cockroachdb/cockroach/pkg/util/timetz"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil/pgdate"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
@@ -117,13 +118,13 @@ func Decode(
 		d, err := a.NewDCollatedString(r, valType.Locale())
 		return d, rkey, err
 	case types.JsonFamily:
-		// Don't attempt to decode the JSON value. Instead, just return the
-		// remaining bytes of the key.
-		jsonLen, err := encoding.PeekLength(key)
+		var json json.JSON
+		json, rkey, err = decodeJSONKey(key, dir)
 		if err != nil {
 			return nil, nil, err
 		}
-		return tree.DNull, key[jsonLen:], nil
+		d := a.NewDJSON(tree.DJSON{JSON: json})
+		return d, rkey, err
 	case types.BytesFamily:
 		var r []byte
 		if dir == encoding.Ascending {

pkg/sql/rowenc/keyside/encode.go

@@ -13,7 +13,6 @@ package keyside
 import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/util/encoding"
-	"github.com/cockroachdb/cockroach/pkg/util/errorutil/unimplemented"
 	"github.com/cockroachdb/errors"
 )
 
@@ -174,7 +173,7 @@ func Encode(b []byte, val tree.Datum, dir encoding.Direction) ([]byte, error) {
 		// DEncodedKey carries an already encoded key.
 		return append(b, []byte(*t)...), nil
 	case *tree.DJSON:
-		return nil, unimplemented.NewWithIssue(35706, "unable to encode JSON as a table key")
+		return encodeJSONKey(b, t, dir)
 	}
 	return nil, errors.Errorf("unable to encode table key: %T", val)
 }

pkg/sql/rowenc/keyside/json.go

@@ -0,0 +1,191 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package keyside
+
+import (
+	"github.com/cockroachdb/apd/v3"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/util/encoding"
+	"github.com/cockroachdb/cockroach/pkg/util/json"
+	"github.com/cockroachdb/errors"
+)
+
+// encodeJSONKey is responsible for encoding the different JSON
+// values.
+func encodeJSONKey(b []byte, json *tree.DJSON, dir encoding.Direction) ([]byte, error) {
+	return json.JSON.EncodeForwardIndex(b, dir)
+}
+
+// decodeJSONKey is responsible for decoding the different JSON
+// values.
+func decodeJSONKey(buf []byte, dir encoding.Direction) (json.JSON, []byte, error) {
+	var err error
+	var typ encoding.Type
+	var jsonVal json.JSON
+
+	buf, typ, err = encoding.ValidateAndConsumeJSONKeyMarker(buf, dir)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	switch typ {
+	case encoding.JSONNull, encoding.JSONNullDesc:
+		jsonVal, err = json.MakeJSON(nil)
+		if err != nil {
+			return nil, nil, errors.NewAssertionErrorWithWrappedErrf(err, "could not decode JSON Null")
+		}
+	case encoding.JSONFalse, encoding.JSONFalseDesc:
+		jsonVal, err = json.MakeJSON(false)
+		if err != nil {
+			return nil, nil, errors.NewAssertionErrorWithWrappedErrf(err, "could not decode JSON False")
+		}
+	case encoding.JSONTrue, encoding.JSONTrueDesc:
+		jsonVal, err = json.MakeJSON(true)
+		if err != nil {
+			return nil, nil, errors.NewAssertionErrorWithWrappedErrf(err, "could not decode JSON True")
+		}
+	case encoding.JSONString, encoding.JSONStringDesc:
+		jsonVal, buf, err = decodeJSONString(buf, dir)
+		if err != nil {
+			return nil, nil, errors.NewAssertionErrorWithWrappedErrf(err, "could not decode JSON String")
+		}
+	case encoding.JSONNumber:
+		var dec apd.Decimal
+		buf, dec, err = encoding.DecodeDecimalAscending(buf, nil)
+		if err != nil {
+			return nil, nil, errors.NewAssertionErrorWithWrappedErrf(err, "could not decode JSON Number")
+		}
+		if len(buf) == 0 || !encoding.IsJSONKeyDone(buf, dir) {
+			return nil, nil, errors.New("cannot find JSON terminator")
+		}
+		buf = buf[1:] // removing the terminator
+		jsonVal = json.FromDecimal(dec)
+	case encoding.JSONNumberDesc:
+		var dec apd.Decimal
+		buf, dec, err = encoding.DecodeDecimalDescending(buf, nil)
+		if err != nil {
+			return nil, nil, errors.NewAssertionErrorWithWrappedErrf(err, "could not decode JSON Number")
+		}
+		if len(buf) == 0 || !encoding.IsJSONKeyDone(buf, dir) {
+			return nil, nil, errors.New("cannot find JSON terminator")
+		}
+		buf = buf[1:] // removing the terminator
+		jsonVal = json.FromDecimal(dec)
+	case encoding.JSONArray, encoding.JSONArrayDesc:
+		jsonVal, buf, err = decodeJSONArray(buf, dir)
+		if err != nil {
+			return nil, nil, errors.NewAssertionErrorWithWrappedErrf(err, "could not decode JSON Array")
+		}
+	case encoding.JSONObject, encoding.JSONObjectDesc:
+		jsonVal, buf, err = decodeJSONObject(buf, dir)
+		if err != nil {
+			return nil, nil, errors.NewAssertionErrorWithWrappedErrf(err, "could not decode JSON Object")
+		}
+	}
+
+	return jsonVal, buf, nil
+}
+
+func decodeJSONString(buf []byte, dir encoding.Direction) (json.JSON, []byte, error) {
+	var err error
+	var str string
+
+	switch dir {
+	case encoding.Ascending:
+		buf, str, err = encoding.DecodeUnsafeStringAscendingDeepCopy(buf, nil)
+	case encoding.Descending:
+		buf, str, err = encoding.DecodeUnsafeStringDescending(buf, nil)
+	}
+	if err != nil {
+		return nil, nil, errors.NewAssertionErrorWithWrappedErrf(err, "could not decode"+
+			"the JSON String")
+	}
+	if len(buf) == 0 || !encoding.IsJSONKeyDone(buf, dir) {
+		return nil, nil, errors.New("cannot find JSON terminator")
+	}
+	buf = buf[1:] // removing the terminator
+	jsonVal, err := json.MakeJSON(str)
+	if err != nil {
+		return nil, nil, errors.NewAssertionErrorWithWrappedErrf(err,
+			"could not make a JSON String from the input string")
+	}
+	return jsonVal, buf, nil
+}
+
+func decodeJSONArray(buf []byte, dir encoding.Direction) (json.JSON, []byte, error) {
+	// Extracting the total number of elements in the json array.
+	var err error
+	buf, length, err := encoding.DecodeJSONValueLength(buf, dir)
+	if err != nil {
+		return nil, nil, errors.AssertionFailedf("could not decode the number of elements in the JSON Array")
+	}
+	// Pre-allocate the array builder with `length` number
+	// of JSON elements.
+	jsonArray := json.NewArrayBuilder(int(length))
+
+	var childElem json.JSON
+	for {
+		if len(buf) == 0 {
+			return nil, nil, errors.AssertionFailedf("invalid JSON array encoding (unterminated)")
+		}
+		if encoding.IsJSONKeyDone(buf, dir) {
+			buf = buf[1:]
+			return jsonArray.Build(), buf, nil
+		}
+		childElem, buf, err = decodeJSONKey(buf, dir)
+		if err != nil {
+			return nil, buf, err
+		}
+		jsonArray.Add(childElem)
+	}
+}
+
+func decodeJSONObject(buf []byte, dir encoding.Direction) (json.JSON, []byte, error) {
+	// Extracting the total number of elements in the json object.
+	var err error
+	buf, length, err := encoding.DecodeJSONValueLength(buf, dir)
+	if err != nil {
+		return nil, nil, errors.AssertionFailedf("could not decode the number of elements in the JSON Object")
+	}
+
+	jsonObject := json.NewObjectBuilder(int(length))
+	var jsonKey, value json.JSON
+	for {
+		if len(buf) == 0 {
+			return nil, nil, errors.AssertionFailedf("invalid JSON Object encoding (unterminated)")
+		}
+		if encoding.IsJSONKeyDone(buf, dir) {
+			// JSONB Objects will have a terminator byte.
+			buf = buf[1:]
+			return jsonObject.Build(), buf, nil
+		}
+
+		// Assumption: The byte array given to us can be decoded into a
+		// valid JSON Object. In other words, for each JSON key there
+		// should be a valid JSON value.
+		jsonKey, buf, err = decodeJSONKey(buf, dir)
+		if err != nil {
+			return nil, buf, err
+		}
+
+		key, err := jsonKey.AsText()
+		if err != nil {
+			return nil, buf, err
+		}
+
+		value, buf, err = decodeJSONKey(buf, dir)
+		if err != nil {
+			return nil, buf, err
+		}
+
+		jsonObject.Add(*key, value)
+	}
+}

pkg/sql/rowenc/keyside/keyside_test.go

@@ -239,7 +239,7 @@ func genEncodingDirection() gopter.Gen {
 func hasKeyEncoding(typ *types.T) bool {
 	// Only some types are round-trip key encodable.
 	switch typ.Family() {
-	case types.JsonFamily, types.CollatedStringFamily, types.TupleFamily, types.DecimalFamily,
+	case types.CollatedStringFamily, types.TupleFamily, types.DecimalFamily,
 		types.GeographyFamily, types.GeometryFamily, types.TSVectorFamily, types.TSQueryFamily:
 		return false
 	case types.ArrayFamily:

pkg/sql/sem/builtins/datums_to_bytes_builtin_test.go

@@ -151,7 +151,7 @@ func TestCrdbInternalDatumsToBytesIllegalType(t *testing.T) {
 	defer s.Stopper().Stop(ctx)
 	tdb := sqlutils.MakeSQLRunner(sqlDB)
 	for _, val := range []string{
-		"'{\"a\": 1}'::JSONB",
+		"'foo:1,2 bar:3'::tsvector",
 	} {
 		t.Run(val, func(t *testing.T) {
 			tdb.ExpectErr(t, ".*illegal argument.*",