sql: support in-memory JSON arrays

This commit adds support for SQL arrays of JSON in memory. This allows
functions like `array_agg` to work against JSON objects, which is useful
for compatibility.

Notably, this commit does not add support for storing JSON arrays in
table columns.

Release note (sql change): add support for SQL arrays containing JSON
for in-memory processing. This does not add support for storing SQL
arrays of JSON in tables.

docs/generated/sql/aggregates.md

@@ -33,6 +33,8 @@
 </span></td></tr>
 <tr><td><a name="array_agg"></a><code>array_agg(arg1: geometry) &rarr; geometry[]</code></td><td><span class="funcdesc"><p>Aggregates the selected values into an array.</p>
 </span></td></tr>
+<tr><td><a name="array_agg"></a><code>array_agg(arg1: jsonb) &rarr; jsonb[]</code></td><td><span class="funcdesc"><p>Aggregates the selected values into an array.</p>
+</span></td></tr>
 <tr><td><a name="array_agg"></a><code>array_agg(arg1: oid) &rarr; oid[]</code></td><td><span class="funcdesc"><p>Aggregates the selected values into an array.</p>
 </span></td></tr>
 <tr><td><a name="array_agg"></a><code>array_agg(arg1: timetz) &rarr; timetz[]</code></td><td><span class="funcdesc"><p>Aggregates the selected values into an array.</p>

docs/generated/sql/functions.md

@@ -35,6 +35,8 @@
 </span></td></tr>
 <tr><td><a name="array_append"></a><code>array_append(array: geometry[], elem: geometry) &rarr; geometry[]</code></td><td><span class="funcdesc"><p>Appends <code>elem</code> to <code>array</code>, returning the result.</p>
 </span></td></tr>
+<tr><td><a name="array_append"></a><code>array_append(array: jsonb[], elem: jsonb) &rarr; jsonb[]</code></td><td><span class="funcdesc"><p>Appends <code>elem</code> to <code>array</code>, returning the result.</p>
+</span></td></tr>
 <tr><td><a name="array_append"></a><code>array_append(array: oid[], elem: oid) &rarr; oid[]</code></td><td><span class="funcdesc"><p>Appends <code>elem</code> to <code>array</code>, returning the result.</p>
 </span></td></tr>
 <tr><td><a name="array_append"></a><code>array_append(array: timetz[], elem: timetz) &rarr; timetz[]</code></td><td><span class="funcdesc"><p>Appends <code>elem</code> to <code>array</code>, returning the result.</p>
@@ -73,6 +75,8 @@
 </span></td></tr>
 <tr><td><a name="array_cat"></a><code>array_cat(left: geometry[], right: geometry[]) &rarr; geometry[]</code></td><td><span class="funcdesc"><p>Appends two arrays.</p>
 </span></td></tr>
+<tr><td><a name="array_cat"></a><code>array_cat(left: jsonb[], right: jsonb[]) &rarr; jsonb[]</code></td><td><span class="funcdesc"><p>Appends two arrays.</p>
+</span></td></tr>
 <tr><td><a name="array_cat"></a><code>array_cat(left: oid[], right: oid[]) &rarr; oid[]</code></td><td><span class="funcdesc"><p>Appends two arrays.</p>
 </span></td></tr>
 <tr><td><a name="array_cat"></a><code>array_cat(left: timetz[], right: timetz[]) &rarr; timetz[]</code></td><td><span class="funcdesc"><p>Appends two arrays.</p>
@@ -115,6 +119,8 @@
 </span></td></tr>
 <tr><td><a name="array_position"></a><code>array_position(array: geometry[], elem: geometry) &rarr; <a href="int.html">int</a></code></td><td><span class="funcdesc"><p>Return the index of the first occurrence of <code>elem</code> in <code>array</code>.</p>
 </span></td></tr>
