sql/schemachanger: execute validation operations in batches for indexes

Previously, when ALTER PRIMARY KEY was issued or multiple indexes were
created in a single transaction, we would end up recounting the table
for each index. This patch modifies the validation logic to batch counts
of indexes together on the same table.

Fixes: #156590

Release note (performance improvement): Optimize validation queries
during ALTER PRIMARY KEY to avoid counting the primary key multiple
times.

Author: fqazi

pkg/ccl/schemachangerccl/testdata/end_to_end/add_column_multiple_regional_by_row/add_column_multiple_regional_by_row.side_effects

@@ -1389,8 +1389,7 @@ update progress of schema change job #1: "Pending: Validating index (2 operation
 commit transaction #24
 begin transaction #25
 ## PostCommitPhase stage 23 of 23 with 2 ValidationType ops
-validate forward indexes [2] in table #108
-validate forward indexes [4] in table #108
+validate forward indexes [2 4] in table #108
 commit transaction #25
 begin transaction #26
 ## PostCommitNonRevertiblePhase stage 1 of 5 with 33 MutationType ops

pkg/ccl/schemachangerccl/testdata/end_to_end/alter_table_alter_primary_key_rbr/alter_table_alter_primary_key_rbr.side_effects

@@ -968,9 +968,8 @@ update progress of schema change job #1: "Pending: Validating NOT NULL constrain
 commit transaction #16
 begin transaction #17
 ## PostCommitPhase stage 15 of 24 with 3 ValidationType ops
-validate forward indexes [10] in table #108
+validate forward indexes [10 4] in table #108
 validate CHECK constraint crdb_internal_k_shard_16_auto_not_null in table #108
-validate forward indexes [4] in table #108
 commit transaction #17
 begin transaction #18
 ## PostCommitPhase stage 16 of 24 with 27 MutationType ops

pkg/sql/schemachanger/scexec/BUILD.bazel

