opt: Normalization rule to normalize json subscripts [...] to fetch…

… value operators `->`

Previously, jsonb subscripts were not normalized to fetch value operators.
This didn't allow queries with filters like `json_col['a'] = '1'` to
scan inverted indexes, resulting in less efficient query plans.
With this rule, these types of queries can be index accelerated.
Corresponding tests have been added as well.

Resolves: cockroachdb#83441

Release note: None

Release justification:

pkg/sql/opt/norm/general_funcs.go

@@ -64,6 +64,12 @@ func (c *CustomFuncs) IsTimestampTZ(scalar opt.ScalarExpr) bool {
 	return scalar.DataType().Family() == types.TimestampTZFamily
 }
 
+// IsJSON returns true if the given scalar expression is of type
+// JsonFamily
+func (c *CustomFuncs) IsJSON(scalar opt.ScalarExpr) bool {
+	return scalar.DataType().Family() == types.JsonFamily
+}
+
 // BoolType returns the boolean SQL type.
 func (c *CustomFuncs) BoolType() *types.T {
 	return types.Bool

pkg/sql/opt/norm/rules/scalar.opt

@@ -390,3 +390,11 @@ $input
 (Not (Function $args:* $private:(FunctionPrivate "st_disjoint")))
 =>
 (MakeIntersectionFunction $args)
+
+# TransformJSONBSubscriptToFetchValue converts jsonb subscripting
+# `[...]` into a fetch value operator `->`
+
+[TransformJSONBSubscriptToFetchValue, Normalize]
+(Indirection $input:* $index:* & (IsJSON $input))
+=>
+(FetchVal $input $index)

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

@@ -1214,10 +1214,11 @@ SELECT j['field'] FROM a
 ----
 project
  ├── columns: j:9
+ ├── immutable
  ├── scan a
  │    └── columns: a.j:5
  └── projections
-      └── a.j:5['field'] [as=j:9, outer=(5)]
+      └── a.j:5->'field' [as=j:9, outer=(5), immutable]
 
 # Regression test for #40404.
 norm expect=FoldIndirection

pkg/sql/opt/norm/testdata/rules/scalar

@@ -18,6 +18,24 @@ exec-ddl
 CREATE TABLE c (c CHAR PRIMARY KEY)
 ----
 
+exec-ddl
+CREATE TABLE b
+(
+    k INT PRIMARY KEY,
+    u INT,
+    v INT,
+    j JSONB,
+    arr STRING[],
+    INDEX u(u),
+    UNIQUE INDEX v(v),
+    INVERTED INDEX j_inv_idx(j)
+)
+----
+
+exec-ddl
+CREATE TABLE d (j INT[])
+----
+
 # --------------------------------------------------
 # CommuteVar
 # --------------------------------------------------
@@ -2130,3 +2148,77 @@ project
  │    └── columns: geom1:2 geom2:3
  └── projections
       └── NOT st_intersects(geom1:2, geom2:3) [as="?column?":8, outer=(2,3), immutable]
+
+# --------------------------------------------------
+# TransformJSONBSubscriptToFetchValue
+# --------------------------------------------------
+
+norm expect=TransformJSONBSubscriptToFetchValue
+SELECT j['c'] FROM b WHERE j['a'] = '"b"'
+----
+project
+ ├── columns: j:9
+ ├── immutable
+ ├── select
+ │    ├── columns: b.j:4
+ │    ├── immutable
+ │    ├── scan b
+ │    │    └── columns: b.j:4
+ │    └── filters
+ │         └── (b.j:4->'a') = '"b"' [outer=(4), immutable]
+ └── projections
+      └── b.j:4->'c' [as=j:9, outer=(4), immutable]
+
+norm expect=TransformJSONBSubscriptToFetchValue
+SELECT j->'a'->'b', j['a']['b'] FROM b WHERE j['c'] = '1'
+----
+project
+ ├── columns: "?column?":9 j:10
+ ├── immutable
+ ├── select
+ │    ├── columns: b.j:4
+ │    ├── immutable
+ │    ├── scan b
+ │    │    └── columns: b.j:4
+ │    └── filters
+ │         └── (b.j:4->'c') = '1' [outer=(4), immutable]
+ └── projections
+      ├── (b.j:4->'a')->'b' [as="?column?":9, outer=(4), immutable]
+      └── (b.j:4->'a')->'b' [as=j:10, outer=(4), immutable]
+
+norm expect=TransformJSONBSubscriptToFetchValue
+SELECT arr[1], j['e'] FROM b
+----
+project
+ ├── columns: arr:9 j:10
+ ├── immutable
+ ├── scan b
+ │    └── columns: b.j:4 b.arr:5
+ └── projections
+      ├── b.arr:5[1] [as=arr:9, outer=(5)]
+      └── b.j:4->'e' [as=j:10, outer=(4), immutable]
+
+norm expect-not=TransformJSONBSubscriptToFetchValue
+SELECT arr[1] FROM b WHERE arr[2] = 'a'
+----
+project
+ ├── columns: arr:9
+ ├── select
+ │    ├── columns: b.arr:5
+ │    ├── scan b
+ │    │    └── columns: b.arr:5
+ │    └── filters
+ │         └── b.arr:5[2] = 'a' [outer=(5)]
+ └── projections
+      └── b.arr:5[1] [as=arr:9, outer=(5)]
+
+norm expect-not=TransformJSONBSubscriptToFetchValue
+SELECT arr[1],arr[2] FROM b;
+----
+project
+ ├── columns: arr:9 arr:10
+ ├── scan b
+ │    └── columns: b.arr:5
+ └── projections
+      ├── b.arr:5[1] [as=arr:9, outer=(5)]
+      └── b.arr:5[2] [as=arr:10, outer=(5)]

pkg/sql/opt/xform/testdata/rules/select

@@ -2791,6 +2791,19 @@ project
       │    └── spans: ["7a\x00\x01\x12b\x00\x01", "7a\x00\x01\x12b\x00\x01"]
       └── key: (1)
 
+opt expect=TransformJSONBSubscriptToFetchValue
+SELECT k FROM b WHERE j['a'] = '"b"'
+----
+project
+ ├── columns: k:1!null
+ ├── immutable
+ ├── key: (1)
+ └── scan b@j_inv_idx
+      ├── columns: k:1!null
+      ├── inverted constraint: /7/1
+      │    └── spans: ["7a\x00\x01\x12b\x00\x01", "7a\x00\x01\x12b\x00\x01"]
+      └── key: (1)
+
 # Chained fetch val operators.
 opt expect=GenerateInvertedIndexScans
 SELECT k FROM b WHERE j->'a'->'b' = '"c"'
@@ -3066,6 +3079,30 @@ project
            ├── key: (1)
            └── fd: (1)-->(7)
 
+opt expect=TransformJSONBSubscriptToFetchValue
+SELECT k FROM b WHERE j['a']['b'] @> '"c"'
+----
+project
+ ├── columns: k:1!null
+ ├── immutable
+ ├── key: (1)
+ └── inverted-filter
+      ├── columns: k:1!null
+      ├── inverted expression: /7
+      │    ├── tight: true, unique: false
+      │    └── union spans
+      │         ├── ["7a\x00\x02b\x00\x01\x12c\x00\x01", "7a\x00\x02b\x00\x01\x12c\x00\x01"]
+      │         └── ["7a\x00\x02b\x00\x02\x00\x03\x00\x01\x12c\x00\x01", "7a\x00\x02b\x00\x02\x00\x03\x00\x01\x12c\x00\x01"]
+      ├── key: (1)
+      └── scan b@j_inv_idx
+           ├── columns: k:1!null j_inverted_key:7!null
+           ├── inverted constraint: /7/1
+           │    └── spans
+           │         ├── ["7a\x00\x02b\x00\x01\x12c\x00\x01", "7a\x00\x02b\x00\x01\x12c\x00\x01"]
+           │         └── ["7a\x00\x02b\x00\x02\x00\x03\x00\x01\x12c\x00\x01", "7a\x00\x02b\x00\x02\x00\x03\x00\x01\x12c\x00\x01"]
+           ├── key: (1)
+           └── fd: (1)-->(7)
+
 opt expect=GenerateInvertedIndexScans
 SELECT k FROM b WHERE j->'a'->'b' <@ '"c"'
 ----