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 forwards indexes.

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.

Author: Uzair5162

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
@@ -249,7 +250,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
 		}
@@ -366,13 +373,17 @@ 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 forwards 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) {
@@ -472,6 +483,23 @@ func createStatsDefaultColumns(
 		return nil
 	}
 
+	// Only collect statistics on single columns that are prefixes of forward
+	// indexes for partial statistics.
+	if partialStats {
+		for _, idx := range desc.ActiveIndexes() {
+			if idx.GetType() != descpb.IndexDescriptor_FORWARD || idx.IsPartial() {
+				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++ {
@@ -690,13 +718,33 @@ 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 _, 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,
+				); 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,
+			)
+		}
+
+		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_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"
@@ -274,13 +273,10 @@ func (dsp *DistSQLPlanner) createPartialStatsPlan(
 	jobID jobspb.JobID,
 	details jobspb.CreateStatsDetails,
 ) (*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]
@@ -729,13 +725,13 @@ func (dsp *DistSQLPlanner) planAndRunCreateStats(
 	planCtx *PlanningCtx,
 	semaCtx *tree.SemaContext,
 	txn *kv.Txn,
-	job *jobs.Job,
 	resultWriter *RowResultWriter,
+	jobId jobspb.JobID,
+	details jobspb.CreateStatsDetails,
 ) 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)
 	if err != nil {
 		return err
 	}

pkg/sql/logictest/testdata/logic_test/distsql_stats

@@ -2503,9 +2503,6 @@ sdn              NULL                                               8          2
 sdnp             (a IS NULL) OR ((a < 0:::INT8) OR (a > 5:::INT8))  2          2
 
 # Verify errors.
-statement error pq: cannot process multiple partial statistics at once
-CREATE STATISTICS abcd_defaults FROM abcd USING EXTREMES;
-
 statement error pq: multi-column partial statistics are not currently supported
 CREATE STATISTICS abcd_a_b ON a, c FROM abcd USING EXTREMES;
 
@@ -3354,3 +3351,68 @@ upper_bound   range_rows  distinct_range_rows  equal_rows
 'paid'        0           0                    1
 'dispatched'  0           0                    1
 'delivered'   0           0                    1
+
+# Verify that partial stats are collected on single column prefixes of forward
+# indexes when columns are unspecified.
+statement ok
+SET enable_create_stats_using_extremes = on
+
+statement ok
+CREATE TABLE pstat_allindex (
+  a INT,
+  b INT,
+  c INT,
+  d INT,
+  j JSONB,
+  PRIMARY KEY (a),
+  INDEX(b, c),
+  INDEX(b, c, d),
+  INDEX(d),
+  INVERTED INDEX (j)
+);
+
+statement ok
+INSERT INTO pstat_allindex VALUES
+  (1, 1, 1, 1, '{"1": "1"}'),
+  (2, 2, 2, 2, '{"2": "2"}'),
+  (3, 3, 3, 3, '{"3": "3"}'),
+  (4, 4, 4, 4, '{"4": "4"}');
+
+statement ok
+CREATE STATISTICS pstat_allindex_full FROM pstat_allindex;
+
+statement ok
+INSERT INTO pstat_allindex VALUES
+  (5, 5, 5, 5, '{"5": "5"}'),
+  (6, 6, 6, 6, '{"6": "6"}'),
+  (7, 7, 7, 7, '{"7": "7"}'),
+  (8, 8, 8, 8, '{"8": "8"}');
+
+statement ok
+CREATE STATISTICS pstat_allindex_partial FROM pstat_allindex USING EXTREMES;
+
+query TTIII colnames
+SELECT
+	statistics_name,
+	column_names,
+	row_count,
+	distinct_count,
+	null_count
+FROM
+	[SHOW STATISTICS FOR TABLE pstat_allindex]
+ORDER BY statistics_name, column_names::STRING
+----
+statistics_name         column_names  row_count  distinct_count  null_count
+pstat_allindex_full     {a}           4          4               0
+pstat_allindex_full     {b,c,d}       4          4               0
+pstat_allindex_full     {b,c}         4          4               0
+pstat_allindex_full     {b}           4          4               0
+pstat_allindex_full     {c}           4          4               0
+pstat_allindex_full     {d}           4          4               0
+pstat_allindex_full     {j}           4          4               0
+pstat_allindex_partial  {a}           4          4               0
+pstat_allindex_partial  {b}           4          4               0
+pstat_allindex_partial  {d}           4          4               0
+
+statement ok
+RESET enable_create_stats_using_extremes