sql: support JSONB subscripting for SELECT cases

[otan]

I've kept the ArraySubscript class names and references as PG does
this too. We can choose to change this in later iterations.

Refs #77434

Release note (sql change): Added support for JSONB subscripting in
SELECT-style cases, e.g. SELECT json_field['a'] ... WHERE
json_field['b'] = ...

[cockroach-teamcity]

This change is 

[rafiss]

just had small comments; maybe someone from sql-queries should look too

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @otan and @rafiss)

---

pkg/sql/opt/memo/typing.go line 255 at r1 (raw file):

	switch t.Family() {
	case types.JsonFamily:
		return t

i don't follow this part. do you think we should update the comment on the function?

---

pkg/sql/opt/norm/fold_constants_funcs.go line 506 at r1 (raw file):

		if resolvedType.Family() == types.ArrayFamily {
			resolvedType = resolvedType.ArrayContents()
		}

nit: should it return an error if it's not an array or json?

---

pkg/sql/opt/optbuilder/scalar.go line 176 at r1 (raw file):

		if len(t.Indirection) != 1 {
			panic(unimplementedWithIssueDetailf(32552, "ind", "multidimensional indexing is not supported"))

would we still want this error message to be shown for arrays? otherwise, it'd hit the assertion error in EvalIndirection, right?

[otan]

maybe @mgartner ?

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @rafiss)

---

pkg/sql/opt/memo/typing.go line 255 at r1 (raw file):

Previously, rafiss (Rafi Shamim) wrote…

i don't follow this part. do you think we should update the comment on the function?

comment updated and hopefully it makes more sense.

---

pkg/sql/opt/optbuilder/scalar.go line 176 at r1 (raw file):

Previously, rafiss (Rafi Shamim) wrote…

would we still want this error message to be shown for arrays? otherwise, it'd hit the assertion error in EvalIndirection, right?

i think it gets caught elsewhere, but added it just in case.

[otan]

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @rafiss)

---

pkg/sql/opt/optbuilder/scalar.go line 176 at r1 (raw file):

Previously, otan (Oliver Tan) wrote…

i think it gets caught elsewhere, but added it just in case.

actually, there is no easy way to check for this. it does have checks further down the stack though:

[email protected]:26257/defaultdb> select array['1', '2'][1][2];
ERROR: unimplemented: multidimensional indexing: ARRAY['1', '2'][1][2]
SQLSTATE: 0A000
HINT: You have attempted to use a feature that is not yet implemented.
See: https://go.crdb.dev/issue-v/32552/v22.2

[email protected]:26257/defaultdb> create table t (a string[]);
CREATE TABLE


Time: 18ms total (execution 18ms / network 0ms)

[email protected]:26257/defaultdb> insert into t values (array['1', '2']);
INSERT 0 1


Time: 12ms total (execution 12ms / network 0ms)

[email protected]:26257/defaultdb> select t[1][2] from t;
ERROR: cannot subscript type tuple{string[] AS a} because it is not an array or json object
SQLSTATE: 42804
[email protected]:26257/defaultdb> select a[1][2] from t;
ERROR: unimplemented: multidimensional indexing: a[1][2]
SQLSTATE: 0A000
HINT: You have attempted to use a feature that is not yet implemented.
See: https://go.crdb.dev/issue-v/32552/v22.2

[otan]

pkg/sql/opt/optbuilder/scalar.go line 176 at r1 (raw file):

Previously, otan (Oliver Tan) wrote…

actually, there is no easy way to check for this. it does have checks further down the stack though:

[email protected]:26257/defaultdb> select array['1', '2'][1][2];
ERROR: unimplemented: multidimensional indexing: ARRAY['1', '2'][1][2]
SQLSTATE: 0A000
HINT: You have attempted to use a feature that is not yet implemented.
See: https://go.crdb.dev/issue-v/32552/v22.2

[email protected]:26257/defaultdb> create table t (a string[]);
CREATE TABLE


Time: 18ms total (execution 18ms / network 0ms)

[email protected]:26257/defaultdb> insert into t values (array['1', '2']);
INSERT 0 1


Time: 12ms total (execution 12ms / network 0ms)

[email protected]:26257/defaultdb> select t[1][2] from t;
ERROR: cannot subscript type tuple{string[] AS a} because it is not an array or json object
SQLSTATE: 42804
[email protected]:26257/defaultdb> select a[1][2] from t;
ERROR: unimplemented: multidimensional indexing: a[1][2]
SQLSTATE: 0A000
HINT: You have attempted to use a feature that is not yet implemented.
See: https://go.crdb.dev/issue-v/32552/v22.2

