sql: support partial stats at extremes without specifying columns

This commit adds support to collect partial statistics at extremes
without specifying a column, such as:

`CREATE STATISTICS my_stat FROM t USING EXTREMES`

This will collect partial stats on a default set of columns, which
are all single columns that are prefixes of forward indexes, excluding
partial, sharded, and implicitly partitioned indexes.

Modifies create stats job progress reporting to work when multiple
indexes will be scanned as part of the stat collection. The i'th index
scan will only report progress between `(i-1)/numIndexes`% and
`i/numIndexes`%. For example, if we expect to scan 4 indexes, each index
will report progress between 0%-25%, 25%-50%, 50%-75%, and 75%-100%
respectively. Also changes the `rowsExpected` count used by the
SampleAggregator and sample size computation to be the estimated number
of stale rows computed as `estimated row count * fraction of stale
rows` for partial stat collections.

Closes: #127832

Release note (sql): Partial stats at extremes can now be collected on
all valid columns of a table using the `CREATE STATISTICS <stat_name>`
`FROM <table_name> USING EXTREMES` syntax, without an `ON <col_name>`
clause. Valid columns are all single column prefixes of a forward index
excluding partial, sharded, and implicitly partitioned indexes.

Author: Uzair5162

pkg/ccl/logictestccl/testdata/logic_test/multi_region_stats

@@ -125,3 +125,38 @@ WHERE stat->>'columns' = '["region"]'
 # Make sure we can still use the histogram in statistics_builder.
 statement ok
 INSERT INTO t124181 (region, a) VALUES ('ca-central-1', 2)
+
+# Verify that we don't collect partial stats on LOCALITY REGIONAL BY ROW table
+# indexes when columns are unspecified since they are implicitly partitioned
+# by the region column.
+statement ok
+CREATE TABLE multiregion_pstats (a INT PRIMARY KEY, b INT) LOCALITY REGIONAL BY ROW
+
+statement ok
+INSERT INTO multiregion_pstats (a, b, crdb_region) VALUES (1, 1, 'ap-southeast-2'), (2, 2, 'ca-central-1'), (3, 3, 'us-east-1')
+
+statement ok
+CREATE STATISTICS multiregion_full FROM multiregion_pstats
+
+statement ok
+INSERT INTO multiregion_pstats (a, b, crdb_region) VALUES (4, 4, 'ap-southeast-2'), (5, 5, 'ca-central-1')
+
+statement ok
+CREATE STATISTICS multiregion_partial FROM multiregion_pstats USING EXTREMES
+
+query TTIII colnames
+SELECT
+	statistics_name,
+	column_names,
+	row_count,
+	distinct_count,
+	null_count
+FROM
+	[SHOW STATISTICS FOR TABLE multiregion_pstats]
+ORDER BY statistics_name, column_names::STRING
+----
+statistics_name   column_names     row_count  distinct_count  null_count
+multiregion_full  {a,crdb_region}  3          3               0
+multiregion_full  {a}              3          3               0
+multiregion_full  {b}              3          3               0
+multiregion_full  {crdb_region}    3          3               0

pkg/sql/create_stats.go

@@ -38,6 +38,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/log/eventpb"
+	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/redact"
 )
