Merge #154892

154892: sql/stats: support partial stats over indexes implied by the filter r=Uzair5162 a=Uzair5162

This commit adds support to collect arbitrary partial statistics when a partial index with a predicate implied by the stat collection's filters exists.

Fixes: #154891

Release note (sql change): Added support for collecting partial statistics when the given `WHERE` clause implies the predicate of a partial index with the requested column as the first key column. For example:

```
CREATE TABLE t (a INT, INDEX idx_partial (a) WHERE a > 5);
CREATE STATISTICS pstat ON a FROM t WHERE a > 7;
```

Co-authored-by: Uzair Ahmad <[email protected]>

Author: craig[bot]

pkg/sql/logictest/testdata/logic_test/constrained_stats

@@ -99,7 +99,7 @@ CREATE STATISTICS stat_wrong_col ON price FROM products WHERE category_id = 1
 statement error pq: filter cannot be a contradiction
 CREATE STATISTICS stat_contradiction ON price FROM products WHERE price < 10 AND price > 10
 
-statement error pq: table products does not contain a non-partial forward index with created_at as a prefix column
+statement error pq: table products does not contain a suitable index with created_at as a prefix column
 CREATE STATISTICS stat_created_at ON created_at FROM products WHERE created_at > '2025-01-05'
 
 statement error column "nonexistent" does not exist
@@ -205,5 +205,91 @@ stats_from_table_id  {price}       10         10              0           price
 statement ok
 CREATE TABLE infer (x INT NOT NULL, y INT NOT NULL AS (x % 5) STORED, PRIMARY KEY (y, x));
 
-statement error pq: table infer does not contain a non-partial forward index with x as a prefix column
+statement error pq: table infer does not contain a suitable index with x as a prefix column
 CREATE STATISTICS inferred_stat ON x FROM infer WHERE x = 100;
+
+# Test partial statistics with a WHERE clause on tables with partial indexes.
+statement ok
+CREATE TABLE partial_indexes (
+    a INT,
+    b INT,
+    INDEX idx_partial (a) WHERE a > 5
+);
+
+statement ok
+INSERT INTO partial_indexes SELECT g, g FROM generate_series(1, 10) AS g;
+
+statement ok
+CREATE STATISTICS fullstat FROM partial_indexes;
+
+# Clear the stat cache so that creating partial statistics has access to the
+# latest full statistic.
+statement ok
+SELECT crdb_internal.clear_table_stats_cache();
+
+# Filter implies the partial index predicate, so it can be used.
+statement ok
+CREATE STATISTICS partial_stat_implied ON a FROM partial_indexes WHERE a > 7;
+
+query TTIIIT colnames
+SELECT statistics_name, column_names, row_count, distinct_count, null_count, partial_predicate
+FROM [SHOW STATISTICS FOR TABLE partial_indexes]
+WHERE statistics_name = 'partial_stat_implied';
+----
+statistics_name       column_names  row_count  distinct_count  null_count  partial_predicate
+partial_stat_implied  {a}           3          3               0           a > 7
+
+# Filter equals partial index predicate, so it can be used.
+statement ok
+CREATE STATISTICS partial_stat_exact ON a FROM partial_indexes WHERE a > 5;
+
+query TTIIIT colnames
+SELECT statistics_name, column_names, row_count, distinct_count, null_count, partial_predicate
+FROM [SHOW STATISTICS FOR TABLE partial_indexes]
+WHERE statistics_name = 'partial_stat_exact';
+----
+statistics_name     column_names  row_count  distinct_count  null_count  partial_predicate
+partial_stat_exact  {a}           5          5               0           a > 5
+
+# Filter does not imply the partial index predicate, so it can't be used.
+statement error pq: table partial_indexes does not contain a suitable index with a as a prefix column
+CREATE STATISTICS partial_stat_not_implied ON a FROM partial_indexes WHERE a > 3;
+
+# Ensure that we use suitable full indexes when a partial index is not implied.
+statement ok
+CREATE INDEX idx_a ON partial_indexes (a);
+
+statement ok
+CREATE STATISTICS partial_stat_not_implied ON a FROM partial_indexes WHERE a > 3;
+
+query TTIIIT colnames
+SELECT statistics_name, column_names, row_count, distinct_count, null_count, partial_predicate
+FROM [SHOW STATISTICS FOR TABLE partial_indexes]
+WHERE statistics_name = 'partial_stat_not_implied';
+----
+statistics_name          column_names  row_count  distinct_count  null_count  partial_predicate
+partial_stat_not_implied  {a}           7          7               0           a > 3
+
+# Ensure that we only use partial indexes with the required prefix columns.
+statement ok
+DROP INDEX partial_indexes@idx_partial;
+
+statement ok
+DROP INDEX partial_indexes@idx_a;
+
+statement ok
+CREATE INDEX idx_b ON partial_indexes (b) WHERE a > 5;
+
+statement error pq: table partial_indexes does not contain a suitable index with a as a prefix column
+CREATE STATISTICS partial_stat_implied ON a FROM partial_indexes WHERE a > 7;
+
+statement ok
+DROP INDEX partial_indexes@idx_b;
+
+# Ensure that we can't use an implied partial index when the filter doesn't
+# create a single constrained span.
+statement ok
+CREATE INDEX idx_b ON partial_indexes (b) WHERE b > 1 AND b % 2 = 0;
+
+statement error pq: predicate could not become a constrained scan of an index
+CREATE STATISTICS partial_stat_implied ON b FROM partial_indexes WHERE b > 1 AND b % 2 = 0;