uhh, nvm

[mgartner]

maybe @mgartner ?

Yes, I'd like to take a look. Apologies for the delay.

[mgartner]

Very cool! We should probably normalize these to -> expressions when possible so that they can be index accelerated. I'll create an issue to track.

Reviewed 7 of 10 files at r1, 3 of 3 files at r2, all commit messages.
Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @otan and @rafiss)

---

pkg/sql/opt/norm/fold_constants_funcs.go line 506 at r2 (raw file):

		switch input.DataType().Family() {
		case types.JsonFamily:
			resolvedType = input.DataType()

You should add some tests for folding JSON indirections in this section:

cockroach/pkg/sql/opt/norm/testdata/rules/fold_constants

Lines 1122 to 1216 in a134352

	# --------------------------------------------------
	# FoldIndirection
	# --------------------------------------------------
	# Fold when input is a static array constructor (but elements are not constant).
	norm expect=FoldIndirection
	SELECT ARRAY[i, i + 1][1] FROM a
	----
	project
	├── columns: array:9
	├── scan a
	│ └── columns: i:2
	└── projections
	└── i:2 [as=array:9, outer=(2)]
	norm expect=FoldIndirection
	SELECT ARRAY[i, i + 1][2] FROM a
	----
	project
	├── columns: array:9
	├── immutable
	├── scan a
	│ └── columns: i:2
	└── projections
	└── i:2 + 1 [as=array:9, outer=(2), immutable]
	# Fold when input is a DArray constant.
	norm expect=FoldIndirection
	SELECT ARRAY[4, 5, 6][2] FROM a
	----
	project
	├── columns: array:9!null
	├── fd: ()-->(9)
	├── scan a
	└── projections
	└── 5 [as=array:9]
	# Array bounds are out-of-range.
	norm expect=FoldIndirection
	SELECT ARRAY[s, 'foo'][0] FROM a
	----
	project
	├── columns: array:9
	├── fd: ()-->(9)
	├── scan a
	└── projections
	└── CAST(NULL AS STRING) [as=array:9]
	norm expect=FoldIndirection
	SELECT ARRAY[i, i + 1][3] FROM a
	----
	project
	├── columns: array:9
	├── fd: ()-->(9)
	├── scan a
	└── projections
	└── CAST(NULL AS INT8) [as=array:9]
	norm expect=FoldIndirection
	SELECT ARRAY[4, 5, 6][0] FROM a
	----
	project
	├── columns: array:9
	├── fd: ()-->(9)
	├── scan a
	└── projections
	└── CAST(NULL AS INT8) [as=array:9]
	# Array is dynamically constructed.
	norm expect-not=FoldIndirection
	SELECT arr[0] FROM a
	----
	project
	├── columns: arr:9
	├── scan a
	│ └── columns: a.arr:6
	└── projections
	└── a.arr:6[0] [as=arr:9, outer=(6)]
	# Regression test for #40404.
	norm expect=FoldIndirection
	SELECT (SELECT x[1]) FROM (VALUES(null::oid[])) v(x)
	----
	values
	├── columns: x:3
	├── cardinality: [1 - 1]
	├── key: ()
	├── fd: ()-->(3)
	└── tuple
	└── subquery
	└── values
	├── columns: x:2
	├── cardinality: [1 - 1]
	├── key: ()
	├── fd: ()-->(2)
	└── (NULL,)

---

pkg/sql/sem/tree/type_check.go line 636 at r2 (raw file):

		for _, t := range expr.Indirection {
			if t.Slice {
				return nil, pgerror.Newf(pgcode.DatatypeMismatch, "cannot reference a slice with JSON")

nit: is "reference" the right word in this error message? Is "index" is more appropriate?

---

pkg/sql/logictest/testdata/logic_test/json line 911 at r2 (raw file):

NULL  NULL  NULL  NULL

# Error cases.

nit: add a test case with slice indexing

[mgartner]

... We should probably normalize these to -> expressions when possible so that they can be index accelerated. I'll create an issue to track.

Done: #83441

[otan]

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @mgartner, @otan, and @rafiss)

---

pkg/sql/sem/tree/type_check.go line 636 at r2 (raw file):

Previously, mgartner (Marcus Gartner) wrote…

nit: is "reference" the right word in this error message? Is "index" is more appropriate?

replaced with jsonb subscript does not support slices a la PG

[otan]

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @mgartner and @rafiss)

---

pkg/sql/opt/norm/fold_constants_funcs.go line 506 at r2 (raw file):

Previously, mgartner (Marcus Gartner) wrote…

You should add some tests for folding JSON indirections in this section:

cockroach/pkg/sql/opt/norm/testdata/rules/fold_constants

Lines 1122 to 1216 in a134352

	# --------------------------------------------------
	# FoldIndirection
	# --------------------------------------------------
	# Fold when input is a static array constructor (but elements are not constant).
	norm expect=FoldIndirection
	SELECT ARRAY[i, i + 1][1] FROM a
	----
	project
	├── columns: array:9
	├── scan a
	│ └── columns: i:2
	└── projections
	└── i:2 [as=array:9, outer=(2)]
	norm expect=FoldIndirection
	SELECT ARRAY[i, i + 1][2] FROM a
	----
	project
	├── columns: array:9
	├── immutable
	├── scan a
	│ └── columns: i:2
	└── projections
	└── i:2 + 1 [as=array:9, outer=(2), immutable]
	# Fold when input is a DArray constant.
	norm expect=FoldIndirection
	SELECT ARRAY[4, 5, 6][2] FROM a
	----
	project
	├── columns: array:9!null
	├── fd: ()-->(9)
	├── scan a
	└── projections
	└── 5 [as=array:9]
	# Array bounds are out-of-range.
	norm expect=FoldIndirection
	SELECT ARRAY[s, 'foo'][0] FROM a
	----
	project
	├── columns: array:9
	├── fd: ()-->(9)
	├── scan a
	└── projections
	└── CAST(NULL AS STRING) [as=array:9]
	norm expect=FoldIndirection
	SELECT ARRAY[i, i + 1][3] FROM a
	----
	project
	├── columns: array:9
	├── fd: ()-->(9)
	├── scan a
	└── projections
	└── CAST(NULL AS INT8) [as=array:9]
	norm expect=FoldIndirection
	SELECT ARRAY[4, 5, 6][0] FROM a
	----
	project
	├── columns: array:9
	├── fd: ()-->(9)
	├── scan a
	└── projections
	└── CAST(NULL AS INT8) [as=array:9]
	# Array is dynamically constructed.
	norm expect-not=FoldIndirection
	SELECT arr[0] FROM a
	----
	project
	├── columns: arr:9
	├── scan a
	│ └── columns: a.arr:6
	└── projections
	└── a.arr:6[0] [as=arr:9, outer=(6)]
	# Regression test for #40404.
	norm expect=FoldIndirection
	SELECT (SELECT x[1]) FROM (VALUES(null::oid[])) v(x)
	----
	values
	├── columns: x:3
	├── cardinality: [1 - 1]
	├── key: ()
	├── fd: ()-->(3)
	└── tuple
	└── subquery
	└── values
	├── columns: x:2
	├── cardinality: [1 - 1]
	├── key: ()
	├── fd: ()-->(2)
	└── (NULL,)

Done.

[otan]

pkg/sql/opt/norm/testdata/rules/cycle line 6 at r4 (raw file):

(NormCycleTestRel (True))
----
error: optimizer factory constructor call stack exceeded max depth of 10000

(this -rewrites for me on my M1 macbook, reverted)

[mgartner]

pkg/sql/opt/norm/testdata/rules/cycle line 6 at r4 (raw file):

Previously, otan (Oliver Tan) wrote…

(this -rewrites for me on my M1 macbook, reverted)

That's odd and unsettling. Is it somehow optimizing the infinite recursion away? Is that possible?

[mgartner]

pkg/sql/opt/norm/testdata/rules/cycle line 6 at r4 (raw file):

Previously, mgartner (Marcus Gartner) wrote…

That's odd and unsettling. Is it somehow optimizing the infinite recursion away? Is that possible?

I created an issue to track: #83512

[mgartner]

Reviewed 3 of 4 files at r3, 1 of 1 files at r4, all commit messages.
Reviewable status: complete! 1 of 0 LGTMs obtained (waiting on @otan and @rafiss)

---

pkg/sql/opt/norm/fold_constants_funcs.go line 506 at r2 (raw file):

Previously, otan (Oliver Tan) wrote…

Done.

Thanks!

[otan]

bors r=mgartner

[craig]

Build failed (retrying...):

• GitHub CI (Cockroach)

[craig]

Build succeeded:

• GitHub CI (Cockroach)