sql/stats: support collecting partial statistics on arbitrary constraints

This commit adds support to manually create partial statistics with a
`WHERE` clause, where the predicate fully constrains the first column of
an index. Additionally, the predicate must reference a single outer
column, which must be the same as the single column that stats are being
collected on. These stats are stored in `system.table_statistics` with
their predicate.

Part of: #93998

Release note (sql change): Users can now manually create single-column
partial statistics on boolean predicate expressions that can become
simple index scans. These statistics can be created by adding a
constraining `WHERE` expression to `CREATE STATISTICS`.

For example:

```
CREATE TABLE t (a INT PRIMARY KEY);
INSERT INTO t VALUES (1), (2), (3), (4), (5);
CREATE STATISTICS constrained_stat ON a FROM t WHERE a > 2;
```

Author: Uzair5162

pkg/jobs/jobspb/jobs.proto

@@ -1264,6 +1264,21 @@ message CreateStatsDetails {
 
   // If true, will collect partial table statistics at extreme values.
   bool using_extremes = 9;
+
+  // WHERE clause for partial statistics collection. This field is only used
+  // to populate the predicate in the system.table_statistics table and to
+  // determine if this is a constrained stats collection. The actual constrained
+  // scan is done over the spans in the where_spans field.
+  string where_clause = 10;
+
+  // Spans over which to collect partial statistics with a WHERE clause.
+  repeated roachpb.Span where_spans = 11 [(gogoproto.nullable) = false];
+
+  // The ID of the index used to collect partial statistics with a WHERE clause.
+  int64 where_index_id = 12 [
+    (gogoproto.customname) = "WhereIndexID",
+    (gogoproto.casttype) = "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb.IndexID"
+  ];
 }
 
 message CreateStatsProgress {

pkg/sql/create_stats.go

@@ -12,6 +12,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/featureflag"
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
@@ -139,6 +140,12 @@ type createStatsNode struct {
 	// If it is false, the flow for create statistics is planned directly; this
 	// is used when the statement is under EXPLAIN or EXPLAIN ANALYZE.
 	runAsJob bool
+
+	// whereSpans are the spans corresponding to the WHERE clause, if any.
+	whereSpans roachpb.Spans
+
+	// whereIndexID is the index to use to collect statistics with a WHERE clause.
+	whereIndexID descpb.IndexID
 }
 
 func (n *createStatsNode) startExec(params runParams) error {
@@ -280,11 +287,15 @@ func (n *createStatsNode) makeJobRecord(ctx context.Context) (*jobs.Record, erro
 		return nil, errors.Errorf(`creating partial statistics at extremes is disabled`)
 	}
 
-	// TODO(93998): Add support for WHERE.
+	var whereClause string
 	if n.Options.Where != nil {
-		return nil, pgerror.New(pgcode.FeatureNotSupported,
-			"creating partial statistics with a WHERE clause is not yet supported",
-		)
+		if n.whereSpans == nil {
+			return nil, errors.AssertionFailedf(
+				"expected whereSpans to be set for statistics with a WHERE clause")
+		}
+		// Safe to use AsString since whereClause is only used to populate the
+		// predicate in system.table_statistics.
+		whereClause = tree.AsString(n.Options.Where.Expr)
 	}
 
 	if err := n.p.CheckPrivilege(ctx, tableDesc, privilege.SELECT); err != nil {
@@ -409,6 +420,9 @@ func (n *createStatsNode) makeJobRecord(ctx context.Context) (*jobs.Record, erro
 			MaxFractionIdle:  n.Options.Throttling,
 			DeleteOtherStats: deleteOtherStats,
 			UsingExtremes:    n.Options.UsingExtremes,
+			WhereClause:      whereClause,
+			WhereSpans:       n.whereSpans,
+			WhereIndexID:     n.whereIndexID,
 		},
 		Progress: jobspb.CreateStatsProgress{},
 	}, nil

pkg/sql/distsql_plan_stats.go

@@ -329,7 +329,11 @@ func (dsp *DistSQLPlanner) createPartialStatsPlan(
 	// Initialize a dummy scanNode for the requested statistic.
 	var scan scanNode
 	scan.desc = desc
-	err = scan.initDescSpecificCol(colCfg, column)
+	if details.UsingExtremes {
+		err = scan.initDescSpecificCol(colCfg, column)
+	} else if details.WhereClause != "" {
+		err = scan.initDescSpecificIndex(colCfg, column, details.WhereIndexID)
+	}
 	if err != nil {
 		return nil, err
 	}
@@ -342,11 +346,6 @@ func (dsp *DistSQLPlanner) createPartialStatsPlan(
 		colIdxMap.Set(c.GetID(), i)
 	}
 
-	var sb span.Builder
-	sb.InitAllowingExternalRowData(
-		planCtx.EvalContext(), planCtx.ExtendedEvalCtx.Codec, desc, scan.index,
-	)
-
 	var stat *stats.TableStatistic
 	// Find the statistic from the newest table statistic for our column that is
 	// not partial and not forecasted. The first one we find will be the latest
@@ -379,27 +378,47 @@ func (dsp *DistSQLPlanner) createPartialStatsPlan(
 			"column %s does not have a prior statistic",
 			column.GetName())
 	}
-	lowerBound, upperBound, err := bounds.GetUsingExtremesBounds(ctx, planCtx.EvalContext(), stat.Histogram)
-	if err != nil {
-		return nil, err
-	}
-	if lowerBound == nil {
-		return nil, pgerror.Newf(
-			pgcode.ObjectNotInPrerequisiteState,
-			"only outer or NULL bounded buckets exist in %s@%s (table ID %d, column IDs %v), "+
-				"so partial stats cannot be collected",
-			scan.desc.GetName(), scan.index.GetName(), stat.TableID, stat.ColumnIDs,
+
+	var predicate string
+	var prevLowerBound tree.Datum
+	if details.UsingExtremes {
+		var sb span.Builder
+		sb.InitAllowingExternalRowData(
+			planCtx.EvalContext(), planCtx.ExtendedEvalCtx.Codec, desc, scan.index,
 		)
-	}
-	extremesSpans, err := bounds.ConstructUsingExtremesSpans(lowerBound, upperBound, scan.index)
-	if err != nil {
-		return nil, err
-	}
-	extremesPredicate := bounds.ConstructUsingExtremesPredicate(lowerBound, upperBound, column.GetName())
-	// Get roachpb.Spans from constraint.Spans
-	scan.spans, err = sb.SpansFromConstraintSpan(&extremesSpans, span.NoopSplitter())
-	if err != nil {
-		return nil, err
+
+		lowerBound, upperBound, err := bounds.GetUsingExtremesBounds(ctx,
+			planCtx.EvalContext(), stat.Histogram)
+		if err != nil {
+			return nil, err
+		}
+		if lowerBound == nil {
+			return nil, pgerror.Newf(
+				pgcode.ObjectNotInPrerequisiteState,
+				"only outer or NULL bounded buckets exist in %s@%s (table ID %d, column IDs %v), "+
+					"so partial stats cannot be collected",
+				scan.desc.GetName(), scan.index.GetName(), stat.TableID, stat.ColumnIDs,
+			)
+		}
+		prevLowerBound = lowerBound
+
+		extremesSpans, err := bounds.ConstructUsingExtremesSpans(lowerBound,
+			upperBound, scan.index)
+		if err != nil {
+			return nil, err
+		}
+		predicate = bounds.ConstructUsingExtremesPredicate(lowerBound, upperBound, column.GetName())
+		// Get roachpb.Spans from constraint.Spans
+		scan.spans, err = sb.SpansFromConstraintSpan(&extremesSpans, span.NoopSplitter())
+		if err != nil {
+			return nil, err
+		}
+	} else if details.WhereClause != "" {
+		predicate = details.WhereClause
+		scan.spans = details.WhereSpans
+	} else {
+		return nil, errors.AssertionFailedf(
+			"partial stats require either USING EXTREMES or a WHERE clause")
 	}
 	p, err := dsp.createTableReaders(ctx, planCtx, &scan)
 	if err != nil {
@@ -419,10 +438,13 @@ func (dsp *DistSQLPlanner) createPartialStatsPlan(
 		HistogramMaxBuckets: reqStat.histogramMaxBuckets,
 		Columns:             make([]uint32, len(reqStat.columns)),
 		StatName:            reqStat.name,
-		PartialPredicate:    extremesPredicate,
-		FullStatisticID:     stat.StatisticID,
-		PrevLowerBound:      tree.Serialize(lowerBound),
+		PartialPredicate:    predicate,
+	}
+	if details.UsingExtremes && prevLowerBound != nil {
+		spec.PrevLowerBound = tree.Serialize(prevLowerBound)
+		spec.FullStatisticID = stat.StatisticID
 	}
+
 	// For now, this loop should iterate only once, as we only
 	// handle single-column partial statistics.
 	// TODO(faizaanmadhani): Add support for multi-column partial stats next
@@ -767,7 +789,7 @@ func (dsp *DistSQLPlanner) createPlanForCreateStats(
 		return nil, errors.New("no stats requested")
 	}
 
-	if details.UsingExtremes {
+	if details.UsingExtremes || details.WhereClause != "" {
 		return dsp.createPartialStatsPlan(ctx, planCtx, tableDesc, reqStats, jobID, details, numIndexes, curIndex)
 	}
 	return dsp.createStatsPlan(ctx, planCtx, semaCtx, tableDesc, reqStats, jobID, details, numIndexes, curIndex)

pkg/sql/distsql_spec_exec_factory.go

@@ -1707,7 +1707,7 @@ func (e *distSQLSpecExecFactory) ConstructCancelSessions(
 }
 
 func (e *distSQLSpecExecFactory) ConstructCreateStatistics(
-	cs *tree.CreateStats,
+	cs *tree.CreateStats, table cat.Table, index cat.Index, whereConstraint *constraint.Constraint,
 ) (exec.Node, error) {
 	return nil, unimplemented.NewWithIssue(47473, "experimental opt-driven distsql planning: create statistics")
 }

pkg/sql/execinfrapb/processors_table_stats.proto

@@ -43,9 +43,9 @@ message SketchSpec {
   // and is the empty string for a full table statistic.
   optional string partial_predicate = 7 [(gogoproto.nullable) = false];
 
-  // FullStatisticID is non-zero for partial statistics and 0 for full
-  // statistics. It is the statistic id of the full statistic that this partial
-  // statistic was derived from.
+  // FullStatisticID is non-zero for partial statistics USING EXTREMES and 0
+  // otherwise. It is the statistic id of the full statistic that this partial
+  // statistic on extreme values was derived from.
   optional uint64 full_statistic_id = 8 [(gogoproto.customname) = "FullStatisticID", (gogoproto.nullable) = false];
 
   // PrevLowerBound is a tree.Datum serialized into a string representing