pkg/sql/opt/optbuilder/misc_statements.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/constraint"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/idxconstraint"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/partialidx"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
@@ -265,17 +266,31 @@ func (b *Builder) buildWhereForStatistics(
 			"WHERE filter must be on the same column as the one specified in the column list"))
 	}
 
-	// Find a non-partial forward index with the filter column as the first key
-	// column.
+	// Add the partial index predicate expressions to the table metadata so we
+	// can check if the filter implies any partial index predicates.
+	b.addPartialIndexPredicatesForTable(tabMeta, nil /* scan */)
+
+	// Find a forward index with the filter column as the first key column.
 	foundIndex := false
 	var orderingCol opt.OrderingColumn
 	var notNullCols opt.ColSet
+
+	var im partialidx.Implicator
+	im.Init(b.ctx, b.factory, b.factory.Metadata(), b.evalCtx)
 	for i := 0; i < tab.IndexCount(); i++ {
 		idx := tab.Index(i)
-		// TODO (uzair): Allow partial indexes that are implied by the filter.
-		if _, isPartial := idx.Predicate(); isPartial || idx.Type() != idxtype.FORWARD {
+		if idx.Type() != idxtype.FORWARD {
 			continue
 		}
+
+		// If this is a partial index, check if the filter implies its predicate.
+		if pred, isPartialIndex := tabMeta.PartialIndexPredicate(i); isPartialIndex {
+			predFilters := *pred.(*memo.FiltersExpr)
+			if _, ok := im.FiltersImplyPredicate(fe, predFilters, tabMeta.ComputedCols); !ok {
+				continue
+			}
+		}
+
 		if idx.KeyColumnCount() > 0 {
 			col := idx.Column(0)
 			if tabID.IndexColumnID(idx, 0) == filterColID {
@@ -292,7 +307,7 @@ func (b *Builder) buildWhereForStatistics(
 	if !foundIndex {
 		panic(
 			pgerror.Newf(pgcode.InvalidColumnReference,
-				"table %s does not contain a non-partial forward index with %s as a prefix column",
+				"table %s does not contain a suitable index with %s as a prefix column",
 				tab.Name(), tab.Column(tabID.ColumnOrdinal(filterColID)).ColName(),
 			),
 		)