+<tr><td><a name="array_position"></a><code>array_position(array: jsonb[], elem: jsonb) &rarr; <a href="int.html">int</a></code></td><td><span class="funcdesc"><p>Return the index of the first occurrence of <code>elem</code> in <code>array</code>.</p>
+</span></td></tr>
 <tr><td><a name="array_position"></a><code>array_position(array: oid[], elem: oid) &rarr; <a href="int.html">int</a></code></td><td><span class="funcdesc"><p>Return the index of the first occurrence of <code>elem</code> in <code>array</code>.</p>
 </span></td></tr>
 <tr><td><a name="array_position"></a><code>array_position(array: timetz[], elem: timetz) &rarr; <a href="int.html">int</a></code></td><td><span class="funcdesc"><p>Return the index of the first occurrence of <code>elem</code> in <code>array</code>.</p>
@@ -153,6 +159,8 @@
 </span></td></tr>
 <tr><td><a name="array_positions"></a><code>array_positions(array: geometry[], elem: geometry) &rarr; <a href="int.html">int</a>[]</code></td><td><span class="funcdesc"><p>Returns and array of indexes of all occurrences of <code>elem</code> in <code>array</code>.</p>
 </span></td></tr>
+<tr><td><a name="array_positions"></a><code>array_positions(array: jsonb[], elem: jsonb) &rarr; <a href="int.html">int</a>[]</code></td><td><span class="funcdesc"><p>Returns and array of indexes of all occurrences of <code>elem</code> in <code>array</code>.</p>
+</span></td></tr>
 <tr><td><a name="array_positions"></a><code>array_positions(array: oid[], elem: oid) &rarr; <a href="int.html">int</a>[]</code></td><td><span class="funcdesc"><p>Returns and array of indexes of all occurrences of <code>elem</code> in <code>array</code>.</p>
 </span></td></tr>
 <tr><td><a name="array_positions"></a><code>array_positions(array: timetz[], elem: timetz) &rarr; <a href="int.html">int</a>[]</code></td><td><span class="funcdesc"><p>Returns and array of indexes of all occurrences of <code>elem</code> in <code>array</code>.</p>
@@ -191,6 +199,8 @@
 </span></td></tr>
 <tr><td><a name="array_prepend"></a><code>array_prepend(elem: geometry, array: geometry[]) &rarr; geometry[]</code></td><td><span class="funcdesc"><p>Prepends <code>elem</code> to <code>array</code>, returning the result.</p>
 </span></td></tr>
+<tr><td><a name="array_prepend"></a><code>array_prepend(elem: jsonb, array: jsonb[]) &rarr; jsonb[]</code></td><td><span class="funcdesc"><p>Prepends <code>elem</code> to <code>array</code>, returning the result.</p>
+</span></td></tr>
 <tr><td><a name="array_prepend"></a><code>array_prepend(elem: oid, array: oid[]) &rarr; oid[]</code></td><td><span class="funcdesc"><p>Prepends <code>elem</code> to <code>array</code>, returning the result.</p>
 </span></td></tr>
 <tr><td><a name="array_prepend"></a><code>array_prepend(elem: timetz, array: timetz[]) &rarr; timetz[]</code></td><td><span class="funcdesc"><p>Prepends <code>elem</code> to <code>array</code>, returning the result.</p>
@@ -229,6 +239,8 @@
 </span></td></tr>
 <tr><td><a name="array_remove"></a><code>array_remove(array: geometry[], elem: geometry) &rarr; geometry[]</code></td><td><span class="funcdesc"><p>Remove from <code>array</code> all elements equal to <code>elem</code>.</p>
 </span></td></tr>
+<tr><td><a name="array_remove"></a><code>array_remove(array: jsonb[], elem: jsonb) &rarr; jsonb[]</code></td><td><span class="funcdesc"><p>Remove from <code>array</code> all elements equal to <code>elem</code>.</p>
+</span></td></tr>
 <tr><td><a name="array_remove"></a><code>array_remove(array: oid[], elem: oid) &rarr; oid[]</code></td><td><span class="funcdesc"><p>Remove from <code>array</code> all elements equal to <code>elem</code>.</p>
 </span></td></tr>
 <tr><td><a name="array_remove"></a><code>array_remove(array: timetz[], elem: timetz) &rarr; timetz[]</code></td><td><span class="funcdesc"><p>Remove from <code>array</code> all elements equal to <code>elem</code>.</p>
