From 4e2707a2fb8428cb52c78018206bea49cc9ea6f2 Mon Sep 17 00:00:00 2001 From: Faizaan Madhani Date: Wed, 14 Sep 2022 12:07:18 -0400 Subject: [PATCH] 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: #83441 Release note: None Release justification: --- pkg/sql/opt/norm/general_funcs.go | 6 ++ pkg/sql/opt/norm/rules/scalar.opt | 8 ++ .../opt/norm/testdata/rules/fold_constants | 3 +- pkg/sql/opt/norm/testdata/rules/scalar | 92 +++++++++++++++++++ pkg/sql/opt/xform/testdata/rules/select | 37 ++++++++ 5 files changed, 145 insertions(+), 1 deletion(-) diff --git a/pkg/sql/opt/norm/general_funcs.go b/pkg/sql/opt/norm/general_funcs.go index 4ed4e4bcf22c..8127e82ca3ed 100644 --- a/pkg/sql/opt/norm/general_funcs.go +++ b/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 diff --git a/pkg/sql/opt/norm/rules/scalar.opt b/pkg/sql/opt/norm/rules/scalar.opt index 5d829093f294..1e1f34dde250 100644 --- a/pkg/sql/opt/norm/rules/scalar.opt +++ b/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) diff --git a/pkg/sql/opt/norm/testdata/rules/fold_constants b/pkg/sql/opt/norm/testdata/rules/fold_constants index be1b49840066..bbe5ac8a0f79 100644 --- a/pkg/sql/opt/norm/testdata/rules/fold_constants +++ b/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 diff --git a/pkg/sql/opt/norm/testdata/rules/scalar b/pkg/sql/opt/norm/testdata/rules/scalar index ca8d5dd7f9f4..73535c4572ee 100644 --- a/pkg/sql/opt/norm/testdata/rules/scalar +++ b/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)] diff --git a/pkg/sql/opt/xform/testdata/rules/select b/pkg/sql/opt/xform/testdata/rules/select index c40c116ee199..159401eb84e9 100644 --- a/pkg/sql/opt/xform/testdata/rules/select +++ b/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"' ----