diff --git a/pkg/sql/logictest/testdata/logic_test/partial_index b/pkg/sql/logictest/testdata/logic_test/partial_index index c838197fdc9b..b3b5b152ec6a 100644 --- a/pkg/sql/logictest/testdata/logic_test/partial_index +++ b/pkg/sql/logictest/testdata/logic_test/partial_index @@ -1281,6 +1281,54 @@ SELECT * FROM inv_c@i WHERE j @> '{"x": "y"}' AND s IN ('foo', 'bar') ORDER BY k 1 {"num": 1, "x": "y"} foo 3 {"num": 3, "x": "y"} bar +# Updates and Upserts with fetch columns pruned. + +statement ok +CREATE TABLE prune ( + a INT PRIMARY KEY, + b INT, + c INT, + d INT, + INDEX idx (b) WHERE c > 0, + FAMILY (a), + FAMILY (b), + FAMILY (c), + FAMILY (d) +) + +statement ok +INSERT INTO prune (a, b, c, d) VALUES (1, 2, 3, 4) + +# Test that an update is successful when fetch columns b and c are pruned +# because an update to idx is not required. +statement ok +UPDATE prune SET d = d + 1 WHERE a = 1 + +query IIII rowsort +SELECT * FROM prune@idx WHERE c > 0 +---- +1 2 3 5 + +# Test that an upsert is successful when fetch columns b and c are pruned +# because an update to idx is not required. +statement ok +UPSERT INTO prune (a, d) VALUES (1, 6) + +query IIII rowsort +SELECT * FROM prune@idx WHERE c > 0 +---- +1 2 3 6 + +# Test that an upsert is successful when fetch columns b and c are pruned +# because an update to idx is not required. +statement ok +INSERT INTO prune (a, d) VALUES (1, 6) ON CONFLICT (a) DO UPDATE SET d = 7 + +query IIII rowsort +SELECT * FROM prune@idx WHERE c > 0 +---- +1 2 3 7 + # Regression tests for #52318. Mutations on partial indexes in the # DELETE_AND_WRITE_ONLY state should update the indexes correctly. subtest regression_52318 diff --git a/pkg/sql/opt/norm/prune_cols_funcs.go b/pkg/sql/opt/norm/prune_cols_funcs.go index 831f982ba586..f65672f32d5d 100644 --- a/pkg/sql/opt/norm/prune_cols_funcs.go +++ b/pkg/sql/opt/norm/prune_cols_funcs.go @@ -144,26 +144,26 @@ func (c *CustomFuncs) NeededMutationFetchCols( // Make sure to consider indexes that are being added or dropped. for i, n := 0, tabMeta.Table.DeletableIndexCount(); i < n; i++ { - // If the columns being updated are not part of the index and the - // index is not a partial index, then the update does not require - // changes to the index. Partial indexes may be updated (even when a - // column in the index is not changing) when rows that were not - // previously in the index must be added to the index because they - // now satisfy the partial index predicate. + // If the columns being updated are not part of the index, then the + // update does not require changes to the index. Partial indexes may + // be updated (even when a column in the index is not changing) when + // the predicate references columns that are being updated. For + // example, rows that were not previously in the index must be added + // to the index because they now satisfy the partial index + // predicate, requiring the index columns to be fetched. // // Note that we use the set of index columns where the virtual // columns have been mapped to their source columns. Virtual columns // are never part of the updated columns. Updates to source columns // trigger index changes. - // - // TODO(mgartner): Index columns are not necessary when neither the - // index columns nor the columns referenced in the partial index - // predicate are being updated. We should prune mutation fetch - // columns when this is the case, rather than always marking index - // columns of partial indexes as "needed". indexCols := tabMeta.IndexColumnsMapVirtual(i) - _, isPartialIndex := tabMeta.Table.Index(i).Predicate() - if !indexCols.Intersects(updateCols) && !isPartialIndex { + pred, isPartialIndex := tabMeta.PartialIndexPredicate(i) + indexAndPredCols := indexCols.Copy() + if isPartialIndex { + predFilters := *pred.(*memo.FiltersExpr) + indexAndPredCols.UnionWith(predFilters.OuterCols()) + } + if !indexAndPredCols.Intersects(updateCols) { continue } diff --git a/pkg/sql/opt/norm/testdata/rules/prune_cols b/pkg/sql/opt/norm/testdata/rules/prune_cols index bde5476683f1..ef54a006d307 100644 --- a/pkg/sql/opt/norm/testdata/rules/prune_cols +++ b/pkg/sql/opt/norm/testdata/rules/prune_cols @@ -48,9 +48,11 @@ CREATE TABLE partial_indexes ( a INT PRIMARY KEY, b INT, c STRING, + d INT, FAMILY (a), FAMILY (b), FAMILY (c), + FAMILY (d), INDEX (c) WHERE b > 1 ) ---- @@ -2058,51 +2060,404 @@ update computed ├── c_new:13 + 1 [as=column14:14, outer=(13), immutable] └── c_new:13 + 10 [as=column15:15, outer=(13), immutable] +# Prune UPDATE fetch columns when the partial index indexes the column but +# neither the column nor the columns referenced in the partial index predicate +# are mutating. +norm expect=(PruneMutationFetchCols,PruneMutationInputCols) +UPDATE partial_indexes SET d = d + 1 WHERE a = 1 +---- +update partial_indexes + ├── columns: + ├── fetch columns: a:6 d:9 + ├── update-mapping: + │ └── d_new:11 => d:4 + ├── partial index put columns: partial_index_put1:12 + ├── partial index del columns: partial_index_put1:12 + ├── cardinality: [0 - 0] + ├── volatile, mutations + └── project + ├── columns: partial_index_put1:12 d_new:11 a:6!null d:9 + ├── cardinality: [0 - 1] + ├── immutable + ├── key: () + ├── fd: ()-->(6,9,11,12) + ├── select + │ ├── columns: a:6!null b:7 d:9 + │ ├── cardinality: [0 - 1] + │ ├── key: () + │ ├── fd: ()-->(6,7,9) + │ ├── scan partial_indexes + │ │ ├── columns: a:6!null b:7 d:9 + │ │ ├── partial index predicates + │ │ │ └── secondary: filters + │ │ │ └── b:7 > 1 [outer=(7), constraints=(/7: [/2 - ]; tight)] + │ │ ├── key: (6) + │ │ └── fd: (6)-->(7,9) + │ └── filters + │ └── a:6 = 1 [outer=(6), constraints=(/6: [/1 - /1]; tight), fd=()-->(6)] + └── projections + ├── b:7 > 1 [as=partial_index_put1:12, outer=(7)] + └── d:9 + 1 [as=d_new:11, outer=(9), immutable] + +# Do not prune the indexed column c when a column in the partial index +# predicate, b, is being updated. +norm expect-not=PruneMutationFetchCols +UPDATE partial_indexes SET d = d + 1, b = 2 WHERE a = 1 +---- +update partial_indexes + ├── columns: + ├── fetch columns: a:6 b:7 c:8 d:9 + ├── update-mapping: + │ ├── b_new:12 => b:2 + │ └── d_new:11 => d:4 + ├── partial index put columns: partial_index_put1:13 + ├── partial index del columns: partial_index_del1:14 + ├── cardinality: [0 - 0] + ├── volatile, mutations + └── project + ├── columns: partial_index_put1:13!null partial_index_del1:14 d_new:11 b_new:12!null a:6!null b:7 c:8 d:9 + ├── cardinality: [0 - 1] + ├── immutable + ├── key: () + ├── fd: ()-->(6-9,11-14) + ├── select + │ ├── columns: a:6!null b:7 c:8 d:9 + │ ├── cardinality: [0 - 1] + │ ├── key: () + │ ├── fd: ()-->(6-9) + │ ├── scan partial_indexes + │ │ ├── columns: a:6!null b:7 c:8 d:9 + │ │ ├── partial index predicates + │ │ │ └── secondary: filters + │ │ │ └── b:7 > 1 [outer=(7), constraints=(/7: [/2 - ]; tight)] + │ │ ├── key: (6) + │ │ └── fd: (6)-->(7-9) + │ └── filters + │ └── a:6 = 1 [outer=(6), constraints=(/6: [/1 - /1]; tight), fd=()-->(6)] + └── projections + ├── true [as=partial_index_put1:13] + ├── b:7 > 1 [as=partial_index_del1:14, outer=(7)] + ├── d:9 + 1 [as=d_new:11, outer=(9), immutable] + └── 2 [as=b_new:12] + +# Prune UPSERT fetch columns when the partial index indexes the column but +# neither the column nor the columns referenced in the partial index predicate +# are mutating. +norm expect=(PruneMutationFetchCols,PruneMutationInputCols) +UPSERT INTO partial_indexes (a, d) VALUES (1, 2) +---- +upsert partial_indexes + ├── columns: + ├── arbiter indexes: primary + ├── canary column: a:10 + ├── fetch columns: a:10 d:13 + ├── insert-mapping: + │ ├── column1:6 => a:1 + │ ├── column8:8 => b:2 + │ ├── column9:9 => c:3 + │ └── column2:7 => d:4 + ├── update-mapping: + │ └── column2:7 => d:4 + ├── partial index put columns: partial_index_put1:18 + ├── partial index del columns: partial_index_del1:19 + ├── cardinality: [0 - 0] + ├── volatile, mutations + └── project + ├── columns: partial_index_put1:18 partial_index_del1:19 column1:6!null column2:7!null column8:8 column9:9 a:10 d:13 + ├── cardinality: [1 - 1] + ├── key: () + ├── fd: ()-->(6-10,13,18,19) + ├── left-join (cross) + │ ├── columns: column1:6!null column2:7!null column8:8 column9:9 a:10 b:11 d:13 + │ ├── cardinality: [1 - 1] + │ ├── multiplicity: left-rows(exactly-one), right-rows(exactly-one) + │ ├── key: () + │ ├── fd: ()-->(6-11,13) + │ ├── values + │ │ ├── columns: column1:6!null column2:7!null column8:8 column9:9 + │ │ ├── cardinality: [1 - 1] + │ │ ├── key: () + │ │ ├── fd: ()-->(6-9) + │ │ └── (1, 2, NULL, NULL) + │ ├── select + │ │ ├── columns: a:10!null b:11 d:13 + │ │ ├── cardinality: [0 - 1] + │ │ ├── key: () + │ │ ├── fd: ()-->(10,11,13) + │ │ ├── scan partial_indexes + │ │ │ ├── columns: a:10!null b:11 d:13 + │ │ │ ├── partial index predicates + │ │ │ │ └── secondary: filters + │ │ │ │ └── b:11 > 1 [outer=(11), constraints=(/11: [/2 - ]; tight)] + │ │ │ ├── key: (10) + │ │ │ └── fd: (10)-->(11,13) + │ │ └── filters + │ │ └── a:10 = 1 [outer=(10), constraints=(/10: [/1 - /1]; tight), fd=()-->(10)] + │ └── filters (true) + └── projections + ├── CASE WHEN a:10 IS NULL THEN column8:8 ELSE b:11 END > 1 [as=partial_index_put1:18, outer=(8,10,11)] + └── b:11 > 1 [as=partial_index_del1:19, outer=(11)] + +# Do not prune the indexed column c when a column in the partial index +# predicate, b, is being updated. +norm expect-not=PruneMutationFetchCols +UPSERT INTO partial_indexes (a, b, d) VALUES (1, 2, 3) +---- +upsert partial_indexes + ├── columns: + ├── arbiter indexes: primary + ├── canary column: a:10 + ├── fetch columns: a:10 b:11 c:12 d:13 + ├── insert-mapping: + │ ├── column1:6 => a:1 + │ ├── column2:7 => b:2 + │ ├── column9:9 => c:3 + │ └── column3:8 => d:4 + ├── update-mapping: + │ ├── column2:7 => b:2 + │ └── column3:8 => d:4 + ├── partial index put columns: partial_index_put1:17 + ├── partial index del columns: partial_index_del1:18 + ├── cardinality: [0 - 0] + ├── volatile, mutations + └── project + ├── columns: partial_index_put1:17!null partial_index_del1:18 column1:6!null column2:7!null column3:8!null column9:9 a:10 b:11 c:12 d:13 + ├── cardinality: [1 - 1] + ├── key: () + ├── fd: ()-->(6-13,17,18) + ├── left-join (cross) + │ ├── columns: column1:6!null column2:7!null column3:8!null column9:9 a:10 b:11 c:12 d:13 + │ ├── cardinality: [1 - 1] + │ ├── multiplicity: left-rows(exactly-one), right-rows(exactly-one) + │ ├── key: () + │ ├── fd: ()-->(6-13) + │ ├── values + │ │ ├── columns: column1:6!null column2:7!null column3:8!null column9:9 + │ │ ├── cardinality: [1 - 1] + │ │ ├── key: () + │ │ ├── fd: ()-->(6-9) + │ │ └── (1, 2, 3, NULL) + │ ├── select + │ │ ├── columns: a:10!null b:11 c:12 d:13 + │ │ ├── cardinality: [0 - 1] + │ │ ├── key: () + │ │ ├── fd: ()-->(10-13) + │ │ ├── scan partial_indexes + │ │ │ ├── columns: a:10!null b:11 c:12 d:13 + │ │ │ ├── partial index predicates + │ │ │ │ └── secondary: filters + │ │ │ │ └── b:11 > 1 [outer=(11), constraints=(/11: [/2 - ]; tight)] + │ │ │ ├── key: (10) + │ │ │ └── fd: (10)-->(11-13) + │ │ └── filters + │ │ └── a:10 = 1 [outer=(10), constraints=(/10: [/1 - /1]; tight), fd=()-->(10)] + │ └── filters (true) + └── projections + ├── column2:7 > 1 [as=partial_index_put1:17, outer=(7)] + └── b:11 > 1 [as=partial_index_del1:18, outer=(11)] + +# Prune INSERT ON CONFLICT DO UPDATE fetch columns when the partial index +# indexes the column but neither the column nor the columns referenced in the +# partial index predicate are mutating. +norm expect=(PruneMutationFetchCols,PruneMutationInputCols) +INSERT INTO partial_indexes (a, d) VALUES (1, 2) ON CONFLICT (a) DO UPDATE SET d = 3 +---- +upsert partial_indexes + ├── columns: + ├── arbiter indexes: primary + ├── canary column: a:10 + ├── fetch columns: a:10 d:13 + ├── insert-mapping: + │ ├── column1:6 => a:1 + │ ├── column8:8 => b:2 + │ ├── column9:9 => c:3 + │ └── column2:7 => d:4 + ├── update-mapping: + │ └── upsert_d:19 => d:4 + ├── partial index put columns: partial_index_put1:20 + ├── partial index del columns: partial_index_del1:21 + ├── cardinality: [0 - 0] + ├── volatile, mutations + └── project + ├── columns: partial_index_put1:20 partial_index_del1:21 upsert_d:19!null column1:6!null column2:7!null column8:8 column9:9 a:10 d:13 + ├── cardinality: [1 - 1] + ├── key: () + ├── fd: ()-->(6-10,13,19-21) + ├── left-join (cross) + │ ├── columns: column1:6!null column2:7!null column8:8 column9:9 a:10 b:11 d:13 + │ ├── cardinality: [1 - 1] + │ ├── multiplicity: left-rows(exactly-one), right-rows(exactly-one) + │ ├── key: () + │ ├── fd: ()-->(6-11,13) + │ ├── values + │ │ ├── columns: column1:6!null column2:7!null column8:8 column9:9 + │ │ ├── cardinality: [1 - 1] + │ │ ├── key: () + │ │ ├── fd: ()-->(6-9) + │ │ └── (1, 2, NULL, NULL) + │ ├── select + │ │ ├── columns: a:10!null b:11 d:13 + │ │ ├── cardinality: [0 - 1] + │ │ ├── key: () + │ │ ├── fd: ()-->(10,11,13) + │ │ ├── scan partial_indexes + │ │ │ ├── columns: a:10!null b:11 d:13 + │ │ │ ├── partial index predicates + │ │ │ │ └── secondary: filters + │ │ │ │ └── b:11 > 1 [outer=(11), constraints=(/11: [/2 - ]; tight)] + │ │ │ ├── key: (10) + │ │ │ └── fd: (10)-->(11,13) + │ │ └── filters + │ │ └── a:10 = 1 [outer=(10), constraints=(/10: [/1 - /1]; tight), fd=()-->(10)] + │ └── filters (true) + └── projections + ├── CASE WHEN a:10 IS NULL THEN column8:8 ELSE b:11 END > 1 [as=partial_index_put1:20, outer=(8,10,11)] + ├── b:11 > 1 [as=partial_index_del1:21, outer=(11)] + └── CASE WHEN a:10 IS NULL THEN column2:7 ELSE 3 END [as=upsert_d:19, outer=(7,10)] + +# Do not prune the indexed column c when a column in the partial index +# predicate, b, is being updated. +norm expect-not=PruneMutationFetchCols +INSERT INTO partial_indexes (a, d) VALUES (1, 2) ON CONFLICT (a) DO UPDATE SET b = 3, d = 4 +---- +upsert partial_indexes + ├── columns: + ├── arbiter indexes: primary + ├── canary column: a:10 + ├── fetch columns: a:10 b:11 c:12 d:13 + ├── insert-mapping: + │ ├── column1:6 => a:1 + │ ├── column8:8 => b:2 + │ ├── column9:9 => c:3 + │ └── column2:7 => d:4 + ├── update-mapping: + │ ├── upsert_b:18 => b:2 + │ └── upsert_d:20 => d:4 + ├── partial index put columns: partial_index_put1:21 + ├── partial index del columns: partial_index_del1:22 + ├── cardinality: [0 - 0] + ├── volatile, mutations + └── project + ├── columns: partial_index_put1:21 partial_index_del1:22 column1:6!null column2:7!null column8:8 column9:9 a:10 b:11 c:12 d:13 upsert_b:18 upsert_d:20!null + ├── cardinality: [1 - 1] + ├── key: () + ├── fd: ()-->(6-13,18,20-22) + ├── project + │ ├── columns: upsert_b:18 upsert_d:20!null column1:6!null column2:7!null column8:8 column9:9 a:10 b:11 c:12 d:13 + │ ├── cardinality: [1 - 1] + │ ├── key: () + │ ├── fd: ()-->(6-13,18,20) + │ ├── left-join (cross) + │ │ ├── columns: column1:6!null column2:7!null column8:8 column9:9 a:10 b:11 c:12 d:13 + │ │ ├── cardinality: [1 - 1] + │ │ ├── multiplicity: left-rows(exactly-one), right-rows(exactly-one) + │ │ ├── key: () + │ │ ├── fd: ()-->(6-13) + │ │ ├── values + │ │ │ ├── columns: column1:6!null column2:7!null column8:8 column9:9 + │ │ │ ├── cardinality: [1 - 1] + │ │ │ ├── key: () + │ │ │ ├── fd: ()-->(6-9) + │ │ │ └── (1, 2, NULL, NULL) + │ │ ├── select + │ │ │ ├── columns: a:10!null b:11 c:12 d:13 + │ │ │ ├── cardinality: [0 - 1] + │ │ │ ├── key: () + │ │ │ ├── fd: ()-->(10-13) + │ │ │ ├── scan partial_indexes + │ │ │ │ ├── columns: a:10!null b:11 c:12 d:13 + │ │ │ │ ├── partial index predicates + │ │ │ │ │ └── secondary: filters + │ │ │ │ │ └── b:11 > 1 [outer=(11), constraints=(/11: [/2 - ]; tight)] + │ │ │ │ ├── key: (10) + │ │ │ │ └── fd: (10)-->(11-13) + │ │ │ └── filters + │ │ │ └── a:10 = 1 [outer=(10), constraints=(/10: [/1 - /1]; tight), fd=()-->(10)] + │ │ └── filters (true) + │ └── projections + │ ├── CASE WHEN a:10 IS NULL THEN column8:8 ELSE 3 END [as=upsert_b:18, outer=(8,10)] + │ └── CASE WHEN a:10 IS NULL THEN column2:7 ELSE 4 END [as=upsert_d:20, outer=(7,10)] + └── projections + ├── upsert_b:18 > 1 [as=partial_index_put1:21, outer=(18)] + └── b:11 > 1 [as=partial_index_del1:22, outer=(11)] + +# Do not prune DELETE fetch columns. +norm +DELETE FROM partial_indexes WHERE a = 1 +---- +delete partial_indexes + ├── columns: + ├── fetch columns: a:6 c:8 + ├── partial index del columns: partial_index_del1:11 + ├── cardinality: [0 - 0] + ├── volatile, mutations + └── project + ├── columns: partial_index_del1:11 a:6!null c:8 + ├── cardinality: [0 - 1] + ├── key: () + ├── fd: ()-->(6,8,11) + ├── select + │ ├── columns: a:6!null b:7 c:8 + │ ├── cardinality: [0 - 1] + │ ├── key: () + │ ├── fd: ()-->(6-8) + │ ├── scan partial_indexes + │ │ ├── columns: a:6!null b:7 c:8 + │ │ ├── partial index predicates + │ │ │ └── secondary: filters + │ │ │ └── b:7 > 1 [outer=(7), constraints=(/7: [/2 - ]; tight)] + │ │ ├── key: (6) + │ │ └── fd: (6)-->(7,8) + │ └── filters + │ └── a:6 = 1 [outer=(6), constraints=(/6: [/1 - /1]; tight), fd=()-->(6)] + └── projections + └── b:7 > 1 [as=partial_index_del1:11, outer=(7)] + # Do not prune columns that are required for evaluating partial index # predicates. -norm expect-not=PruneMutationFetchCols +norm UPDATE partial_indexes SET b = b + 1 WHERE a = 1 ---- update partial_indexes ├── columns: - ├── fetch columns: a:5 b:6 c:7 + ├── fetch columns: a:6 b:7 c:8 ├── update-mapping: - │ └── b_new:9 => b:2 - ├── partial index put columns: partial_index_put1:10 - ├── partial index del columns: partial_index_del1:11 + │ └── b_new:11 => b:2 + ├── partial index put columns: partial_index_put1:12 + ├── partial index del columns: partial_index_del1:13 ├── cardinality: [0 - 0] ├── volatile, mutations └── project - ├── columns: partial_index_put1:10 partial_index_del1:11 a:5!null b:6 c:7 b_new:9 + ├── columns: partial_index_put1:12 partial_index_del1:13 a:6!null b:7 c:8 b_new:11 ├── cardinality: [0 - 1] ├── immutable ├── key: () - ├── fd: ()-->(5-7,9-11) + ├── fd: ()-->(6-8,11-13) ├── project - │ ├── columns: b_new:9 a:5!null b:6 c:7 + │ ├── columns: b_new:11 a:6!null b:7 c:8 │ ├── cardinality: [0 - 1] │ ├── immutable │ ├── key: () - │ ├── fd: ()-->(5-7,9) + │ ├── fd: ()-->(6-8,11) │ ├── select - │ │ ├── columns: a:5!null b:6 c:7 + │ │ ├── columns: a:6!null b:7 c:8 │ │ ├── cardinality: [0 - 1] │ │ ├── key: () - │ │ ├── fd: ()-->(5-7) + │ │ ├── fd: ()-->(6-8) │ │ ├── scan partial_indexes - │ │ │ ├── columns: a:5!null b:6 c:7 + │ │ │ ├── columns: a:6!null b:7 c:8 │ │ │ ├── partial index predicates │ │ │ │ └── secondary: filters - │ │ │ │ └── b:6 > 1 [outer=(6), constraints=(/6: [/2 - ]; tight)] - │ │ │ ├── key: (5) - │ │ │ └── fd: (5)-->(6,7) + │ │ │ │ └── b:7 > 1 [outer=(7), constraints=(/7: [/2 - ]; tight)] + │ │ │ ├── key: (6) + │ │ │ └── fd: (6)-->(7,8) │ │ └── filters - │ │ └── a:5 = 1 [outer=(5), constraints=(/5: [/1 - /1]; tight), fd=()-->(5)] + │ │ └── a:6 = 1 [outer=(6), constraints=(/6: [/1 - /1]; tight), fd=()-->(6)] │ └── projections - │ └── b:6 + 1 [as=b_new:9, outer=(6), immutable] + │ └── b:7 + 1 [as=b_new:11, outer=(7), immutable] └── projections - ├── b_new:9 > 1 [as=partial_index_put1:10, outer=(9)] - └── b:6 > 1 [as=partial_index_del1:11, outer=(6)] + ├── b_new:11 > 1 [as=partial_index_put1:12, outer=(11)] + └── b:7 > 1 [as=partial_index_del1:13, outer=(7)] # Prune secondary family column not needed for the update. norm expect=(PruneMutationFetchCols,PruneMutationInputCols) diff --git a/pkg/sql/opt/optbuilder/insert.go b/pkg/sql/opt/optbuilder/insert.go index aa929e258714..cdf095caad30 100644 --- a/pkg/sql/opt/optbuilder/insert.go +++ b/pkg/sql/opt/optbuilder/insert.go @@ -1061,6 +1061,11 @@ func (mb *mutationBuilder) buildUpsert(returning tree.ReturningExprs) { // Add any check constraint boolean columns to the input. mb.addCheckConstraintCols() + // Add the partial index predicate expressions to the table metadata. + // These expressions are used to prune fetch columns during + // normalization. + mb.b.addPartialIndexPredicatesForTable(mb.md.TableMeta(mb.tabID), nil /* scan */, true /* includeDeletable */) + // Project partial index PUT and DEL boolean columns. // // In some cases existing rows may not be fetched for an UPSERT (see diff --git a/pkg/sql/opt/optbuilder/partial_index.go b/pkg/sql/opt/optbuilder/partial_index.go index 413bae2e302f..c60cc64136c7 100644 --- a/pkg/sql/opt/optbuilder/partial_index.go +++ b/pkg/sql/opt/optbuilder/partial_index.go @@ -31,11 +31,18 @@ import ( // outputs all the ordinary columns in the table, we avoid constructing a new // scan. A scan and its logical properties are required in order to fully // normalize the partial index predicates. -func (b *Builder) addPartialIndexPredicatesForTable(tabMeta *opt.TableMeta, scan memo.RelExpr) { +func (b *Builder) addPartialIndexPredicatesForTable( + tabMeta *opt.TableMeta, scan memo.RelExpr, includeDeletable bool, +) { tab := tabMeta.Table + var numIndexes int + if includeDeletable { + numIndexes = tab.DeletableIndexCount() + } else { + numIndexes = tab.IndexCount() + } // Find the first partial index. - numIndexes := tab.IndexCount() indexOrd := 0 for ; indexOrd < numIndexes; indexOrd++ { if _, ok := tab.Index(indexOrd).Predicate(); ok { diff --git a/pkg/sql/opt/optbuilder/select.go b/pkg/sql/opt/optbuilder/select.go index d155bfe68c42..3a74a991df95 100644 --- a/pkg/sql/opt/optbuilder/select.go +++ b/pkg/sql/opt/optbuilder/select.go @@ -536,9 +536,10 @@ func (b *Builder) buildScan( outScope.expr = b.factory.ConstructScan(&private) // Add the partial indexes after constructing the scan so we can use the - // logical properties of the scan to fully normalize the index - // predicates. - b.addPartialIndexPredicatesForTable(tabMeta, outScope.expr) + // logical properties of the scan to fully normalize the index predicates. + // We don't need to add deletable partial index predicates in the context of + // a scan. + b.addPartialIndexPredicatesForTable(tabMeta, outScope.expr, false /* includeDeletable */) if !virtualColIDs.Empty() { // Project the expressions for the virtual columns (and pass through all diff --git a/pkg/sql/opt/optbuilder/update.go b/pkg/sql/opt/optbuilder/update.go index e91158ef185a..a38f3370a795 100644 --- a/pkg/sql/opt/optbuilder/update.go +++ b/pkg/sql/opt/optbuilder/update.go @@ -336,6 +336,11 @@ func (mb *mutationBuilder) buildUpdate(returning tree.ReturningExprs) { mb.addCheckConstraintCols() + // Add the partial index predicate expressions to the table metadata. + // These expressions are used to prune fetch columns during + // normalization. + mb.b.addPartialIndexPredicatesForTable(mb.md.TableMeta(mb.tabID), nil /* scan */, true /* includeDeletable */) + // Project partial index PUT and DEL boolean columns. mb.projectPartialIndexPutAndDelCols(preCheckScope, mb.fetchScope)