@@ -39,6 +39,7 @@ go_library(
         "//pkg/sql/sem/catid",
         "//pkg/sql/sem/idxtype",
         "//pkg/sql/sessiondata",
+        "//pkg/util/buildutil",
         "//pkg/util/ctxgroup",
         "//pkg/util/hlc",
         "//pkg/util/intsets",

pkg/sql/schemachanger/scexec/exec_validation.go

@@ -9,15 +9,32 @@ import (
 	"context"
 
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scerrors"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scop"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/idxtype"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
+	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/errors"
 )
 
-func executeValidateIndex(ctx context.Context, deps Dependencies, op *scop.ValidateIndex) error {
-	descs, err := deps.Catalog().MustReadImmutableDescriptors(ctx, op.TableID)
+func executeValidateIndexes(
+	ctx context.Context, deps Dependencies, ops []*scop.ValidateIndex,
+) error {
+	// Nothing to validate.
+	if len(ops) == 0 {
+		return nil
+	}
+	// Validate all validation operations are against the same table.
+	if buildutil.CrdbTestBuild {
+		for _, op := range ops {
+			if op.TableID != ops[0].TableID {
+				return errors.AssertionFailedf("all validation operations must be against the same table")
+			}
+		}
+	}
+
+	descs, err := deps.Catalog().MustReadImmutableDescriptors(ctx, ops[0].TableID)
 	if err != nil {
 		return err
 	}
@@ -26,25 +43,35 @@ func executeValidateIndex(ctx context.Context, deps Dependencies, op *scop.Valid
 	if !ok {
 		return catalog.WrapTableDescRefErr(desc.GetID(), catalog.NewDescriptorTypeError(desc))
 	}
-	index, err := catalog.MustFindIndexByID(table, op.IndexID)
-	if err != nil {
-		return err
+	indexTypes := make(map[idxtype.T][]catalog.Index)
+	for _, validateIndex := range ops {
+		index, err := catalog.MustFindIndexByID(table, validateIndex.IndexID)
+		if err != nil {
+			return err
+		}
+		indexTypes[index.GetType()] = append(indexTypes[index.GetType()], index)
 	}
 	// Execute the validation operation as a node user.
 	execOverride := sessiondata.NodeUserSessionDataOverride
-	switch index.GetType() {
-	case idxtype.FORWARD:
-		err = deps.Validator().ValidateForwardIndexes(ctx, deps.TransactionalJobRegistry().CurrentJob(), table, []catalog.Index{index}, execOverride)
-	case idxtype.INVERTED:
-		err = deps.Validator().ValidateInvertedIndexes(ctx, deps.TransactionalJobRegistry().CurrentJob(), table, []catalog.Index{index}, execOverride)
-	case idxtype.VECTOR:
-		// TODO(drewk): consider whether we can perform useful validation for
-		// vector indexes.
-	default:
-		return errors.AssertionFailedf("unexpected index type %v", index.GetType())
-	}
-	if err != nil {
-		return scerrors.SchemaChangerUserError(err)
+	// Execute each type of index together, so that the table counts are only
+	// fetched once.
+	for typ, indexes := range indexTypes {
+		var err error
+		switch typ {
+		case idxtype.FORWARD:
+			err = deps.Validator().ValidateForwardIndexes(ctx, deps.TransactionalJobRegistry().CurrentJob(), table, indexes, execOverride)
+		case idxtype.INVERTED:
+			err = deps.Validator().ValidateInvertedIndexes(ctx, deps.TransactionalJobRegistry().CurrentJob(), table, indexes, execOverride)
+		case idxtype.VECTOR:
+			// TODO(drewk): consider whether we can perform useful validation for
+			// vector indexes.
+			continue
+		default:
+			return errors.AssertionFailedf("unexpected index type %v", typ)
+		}
+		if err != nil {
+			return scerrors.SchemaChangerUserError(err)
+		}
 	}
 	return nil
 }
@@ -105,40 +132,64 @@ func executeValidateColumnNotNull(
 }
 
 func executeValidationOps(ctx context.Context, deps Dependencies, ops []scop.Op) (err error) {
+	v := makeValidationAccumulator(ops)
+	return v.validate(ctx, deps)
+}
+
+// validationAccumulator batches validation operations on a per-table basis, so
+// that index validations can be batched together.
+type validationAccumulator struct {
+	indexes     map[descpb.ID][]*scop.ValidateIndex
+	constraints map[descpb.ID][]*scop.ValidateConstraint
+	notNulls    map[descpb.ID][]*scop.ValidateColumnNotNull
+}
+
+// makeValidationAccumulator creates a validationAccumulator from a list of
+// validation operations.
+func makeValidationAccumulator(ops []scop.Op) validationAccumulator {
+	v := validationAccumulator{
+		indexes:     make(map[descpb.ID][]*scop.ValidateIndex),
+		constraints: make(map[descpb.ID][]*scop.ValidateConstraint),
+		notNulls:    make(map[descpb.ID][]*scop.ValidateColumnNotNull),
+	}
 	for _, op := range ops {
-		if err = executeValidationOp(ctx, deps, op); err != nil {
-			return err
+		switch op := op.(type) {
+		case *scop.ValidateIndex:
+			v.indexes[op.TableID] = append(v.indexes[op.TableID], op)
+		case *scop.ValidateConstraint:
+			v.constraints[op.TableID] = append(v.constraints[op.TableID], op)
+		case *scop.ValidateColumnNotNull:
+			v.notNulls[op.TableID] = append(v.notNulls[op.TableID], op)
+		default:
+			panic("unimplemented")
 		}
 	}
-	return nil
+	return v
 }
 
-func executeValidationOp(ctx context.Context, deps Dependencies, op scop.Op) (err error) {
-	switch op := op.(type) {
-	case *scop.ValidateIndex:
-		if err = executeValidateIndex(ctx, deps, op); err != nil {
-			if !scerrors.HasSchemaChangerUserError(err) {
-				return errors.Wrapf(err, "%T: %v", op, op)
-			}
-			return err
+// validate executes all validation operations in the accumulator.
+func (v validationAccumulator) validate(ctx context.Context, deps Dependencies) error {
+	// Batch all index operations per table for efficient execution.
+	for _, tableIndexes := range v.indexes {
+		if err := executeValidateIndexes(ctx, deps, tableIndexes); err != nil {
+			// Error is wrapped with the operation inside the validation method.
+			return errors.Wrapf(err, "%T: %v", tableIndexes, tableIndexes)
 		}
-	case *scop.ValidateConstraint:
-		if err = executeValidateConstraint(ctx, deps, op); err != nil {
-			if !scerrors.HasSchemaChangerUserError(err) {
-				return errors.Wrapf(err, "%T: %v", op, op)
+	}
+	// These operations do not support any type of batching.
+	for _, tableNotNulls := range v.notNulls {
+		for _, notNull := range tableNotNulls {
+			if err := executeValidateColumnNotNull(ctx, deps, notNull); err != nil {
+				return errors.Wrapf(err, "%T: %v", notNull, notNull)
 			}
-			return err
 		}
-	case *scop.ValidateColumnNotNull:
-		if err = executeValidateColumnNotNull(ctx, deps, op); err != nil {
-			if !scerrors.HasSchemaChangerUserError(err) {
-				return errors.Wrapf(err, "%T: %v", op, op)
+	}
+	for _, tableConstraints := range v.constraints {
+		for _, constraint := range tableConstraints {
+			if err := executeValidateConstraint(ctx, deps, constraint); err != nil {
+				return errors.Wrapf(err, "%T: %v", constraint, constraint)
 			}
-			return err
 		}
-
-	default:
-		panic("unimplemented")
 	}
 	return nil
 }

pkg/sql/schemachanger/testdata/end_to_end/alter_table_add_primary_key_drop_rowid_with_secondary_idx/alter_table_add_primary_key_drop_rowid_with_secondary_idx.side_effects

@@ -887,9 +887,7 @@ update progress of schema change job #1: "Pending: Validating index (3 operation
 commit transaction #8
 begin transaction #9
 ## PostCommitPhase stage 7 of 16 with 3 ValidationType ops
-validate forward indexes [10] in table #104
-validate forward indexes [6] in table #104
-validate forward indexes [8] in table #104
+validate forward indexes [10 6 8] in table #104
 commit transaction #9
 begin transaction #10
 ## PostCommitPhase stage 8 of 16 with 23 MutationType ops

pkg/sql/schemachanger/testdata/end_to_end/alter_table_alter_primary_key_drop_rowid_with_idx/alter_table_alter_primary_key_drop_rowid_with_idx.side_effects

@@ -814,8 +814,7 @@ update progress of schema change job #1: "Pending: Validating index (2 operation
 commit transaction #16
 begin transaction #17
 ## PostCommitPhase stage 15 of 23 with 2 ValidationType ops
-validate forward indexes [8] in table #104
-validate forward indexes [4] in table #104
+validate forward indexes [8 4] in table #104
 commit transaction #17
 begin transaction #18
 ## PostCommitPhase stage 16 of 23 with 17 MutationType ops

pkg/sql/schemachanger/testdata/end_to_end/alter_table_alter_primary_key_using_hash/alter_table_alter_primary_key_using_hash.side_effects

@@ -616,9 +616,8 @@ update progress of schema change job #1: "Pending: Validating NOT NULL constrain
 commit transaction #8
 begin transaction #9
 ## PostCommitPhase stage 7 of 16 with 3 ValidationType ops
-validate forward indexes [8] in table #104
+validate forward indexes [8 4] in table #104
 validate CHECK constraint crdb_internal_j_shard_3_auto_not_null in table #104
-validate forward indexes [4] in table #104
 commit transaction #9
 begin transaction #10
 ## PostCommitPhase stage 8 of 16 with 20 MutationType ops

pkg/sql/schemachanger/testdata/end_to_end/alter_table_alter_primary_key_vanilla/alter_table_alter_primary_key_vanilla.side_effects

@@ -739,10 +739,7 @@ update progress of schema change job #1: "Pending: Validating index (4 operation
 commit transaction #8
 begin transaction #9
 ## PostCommitPhase stage 7 of 15 with 4 ValidationType ops
-validate forward indexes [16] in table #104
-validate forward indexes [8] in table #104
-validate forward indexes [10] in table #104
-validate forward indexes [12] in table #104
+validate forward indexes [16 8 10 12] in table #104
 commit transaction #9
 begin transaction #10
 ## PostCommitPhase stage 8 of 15 with 23 MutationType ops