diff --git a/pkg/sql/logictest/testdata/logic_test/inverted_index b/pkg/sql/logictest/testdata/logic_test/inverted_index index a17ad13723fc..6a499bb2cd24 100644 --- a/pkg/sql/logictest/testdata/logic_test/inverted_index +++ b/pkg/sql/logictest/testdata/logic_test/inverted_index @@ -753,6 +753,22 @@ SELECT j FROM f@i WHERE j->'a' = '1' ORDER BY k {"a": 1, "b": 2} {"a": 1, "c": 3} +query T +SELECT j FROM f@i WHERE j->'a' = '1' OR j->'b' = '2' ORDER BY k +---- +{"a": 1} +{"b": 2} +{"a": 1, "b": 2} +{"a": 1, "c": 3} + +query T +SELECT j FROM f@i WHERE j->'a' = '1' OR j @> '{"b": 2}' ORDER BY k +---- +{"a": 1} +{"b": 2} +{"a": 1, "b": 2} +{"a": 1, "c": 3} + subtest arrays statement ok diff --git a/pkg/sql/opt/exec/execbuilder/testdata/inverted_index b/pkg/sql/opt/exec/execbuilder/testdata/inverted_index index e5880d076742..27db8fc893c3 100644 --- a/pkg/sql/opt/exec/execbuilder/testdata/inverted_index +++ b/pkg/sql/opt/exec/execbuilder/testdata/inverted_index @@ -350,7 +350,7 @@ vectorized: true │ └── • scan columns: (a) - estimated row count: 110 (missing stats) + estimated row count: 111 (missing stats) table: d@foo_inv spans: /"a"/"b"-/"a"/"b"/PrefixEnd @@ -399,7 +399,7 @@ vectorized: true │ └── • scan columns: (a) - estimated row count: 110 (missing stats) + estimated row count: 111 (missing stats) table: d@foo_inv spans: /"a"/"b"-/"a"/"b"/PrefixEnd diff --git a/pkg/sql/opt/invertedidx/BUILD.bazel b/pkg/sql/opt/invertedidx/BUILD.bazel index e7d90cd74836..c9b1e4f654a4 100644 --- a/pkg/sql/opt/invertedidx/BUILD.bazel +++ b/pkg/sql/opt/invertedidx/BUILD.bazel @@ -28,6 +28,7 @@ go_library( "//pkg/sql/sem/tree", "//pkg/sql/types", "//pkg/util/encoding", + "//pkg/util/json", "@com_github_cockroachdb_errors//:errors", "@com_github_golang_geo//r1", "@com_github_golang_geo//s1", diff --git a/pkg/sql/opt/invertedidx/json_array.go b/pkg/sql/opt/invertedidx/json_array.go index 72cd86280089..34cab66a8fc9 100644 --- a/pkg/sql/opt/invertedidx/json_array.go +++ b/pkg/sql/opt/invertedidx/json_array.go @@ -23,6 +23,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/rowenc" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" "github.com/cockroachdb/cockroach/pkg/sql/types" + "github.com/cockroachdb/cockroach/pkg/util/json" "github.com/cockroachdb/errors" ) @@ -270,24 +271,32 @@ func (j *jsonOrArrayFilterPlanner) extractInvertedFilterConditionFromLeaf( _ *invertedexpr.PreFiltererStateForInvertedFilterer, ) { switch t := expr.(type) { - // TODO(rytaft): Support JSON fetch val operator (->). case *memo.ContainsExpr: - invertedExpr := j.extractJSONOrArrayContainsCondition(evalCtx, t.Left, t.Right) - if !invertedExpr.IsTight() { - remainingFilters = expr + invertedExpr = j.extractJSONOrArrayContainsCondition(evalCtx, t.Left, t.Right) + case *memo.EqExpr: + if fetch, ok := t.Left.(*memo.FetchValExpr); ok { + invertedExpr = j.extractJSONFetchValEqCondition(evalCtx, fetch, t.Right) } + } - // We do not currently support pre-filtering for JSON and Array indexes, so - // the returned pre-filter state is nil. - return invertedExpr, remainingFilters, nil - - default: + if invertedExpr == nil { + // An inverted expression could not be extracted. return invertedexpr.NonInvertedColExpression{}, expr, nil } + + // If the extracted inverted expression is not tight then remaining filters + // must be applied after the inverted index scan. + if !invertedExpr.IsTight() { + remainingFilters = expr + } + + // We do not currently support pre-filtering for JSON and Array indexes, so + // the returned pre-filter state is nil. + return invertedExpr, remainingFilters, nil } // extractJSONOrArrayContainsCondition extracts an InvertedExpression -// representing an inverted filter over the given inverted index, based +// representing an inverted filter over the planner's inverted index, based // on the given left and right expression arguments. Returns an empty // InvertedExpression if no inverted filter could be extracted. func (j *jsonOrArrayFilterPlanner) extractJSONOrArrayContainsCondition( @@ -316,3 +325,57 @@ func (j *jsonOrArrayFilterPlanner) extractJSONOrArrayContainsCondition( return getSpanExprForJSONOrArrayIndex(evalCtx, d) } + +// extractJSONFetchValEqCondition extracts an InvertedExpression representing an +// inverted filter over the planner's inverted index, based on equality between +// a fetch val expression and a right scalar expression. If the following criteria +// are not met, an empty InvertedExpression is returned. +// +// 1. The fetch value operator's left expression must be a variable +// referencing the inverted column in the index. +// 2. The fetch value operator's right expression must be a constant string. +// 3. The right expression in the equality expression must be a constant JSON +// value that is not an object or an array. +// +// TODO(mgartner): Support chained fetch val operators, e.g., j->'a'->'b' = '1'. +func (j *jsonOrArrayFilterPlanner) extractJSONFetchValEqCondition( + evalCtx *tree.EvalContext, fetch *memo.FetchValExpr, right opt.ScalarExpr, +) invertedexpr.InvertedExpression { + // The left side of the fetch val expression, the Json field, should be a + // variable corresponding to the index column. + variable, ok := indexColumnVariable(j.tabID, j.index, fetch.Json) + if !ok { + return invertedexpr.NonInvertedColExpression{} + } + + // The right side of the fetch val expression, the Index field, should be a + // constant string. + if !memo.CanExtractConstDatum(fetch.Index) { + return invertedexpr.NonInvertedColExpression{} + } + key, ok := memo.ExtractConstDatum(fetch.Index).(*tree.DString) + if !ok { + return invertedexpr.NonInvertedColExpression{} + } + + // The right side of the equals expression should be a constant JSON value + // that is not an object or array. + if !memo.CanExtractConstDatum(right) { + return invertedexpr.NonInvertedColExpression{} + } + val, ok := memo.ExtractConstDatum(right).(*tree.DJSON) + if !ok { + return invertedexpr.NonInvertedColExpression{} + } + typ := val.JSON.Type() + if typ == json.ObjectJSONType || typ == json.ArrayJSONType { + return invertedexpr.NonInvertedColExpression{} + } + + // Build a new JSON object of the form: {: }. + b := json.NewObjectBuilder(1) + b.Add(string(*key), val.JSON) + obj := tree.NewDJSON(b.Build()) + + return getSpanExprForJSONOrArrayIndex(evalCtx, obj) +} diff --git a/pkg/sql/opt/invertedidx/json_array_test.go b/pkg/sql/opt/invertedidx/json_array_test.go index e6bd6ac6d674..031657fd7ebe 100644 --- a/pkg/sql/opt/invertedidx/json_array_test.go +++ b/pkg/sql/opt/invertedidx/json_array_test.go @@ -391,6 +391,73 @@ func TestTryFilterJsonOrArrayIndex(t *testing.T) { unique: false, remainingFilters: "j @> '[[1, 2]]'", }, + { + filters: "j->'a' = '1'", + indexOrd: jsonOrd, + ok: true, + tight: true, + unique: true, + }, + { + // Integer indexes are not yet supported. + filters: "j->0 = '1'", + indexOrd: jsonOrd, + ok: false, + }, + { + // Arrays on the right side of the equality are not yet supported. + filters: "j->'a' = '[1]'", + indexOrd: jsonOrd, + ok: false, + }, + { + // Objects on the right side of the equality are not yet supported. + filters: `j->'a' = '{"b": "c"}'`, + indexOrd: jsonOrd, + ok: false, + }, + { + // Wrong index ordinal. + filters: "j->'a' = '1'", + indexOrd: arrayOrd, + ok: false, + }, + { + filters: "j->'a' = '1' AND j->'b' = '2'", + indexOrd: jsonOrd, + ok: true, + tight: true, + unique: false, + }, + { + filters: "j->'a' = '1' OR j->'b' = '2'", + indexOrd: jsonOrd, + ok: true, + tight: true, + unique: false, + }, + { + filters: `j->'a' = '1' AND j @> '{"b": "c"}'`, + indexOrd: jsonOrd, + ok: true, + tight: true, + unique: false, + }, + { + filters: `j->'a' = '1' OR j @> '{"b": "c"}'`, + indexOrd: jsonOrd, + ok: true, + tight: true, + unique: false, + }, + { + filters: `j->'a' = '1' AND j @> '[[1, 2]]'`, + indexOrd: jsonOrd, + ok: true, + tight: false, + unique: false, + remainingFilters: "j @> '[[1, 2]]'", + }, } for _, tc := range testCases { diff --git a/pkg/sql/opt/xform/testdata/rules/select b/pkg/sql/opt/xform/testdata/rules/select index 216c59614b6e..1de1c2f20c2e 100644 --- a/pkg/sql/opt/xform/testdata/rules/select +++ b/pkg/sql/opt/xform/testdata/rules/select @@ -1937,6 +1937,162 @@ index-join b │ └── spans: ["7a\x00\x02\x00\x03\x00\x03b\x00\x02c\x00\x02\x00\x03d\x00\x01\x12e\x00\x01", "7a\x00\x02\x00\x03\x00\x03b\x00\x02c\x00\x02\x00\x03d\x00\x01\x12e\x00\x01"] └── key: (1) +# Query using the fetch val and equality operators. +opt expect=GenerateInvertedIndexScans +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: /6/1 + │ └── spans: ["7a\x00\x01\x12b\x00\x01", "7a\x00\x01\x12b\x00\x01"] + └── key: (1) + +# Do not generate an inverted scan when the index of the fetch val operator is +# not a string. +opt expect-not=GenerateInvertedIndexScans +SELECT k FROM b WHERE j->0 = '"b"' +---- +project + ├── columns: k:1!null + ├── immutable + ├── key: (1) + └── select + ├── columns: k:1!null j:4 + ├── immutable + ├── key: (1) + ├── fd: (1)-->(4) + ├── scan b + │ ├── columns: k:1!null j:4 + │ ├── key: (1) + │ └── fd: (1)-->(4) + └── filters + └── (j:4->0) = '"b"' [outer=(4), immutable] + +# Do not generate an inverted scan when right side of the equality is an array. +opt expect-not=GenerateInvertedIndexScans +SELECT k FROM b WHERE j->'a' = '["b"]' +---- +project + ├── columns: k:1!null + ├── immutable + ├── key: (1) + └── select + ├── columns: k:1!null j:4 + ├── immutable + ├── key: (1) + ├── fd: (1)-->(4) + ├── scan b + │ ├── columns: k:1!null j:4 + │ ├── key: (1) + │ └── fd: (1)-->(4) + └── filters + └── (j:4->'a') = '["b"]' [outer=(4), immutable] + +# Do not generate an inverted scan when right side of the equality is an object. +opt expect-not=GenerateInvertedIndexScans +SELECT k FROM b WHERE j->'a' = '{"b": "c"}' +---- +project + ├── columns: k:1!null + ├── immutable + ├── key: (1) + └── select + ├── columns: k:1!null j:4 + ├── immutable + ├── key: (1) + ├── fd: (1)-->(4) + ├── scan b + │ ├── columns: k:1!null j:4 + │ ├── key: (1) + │ └── fd: (1)-->(4) + └── filters + └── (j:4->'a') = '{"b": "c"}' [outer=(4), immutable] + +# Query using the fetch val and equality operators in a conjunction. +opt expect=GenerateInvertedIndexScans disable=GenerateInvertedIndexZigzagJoins +SELECT k FROM b WHERE j->'a' = '"b"' AND j->'c' = '"d"' +---- +project + ├── columns: k:1!null + ├── immutable + ├── key: (1) + └── inverted-filter + ├── columns: k:1!null + ├── inverted expression: /6 + │ ├── tight: true, unique: false + │ ├── union spans: empty + │ └── INTERSECTION + │ ├── span expression + │ │ ├── tight: true, unique: true + │ │ └── union spans: ["7a\x00\x01\x12b\x00\x01", "7a\x00\x01\x12b\x00\x01"] + │ └── span expression + │ ├── tight: true, unique: true + │ └── union spans: ["7c\x00\x01\x12d\x00\x01", "7c\x00\x01\x12d\x00\x01"] + ├── key: (1) + └── scan b@j_inv_idx + ├── columns: k:1!null j_inverted_key:6!null + ├── inverted constraint: /6/1 + │ └── spans + │ ├── ["7a\x00\x01\x12b\x00\x01", "7a\x00\x01\x12b\x00\x01"] + │ └── ["7c\x00\x01\x12d\x00\x01", "7c\x00\x01\x12d\x00\x01"] + ├── key: (1) + └── fd: (1)-->(6) + +# Query using the fetch val and equality operators in a disjunction. +opt expect=GenerateInvertedIndexScans +SELECT k FROM b WHERE j->'a' = '"b"' OR j->'c' = '"d"' +---- +project + ├── columns: k:1!null + ├── immutable + ├── key: (1) + └── inverted-filter + ├── columns: k:1!null + ├── inverted expression: /6 + │ ├── tight: true, unique: false + │ └── union spans + │ ├── ["7a\x00\x01\x12b\x00\x01", "7a\x00\x01\x12b\x00\x01"] + │ └── ["7c\x00\x01\x12d\x00\x01", "7c\x00\x01\x12d\x00\x01"] + ├── key: (1) + └── scan b@j_inv_idx + ├── columns: k:1!null j_inverted_key:6!null + ├── inverted constraint: /6/1 + │ └── spans + │ ├── ["7a\x00\x01\x12b\x00\x01", "7a\x00\x01\x12b\x00\x01"] + │ └── ["7c\x00\x01\x12d\x00\x01", "7c\x00\x01\x12d\x00\x01"] + ├── key: (1) + └── fd: (1)-->(6) + +# Query using the fetch val and equality operators in a disjunction with a +# contains operator. +opt expect=GenerateInvertedIndexScans +SELECT k FROM b WHERE j->'a' = '"b"' OR j @> '{"c": "d"}' +---- +project + ├── columns: k:1!null + ├── immutable + ├── key: (1) + └── inverted-filter + ├── columns: k:1!null + ├── inverted expression: /6 + │ ├── tight: true, unique: false + │ └── union spans + │ ├── ["7a\x00\x01\x12b\x00\x01", "7a\x00\x01\x12b\x00\x01"] + │ └── ["7c\x00\x01\x12d\x00\x01", "7c\x00\x01\x12d\x00\x01"] + ├── key: (1) + └── scan b@j_inv_idx + ├── columns: k:1!null j_inverted_key:6!null + ├── inverted constraint: /6/1 + │ └── spans + │ ├── ["7a\x00\x01\x12b\x00\x01", "7a\x00\x01\x12b\x00\x01"] + │ └── ["7c\x00\x01\x12d\x00\x01", "7c\x00\x01\x12d\x00\x01"] + ├── key: (1) + └── fd: (1)-->(6) + # GenerateInvertedIndexScans propagates row-level locking information. opt expect=GenerateInvertedIndexScans SELECT k FROM b WHERE j @> '{"a": "b"}' FOR UPDATE @@ -4783,7 +4939,8 @@ SELECT k FROM inv_zz_partial WHERE j->'a' = '1' AND j->'b' = '2' ---- project └── scan inv_zz_partial@zz_idx,partial - └── constraint: /2/1: [/'{"b": 2}' - /'{"b": 2}'] + └── inverted constraint: /7/1 + └── spans: ["7b\x00\x01*\x04\x00", "7b\x00\x01*\x04\x00"] # -------------------------------------------------- # SplitDisjunction