opt: 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: #83441

Release note (performance improvement): The optimizer will now plan
inverted index scans for queries with JSON subscripting filters, like
json_col['field'] = '"value".

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
+// JSON.
+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,10 @@ $input
 (Not (Function $args:* $private:(FunctionPrivate "st_disjoint")))
 =>
 (MakeIntersectionFunction $args)
+
+# ConvertJSONSubscriptToFetchValue converts json subscripting
+# `[...]` into a fetch value operator `->`.
+[ConvertJSONSubscriptToFetchValue, 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,17 @@ exec-ddl
 CREATE TABLE c (c CHAR PRIMARY KEY)
 ----
 
+exec-ddl
+CREATE TABLE b
+(
+    k INT PRIMARY KEY,
+    j JSONB,
+    i INT,
+    s STRING,
+    arr STRING[]
+)
+----
+
 # --------------------------------------------------
 # CommuteVar
 # --------------------------------------------------
@@ -2130,3 +2141,92 @@ project
  │    └── columns: geom1:2 geom2:3
  └── projections
       └── NOT st_intersects(geom1:2, geom2:3) [as="?column?":8, outer=(2,3), immutable]
+
+# --------------------------------------------------
+# ConvertJSONSubscriptToFetchValue
+# --------------------------------------------------
+
+norm expect=ConvertJSONSubscriptToFetchValue
+SELECT j['c'] FROM b WHERE j['a'] = '"b"'
+----
+project
+ ├── columns: j:8
+ ├── immutable
+ ├── select
+ │    ├── columns: b.j:2
+ │    ├── immutable
+ │    ├── scan b
+ │    │    └── columns: b.j:2
+ │    └── filters
+ │         └── (b.j:2->'a') = '"b"' [outer=(2), immutable]
+ └── projections
+      └── b.j:2->'c' [as=j:8, outer=(2), immutable]
+
+norm expect=ConvertJSONSubscriptToFetchValue
+SELECT j['a']['b'] FROM b WHERE j['c'] = '1'
+----
+project
+ ├── columns: j:8
+ ├── immutable
+ ├── select
+ │    ├── columns: b.j:2
+ │    ├── immutable
+ │    ├── scan b
+ │    │    └── columns: b.j:2
+ │    └── filters
+ │         └── (b.j:2->'c') = '1' [outer=(2), immutable]
+ └── projections
+      └── (b.j:2->'a')->'b' [as=j:8, outer=(2), immutable]
+
+norm expect=ConvertJSONSubscriptToFetchValue
+SELECT j[0] FROM b WHERE j[0][1] = '1'
+----
+project
+ ├── columns: j:8
+ ├── immutable
+ ├── select
+ │    ├── columns: b.j:2
+ │    ├── immutable
+ │    ├── scan b
+ │    │    └── columns: b.j:2
+ │    └── filters
+ │         └── ((b.j:2->0)->1) = '1' [outer=(2), immutable]
+ └── projections
+      └── b.j:2->0 [as=j:8, outer=(2), immutable]
+
+norm expect=ConvertJSONSubscriptToFetchValue
+SELECT j[i], j[s] FROM b
+----
+project
+ ├── columns: j:8 j:9
+ ├── immutable
+ ├── scan b
+ │    └── columns: b.j:2 i:3 s:4
+ └── projections
+      ├── b.j:2->i:3 [as=j:8, outer=(2,3), immutable]
+      └── b.j:2->s:4 [as=j:9, outer=(2,4), immutable]
+
+norm expect-not=ConvertJSONSubscriptToFetchValue
+SELECT arr[1] FROM b WHERE arr[2] = 'a'
+----
+project
+ ├── columns: arr:8
+ ├── 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:8, outer=(5)]
+
+norm expect-not=ConvertJSONSubscriptToFetchValue
+SELECT arr[1], arr[2] FROM b;
+----
+project
+ ├── columns: arr:8 arr:9
+ ├── scan b
+ │    └── columns: b.arr:5
+ └── projections
+      ├── b.arr:5[1] [as=arr:8, outer=(5)]
+      └── b.arr:5[2] [as=arr:9, 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=ConvertJSONSubscriptToFetchValue
+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=ConvertJSONSubscriptToFetchValue
+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"'
 ----