@@ -267,6 +279,8 @@
 </span></td></tr>
 <tr><td><a name="array_replace"></a><code>array_replace(array: geometry[], toreplace: geometry, replacewith: geometry) &rarr; geometry[]</code></td><td><span class="funcdesc"><p>Replace all occurrences of <code>toreplace</code> in <code>array</code> with <code>replacewith</code>.</p>
 </span></td></tr>
+<tr><td><a name="array_replace"></a><code>array_replace(array: jsonb[], toreplace: jsonb, replacewith: jsonb) &rarr; jsonb[]</code></td><td><span class="funcdesc"><p>Replace all occurrences of <code>toreplace</code> in <code>array</code> with <code>replacewith</code>.</p>
+</span></td></tr>
 <tr><td><a name="array_replace"></a><code>array_replace(array: oid[], toreplace: oid, replacewith: oid) &rarr; oid[]</code></td><td><span class="funcdesc"><p>Replace all occurrences of <code>toreplace</code> in <code>array</code> with <code>replacewith</code>.</p>
 </span></td></tr>
 <tr><td><a name="array_replace"></a><code>array_replace(array: timetz[], toreplace: timetz, replacewith: timetz) &rarr; timetz[]</code></td><td><span class="funcdesc"><p>Replace all occurrences of <code>toreplace</code> in <code>array</code> with <code>replacewith</code>.</p>

pkg/sql/catalog/colinfo/BUILD.bazel