@@ -77,7 +78,7 @@ func StubTableStats(
 ) ([]*stats.TableStatisticProto, error) {
 	colStats, err := createStatsDefaultColumns(
 		context.Background(), desc, false /* virtColEnabled */, false, /* multiColEnabled */
-		nonIndexColHistogramBuckets, nil, /* evalCtx */
+		false /* partialStats */, nonIndexColHistogramBuckets, nil, /* evalCtx */
 	)
 	if err != nil {
 		return nil, err
@@ -247,7 +248,13 @@ func (n *createStatsNode) makeJobRecord(ctx context.Context) (*jobs.Record, erro
 		}
 		defaultHistogramBuckets := stats.GetDefaultHistogramBuckets(n.p.ExecCfg().SV(), tableDesc)
 		if colStats, err = createStatsDefaultColumns(
-			ctx, tableDesc, virtColEnabled, multiColEnabled, defaultHistogramBuckets, n.p.EvalContext(),
+			ctx,
+			tableDesc,
+			virtColEnabled,
+			multiColEnabled,
+			n.Options.UsingExtremes,
+			defaultHistogramBuckets,
+			n.p.EvalContext(),
 		); err != nil {
 			return nil, err
 		}
@@ -364,13 +371,18 @@ const maxNonIndexCols = 100
 // predicate expressions are also likely to appear in query filters, so stats
 // are collected for those columns as well.
 //
+// If partialStats is true, we only collect statistics on single columns that
+// are prefixes of forward indexes, and skip over partial, sharded, and
+// implicitly partitioned indexes. Partial statistic creation only supports
+// these columns.
+//
 // In addition to the index columns, we collect stats on up to maxNonIndexCols
 // other columns from the table. We only collect histograms for index columns,
 // plus any other boolean or enum columns (where the "histogram" is tiny).
 func createStatsDefaultColumns(
 	ctx context.Context,
 	desc catalog.TableDescriptor,
-	virtColEnabled, multiColEnabled bool,
+	virtColEnabled, multiColEnabled, partialStats bool,
 	defaultHistogramBuckets uint32,
 	evalCtx *eval.Context,
 ) ([]jobspb.CreateStatsDetails_ColStat, error) {
@@ -470,6 +482,27 @@ func createStatsDefaultColumns(
 		return nil
 	}
 
+	// Only collect statistics on single columns that are prefixes of forward
+	// indexes for partial statistics, and skip over partial, sharded, and
+	// implicitly partitioned indexes.
+	if partialStats {
+		for _, idx := range desc.ActiveIndexes() {
+			if idx.GetType() != descpb.IndexDescriptor_FORWARD ||
+				idx.IsPartial() ||
+				idx.IsSharded() ||
+				idx.ImplicitPartitioningColumnCount() > 0 {
+				continue
+			}
+			if idx.NumKeyColumns() != 0 {
+				colID := idx.GetKeyColumnID(0)
+				if err := addIndexColumnStatsIfNotExists(colID, false /* isInverted */); err != nil {
+					return nil, err
+				}
+			}
+		}
+		return colStats, nil
+	}
+
 	// Add column stats for the primary key.
 	primaryIdx := desc.GetPrimaryIndex()
 	for i := 0; i < primaryIdx.NumKeyColumns(); i++ {
@@ -688,13 +721,35 @@ func (r *createStatsResumer) Resume(ctx context.Context, execCtx interface{}) er
 		}
 
 		dsp := innerP.DistSQLPlanner()
-		planCtx := dsp.NewPlanningCtx(ctx, innerEvalCtx, innerP, txn.KV(), FullDistribution)
 		// CREATE STATS flow doesn't produce any rows and only emits the
 		// metadata, so we can use a nil rowContainerHelper.
 		resultWriter := NewRowResultWriter(nil /* rowContainer */)
-		if err := dsp.planAndRunCreateStats(
-			ctx, innerEvalCtx, planCtx, innerP.SemaCtx(), txn.KV(), r.job, resultWriter,
-		); err != nil {
+
+		var err error
+		if details.UsingExtremes {
+			for i, colStat := range details.ColumnStats {
+				// Plan and run partial stats on multiple columns separately since each
+				// partial stat collection will use a different index and have different
+				// plans.
+				singleColDetails := protoutil.Clone(&details).(*jobspb.CreateStatsDetails)
+				singleColDetails.ColumnStats = []jobspb.CreateStatsDetails_ColStat{colStat}
+				planCtx := dsp.NewPlanningCtx(ctx, innerEvalCtx, innerP, txn.KV(), FullDistribution)
+				if err = dsp.planAndRunCreateStats(
+					ctx, innerEvalCtx, planCtx, innerP.SemaCtx(), txn.KV(), resultWriter, r.job.ID(), *singleColDetails,
+					len(details.ColumnStats), i,
+				); err != nil {
+					break
+				}
+			}
+		} else {
+			planCtx := dsp.NewPlanningCtx(ctx, innerEvalCtx, innerP, txn.KV(), FullDistribution)
+			err = dsp.planAndRunCreateStats(
+				ctx, innerEvalCtx, planCtx, innerP.SemaCtx(), txn.KV(), resultWriter, r.job.ID(), details,
+				1 /* numIndexes */, 0, /* curIndex */
+			)
+		}
+
+		if err != nil {
 			// Check if this was a context canceled error and restart if it was.
 			if grpcutil.IsContextCanceled(err) {
 				return jobs.MarkAsRetryJobError(err)

pkg/sql/distsql_physical_planner.go

@@ -3792,8 +3792,16 @@ func (dsp *DistSQLPlanner) createPhysPlanForPlanNode(
 			if err != nil {
 				return nil, err
 			}
+			details := record.Details.(jobspb.CreateStatsDetails)
+
+			numIndexes := 1
+			if details.UsingExtremes {
+				// Partial stats collections scan a different index for each column.
+				numIndexes = len(details.ColumnStats)
+			}
 			plan, err = dsp.createPlanForCreateStats(
-				ctx, planCtx, planCtx.planner.SemaCtx(), 0 /* jobID */, record.Details.(jobspb.CreateStatsDetails),
+				ctx, planCtx, planCtx.planner.SemaCtx(), 0 /* jobID */, details,
+				numIndexes, 0, /* curIndex */
 			)
 		}
 

pkg/sql/distsql_plan_stats.go

@@ -15,7 +15,6 @@ import (
 	"math"
 	"time"
 
-	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/settings"
@@ -160,6 +159,8 @@ func (dsp *DistSQLPlanner) createAndAttachSamplers(
 	jobID jobspb.JobID,
 	reqStats []requestedStat,
 	sketchSpec, invSketchSpec []execinfrapb.SketchSpec,
+	numIndexes int,
+	curIndex int,
 ) *PhysicalPlan {
 	// Estimate the expected number of rows based on existing stats in the cache.
 	var rowsExpected uint64
@@ -169,11 +170,20 @@ func (dsp *DistSQLPlanner) createAndAttachSamplers(
 			overhead = autoStatsFractionStaleRowsForTable
 		}
 		// Convert to a signed integer first to make the linter happy.
-		rowsExpected = uint64(int64(
-			// The total expected number of rows is the same number that was measured
-			// most recently, plus some overhead for possible insertions.
-			float64(tableStats[0].RowCount) * (1 + overhead),
-		))
+		if details.UsingExtremes {
+			rowsExpected = uint64(int64(
+				// The total expected number of rows is the estimated number of stale
+				// rows since we're only collecting stats on rows outside the bounds of
+				// the most recent statistic.
+				float64(tableStats[0].RowCount) * overhead,
+			))
+		} else {
+			rowsExpected = uint64(int64(
+				// The total expected number of rows is the same number that was measured
+				// most recently, plus some overhead for possible insertions.
+				float64(tableStats[0].RowCount) * (1 + overhead),
+			))
+		}
 	}
 
 	// Set up the samplers.
@@ -249,6 +259,8 @@ func (dsp *DistSQLPlanner) createAndAttachSamplers(
 		JobID:            jobID,
 		RowsExpected:     rowsExpected,
 		DeleteOtherStats: details.DeleteOtherStats,
+		NumIndexes:       uint64(numIndexes),
+		CurIndex:         uint64(curIndex),
 	}
 	// Plan the SampleAggregator on the gateway, unless we have a single Sampler.
 	node := dsp.gatewaySQLInstanceID
@@ -273,14 +285,13 @@ func (dsp *DistSQLPlanner) createPartialStatsPlan(
 	reqStats []requestedStat,
 	jobID jobspb.JobID,
 	details jobspb.CreateStatsDetails,
+	numIndexes int,
+	curIndex int,
 ) (*PhysicalPlan, error) {
-
-	// Currently, we limit the number of requests for partial statistics
-	// stats at a given point in time to 1.
-	// TODO (faizaanmadhani): Add support for multiple distinct requested
-	// partial stats in one job.
+	// Partial stats collections on multiple columns create different plans,
+	// so we only support one requested stat at a time here.
 	if len(reqStats) > 1 {
-		return nil, pgerror.Newf(pgcode.FeatureNotSupported, "cannot process multiple partial statistics at once")
+		return nil, errors.AssertionFailedf("only one partial statistic can be requested at a time")
 	}
 
 	reqStat := reqStats[0]
@@ -445,7 +456,8 @@ func (dsp *DistSQLPlanner) createPartialStatsPlan(
 		sampledColumnIDs,
 		jobID,
 		reqStats,
-		sketchSpec, invSketchSpec), nil
+		sketchSpec, invSketchSpec,
+		numIndexes, curIndex), nil
 }
 
 func (dsp *DistSQLPlanner) createStatsPlan(
@@ -456,6 +468,8 @@ func (dsp *DistSQLPlanner) createStatsPlan(
 	reqStats []requestedStat,
 	jobID jobspb.JobID,
 	details jobspb.CreateStatsDetails,
+	numIndexes int,
+	curIndex int,
 ) (*PhysicalPlan, error) {
 	if len(reqStats) == 0 {
 		return nil, errors.New("no stats requested")
@@ -681,7 +695,8 @@ func (dsp *DistSQLPlanner) createStatsPlan(
 		sampledColumnIDs,
 		jobID,
 		reqStats,
-		sketchSpecs, invSketchSpecs), nil
+		sketchSpecs, invSketchSpecs,
+		numIndexes, curIndex), nil
 }
 
 func (dsp *DistSQLPlanner) createPlanForCreateStats(
@@ -690,6 +705,8 @@ func (dsp *DistSQLPlanner) createPlanForCreateStats(
 	semaCtx *tree.SemaContext,
 	jobID jobspb.JobID,
 	details jobspb.CreateStatsDetails,
+	numIndexes int,
+	curIndex int,
 ) (*PhysicalPlan, error) {
 	reqStats := make([]requestedStat, len(details.ColumnStats))
 	histogramCollectionEnabled := stats.HistogramClusterMode.Get(&dsp.st.SV)
@@ -718,9 +735,9 @@ func (dsp *DistSQLPlanner) createPlanForCreateStats(
 	}
 
 	if details.UsingExtremes {
-		return dsp.createPartialStatsPlan(ctx, planCtx, tableDesc, reqStats, jobID, details)
+		return dsp.createPartialStatsPlan(ctx, planCtx, tableDesc, reqStats, jobID, details, numIndexes, curIndex)
 	}
-	return dsp.createStatsPlan(ctx, planCtx, semaCtx, tableDesc, reqStats, jobID, details)
+	return dsp.createStatsPlan(ctx, planCtx, semaCtx, tableDesc, reqStats, jobID, details, numIndexes, curIndex)
 }
 
 func (dsp *DistSQLPlanner) planAndRunCreateStats(
@@ -729,13 +746,15 @@ func (dsp *DistSQLPlanner) planAndRunCreateStats(
 	planCtx *PlanningCtx,
 	semaCtx *tree.SemaContext,
 	txn *kv.Txn,
-	job *jobs.Job,
 	resultWriter *RowResultWriter,
+	jobId jobspb.JobID,
+	details jobspb.CreateStatsDetails,
+	numIndexes int,
+	curIndex int,
 ) error {
 	ctx = logtags.AddTag(ctx, "create-stats-distsql", nil)
 
-	details := job.Details().(jobspb.CreateStatsDetails)
-	physPlan, err := dsp.createPlanForCreateStats(ctx, planCtx, semaCtx, job.ID(), details)
+	physPlan, err := dsp.createPlanForCreateStats(ctx, planCtx, semaCtx, jobId, details, numIndexes, curIndex)
 	if err != nil {
 		return err
 	}

pkg/sql/execinfrapb/processors_table_stats.proto

@@ -194,4 +194,12 @@ message SampleAggregatorSpec {
   // If true, calculate partial table statistics on the extreme values of
   // the previous full table stat.
   optional bool using_extremes = 11 [(gogoproto.nullable) = false];
+
+  // The number of indexes expected to be scanned as part of this stats
+  // collection. Used for progress reporting.
+  optional uint64 num_indexes = 12 [(gogoproto.nullable) = false];
+
+  // The number of indexes that have been scanned so far. Used for progress
+  // reporting.
+  optional uint64 cur_index = 13 [(gogoproto.nullable) = false];
 }