@@ -23,6 +23,7 @@ go_library(
         "//pkg/sql/sem/tree",
         "//pkg/sql/types",
         "//pkg/util/encoding",
+        "//pkg/util/errorutil/unimplemented",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_lib_pq//oid",
         "@org_golang_x_text//language",

pkg/sql/catalog/colinfo/col_type_info.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/cockroachdb/cockroach/pkg/util/errorutil/unimplemented"
 	"github.com/cockroachdb/errors"
 	"github.com/lib/pq/oid"
 	"golang.org/x/text/language"
@@ -92,6 +93,11 @@ func ValidateColumnDefType(t *types.T) error {
 			// Nested arrays are not supported as a column type.
 			return errors.Errorf("nested array unsupported as column type: %s", t.String())
 		}
+		if t.ArrayContents().Family() == types.JsonFamily {
+			// JSON arrays are not supported as a column type.
+			return unimplemented.NewWithIssueDetailf(23468, t.String(),
+				"arrays of JSON unsupported as column type")
+		}
 		if err := types.CheckArrayElementType(t.ArrayContents()); err != nil {
 			return err
 		}

pkg/sql/logictest/testdata/logic_test/json

@@ -84,20 +84,17 @@ SELECT NULL::JSON
 ----
 NULL
 
-statement error arrays of jsonb not allowed.*\nHINT:.*\n.*23468
-SELECT ARRAY['"hello"'::JSON]
-
-statement error arrays of jsonb not allowed.*\nHINT:.*\n.*23468
-SELECT '{}'::JSONB[]
-
-statement error arrays of jsonb not allowed.*\nHINT:.*\n.*23468
+statement error arrays of JSON unsupported as column type.*\nHINT:.*\n.*23468
 CREATE TABLE x (y JSONB[])
 
 statement ok
-CREATE TABLE foo (bar JSON)
+CREATE TABLE foo (pk INT DEFAULT unique_rowid(), bar JSON)
+
+statement error arrays of JSON unsupported as column type.*\nHINT:.*\n.*23468
+CREATE VIEW x AS SELECT array_agg(bar) FROM foo
 
 statement ok
-INSERT INTO foo VALUES
+INSERT INTO foo(bar) VALUES
   ('{"a": "b"}'),
   ('[1, 2, 3]'),
   ('"hello"'),
@@ -201,19 +198,24 @@ NULL
 NULL
 NULL
 
-query T
+query IT
 SELECT * from foo where bar->'x' = '[1]'
 ----
 
-query T
+query IT
 SELECT * from foo where bar->'x' = '{}'
 ----
 
+query T
+SELECT array_agg(bar ORDER BY pk) FROM foo
+----
+{"'{\"a\": \"b\"}'","'[1, 2, 3]'","\"hello\"",1.000,true,false,NULL,"'{\"x\": [1, 2, 3]}'","'{\"x\": {\"y\": \"z\"}}'"}
+
 statement ok
 DELETE FROM foo
 
 statement ok
-INSERT INTO foo VALUES ('{"a": {"c": "d"}}');
+INSERT INTO foo(bar) VALUES ('{"a": {"c": "d"}}');
 
 query TT
 SELECT bar->'a'->'c', bar->'a'->>'c' FROM foo
@@ -840,3 +842,13 @@ SELECT j - s FROM t57165
 ----
 {}
 {}
+
+query T
+SELECT ARRAY['"hello"'::JSON]
+----
+{"\"hello\""}
+
+query T
+SELECT '{}'::JSONB[]
+----
+{}

pkg/sql/opt/optbuilder/testdata/aggregate

@@ -225,6 +225,7 @@ array_agg(uuid) -> uuid[]
 array_agg(inet) -> inet[]
 array_agg(time) -> time[]
 array_agg(timetz) -> timetz[]
+array_agg(jsonb) -> jsonb[]
 array_agg(varbit) -> varbit[]
 array_agg(bool) -> bool[]
 

pkg/sql/opt/optbuilder/testdata/scalar

@@ -896,12 +896,14 @@ concat [type=int[]]
 build-scalar
 ARRAY['"foo"'::jsonb]
 ----
-error: unimplemented: arrays of jsonb not allowed
+array: [type=jsonb[]]
+ └── const: '"foo"' [type=jsonb]
 
 build-scalar
 ARRAY['"foo"'::json]
 ----
-error: unimplemented: arrays of jsonb not allowed
+array: [type=jsonb[]]
+ └── const: '"foo"' [type=jsonb]
 
 opt
 SELECT -((-9223372036854775808):::int)
@@ -1030,7 +1032,15 @@ project
 build
 SELECT ARRAY(VALUES ('{}'::JSONB))
 ----
-error (0A000): unimplemented: arrays of jsonb not allowed
+project
+ ├── columns: array:2
+ ├── values
+ │    └── ()
+ └── projections
+      └── array-flatten [as=array:2]
+           └── values
+                ├── columns: column1:1!null
+                └── ('{}',)
 
 build
 SELECT ARRAY(SELECT 1, 2)

pkg/sql/rowenc/column_type_encoding.go

@@ -1356,6 +1356,8 @@ func DatumTypeToArrayElementEncodingType(t *types.T) (encoding.Type, error) {
 		return encoding.UUID, nil
 	case types.INetFamily:
 		return encoding.IPAddr, nil
+	case types.JsonFamily:
+		return encoding.JSON, nil
 	default:
 		return 0, errors.AssertionFailedf("no known encoding type for %s", t)
 	}
@@ -1427,6 +1429,12 @@ func encodeArrayElement(b []byte, d tree.Datum) ([]byte, error) {
 		return encodeArrayElement(b, t.Wrapped)
 	case *tree.DEnum:
 		return encoding.EncodeUntaggedBytesValue(b, t.PhysicalRep), nil
+	case *tree.DJSON:
+		encoded, err := json.EncodeJSON(nil, t.JSON)
+		if err != nil {
+			return nil, err
+		}
+		return encoding.EncodeUntaggedBytesValue(b, encoded), nil
 	default:
 		return nil, errors.Errorf("don't know how to encode %s (%T)", d, d)
 	}

pkg/sql/types/types.go

@@ -2452,8 +2452,6 @@ func IsStringType(t *T) bool {
 // the issue number should be included in the error report to inform the user.
 func IsValidArrayElementType(t *T) (valid bool, issueNum int) {
 	switch t.Family() {
-	case JsonFamily:
-		return false, 23468
 	default:
 		return true, 0
 	}