From a24d6c215506da67243c454c7645c366c029c4f4 Mon Sep 17 00:00:00 2001 From: Marius Posta Date: Wed, 21 Apr 2021 12:28:30 -0400 Subject: [PATCH] sql: add catalog.Partitioning interface This commit aims at reducing direct usage of descpb.PartitioningDescriptor types and instead using the new catalog.Partitioning interface. This refactoring is in line with recent virtualization work tracked under #56306. In an effort to improve readability and ability to reason about state changes, this commit changes the type signature of sql.CreatePartitioning and, notably, removes its side-effects on the index descriptor. The latter is now modified using a new tabledesc.UpdateIndexPartitioning function. Release note: None --- pkg/ccl/partitionccl/BUILD.bazel | 1 + pkg/ccl/partitionccl/partition.go | 159 +++++++++----------- pkg/ccl/partitionccl/partition_test.go | 11 +- pkg/sql/alter_database.go | 4 +- pkg/sql/alter_index.go | 46 +++--- pkg/sql/alter_primary_key.go | 14 +- pkg/sql/alter_table.go | 41 ++--- pkg/sql/backfill.go | 4 +- pkg/sql/catalog/descpb/structured.go | 33 ---- pkg/sql/catalog/table_elements.go | 53 ++++++- pkg/sql/catalog/tabledesc/index.go | 142 +++++++++++++++-- pkg/sql/catalog/tabledesc/index_test.go | 4 - pkg/sql/catalog/tabledesc/structured.go | 16 -- pkg/sql/catalog/tabledesc/table_desc.go | 35 +++++ pkg/sql/catalog/tabledesc/validate.go | 93 ++++++------ pkg/sql/crdb_internal.go | 50 +++--- pkg/sql/create_index.go | 8 +- pkg/sql/create_table.go | 72 +++++---- pkg/sql/database_region_change_finalizer.go | 19 +-- pkg/sql/opt_catalog.go | 37 ++--- pkg/sql/partition.go | 91 ++++++----- pkg/sql/partition_utils.go | 54 ++++--- pkg/sql/rowenc/partition.go | 9 +- pkg/sql/set_zone_config.go | 5 +- pkg/sql/show_create.go | 4 +- pkg/sql/show_create_clauses.go | 53 ++++--- pkg/sql/zone_config.go | 24 +-- 27 files changed, 625 insertions(+), 457 deletions(-) diff --git a/pkg/ccl/partitionccl/BUILD.bazel b/pkg/ccl/partitionccl/BUILD.bazel index 0ff586a09bac..5e1fbed90c13 100644 --- a/pkg/ccl/partitionccl/BUILD.bazel +++ b/pkg/ccl/partitionccl/BUILD.bazel @@ -22,6 +22,7 @@ go_library( "//pkg/sql/types", "//pkg/util/encoding", "//pkg/util/errorutil/unimplemented", + "//pkg/util/log", "@com_github_cockroachdb_errors//:errors", ], ) diff --git a/pkg/ccl/partitionccl/partition.go b/pkg/ccl/partitionccl/partition.go index 6c1899917fb7..c38b1abd45e4 100644 --- a/pkg/ccl/partitionccl/partition.go +++ b/pkg/ccl/partitionccl/partition.go @@ -28,6 +28,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/types" "github.com/cockroachdb/cockroach/pkg/util/encoding" "github.com/cockroachdb/cockroach/pkg/util/errorutil/unimplemented" + "github.com/cockroachdb/cockroach/pkg/util/log" "github.com/cockroachdb/errors" ) @@ -149,7 +150,7 @@ func createPartitioningImpl( ctx context.Context, evalCtx *tree.EvalContext, tableDesc *tabledesc.Mutable, - indexDesc *descpb.IndexDescriptor, + newIdxColumnNames []string, partBy *tree.PartitionBy, allowedNewColumnNames []tree.Name, numImplicitColumns int, @@ -166,7 +167,7 @@ func createPartitioningImpl( // We don't have the fields for our parent partitions handy, but we can use // the names from the index we're partitioning. They must have matched or we // would have already returned an error. - partCols := append([]string(nil), indexDesc.ColumnNames[:colOffset]...) + partCols := append([]string(nil), newIdxColumnNames[:colOffset]...) for _, p := range partBy.Fields { partCols = append(partCols, string(p)) } @@ -175,16 +176,16 @@ func createPartitioningImpl( var cols []catalog.Column for i := 0; i < len(partBy.Fields); i++ { - if colOffset+i >= len(indexDesc.ColumnNames) { + if colOffset+i >= len(newIdxColumnNames) { return partDesc, pgerror.Newf(pgcode.Syntax, "declared partition columns (%s) exceed the number of columns in index being partitioned (%s)", - partitioningString(), strings.Join(indexDesc.ColumnNames, ", ")) + partitioningString(), strings.Join(newIdxColumnNames, ", ")) } // Search by name because some callsites of this method have not // allocated ids yet (so they are still all the 0 value). col, err := findColumnByNameOnTable( tableDesc, - tree.Name(indexDesc.ColumnNames[colOffset+i]), + tree.Name(newIdxColumnNames[colOffset+i]), allowedNewColumnNames, ) if err != nil { @@ -197,7 +198,7 @@ func createPartitioningImpl( n := colOffset + i + 1 return partDesc, pgerror.Newf(pgcode.Syntax, "declared partition columns (%s) do not match first %d columns in index being partitioned (%s)", - partitioningString(), n, strings.Join(indexDesc.ColumnNames[:n], ", ")) + partitioningString(), n, strings.Join(newIdxColumnNames[:n], ", ")) } } @@ -225,7 +226,7 @@ func createPartitioningImpl( ctx, evalCtx, tableDesc, - indexDesc, + newIdxColumnNames, l.Subpartition, allowedNewColumnNames, 0, /* implicitColumnNames */ @@ -263,25 +264,19 @@ func createPartitioningImpl( return partDesc, nil } -// detectImplicitPartitionColumns detects implicit partitioning columns -// and returns a new index descriptor with the implicit columns modified -// on the index descriptor and the number of implicit columns prepended. -func detectImplicitPartitionColumns( - evalCtx *tree.EvalContext, +// collectImplicitPartitionColumns collects implicit partitioning columns. +func collectImplicitPartitionColumns( tableDesc *tabledesc.Mutable, - indexDesc descpb.IndexDescriptor, + indexFirstColumnName string, partBy *tree.PartitionBy, allowedNewColumnNames []tree.Name, -) (descpb.IndexDescriptor, int, error) { +) (implicitCols []catalog.Column, _ error) { seenImplicitColumnNames := map[string]struct{}{} - var implicitColumnIDs []descpb.ColumnID - var implicitColumns []string - var implicitColumnDirections []descpb.IndexDescriptor_Direction // Iterate over each field in the PARTITION BY until it matches the start // of the actual explicitly indexed columns. for _, field := range partBy.Fields { // As soon as the fields match, we have no implicit columns to add. - if string(field) == indexDesc.ColumnNames[0] { + if string(field) == indexFirstColumnName { break } @@ -291,27 +286,20 @@ func detectImplicitPartitionColumns( allowedNewColumnNames, ) if err != nil { - return indexDesc, 0, err + return nil, err } if _, ok := seenImplicitColumnNames[col.GetName()]; ok { - return indexDesc, 0, pgerror.Newf( + return nil, pgerror.Newf( pgcode.InvalidObjectDefinition, `found multiple definitions in partition using column "%s"`, col.GetName(), ) } seenImplicitColumnNames[col.GetName()] = struct{}{} - implicitColumns = append(implicitColumns, col.GetName()) - implicitColumnIDs = append(implicitColumnIDs, col.GetID()) - implicitColumnDirections = append(implicitColumnDirections, descpb.IndexDescriptor_ASC) + implicitCols = append(implicitCols, col) } - if len(implicitColumns) > 0 { - indexDesc.ColumnNames = append(implicitColumns, indexDesc.ColumnNames...) - indexDesc.ColumnIDs = append(implicitColumnIDs, indexDesc.ColumnIDs...) - indexDesc.ColumnDirections = append(implicitColumnDirections, indexDesc.ColumnDirections...) - } - return indexDesc, len(implicitColumns), nil + return implicitCols, nil } // findColumnByNameOnTable finds the given column from the table. @@ -348,38 +336,37 @@ func createPartitioning( partBy *tree.PartitionBy, allowedNewColumnNames []tree.Name, allowImplicitPartitioning bool, -) (descpb.IndexDescriptor, error) { +) (newImplicitCols []catalog.Column, newPartitioning descpb.PartitioningDescriptor, err error) { org := sql.ClusterOrganization.Get(&st.SV) if err := utilccl.CheckEnterpriseEnabled(st, evalCtx.ClusterID, org, "partitions"); err != nil { - return indexDesc, err + return nil, newPartitioning, err } - // Truncate existing implicitly partitioned columns. + // Truncate existing implicitly partitioned column names. oldNumImplicitColumns := int(indexDesc.Partitioning.NumImplicitColumns) - oldImplicitColumnIDs := indexDesc.ColumnIDs[:oldNumImplicitColumns] + newIdxColumnNames := indexDesc.ColumnNames[oldNumImplicitColumns:] - indexDesc.ColumnIDs = indexDesc.ColumnIDs[oldNumImplicitColumns:] - indexDesc.ColumnNames = indexDesc.ColumnNames[oldNumImplicitColumns:] - indexDesc.ColumnDirections = indexDesc.ColumnDirections[oldNumImplicitColumns:] - - var numImplicitColumns int - var err error if allowImplicitPartitioning { - indexDesc, numImplicitColumns, err = detectImplicitPartitionColumns( - evalCtx, + newImplicitCols, err = collectImplicitPartitionColumns( tableDesc, - indexDesc, + newIdxColumnNames[0], partBy, allowedNewColumnNames, ) if err != nil { - return indexDesc, err + return nil, newPartitioning, err } - if numImplicitColumns > 0 { - if err := checkClusterSupportsImplicitPartitioning(evalCtx); err != nil { - return indexDesc, err - } + } + if len(newImplicitCols) > 0 { + if err := checkClusterSupportsImplicitPartitioning(evalCtx); err != nil { + return nil, newPartitioning, err } + // Prepend with new implicit column names. + newIdxColumnNames = make([]string, len(newImplicitCols), len(newImplicitCols)+len(newIdxColumnNames)) + for i, col := range newImplicitCols { + newIdxColumnNames[i] = col.GetName() + } + newIdxColumnNames = append(newIdxColumnNames, indexDesc.ColumnNames[oldNumImplicitColumns:]...) } // If we had implicit column partitioning beforehand, check we have the @@ -387,35 +374,34 @@ func createPartitioning( // Having different implicitly partitioned columns requires rewrites, // which is outside the scope of createPartitioning. if oldNumImplicitColumns > 0 { - if numImplicitColumns != oldNumImplicitColumns { - return indexDesc, errors.AssertionFailedf( + if len(newImplicitCols) != oldNumImplicitColumns { + return nil, newPartitioning, errors.AssertionFailedf( "mismatching number of implicit columns: old %d vs new %d", oldNumImplicitColumns, - numImplicitColumns, + len(newImplicitCols), ) } - for i, oldColID := range oldImplicitColumnIDs { - if oldColID != indexDesc.ColumnIDs[i] { - return indexDesc, errors.AssertionFailedf("found new implicit partitioning at index %d", i) + for i, col := range newImplicitCols { + if indexDesc.ColumnIDs[i] != col.GetID() { + return nil, newPartitioning, errors.AssertionFailedf("found new implicit partitioning at column ordinal %d", i) } } } - partitioning, err := createPartitioningImpl( + newPartitioning, err = createPartitioningImpl( ctx, evalCtx, tableDesc, - &indexDesc, + newIdxColumnNames, partBy, allowedNewColumnNames, - numImplicitColumns, + len(newImplicitCols), 0, /* colOffset */ ) if err != nil { - return indexDesc, err + return nil, descpb.PartitioningDescriptor{}, err } - indexDesc.Partitioning = partitioning - return indexDesc, err + return newImplicitCols, newPartitioning, err } // selectPartitionExprs constructs an expression for selecting all rows in the @@ -434,7 +420,7 @@ func selectPartitionExprs( AddMutations: true, }, func(idx catalog.Index) error { return selectPartitionExprsByName( - a, evalCtx, tableDesc, idx, &idx.IndexDesc().Partitioning, prefixDatums, exprsByPartName, true /* genExpr */) + a, evalCtx, tableDesc, idx, idx.GetPartitioning(), prefixDatums, exprsByPartName, true /* genExpr */) }); err != nil { return nil, err } @@ -488,38 +474,37 @@ func selectPartitionExprsByName( evalCtx *tree.EvalContext, tableDesc catalog.TableDescriptor, idx catalog.Index, - partDesc *descpb.PartitioningDescriptor, + part catalog.Partitioning, prefixDatums tree.Datums, exprsByPartName map[string]tree.TypedExpr, genExpr bool, ) error { - if partDesc.NumColumns == 0 { + if part.NumColumns() == 0 { return nil } // Setting genExpr to false skips the expression generation and only // registers each descendent partition in the map with a placeholder entry. if !genExpr { - for _, l := range partDesc.List { - exprsByPartName[l.Name] = tree.DBoolFalse + err := part.ForEachList(func(name string, _ [][]byte, subPartitioning catalog.Partitioning) error { + exprsByPartName[name] = tree.DBoolFalse var fakeDatums tree.Datums - if err := selectPartitionExprsByName( - a, evalCtx, tableDesc, idx, &l.Subpartitioning, fakeDatums, exprsByPartName, genExpr, - ); err != nil { - return err - } - } - for _, r := range partDesc.Range { - exprsByPartName[r.Name] = tree.DBoolFalse + return selectPartitionExprsByName(a, evalCtx, tableDesc, idx, subPartitioning, fakeDatums, exprsByPartName, genExpr) + }) + if err != nil { + return err } - return nil + return part.ForEachRange(func(name string, _, _ []byte) error { + exprsByPartName[name] = tree.DBoolFalse + return nil + }) } var colVars tree.Exprs { // The recursive calls of selectPartitionExprsByName don't pass though // the column ordinal references, so reconstruct them here. - colVars = make(tree.Exprs, len(prefixDatums)+int(partDesc.NumColumns)) + colVars = make(tree.Exprs, len(prefixDatums)+part.NumColumns()) for i := range colVars { col, err := tabledesc.FindPublicColumnWithID(tableDesc, idx.GetColumnID(i)) if err != nil { @@ -529,7 +514,7 @@ func selectPartitionExprsByName( } } - if len(partDesc.List) > 0 { + if part.NumLists() > 0 { type exprAndPartName struct { expr tree.TypedExpr name string @@ -539,12 +524,11 @@ func selectPartitionExprsByName( // `(1, 2)`, the expr for the former must exclude the latter. This is // done by bucketing the expression for each partition value by the // number of DEFAULTs it involves. - partValueExprs := make([][]exprAndPartName, int(partDesc.NumColumns)+1) + partValueExprs := make([][]exprAndPartName, part.NumColumns()+1) - for _, l := range partDesc.List { - for _, valueEncBuf := range l.Values { - t, _, err := rowenc.DecodePartitionTuple( - a, evalCtx.Codec, tableDesc, idx, partDesc, valueEncBuf, prefixDatums) + err := part.ForEachList(func(name string, values [][]byte, subPartitioning catalog.Partitioning) error { + for _, valueEncBuf := range values { + t, _, err := rowenc.DecodePartitionTuple(a, evalCtx.Codec, tableDesc, idx, part, valueEncBuf, prefixDatums) if err != nil { return err } @@ -562,11 +546,11 @@ func selectPartitionExprsByName( tree.NewDTuple(tupleTyp, allDatums...)) partValueExprs[len(t.Datums)] = append(partValueExprs[len(t.Datums)], exprAndPartName{ expr: partValueExpr, - name: l.Name, + name: name, }) genExpr := true - if _, ok := exprsByPartName[l.Name]; ok { + if _, ok := exprsByPartName[name]; ok { // Presence of a partition name in the exprsByPartName map // means the caller has expressed an interested in this // partition, which means any subpartitions can be skipped @@ -578,11 +562,15 @@ func selectPartitionExprsByName( genExpr = false } if err := selectPartitionExprsByName( - a, evalCtx, tableDesc, idx, &l.Subpartitioning, allDatums, exprsByPartName, genExpr, + a, evalCtx, tableDesc, idx, subPartitioning, allDatums, exprsByPartName, genExpr, ); err != nil { return err } } + return nil + }) + if err != nil { + return err } // Walk backward through partValueExprs, so partition values with fewest @@ -617,10 +605,9 @@ func selectPartitionExprsByName( } } - for range partDesc.Range { - return errors.New("TODO(dan): unsupported for range partitionings") + if part.NumRanges() > 0 { + log.Fatal(evalCtx.Context, "TODO(dan): unsupported for range partitionings") } - return nil } diff --git a/pkg/ccl/partitionccl/partition_test.go b/pkg/ccl/partitionccl/partition_test.go index 66745c7d1da9..0eb60760a44f 100644 --- a/pkg/ccl/partitionccl/partition_test.go +++ b/pkg/ccl/partitionccl/partition_test.go @@ -1342,12 +1342,12 @@ func TestRepartitioning(t *testing.T) { } else { fmt.Fprintf(&repartition, `ALTER INDEX %s@%s `, test.new.parsed.tableName, testIndex.GetName()) } - if testIndex.GetPartitioning().NumColumns == 0 { + if testIndex.GetPartitioning().NumColumns() == 0 { repartition.WriteString(`PARTITION BY NOTHING`) } else { if err := sql.ShowCreatePartitioning( &rowenc.DatumAlloc{}, keys.SystemSQLCodec, test.new.parsed.tableDesc, testIndex, - &testIndex.IndexDesc().Partitioning, &repartition, 0 /* indent */, 0, /* colOffset */ + testIndex.GetPartitioning(), &repartition, 0 /* indent */, 0, /* colOffset */ ); err != nil { t.Fatalf("%+v", err) } @@ -1357,8 +1357,11 @@ func TestRepartitioning(t *testing.T) { // Verify that repartitioning removes zone configs for partitions that // have been removed. newPartitionNames := map[string]struct{}{} - for _, name := range test.new.parsed.tableDesc.PartitionNames() { - newPartitionNames[name] = struct{}{} + for _, index := range test.new.parsed.tableDesc.NonDropIndexes() { + _ = index.GetPartitioning().ForEachPartitionName(func(name string) error { + newPartitionNames[name] = struct{}{} + return nil + }) } for _, row := range sqlDB.QueryStr( t, "SELECT partition_name FROM crdb_internal.zones WHERE partition_name IS NOT NULL") { diff --git a/pkg/sql/alter_database.go b/pkg/sql/alter_database.go index 1e885a4a40da..a1be78652cbc 100644 --- a/pkg/sql/alter_database.go +++ b/pkg/sql/alter_database.go @@ -758,7 +758,7 @@ func addDefaultLocalityConfigToAllTables( func checkCanConvertTableToMultiRegion( dbDesc catalog.DatabaseDescriptor, tableDesc catalog.TableDescriptor, ) error { - if tableDesc.GetPrimaryIndex().GetPartitioning().NumColumns > 0 { + if tableDesc.GetPrimaryIndex().GetPartitioning().NumColumns() > 0 { return errors.WithDetailf( pgerror.Newf( pgcode.ObjectNotInPrerequisiteState, @@ -770,7 +770,7 @@ func checkCanConvertTableToMultiRegion( ) } for _, idx := range tableDesc.AllIndexes() { - if idx.GetPartitioning().NumColumns > 0 { + if idx.GetPartitioning().NumColumns() > 0 { return errors.WithDetailf( pgerror.Newf( pgcode.ObjectNotInPrerequisiteState, diff --git a/pkg/sql/alter_index.go b/pkg/sql/alter_index.go index 869358cd1633..3a2c9a481a51 100644 --- a/pkg/sql/alter_index.go +++ b/pkg/sql/alter_index.go @@ -79,8 +79,7 @@ func (n *alterIndexNode) startExec(params runParams) error { "cannot change the partitioning of an index if the table has PARTITION ALL BY defined", ) } - existingIndexDesc := n.index.IndexDescDeepCopy() - if existingIndexDesc.Partitioning.NumImplicitColumns > 0 { + if n.index.GetPartitioning().NumImplicitColumns() > 0 { return unimplemented.New( "ALTER INDEX PARTITION BY", "cannot ALTER INDEX PARTITION BY on an index which already has implicit column partitioning", @@ -88,12 +87,13 @@ func (n *alterIndexNode) startExec(params runParams) error { } allowImplicitPartitioning := params.p.EvalContext().SessionData.ImplicitColumnPartitioningEnabled || n.tableDesc.IsLocalityRegionalByRow() - newIndexDesc, err := CreatePartitioning( + alteredIndexDesc := n.index.IndexDescDeepCopy() + newImplicitCols, newPartitioning, err := CreatePartitioning( params.ctx, params.extendedEvalCtx.Settings, params.EvalContext(), n.tableDesc, - existingIndexDesc, + alteredIndexDesc, t.PartitionBy, nil, /* allowedNewColumnNames */ allowImplicitPartitioning, @@ -101,31 +101,39 @@ func (n *alterIndexNode) startExec(params runParams) error { if err != nil { return err } - if newIndexDesc.Partitioning.NumImplicitColumns > 0 { + if newPartitioning.NumImplicitColumns > 0 { return unimplemented.New( "ALTER INDEX PARTITION BY", "cannot ALTER INDEX and change the partitioning to contain implicit columns", ) } - descriptorChanged = !existingIndexDesc.Equal(&newIndexDesc) - if err = deleteRemovedPartitionZoneConfigs( - params.ctx, - params.p.txn, - n.tableDesc, - n.index.GetID(), - &existingIndexDesc.Partitioning, - &newIndexDesc.Partitioning, - params.extendedEvalCtx.ExecCfg, - ); err != nil { - return err + isIndexAltered := tabledesc.UpdateIndexPartitioning(&alteredIndexDesc, newImplicitCols, newPartitioning) + if isIndexAltered { + oldPartitioning := n.index.GetPartitioning().DeepCopy() + if n.index.Primary() { + n.tableDesc.SetPrimaryIndex(alteredIndexDesc) + } else { + n.tableDesc.SetPublicNonPrimaryIndex(n.index.Ordinal(), alteredIndexDesc) + } + n.index = n.tableDesc.ActiveIndexes()[n.index.Ordinal()] + descriptorChanged = true + if err := deleteRemovedPartitionZoneConfigs( + params.ctx, + params.p.txn, + n.tableDesc, + n.index.GetID(), + oldPartitioning, + n.index.GetPartitioning(), + params.extendedEvalCtx.ExecCfg, + ); err != nil { + return err + } } - // Update value to make sure it doesn't become stale. - n.index = n.tableDesc.AllIndexes()[n.index.Ordinal()] - *n.index.IndexDesc() = newIndexDesc default: return errors.AssertionFailedf( "unsupported alter command: %T", cmd) } + } if err := n.tableDesc.AllocateIDs(params.ctx); err != nil { diff --git a/pkg/sql/alter_primary_key.go b/pkg/sql/alter_primary_key.go index 287ea290d315..fae693878d57 100644 --- a/pkg/sql/alter_primary_key.go +++ b/pkg/sql/alter_primary_key.go @@ -366,7 +366,7 @@ func (p *planner) AlterPrimaryKey( } if partitionAllBy != nil { - *newPrimaryIndexDesc, err = CreatePartitioning( + newImplicitCols, newPartitioning, err := CreatePartitioning( ctx, p.ExecCfg().Settings, p.EvalContext(), @@ -379,6 +379,7 @@ func (p *planner) AlterPrimaryKey( if err != nil { return err } + tabledesc.UpdateIndexPartitioning(newPrimaryIndexDesc, newImplicitCols, newPartitioning) } // Create a new index that indexes everything the old primary index @@ -461,10 +462,7 @@ func (p *planner) AlterPrimaryKey( // Drop any PARTITION ALL BY clause. if dropPartitionAllBy { - newIndex.ColumnNames = newIndex.ColumnNames[newIndex.Partitioning.NumImplicitColumns:] - newIndex.ColumnIDs = newIndex.ColumnIDs[newIndex.Partitioning.NumImplicitColumns:] - newIndex.ColumnDirections = newIndex.ColumnDirections[newIndex.Partitioning.NumImplicitColumns:] - newIndex.Partitioning = descpb.PartitioningDescriptor{} + tabledesc.UpdateIndexPartitioning(&newIndex, nil /* newImplicitCols */, descpb.PartitioningDescriptor{}) } newIndex.Name = tabledesc.GenerateUniqueConstraintName(basename, nameExists) @@ -482,7 +480,7 @@ func (p *planner) AlterPrimaryKey( } // Create partitioning if we are newly adding a PARTITION BY ALL statement. if isNewPartitionAllBy { - if newIndex, err = CreatePartitioning( + newImplicitCols, newPartitioning, err := CreatePartitioning( ctx, p.ExecCfg().Settings, p.EvalContext(), @@ -491,9 +489,11 @@ func (p *planner) AlterPrimaryKey( partitionAllBy, allowedNewColumnNames, allowImplicitPartitioning, - ); err != nil { + ) + if err != nil { return err } + tabledesc.UpdateIndexPartitioning(&newIndex, newImplicitCols, newPartitioning) } oldIndexIDs = append(oldIndexIDs, idx.GetID()) newIndexIDs = append(newIndexIDs, newIndex.ID) diff --git a/pkg/sql/alter_table.go b/pkg/sql/alter_table.go index f34bd3894379..0293d31a8207 100644 --- a/pkg/sql/alter_table.go +++ b/pkg/sql/alter_table.go @@ -842,18 +842,19 @@ func (n *alterTableNode) startExec(params runParams) error { if n.tableDesc.IsPartitionAllBy() { return unimplemented.NewWithIssue(58736, "changing partition of table with PARTITION ALL BY not yet implemented") } - oldPartitioning := n.tableDesc.GetPrimaryIndex().GetPartitioning() - if oldPartitioning.NumImplicitColumns > 0 { + oldPartitioning := n.tableDesc.GetPrimaryIndex().GetPartitioning().DeepCopy() + if oldPartitioning.NumImplicitColumns() > 0 { return unimplemented.NewWithIssue( 58731, "cannot ALTER TABLE PARTITION BY on a table which already has implicit column partitioning", ) } - newPrimaryIndex, err := CreatePartitioning( + newPrimaryIndexDesc := n.tableDesc.GetPrimaryIndex().IndexDescDeepCopy() + newImplicitCols, newPartitioning, err := CreatePartitioning( params.ctx, params.p.ExecCfg().Settings, params.EvalContext(), n.tableDesc, - *n.tableDesc.GetPrimaryIndex().IndexDesc(), + newPrimaryIndexDesc, t.PartitionBy, nil, /* allowedNewColumnNames */ params.p.EvalContext().SessionData.ImplicitColumnPartitioningEnabled || @@ -862,27 +863,27 @@ func (n *alterTableNode) startExec(params runParams) error { if err != nil { return err } - if newPrimaryIndex.Partitioning.NumImplicitColumns > 0 { + if newPartitioning.NumImplicitColumns > 0 { return unimplemented.NewWithIssue( 58731, "cannot ALTER TABLE and change the partitioning to contain implicit columns", ) } - descriptorChanged = descriptorChanged || !n.tableDesc.GetPrimaryIndex().IndexDesc().Equal(&newPrimaryIndex) - err = deleteRemovedPartitionZoneConfigs( - params.ctx, - params.p.txn, - n.tableDesc, - n.tableDesc.GetPrimaryIndexID(), - &oldPartitioning, - &newPrimaryIndex.Partitioning, - params.extendedEvalCtx.ExecCfg, - ) - if err != nil { - return err - } - { - n.tableDesc.SetPrimaryIndex(newPrimaryIndex) + isIndexAltered := tabledesc.UpdateIndexPartitioning(&newPrimaryIndexDesc, newImplicitCols, newPartitioning) + if isIndexAltered { + n.tableDesc.SetPrimaryIndex(newPrimaryIndexDesc) + descriptorChanged = true + if err := deleteRemovedPartitionZoneConfigs( + params.ctx, + params.p.txn, + n.tableDesc, + n.tableDesc.GetPrimaryIndexID(), + oldPartitioning, + n.tableDesc.GetPrimaryIndex().GetPartitioning(), + params.extendedEvalCtx.ExecCfg, + ); err != nil { + return err + } } case *tree.AlterTableSetAudit: diff --git a/pkg/sql/backfill.go b/pkg/sql/backfill.go index 3a79f0bdd4d8..1fde23862106 100644 --- a/pkg/sql/backfill.go +++ b/pkg/sql/backfill.go @@ -1781,12 +1781,12 @@ func ValidateForwardIndexes( // For implicitly partitioned unique indexes, we need to independently // validate that the non-implicitly partitioned columns are unique. - if idx.IsUnique() && idx.GetPartitioning().NumImplicitColumns > 0 && !skipUniquenessChecks { + if idx.IsUnique() && idx.GetPartitioning().NumImplicitColumns() > 0 && !skipUniquenessChecks { if err := validateUniqueConstraint( ctx, tableDesc, idx.GetName(), - idx.IndexDesc().ColumnIDs[idx.GetPartitioning().NumImplicitColumns:], + idx.IndexDesc().ColumnIDs[idx.GetPartitioning().NumImplicitColumns():], idx.GetPredicate(), ie, txn, diff --git a/pkg/sql/catalog/descpb/structured.go b/pkg/sql/catalog/descpb/structured.go index cce468088201..0674acf76004 100644 --- a/pkg/sql/catalog/descpb/structured.go +++ b/pkg/sql/catalog/descpb/structured.go @@ -186,39 +186,6 @@ func (f ForeignKeyReference) IsSet() bool { return f.Table != 0 } -// FindPartitionByName searches this partitioning descriptor for a partition -// whose name is the input and returns it, or nil if no match is found. -func (desc *PartitioningDescriptor) FindPartitionByName(name string) *PartitioningDescriptor { - for _, l := range desc.List { - if l.Name == name { - return desc - } - if s := l.Subpartitioning.FindPartitionByName(name); s != nil { - return s - } - } - for _, r := range desc.Range { - if r.Name == name { - return desc - } - } - return nil -} - -// PartitionNames returns a slice containing the name of every partition and -// subpartition in an arbitrary order. -func (desc *PartitioningDescriptor) PartitionNames() []string { - var names []string - for _, l := range desc.List { - names = append(names, l.Name) - names = append(names, l.Subpartitioning.PartitionNames()...) - } - for _, r := range desc.Range { - names = append(names, r.Name) - } - return names -} - // Public implements the Descriptor interface. func (desc *TableDescriptor) Public() bool { return desc.State == DescriptorState_PUBLIC diff --git a/pkg/sql/catalog/table_elements.go b/pkg/sql/catalog/table_elements.go index 7b78e80de649..99f4007fbd02 100644 --- a/pkg/sql/catalog/table_elements.go +++ b/pkg/sql/catalog/table_elements.go @@ -137,9 +137,7 @@ type Index interface { IsValidOriginIndex(originColIDs descpb.ColumnIDs) bool IsValidReferencedUniqueConstraint(referencedColIDs descpb.ColumnIDs) bool - GetPartitioning() descpb.PartitioningDescriptor - FindPartitionByName(name string) descpb.PartitioningDescriptor - PartitionNames() []string + GetPartitioning() Partitioning ExplicitColumnStartIdx() int @@ -364,6 +362,55 @@ type MaterializedViewRefresh interface { TableWithNewIndexes(tbl TableDescriptor) TableDescriptor } +// Partitioning is an interface around an index partitioning. +type Partitioning interface { + + // PartitioningDesc returns the underlying protobuf descriptor. + PartitioningDesc() *descpb.PartitioningDescriptor + + // DeepCopy returns a deep copy of the receiver. + DeepCopy() Partitioning + + // FindPartitionByName recursively searches the partitioning for a partition + // whose name matches the input and returns it, or nil if no match is found. + FindPartitionByName(name string) Partitioning + + // ForEachPartitionName applies fn on each of the partition names in this + // partition and recursively in its subpartitions. + // Supports iterutil.Done. + ForEachPartitionName(fn func(name string) error) error + + // ForEachList applies fn on each list element of the wrapped partitioning. + // Supports iterutil.Done. + ForEachList(fn func(name string, values [][]byte, subPartitioning Partitioning) error) error + + // ForEachRange applies fn on each range element of the wrapped partitioning. + // Supports iterutil.Done. + ForEachRange(fn func(name string, from, to []byte) error) error + + // NumColumns is how large of a prefix of the columns in an index are used in + // the function mapping column values to partitions. If this is a + // subpartition, this is offset to start from the end of the parent + // partition's columns. If NumColumns is 0, then there is no partitioning. + NumColumns() int + + // NumImplicitColumns specifies the number of columns that implicitly prefix a + // given index. This occurs if a user specifies a PARTITION BY which is not a + // prefix of the given index, in which case the ColumnIDs are added in front + // of the index and this field denotes the number of columns added as a + // prefix. + // If NumImplicitColumns is 0, no implicit columns are defined for the index. + NumImplicitColumns() int + + // NumLists returns the number of list elements in the underlying partitioning + // descriptor. + NumLists() int + + // NumRanges returns the number of range elements in the underlying + // partitioning descriptor. + NumRanges() int +} + func isIndexInSearchSet(desc TableDescriptor, opts IndexOpts, idx Index) bool { if !opts.NonPhysicalPrimaryIndex && idx.Primary() && !desc.IsPhysicalTable() { return false diff --git a/pkg/sql/catalog/tabledesc/index.go b/pkg/sql/catalog/tabledesc/index.go index 4bb84e95de8b..262df80ef662 100644 --- a/pkg/sql/catalog/tabledesc/index.go +++ b/pkg/sql/catalog/tabledesc/index.go @@ -14,6 +14,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/geo/geoindex" "github.com/cockroachdb/cockroach/pkg/sql/catalog" "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb" + "github.com/cockroachdb/cockroach/pkg/util/iterutil" "github.com/cockroachdb/cockroach/pkg/util/protoutil" ) @@ -112,20 +113,8 @@ func (w index) GetType() descpb.IndexDescriptor_Type { } // GetPartitioning returns the partitioning descriptor of the index. -func (w index) GetPartitioning() descpb.PartitioningDescriptor { - return w.desc.Partitioning -} - -// FindPartitionByName searches the index's partitioning descriptor for a -// partition whose name is the input and returns it, or nil if no match is found. -func (w index) FindPartitionByName(name string) descpb.PartitioningDescriptor { - return *w.desc.Partitioning.FindPartitionByName(name) -} - -// PartitionNames returns a slice containing the name of every partition and -// subpartition in an arbitrary order. -func (w index) PartitionNames() []string { - return w.desc.Partitioning.PartitionNames() +func (w index) GetPartitioning() catalog.Partitioning { + return &partitioning{desc: &w.desc.Partitioning} } // ExplicitColumnStartIdx returns the first index in which the column is @@ -296,6 +285,131 @@ func (w index) GetCompositeColumnID(compositeColumnOrdinal int) descpb.ColumnID return w.desc.CompositeColumnIDs[compositeColumnOrdinal] } +// partitioning is the backing struct for a catalog.Partitioning interface. +type partitioning struct { + desc *descpb.PartitioningDescriptor +} + +// PartitioningDesc returns the underlying protobuf descriptor. +func (p partitioning) PartitioningDesc() *descpb.PartitioningDescriptor { + return p.desc +} + +// DeepCopy returns a deep copy of the receiver. +func (p partitioning) DeepCopy() catalog.Partitioning { + return &partitioning{desc: protoutil.Clone(p.desc).(*descpb.PartitioningDescriptor)} +} + +// FindPartitionByName recursively searches the partitioning for a partition +// whose name matches the input and returns it, or nil if no match is found. +func (p partitioning) FindPartitionByName(name string) (found catalog.Partitioning) { + _ = p.forEachPartitionName(func(partitioning catalog.Partitioning, currentName string) error { + if name == currentName { + found = partitioning + return iterutil.StopIteration() + } + return nil + }) + return found +} + +// ForEachPartitionName applies fn on each of the partition names in this +// partition and recursively in its subpartitions. +// Supports iterutil.Done. +func (p partitioning) ForEachPartitionName(fn func(name string) error) error { + err := p.forEachPartitionName(func(_ catalog.Partitioning, name string) error { + return fn(name) + }) + if iterutil.Done(err) { + return nil + } + return err +} + +func (p partitioning) forEachPartitionName( + fn func(partitioning catalog.Partitioning, name string) error, +) error { + for _, l := range p.desc.List { + err := fn(p, l.Name) + if err != nil { + return err + } + err = partitioning{desc: &l.Subpartitioning}.forEachPartitionName(fn) + if err != nil { + return err + } + } + for _, r := range p.desc.Range { + err := fn(p, r.Name) + if err != nil { + return err + } + } + return nil +} + +// NumLists returns the number of list elements in the underlying partitioning +// descriptor. +func (p partitioning) NumLists() int { + return len(p.desc.List) +} + +// NumRanges returns the number of range elements in the underlying +// partitioning descriptor. +func (p partitioning) NumRanges() int { + return len(p.desc.Range) +} + +// ForEachList applies fn on each list element of the wrapped partitioning. +// Supports iterutil.Done. +func (p partitioning) ForEachList( + fn func(name string, values [][]byte, subPartitioning catalog.Partitioning) error, +) error { + for _, l := range p.desc.List { + subp := partitioning{desc: &l.Subpartitioning} + err := fn(l.Name, l.Values, subp) + if err != nil { + if iterutil.Done(err) { + return nil + } + return err + } + } + return nil +} + +// ForEachRange applies fn on each range element of the wrapped partitioning. +// Supports iterutil.Done. +func (p partitioning) ForEachRange(fn func(name string, from, to []byte) error) error { + for _, r := range p.desc.Range { + err := fn(r.Name, r.FromInclusive, r.ToExclusive) + if err != nil { + if iterutil.Done(err) { + return nil + } + return err + } + } + return nil +} + +// NumColumns is how large of a prefix of the columns in an index are used in +// the function mapping column values to partitions. If this is a +// subpartition, this is offset to start from the end of the parent +// partition's columns. If NumColumns is 0, then there is no partitioning. +func (p partitioning) NumColumns() int { + return int(p.desc.NumColumns) +} + +// NumImplicitColumns specifies the number of columns that implicitly prefix a +// given index. This occurs if a user specifies a PARTITION BY which is not a +// prefix of the given index, in which case the ColumnIDs are added in front of +// the index and this field denotes the number of columns added as a prefix. +// If NumImplicitColumns is 0, no implicit columns are defined for the index. +func (p partitioning) NumImplicitColumns() int { + return int(p.desc.NumImplicitColumns) +} + // indexCache contains precomputed slices of catalog.Index interfaces. type indexCache struct { primary catalog.Index diff --git a/pkg/sql/catalog/tabledesc/index_test.go b/pkg/sql/catalog/tabledesc/index_test.go index 810d5902a49a..784ad28fd26c 100644 --- a/pkg/sql/catalog/tabledesc/index_test.go +++ b/pkg/sql/catalog/tabledesc/index_test.go @@ -240,10 +240,6 @@ func TestIndexInterface(t *testing.T) { errMsgFmt, "IsCreatedExplicitly", idx.GetName()) require.Equal(t, descpb.IndexDescriptorVersion(0x2), idx.GetVersion(), errMsgFmt, "GetVersion", idx.GetName()) - require.Equal(t, descpb.PartitioningDescriptor{}, idx.GetPartitioning(), - errMsgFmt, "GetPartitioning", idx.GetName()) - require.Equal(t, []string(nil), idx.PartitionNames(), - errMsgFmt, "PartitionNames", idx.GetName()) require.Equal(t, 0, idx.NumInterleaveAncestors(), errMsgFmt, "NumInterleaveAncestors", idx.GetName()) require.Equal(t, 0, idx.NumInterleavedBy(), diff --git a/pkg/sql/catalog/tabledesc/structured.go b/pkg/sql/catalog/tabledesc/structured.go index 08b6cfddfd5f..84d60bc500c1 100644 --- a/pkg/sql/catalog/tabledesc/structured.go +++ b/pkg/sql/catalog/tabledesc/structured.go @@ -85,12 +85,6 @@ type PostDeserializationTableDescriptorChanges struct { UpgradedForeignKeyRepresentation bool } -// FindIndexPartitionByName searches this index descriptor for a partition whose name -// is the input and returns it, or nil if no match is found. -func FindIndexPartitionByName(idx catalog.Index, name string) *descpb.PartitioningDescriptor { - return idx.IndexDesc().Partitioning.FindPartitionByName(name) -} - // DescriptorType returns the type of this descriptor. func (desc *wrapper) DescriptorType() catalog.DescriptorType { return catalog.Table @@ -2193,16 +2187,6 @@ func (desc *wrapper) GetFamilyOfColumn( return nil, errors.Newf("no column family found for column id %v", colID) } -// PartitionNames returns a slice containing the name of every partition and -// subpartition in an arbitrary order. -func (desc *wrapper) PartitionNames() []string { - var names []string - for _, index := range desc.NonDropIndexes() { - names = append(names, index.PartitionNames()...) - } - return names -} - // SetAuditMode configures the audit mode on the descriptor. func (desc *Mutable) SetAuditMode(mode tree.AuditMode) (bool, error) { prev := desc.AuditMode diff --git a/pkg/sql/catalog/tabledesc/table_desc.go b/pkg/sql/catalog/tabledesc/table_desc.go index 9e158cee8f8b..51135095ab62 100644 --- a/pkg/sql/catalog/tabledesc/table_desc.go +++ b/pkg/sql/catalog/tabledesc/table_desc.go @@ -154,6 +154,41 @@ func (desc *Mutable) SetPublicNonPrimaryIndex(indexOrdinal int, index descpb.Ind desc.Indexes[indexOrdinal-1] = index } +// UpdateIndexPartitioning applies the new partition and adjusts the column info +// for the specified index descriptor. Returns false iff this was a no-op. +func UpdateIndexPartitioning( + idx *descpb.IndexDescriptor, + newImplicitCols []catalog.Column, + newPartitioning descpb.PartitioningDescriptor, +) bool { + oldNumImplicitCols := int(idx.Partitioning.NumImplicitColumns) + isNoOp := oldNumImplicitCols == len(newImplicitCols) && idx.Partitioning.Equal(newPartitioning) + numCols := len(idx.ColumnIDs) + newCap := numCols + len(newImplicitCols) - oldNumImplicitCols + newColumnIDs := make([]descpb.ColumnID, len(newImplicitCols), newCap) + newColumnNames := make([]string, len(newImplicitCols), newCap) + newColumnDirections := make([]descpb.IndexDescriptor_Direction, len(newImplicitCols), newCap) + for i, col := range newImplicitCols { + newColumnIDs[i] = col.GetID() + newColumnNames[i] = col.GetName() + newColumnDirections[i] = descpb.IndexDescriptor_ASC + if isNoOp && + (idx.ColumnIDs[i] != newColumnIDs[i] || + idx.ColumnNames[i] != newColumnNames[i] || + idx.ColumnDirections[i] != newColumnDirections[i]) { + isNoOp = false + } + } + if isNoOp { + return false + } + idx.ColumnIDs = append(newColumnIDs, idx.ColumnIDs[oldNumImplicitCols:]...) + idx.ColumnNames = append(newColumnNames, idx.ColumnNames[oldNumImplicitCols:]...) + idx.ColumnDirections = append(newColumnDirections, idx.ColumnDirections[oldNumImplicitCols:]...) + idx.Partitioning = newPartitioning + return true +} + // GetPrimaryIndex returns the primary index in the form of a catalog.Index // interface. func (desc *wrapper) GetPrimaryIndex() catalog.Index { diff --git a/pkg/sql/catalog/tabledesc/validate.go b/pkg/sql/catalog/tabledesc/validate.go index 446af5299e81..75bff0b65063 100644 --- a/pkg/sql/catalog/tabledesc/validate.go +++ b/pkg/sql/catalog/tabledesc/validate.go @@ -407,7 +407,7 @@ func (desc *wrapper) validateInboundFK( func (desc *wrapper) matchingPartitionbyAll(indexI catalog.Index) bool { primaryIndexPartitioning := desc.PrimaryIndex.ColumnIDs[:desc.PrimaryIndex.Partitioning.NumColumns] - indexPartitioning := indexI.IndexDesc().ColumnIDs[:indexI.GetPartitioning().NumColumns] + indexPartitioning := indexI.IndexDesc().ColumnIDs[:indexI.GetPartitioning().NumColumns()] if len(primaryIndexPartitioning) != len(indexPartitioning) { return false } @@ -1032,18 +1032,18 @@ func (desc *wrapper) ensureShardedIndexNotComputed(index *descpb.IndexDescriptor func (desc *wrapper) validatePartitioningDescriptor( a *rowenc.DatumAlloc, idx catalog.Index, - partDesc *descpb.PartitioningDescriptor, + part catalog.Partitioning, colOffset int, partitionNames map[string]string, ) error { - if partDesc.NumImplicitColumns > partDesc.NumColumns { + if part.NumImplicitColumns() > part.NumColumns() { return errors.Newf( "cannot have implicit partitioning columns (%d) > partitioning columns (%d)", - partDesc.NumImplicitColumns, - partDesc.NumColumns, + part.NumImplicitColumns(), + part.NumColumns(), ) } - if partDesc.NumColumns == 0 { + if part.NumColumns() == 0 { return nil } @@ -1066,10 +1066,10 @@ func (desc *wrapper) validatePartitioningDescriptor( fakePrefixDatums[i] = tree.DNull } - if len(partDesc.List) == 0 && len(partDesc.Range) == 0 { + if part.NumLists() == 0 && part.NumRanges() == 0 { return fmt.Errorf("at least one of LIST or RANGE partitioning must be used") } - if len(partDesc.List) > 0 && len(partDesc.Range) > 0 { + if part.NumLists() > 0 && part.NumRanges() > 0 { return fmt.Errorf("only one LIST or RANGE partitioning may used") } @@ -1077,8 +1077,7 @@ func (desc *wrapper) validatePartitioningDescriptor( // This should only happen at read time and descriptors should not become // invalid at read time, only at write time. { - numColumns := int(partDesc.NumColumns) - for i := colOffset; i < colOffset+numColumns; i++ { + for i := colOffset; i < colOffset+part.NumColumns(); i++ { // The partitioning descriptor may be invalid and refer to columns // not stored in the index. In that case, skip this check as the // validation will fail later. @@ -1114,23 +1113,23 @@ func (desc *wrapper) validatePartitioningDescriptor( // so it's fine to ignore the tenant ID prefix. codec := keys.SystemSQLCodec - if len(partDesc.List) > 0 { - listValues := make(map[string]struct{}, len(partDesc.List)) - for _, p := range partDesc.List { - if err := checkName(p.Name); err != nil { + if part.NumLists() > 0 { + listValues := make(map[string]struct{}, part.NumLists()) + err := part.ForEachList(func(name string, values [][]byte, subPartitioning catalog.Partitioning) error { + if err := checkName(name); err != nil { return err } - if len(p.Values) == 0 { - return fmt.Errorf("PARTITION %s: must contain values", p.Name) + if len(values) == 0 { + return fmt.Errorf("PARTITION %s: must contain values", name) } // NB: key encoding is used to check uniqueness because it has // to match the behavior of the value when indexed. - for _, valueEncBuf := range p.Values { + for _, valueEncBuf := range values { tuple, keyPrefix, err := rowenc.DecodePartitionTuple( - a, codec, desc, idx, partDesc, valueEncBuf, fakePrefixDatums) + a, codec, desc, idx, part, valueEncBuf, fakePrefixDatums) if err != nil { - return fmt.Errorf("PARTITION %s: %v", p.Name, err) + return fmt.Errorf("PARTITION %s: %v", name, err) } if _, exists := listValues[string(keyPrefix)]; exists { return fmt.Errorf("%s cannot be present in more than one partition", tuple) @@ -1138,48 +1137,53 @@ func (desc *wrapper) validatePartitioningDescriptor( listValues[string(keyPrefix)] = struct{}{} } - newColOffset := colOffset + int(partDesc.NumColumns) - if err := desc.validatePartitioningDescriptor( - a, idx, &p.Subpartitioning, newColOffset, partitionNames, - ); err != nil { - return err - } + newColOffset := colOffset + part.NumColumns() + return desc.validatePartitioningDescriptor( + a, idx, subPartitioning, newColOffset, partitionNames, + ) + }) + if err != nil { + return err } } - if len(partDesc.Range) > 0 { + if part.NumRanges() > 0 { tree := interval.NewTree(interval.ExclusiveOverlapper) - for _, p := range partDesc.Range { - if err := checkName(p.Name); err != nil { + err := part.ForEachRange(func(name string, from, to []byte) error { + if err := checkName(name); err != nil { return err } // NB: key encoding is used to check uniqueness because it has to match // the behavior of the value when indexed. fromDatums, fromKey, err := rowenc.DecodePartitionTuple( - a, codec, desc, idx, partDesc, p.FromInclusive, fakePrefixDatums) + a, codec, desc, idx, part, from, fakePrefixDatums) if err != nil { - return fmt.Errorf("PARTITION %s: %v", p.Name, err) + return fmt.Errorf("PARTITION %s: %v", name, err) } toDatums, toKey, err := rowenc.DecodePartitionTuple( - a, codec, desc, idx, partDesc, p.ToExclusive, fakePrefixDatums) + a, codec, desc, idx, part, to, fakePrefixDatums) if err != nil { - return fmt.Errorf("PARTITION %s: %v", p.Name, err) + return fmt.Errorf("PARTITION %s: %v", name, err) } - pi := partitionInterval{p.Name, fromKey, toKey} + pi := partitionInterval{name, fromKey, toKey} if overlaps := tree.Get(pi.Range()); len(overlaps) > 0 { return fmt.Errorf("partitions %s and %s overlap", - overlaps[0].(partitionInterval).name, p.Name) + overlaps[0].(partitionInterval).name, name) } if err := tree.Insert(pi, false /* fast */); errors.Is(err, interval.ErrEmptyRange) { return fmt.Errorf("PARTITION %s: empty range: lower bound %s is equal to upper bound %s", - p.Name, fromDatums, toDatums) + name, fromDatums, toDatums) } else if errors.Is(err, interval.ErrInvertedRange) { return fmt.Errorf("PARTITION %s: empty range: lower bound %s is greater than upper bound %s", - p.Name, fromDatums, toDatums) + name, fromDatums, toDatums) } else if err != nil { - return errors.Wrapf(err, "PARTITION %s", p.Name) + return errors.Wrapf(err, "PARTITION %s", name) } + return nil + }) + if err != nil { + return err } } @@ -1210,7 +1214,7 @@ func (desc *wrapper) validatePartitioning() error { a := &rowenc.DatumAlloc{} return catalog.ForEachNonDropIndex(desc, func(idx catalog.Index) error { return desc.validatePartitioningDescriptor( - a, idx, &idx.IndexDesc().Partitioning, 0 /* colOffset */, partitionNames, + a, idx, idx.GetPartitioning(), 0 /* colOffset */, partitionNames, ) }) } @@ -1347,22 +1351,27 @@ func (desc *wrapper) validateTableLocalityConfig( transitioningRegionNames[region] = struct{}{} } - for _, partitioning := range desc.GetPrimaryIndex().GetPartitioning().List { - regionName := descpb.RegionName(partitioning.Name) + part := desc.GetPrimaryIndex().GetPartitioning() + err = part.ForEachList(func(name string, _ [][]byte, _ catalog.Partitioning) error { + regionName := descpb.RegionName(name) // Any transitioning region names may exist. if _, ok := transitioningRegionNames[regionName]; ok { - continue + return nil } // If a region is not found in any of the region names, we have an unknown // partition. if _, ok := regionNames[regionName]; !ok { return errors.AssertionFailedf( "unknown partition %s on PRIMARY INDEX of table %s", - partitioning.Name, + name, desc.GetName(), ) } delete(regionNames, regionName) + return nil + }) + if err != nil { + return err } // Any regions that are not deleted from the above loop is missing. diff --git a/pkg/sql/crdb_internal.go b/pkg/sql/crdb_internal.go index c08b19e425e5..5f73f31db827 100644 --- a/pkg/sql/crdb_internal.go +++ b/pkg/sql/crdb_internal.go @@ -2077,7 +2077,7 @@ CREATE TABLE crdb_internal.create_statements ( createNofk = stmt } hasPartitions := nil != catalog.FindIndex(table, catalog.IndexOpts{}, func(idx catalog.Index) bool { - return idx.GetPartitioning().NumColumns != 0 + return idx.GetPartitioning().NumColumns() != 0 }) return addRow( dbDescID, @@ -3576,7 +3576,7 @@ func addPartitioningRows( database string, table catalog.TableDescriptor, index catalog.Index, - partitioning *descpb.PartitioningDescriptor, + partitioning catalog.Partitioning, parentName tree.Datum, colOffset int, addRow func(...tree.Datum) error, @@ -3591,10 +3591,10 @@ func addPartitioningRows( tableID := tree.NewDInt(tree.DInt(table.GetID())) indexID := tree.NewDInt(tree.DInt(index.GetID())) - numColumns := tree.NewDInt(tree.DInt(partitioning.NumColumns)) + numColumns := tree.NewDInt(tree.DInt(partitioning.NumColumns())) var buf bytes.Buffer - for i := uint32(colOffset); i < uint32(colOffset)+partitioning.NumColumns; i++ { + for i := uint32(colOffset); i < uint32(colOffset+partitioning.NumColumns()); i++ { if i != uint32(colOffset) { buf.WriteString(`, `) } @@ -3612,9 +3612,9 @@ func addPartitioningRows( } // This produces the list_value column. - for _, l := range partitioning.List { + err := partitioning.ForEachList(func(name string, values [][]byte, subPartitioning catalog.Partitioning) error { var buf bytes.Buffer - for j, values := range l.Values { + for j, values := range values { if j != 0 { buf.WriteString(`, `) } @@ -3628,11 +3628,11 @@ func addPartitioningRows( } partitionValue := tree.NewDString(buf.String()) - name := tree.NewDString(l.Name) + nameDString := tree.NewDString(name) // Figure out which zone and subzone this partition should correspond to. zoneID, zone, subzone, err := GetZoneConfigInTxn( - ctx, p.txn, config.SystemTenantObjectID(table.GetID()), index, l.Name, false /* getInheritedDefault */) + ctx, p.txn, config.SystemTenantObjectID(table.GetID()), index, name, false /* getInheritedDefault */) if err != nil { return err } @@ -3649,7 +3649,7 @@ func addPartitioningRows( tableID, indexID, parentName, - name, + nameDString, numColumns, colNames, partitionValue, @@ -3659,18 +3659,18 @@ func addPartitioningRows( ); err != nil { return err } - err = addPartitioningRows(ctx, p, database, table, index, &l.Subpartitioning, name, - colOffset+int(partitioning.NumColumns), addRow) - if err != nil { - return err - } + return addPartitioningRows(ctx, p, database, table, index, subPartitioning, nameDString, + colOffset+partitioning.NumColumns(), addRow) + }) + if err != nil { + return err } // This produces the range_value column. - for _, r := range partitioning.Range { + err = partitioning.ForEachRange(func(name string, from, to []byte) error { var buf bytes.Buffer fromTuple, _, err := rowenc.DecodePartitionTuple( - &datumAlloc, p.ExecCfg().Codec, table, index, partitioning, r.FromInclusive, fakePrefixDatums, + &datumAlloc, p.ExecCfg().Codec, table, index, partitioning, from, fakePrefixDatums, ) if err != nil { return err @@ -3678,7 +3678,7 @@ func addPartitioningRows( buf.WriteString(fromTuple.String()) buf.WriteString(" TO ") toTuple, _, err := rowenc.DecodePartitionTuple( - &datumAlloc, p.ExecCfg().Codec, table, index, partitioning, r.ToExclusive, fakePrefixDatums, + &datumAlloc, p.ExecCfg().Codec, table, index, partitioning, to, fakePrefixDatums, ) if err != nil { return err @@ -3688,7 +3688,7 @@ func addPartitioningRows( // Figure out which zone and subzone this partition should correspond to. zoneID, zone, subzone, err := GetZoneConfigInTxn( - ctx, p.txn, config.SystemTenantObjectID(table.GetID()), index, r.Name, false /* getInheritedDefault */) + ctx, p.txn, config.SystemTenantObjectID(table.GetID()), index, name, false /* getInheritedDefault */) if err != nil { return err } @@ -3701,23 +3701,21 @@ func addPartitioningRows( } } - if err := addRow( + return addRow( tableID, indexID, parentName, - tree.NewDString(r.Name), + tree.NewDString(name), numColumns, colNames, tree.DNull, /* null value for partition list */ partitionRange, tree.NewDInt(tree.DInt(zoneID)), tree.NewDInt(tree.DInt(subzoneID)), - ); err != nil { - return err - } - } + ) + }) - return nil + return err } // crdbInternalPartitionsTable decodes and exposes the partitions of each @@ -3751,7 +3749,7 @@ CREATE TABLE crdb_internal.partitions ( return catalog.ForEachIndex(table, catalog.IndexOpts{ AddMutations: true, }, func(index catalog.Index) error { - return addPartitioningRows(ctx, p, dbName, table, index, &index.IndexDesc().Partitioning, + return addPartitioningRows(ctx, p, dbName, table, index, index.GetPartitioning(), tree.DNull /* parentName */, 0 /* colOffset */, pusher.pushRow) }) }) diff --git a/pkg/sql/create_index.go b/pkg/sql/create_index.go index 8b065e8d9b35..9d8bf3ebec2e 100644 --- a/pkg/sql/create_index.go +++ b/pkg/sql/create_index.go @@ -452,7 +452,7 @@ func (n *createIndexNode) startExec(params runParams) error { // Avoid the warning if we have PARTITION ALL BY as all indexes will implicitly // have relevant partitioning columns prepended at the front. if n.n.PartitionByIndex == nil && - n.tableDesc.GetPrimaryIndex().GetPartitioning().NumColumns > 0 && + n.tableDesc.GetPrimaryIndex().GetPartitioning().NumColumns() > 0 && !n.tableDesc.IsPartitionAllBy() { params.p.BufferClientNotice( params.ctx, @@ -596,7 +596,7 @@ func (p *planner) configureIndexDescForNewIndexPartitioning( allowImplicitPartitioning := p.EvalContext().SessionData.ImplicitColumnPartitioningEnabled || tableDesc.IsLocalityRegionalByRow() if partitionBy != nil { - if indexDesc, err = CreatePartitioning( + newImplicitCols, newPartitioning, err := CreatePartitioning( ctx, p.ExecCfg().Settings, p.EvalContext(), @@ -605,9 +605,11 @@ func (p *planner) configureIndexDescForNewIndexPartitioning( partitionBy, nil, /* allowedNewColumnNames */ allowImplicitPartitioning, - ); err != nil { + ) + if err != nil { return indexDesc, err } + tabledesc.UpdateIndexPartitioning(&indexDesc, newImplicitCols, newPartitioning) } } return indexDesc, nil diff --git a/pkg/sql/create_table.go b/pkg/sql/create_table.go index 0a0bbea55a84..aaf773ec22be 100644 --- a/pkg/sql/create_table.go +++ b/pkg/sql/create_table.go @@ -905,7 +905,7 @@ func ResolveFK( referencedColNames := d.ToCols // If no columns are specified, attempt to default to PK, ignoring implicit columns. if len(referencedColNames) == 0 { - numImplicitCols := int(target.GetPrimaryIndex().GetPartitioning().NumImplicitColumns) + numImplicitCols := target.GetPrimaryIndex().GetPartitioning().NumImplicitColumns() referencedColNames = make( tree.NameList, 0, @@ -1291,8 +1291,9 @@ func (p *planner) finalizeInterleave( return nil } -// CreatePartitioning returns a new index descriptor with partitioning fields -// populated to align with the tree.PartitionBy clause. +// CreatePartitioning returns a set of implicit columns and a new partitioning +// descriptor to build an index with partitioning fields populated to align with +// the tree.PartitionBy clause. func CreatePartitioning( ctx context.Context, st *cluster.Settings, @@ -1302,18 +1303,17 @@ func CreatePartitioning( partBy *tree.PartitionBy, allowedNewColumnNames []tree.Name, allowImplicitPartitioning bool, -) (descpb.IndexDescriptor, error) { +) (newImplicitCols []catalog.Column, newPartitioning descpb.PartitioningDescriptor, err error) { if partBy == nil { if indexDesc.Partitioning.NumImplicitColumns > 0 { - return indexDesc, unimplemented.Newf( + return nil, newPartitioning, unimplemented.Newf( "ALTER ... PARTITION BY NOTHING", "cannot alter to PARTITION BY NOTHING if the object has implicit column partitioning", ) } // No CCL necessary if we're looking at PARTITION BY NOTHING - we can // set the partitioning to nothing. - indexDesc.Partitioning = descpb.PartitioningDescriptor{} - return indexDesc, nil + return nil, newPartitioning, nil } return CreatePartitioningCCL( ctx, st, evalCtx, tableDesc, indexDesc, partBy, allowedNewColumnNames, allowImplicitPartitioning, @@ -1331,8 +1331,8 @@ var CreatePartitioningCCL = func( partBy *tree.PartitionBy, allowedNewColumnNames []tree.Name, allowImplicitPartitioning bool, -) (descpb.IndexDescriptor, error) { - return descpb.IndexDescriptor{}, sqlerrors.NewCCLRequiredError(errors.New( +) (newImplicitCols []catalog.Column, newPartitioning descpb.PartitioningDescriptor, err error) { + return nil, descpb.PartitioningDescriptor{}, sqlerrors.NewCCLRequiredError(errors.New( "creating or manipulating partitions requires a CCL binary")) } @@ -1802,7 +1802,7 @@ func NewTableDesc( // partitioning for PARTITION ALL BY. if desc.PartitionAllBy && !implicitColumnDefIdx.def.PrimaryKey.IsPrimaryKey { var err error - *implicitColumnDefIdx.idx, err = CreatePartitioning( + newImplicitCols, newPartitioning, err := CreatePartitioning( ctx, st, evalCtx, @@ -1815,6 +1815,7 @@ func NewTableDesc( if err != nil { return nil, err } + tabledesc.UpdateIndexPartitioning(implicitColumnDefIdx.idx, newImplicitCols, newPartitioning) } if err := desc.AddIndex(*implicitColumnDefIdx.idx, implicitColumnDefIdx.def.PrimaryKey.IsPrimaryKey); err != nil { @@ -1946,7 +1947,7 @@ func NewTableDesc( if partitionBy != nil { var err error - idx, err = CreatePartitioning( + newImplicitCols, newPartitioning, err := CreatePartitioning( ctx, st, evalCtx, @@ -1959,6 +1960,7 @@ func NewTableDesc( if err != nil { return nil, err } + tabledesc.UpdateIndexPartitioning(&idx, newImplicitCols, newPartitioning) } } if d.Predicate != nil { @@ -2042,7 +2044,7 @@ func NewTableDesc( if partitionBy != nil { var err error - idx, err = CreatePartitioning( + newImplicitCols, newPartitioning, err := CreatePartitioning( ctx, st, evalCtx, @@ -2055,6 +2057,7 @@ func NewTableDesc( if err != nil { return nil, err } + tabledesc.UpdateIndexPartitioning(&idx, newImplicitCols, newPartitioning) } } if d.Predicate != nil { @@ -2159,12 +2162,13 @@ func NewTableDesc( } // At this point, we could have PARTITION ALL BY NOTHING, so check it is != nil. if partitionBy != nil { - newPrimaryIndex, err := CreatePartitioning( + newPrimaryIndex := desc.GetPrimaryIndex().IndexDescDeepCopy() + newImplicitCols, newPartitioning, err := CreatePartitioning( ctx, st, evalCtx, &desc, - *desc.GetPrimaryIndex().IndexDesc(), + newPrimaryIndex, partitionBy, nil, /* allowedNewColumnNames */ allowImplicitPartitioning, @@ -2172,29 +2176,35 @@ func NewTableDesc( if err != nil { return nil, err } - // During CreatePartitioning, implicitly partitioned columns may be - // created. AllocateIDs which allocates column IDs to each index - // needs to be called before CreatePartitioning as CreatePartitioning - // requires IDs to be allocated. - // - // As such, do a post check for implicitly partitioned columns, and - // if they are detected, ensure each index contains the implicitly - // partitioned column. - if numImplicitCols := newPrimaryIndex.Partitioning.NumImplicitColumns; numImplicitCols > 0 { - for _, idx := range desc.AllIndexes() { - if idx.GetEncodingType() == descpb.SecondaryIndexEncoding { + isIndexAltered := tabledesc.UpdateIndexPartitioning(&newPrimaryIndex, newImplicitCols, newPartitioning) + if isIndexAltered { + // During CreatePartitioning, implicitly partitioned columns may be + // created. AllocateIDs which allocates column IDs to each index + // needs to be called before CreatePartitioning as CreatePartitioning + // requires IDs to be allocated. + // + // As such, do a post check for implicitly partitioned columns, and + // if they are detected, ensure each index contains the implicitly + // partitioned column. + if numImplicitCols := newPrimaryIndex.Partitioning.NumImplicitColumns; numImplicitCols > 0 { + for _, idx := range desc.PublicNonPrimaryIndexes() { + missingExtraColumnIDs := make([]descpb.ColumnID, 0, numImplicitCols) for _, implicitPrimaryColID := range newPrimaryIndex.ColumnIDs[:numImplicitCols] { if !idx.ContainsColumnID(implicitPrimaryColID) { - idx.IndexDesc().ExtraColumnIDs = append( - idx.IndexDesc().ExtraColumnIDs, - implicitPrimaryColID, - ) + missingExtraColumnIDs = append(missingExtraColumnIDs, implicitPrimaryColID) + idx.IndexDesc().ExtraColumnIDs = append(idx.IndexDesc().ExtraColumnIDs, implicitPrimaryColID) } } + if len(missingExtraColumnIDs) == 0 { + continue + } + newIdxDesc := idx.IndexDescDeepCopy() + newIdxDesc.ExtraColumnIDs = append(newIdxDesc.ExtraColumnIDs, missingExtraColumnIDs...) + desc.SetPublicNonPrimaryIndex(idx.Ordinal(), newIdxDesc) } } + desc.SetPrimaryIndex(newPrimaryIndex) } - desc.SetPrimaryIndex(newPrimaryIndex) } } @@ -2332,7 +2342,7 @@ func NewTableDesc( if idx.NumColumns() > 1 { telemetry.Inc(sqltelemetry.MultiColumnInvertedIndexCounter) } - if idx.GetPartitioning().NumColumns != 0 { + if idx.GetPartitioning().NumColumns() != 0 { telemetry.Inc(sqltelemetry.PartitionedInvertedIndexCounter) } } diff --git a/pkg/sql/database_region_change_finalizer.go b/pkg/sql/database_region_change_finalizer.go index 70d5ae35b420..1b7560c0f326 100644 --- a/pkg/sql/database_region_change_finalizer.go +++ b/pkg/sql/database_region_change_finalizer.go @@ -187,15 +187,16 @@ func (r *databaseRegionChangeFinalizer) repartitionRegionalByRowTables( } partitionAllBy := partitionByForRegionalByRow(regionConfig, colName) - // oldPartitioningDescs saves the old partitioning descriptors for each + // oldPartitionings saves the old partitionings for each // index that is repartitioned. This is later used to remove zone // configurations from any partitions that are removed. - oldPartitioningDescs := make(map[descpb.IndexID]descpb.PartitioningDescriptor) + oldPartitionings := make(map[descpb.IndexID]catalog.Partitioning) // Update the partitioning on all indexes of the table that aren't being // dropped. for _, index := range tableDesc.NonDropIndexes() { - newIdx, err := CreatePartitioning( + oldPartitionings[index.GetID()] = index.GetPartitioning().DeepCopy() + newImplicitCols, newPartitioning, err := CreatePartitioning( ctx, r.localPlanner.extendedEvalCtx.Settings, r.localPlanner.EvalContext(), @@ -208,11 +209,7 @@ func (r *databaseRegionChangeFinalizer) repartitionRegionalByRowTables( if err != nil { return err } - - oldPartitioningDescs[index.GetID()] = index.IndexDesc().Partitioning - - // Update the index descriptor proto's partitioning. - index.IndexDesc().Partitioning = newIdx.Partitioning + tabledesc.UpdateIndexPartitioning(index.IndexDesc(), newImplicitCols, newPartitioning) } // Remove zone configurations that applied to partitions that were removed @@ -225,8 +222,6 @@ func (r *databaseRegionChangeFinalizer) repartitionRegionalByRowTables( // remove the partition from all indexes before trying to delete zone // configurations. for _, index := range tableDesc.NonDropIndexes() { - oldPartitioning := oldPartitioningDescs[index.GetID()] - // Remove zone configurations that reference partition values we removed // in the previous step. if err = deleteRemovedPartitionZoneConfigs( @@ -234,8 +229,8 @@ func (r *databaseRegionChangeFinalizer) repartitionRegionalByRowTables( txn, tableDesc, index.GetID(), - &oldPartitioning, - &index.IndexDesc().Partitioning, + oldPartitionings[index.GetID()], + index.GetPartitioning(), r.localPlanner.ExecCfg(), ); err != nil { return err diff --git a/pkg/sql/opt_catalog.go b/pkg/sql/opt_catalog.go index cf004268572c..98839b2c2b2f 100644 --- a/pkg/sql/opt_catalog.go +++ b/pkg/sql/opt_catalog.go @@ -38,6 +38,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/stats" "github.com/cockroachdb/cockroach/pkg/sql/types" "github.com/cockroachdb/cockroach/pkg/util" + "github.com/cockroachdb/cockroach/pkg/util/log" "github.com/cockroachdb/errors" "github.com/lib/pq/oid" ) @@ -816,11 +817,11 @@ func newOptTable( } // Add unique constraints for implicitly partitioned unique indexes. - if idx.IsUnique() && idx.GetPartitioning().NumImplicitColumns > 0 { + if idx.IsUnique() && idx.GetPartitioning().NumImplicitColumns() > 0 { ot.uniqueConstraints = append(ot.uniqueConstraints, optUniqueConstraint{ name: idx.GetName(), table: ot.ID(), - columns: idx.IndexDesc().ColumnIDs[idx.GetPartitioning().NumImplicitColumns:], + columns: idx.IndexDesc().ColumnIDs[idx.GetPartitioning().NumImplicitColumns():], withoutIndex: true, predicate: idx.GetPredicate(), // TODO(rytaft): will we ever support an unvalidated unique constraint @@ -1206,37 +1207,39 @@ func (oi *optIndex) init( // Collect information about the partitions. idxPartitioning := idx.GetPartitioning() - oi.partitions = make([]optPartition, len(idxPartitioning.List)) - for i := range idxPartitioning.List { - p := &idxPartitioning.List[i] - oi.partitions[i] = optPartition{ - name: p.Name, + oi.partitions = make([]optPartition, 0, idxPartitioning.NumLists()) + _ = idxPartitioning.ForEachList(func(name string, values [][]byte, subPartitioning catalog.Partitioning) error { + op := optPartition{ + name: name, zone: &zonepb.ZoneConfig{}, - datums: make([]tree.Datums, 0, len(p.Values)), + datums: make([]tree.Datums, 0, len(values)), } // Get the zone. - if zone, ok := partZones[p.Name]; ok { - oi.partitions[i].zone = zone + if zone, ok := partZones[name]; ok { + op.zone = zone } // Get the partition values. var a rowenc.DatumAlloc - for _, valueEncBuf := range p.Values { + for _, valueEncBuf := range values { t, _, err := rowenc.DecodePartitionTuple( - &a, oi.tab.codec, oi.tab.desc, oi.idx, &oi.idx.IndexDesc().Partitioning, + &a, oi.tab.codec, oi.tab.desc, oi.idx, oi.idx.GetPartitioning(), valueEncBuf, nil, /* prefixDatums */ ) if err != nil { - panic(errors.NewAssertionErrorWithWrappedErrf(err, - "while decoding partition tuple: %+v %+v", oi.tab.desc, oi.tab.desc.GetDependsOnTypes())) + log.Fatalf(context.TODO(), "error while decoding partition tuple: %+v %+v", + oi.tab.desc, oi.tab.desc.GetDependsOnTypes()) } - oi.partitions[i].datums = append(oi.partitions[i].datums, t.Datums) + op.datums = append(op.datums, t.Datums) // TODO(radu): split into multiple prefixes if Subpartition is also by list. // Note that this functionality should be kept in sync with the test catalog // implementation (test_catalog.go). } - } + + oi.partitions = append(oi.partitions, op) + return nil + }) if idx.IsUnique() { notNull := true @@ -1385,7 +1388,7 @@ func (oi *optIndex) Ordinal() int { // ImplicitPartitioningColumnCount is part of the cat.Index interface. func (oi *optIndex) ImplicitPartitioningColumnCount() int { - return int(oi.idx.GetPartitioning().NumImplicitColumns) + return oi.idx.GetPartitioning().NumImplicitColumns() } // InterleaveAncestorCount is part of the cat.Index interface. diff --git a/pkg/sql/partition.go b/pkg/sql/partition.go index 567e260d4b7a..480b83f6bb96 100644 --- a/pkg/sql/partition.go +++ b/pkg/sql/partition.go @@ -13,7 +13,6 @@ package sql import ( "github.com/cockroachdb/cockroach/pkg/keys" "github.com/cockroachdb/cockroach/pkg/sql/catalog" - "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb" "github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc" "github.com/cockroachdb/cockroach/pkg/sql/rowenc" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" @@ -25,7 +24,7 @@ func partitionByFromTableDesc( codec keys.SQLCodec, tableDesc *tabledesc.Mutable, ) (*tree.PartitionBy, error) { idx := tableDesc.GetPrimaryIndex() - return partitionByFromTableDescImpl(codec, tableDesc, idx, &idx.IndexDesc().Partitioning, 0) + return partitionByFromTableDescImpl(codec, tableDesc, idx, idx.GetPartitioning(), 0) } // partitionByFromTableDescImpl contains the inner logic of partitionByFromTableDesc. @@ -35,10 +34,10 @@ func partitionByFromTableDescImpl( codec keys.SQLCodec, tableDesc *tabledesc.Mutable, idx catalog.Index, - partDesc *descpb.PartitioningDescriptor, + part catalog.Partitioning, colOffset int, ) (*tree.PartitionBy, error) { - if partDesc.NumColumns == 0 { + if part.NumColumns() == 0 { return nil, nil } @@ -50,87 +49,81 @@ func partitionByFromTableDescImpl( } partitionBy := &tree.PartitionBy{ - Fields: make(tree.NameList, partDesc.NumColumns), - List: make([]tree.ListPartition, len(partDesc.List)), - Range: make([]tree.RangePartition, len(partDesc.Range)), + Fields: make(tree.NameList, part.NumColumns()), + List: make([]tree.ListPartition, 0, part.NumLists()), + Range: make([]tree.RangePartition, 0, part.NumRanges()), } - for i := 0; i < int(partDesc.NumColumns); i++ { + for i := 0; i < part.NumColumns(); i++ { partitionBy.Fields[i] = tree.Name(idx.GetColumnName(colOffset + i)) } // Copy the LIST of the PARTITION BY clause. a := &rowenc.DatumAlloc{} - for i := range partDesc.List { - part := &partDesc.List[i] - partitionBy.List[i].Name = tree.UnrestrictedName(part.Name) - partitionBy.List[i].Exprs = make(tree.Exprs, len(part.Values)) - for j, values := range part.Values { + err := part.ForEachList(func(name string, values [][]byte, subPartitioning catalog.Partitioning) (err error) { + lp := tree.ListPartition{ + Name: tree.UnrestrictedName(name), + Exprs: make(tree.Exprs, len(values)), + } + for j, values := range values { tuple, _, err := rowenc.DecodePartitionTuple( a, codec, tableDesc, idx, - partDesc, + part, values, fakePrefixDatums, ) if err != nil { - return nil, err + return err } exprs, err := partitionTupleToExprs(tuple) if err != nil { - return nil, err + return err } - partitionBy.List[i].Exprs[j] = &tree.Tuple{ + lp.Exprs[j] = &tree.Tuple{ Exprs: exprs, } } - var err error - if partitionBy.List[i].Subpartition, err = partitionByFromTableDescImpl( + lp.Subpartition, err = partitionByFromTableDescImpl( codec, tableDesc, idx, - &part.Subpartitioning, - colOffset+int(partDesc.NumColumns), - ); err != nil { - return nil, err - } + subPartitioning, + colOffset+part.NumColumns(), + ) + partitionBy.List = append(partitionBy.List, lp) + return err + }) + if err != nil { + return nil, err } // Copy the RANGE of the PARTITION BY clause. - for i, part := range partDesc.Range { - partitionBy.Range[i].Name = tree.UnrestrictedName(part.Name) + err = part.ForEachRange(func(name string, from, to []byte) error { + rp := tree.RangePartition{Name: tree.UnrestrictedName(name)} fromTuple, _, err := rowenc.DecodePartitionTuple( - a, - codec, - tableDesc, - idx, - partDesc, - part.FromInclusive, - fakePrefixDatums, - ) + a, codec, tableDesc, idx, part, from, fakePrefixDatums) if err != nil { - return nil, err + return err } - if partitionBy.Range[i].From, err = partitionTupleToExprs(fromTuple); err != nil { - return nil, err + rp.From, err = partitionTupleToExprs(fromTuple) + if err != nil { + return err } toTuple, _, err := rowenc.DecodePartitionTuple( - a, - codec, - tableDesc, - idx, - partDesc, - part.ToExclusive, - fakePrefixDatums, - ) + a, codec, tableDesc, idx, part, to, fakePrefixDatums) if err != nil { - return nil, err - } - if partitionBy.Range[i].To, err = partitionTupleToExprs(toTuple); err != nil { - return nil, err + return err } + rp.To, err = partitionTupleToExprs(toTuple) + partitionBy.Range = append(partitionBy.Range, rp) + return err + }) + if err != nil { + return nil, err } + return partitionBy, nil } diff --git a/pkg/sql/partition_utils.go b/pkg/sql/partition_utils.go index fec2ad869334..5f18abe0f063 100644 --- a/pkg/sql/partition_utils.go +++ b/pkg/sql/partition_utils.go @@ -125,7 +125,7 @@ func GenerateSubzoneSpans( var emptyPrefix []tree.Datum indexPartitionCoverings, err := indexCoveringsForPartitioning( - a, codec, tableDesc, idx, &idx.IndexDesc().Partitioning, subzoneIndexByPartition, emptyPrefix) + a, codec, tableDesc, idx, idx.GetPartitioning(), subzoneIndexByPartition, emptyPrefix) if err != nil { return err } @@ -186,18 +186,18 @@ func indexCoveringsForPartitioning( codec keys.SQLCodec, tableDesc catalog.TableDescriptor, idx catalog.Index, - partDesc *descpb.PartitioningDescriptor, + part catalog.Partitioning, relevantPartitions map[string]int32, prefixDatums []tree.Datum, ) ([]covering.Covering, error) { - if partDesc.NumColumns == 0 { + if part.NumColumns() == 0 { return nil, nil } var coverings []covering.Covering var descendentCoverings []covering.Covering - if len(partDesc.List) > 0 { + if part.NumLists() > 0 { // The returned spans are required to be ordered with highest precedence // first. The span for (1, DEFAULT) overlaps with (1, 2) and needs to be // returned at a lower precedence. Luckily, because of the partitioning @@ -205,28 +205,32 @@ func indexCoveringsForPartitioning( // with the same number of DEFAULTs are non-overlapping. So, bucket the // `interval.Range`s by the number of non-DEFAULT columns and return // them ordered from least # of DEFAULTs to most. - listCoverings := make([]covering.Covering, int(partDesc.NumColumns)+1) - for _, p := range partDesc.List { - for _, valueEncBuf := range p.Values { + listCoverings := make([]covering.Covering, part.NumColumns()+1) + err := part.ForEachList(func(name string, values [][]byte, subPartitioning catalog.Partitioning) error { + for _, valueEncBuf := range values { t, keyPrefix, err := rowenc.DecodePartitionTuple( - a, codec, tableDesc, idx, partDesc, valueEncBuf, prefixDatums) + a, codec, tableDesc, idx, part, valueEncBuf, prefixDatums) if err != nil { - return nil, err + return err } - if _, ok := relevantPartitions[p.Name]; ok { + if _, ok := relevantPartitions[name]; ok { listCoverings[len(t.Datums)] = append(listCoverings[len(t.Datums)], covering.Range{ Start: keyPrefix, End: roachpb.Key(keyPrefix).PrefixEnd(), - Payload: zonepb.Subzone{PartitionName: p.Name}, + Payload: zonepb.Subzone{PartitionName: name}, }) } newPrefixDatums := append(prefixDatums, t.Datums...) subpartitionCoverings, err := indexCoveringsForPartitioning( - a, codec, tableDesc, idx, &p.Subpartitioning, relevantPartitions, newPrefixDatums) + a, codec, tableDesc, idx, subPartitioning, relevantPartitions, newPrefixDatums) if err != nil { - return nil, err + return err } descendentCoverings = append(descendentCoverings, subpartitionCoverings...) } + return nil + }) + if err != nil { + return nil, err } for i := range listCoverings { if covering := listCoverings[len(listCoverings)-i-1]; len(covering) > 0 { @@ -235,27 +239,31 @@ func indexCoveringsForPartitioning( } } - if len(partDesc.Range) > 0 { - for _, p := range partDesc.Range { - if _, ok := relevantPartitions[p.Name]; !ok { - continue + if part.NumRanges() > 0 { + err := part.ForEachRange(func(name string, from, to []byte) error { + if _, ok := relevantPartitions[name]; !ok { + return nil } _, fromKey, err := rowenc.DecodePartitionTuple( - a, codec, tableDesc, idx, partDesc, p.FromInclusive, prefixDatums) + a, codec, tableDesc, idx, part, from, prefixDatums) if err != nil { - return nil, err + return err } _, toKey, err := rowenc.DecodePartitionTuple( - a, codec, tableDesc, idx, partDesc, p.ToExclusive, prefixDatums) + a, codec, tableDesc, idx, part, to, prefixDatums) if err != nil { - return nil, err + return err } - if _, ok := relevantPartitions[p.Name]; ok { + if _, ok := relevantPartitions[name]; ok { coverings = append(coverings, covering.Covering{{ Start: fromKey, End: toKey, - Payload: zonepb.Subzone{PartitionName: p.Name}, + Payload: zonepb.Subzone{PartitionName: name}, }}) } + return nil + }) + if err != nil { + return nil, err } } diff --git a/pkg/sql/rowenc/partition.go b/pkg/sql/rowenc/partition.go index 99a28c5ccd59..2658e87f0409 100644 --- a/pkg/sql/rowenc/partition.go +++ b/pkg/sql/rowenc/partition.go @@ -16,7 +16,6 @@ import ( "github.com/cockroachdb/cockroach/pkg/keys" "github.com/cockroachdb/cockroach/pkg/roachpb" "github.com/cockroachdb/cockroach/pkg/sql/catalog" - "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" "github.com/cockroachdb/cockroach/pkg/util/encoding" "github.com/cockroachdb/errors" @@ -107,19 +106,19 @@ func DecodePartitionTuple( codec keys.SQLCodec, tableDesc catalog.TableDescriptor, index catalog.Index, - partDesc *descpb.PartitioningDescriptor, + part catalog.Partitioning, valueEncBuf []byte, prefixDatums tree.Datums, ) (*PartitionTuple, []byte, error) { - if len(prefixDatums)+int(partDesc.NumColumns) > index.NumColumns() { + if len(prefixDatums)+part.NumColumns() > index.NumColumns() { return nil, nil, fmt.Errorf("not enough columns in index for this partitioning") } t := &PartitionTuple{ - Datums: make(tree.Datums, 0, int(partDesc.NumColumns)), + Datums: make(tree.Datums, 0, part.NumColumns()), } - for i := len(prefixDatums); i < index.NumColumns() && i < len(prefixDatums)+int(partDesc.NumColumns); i++ { + for i := len(prefixDatums); i < index.NumColumns() && i < len(prefixDatums)+part.NumColumns(); i++ { colID := index.GetColumnID(i) col, err := tableDesc.FindColumnWithID(colID) if err != nil { diff --git a/pkg/sql/set_zone_config.go b/pkg/sql/set_zone_config.go index 1085942409d5..c6b4dce05572 100644 --- a/pkg/sql/set_zone_config.go +++ b/pkg/sql/set_zone_config.go @@ -27,7 +27,6 @@ import ( "github.com/cockroachdb/cockroach/pkg/server/telemetry" "github.com/cockroachdb/cockroach/pkg/sql/catalog" "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb" - "github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc" "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode" "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror" "github.com/cockroachdb/cockroach/pkg/sql/privilege" @@ -426,7 +425,7 @@ func (n *setZoneConfigNode) startExec(params runParams) error { var indexes []catalog.Index for _, idx := range table.NonDropIndexes() { - if tabledesc.FindIndexPartitionByName(idx, partitionName) != nil { + if idx.GetPartitioning().FindPartitionByName(partitionName) != nil { indexes = append(indexes, idx) } } @@ -452,7 +451,7 @@ func (n *setZoneConfigNode) startExec(params runParams) error { if n.zoneSpecifier.TargetsPartition() && n.allIndexes { sqltelemetry.IncrementPartitioningCounter(sqltelemetry.AlterAllPartitions) for _, idx := range table.NonDropIndexes() { - if p := tabledesc.FindIndexPartitionByName(idx, string(n.zoneSpecifier.Partition)); p != nil { + if idx.GetPartitioning().FindPartitionByName(string(n.zoneSpecifier.Partition)) != nil { zs := n.zoneSpecifier zs.TableOrIndex.Index = tree.UnrestrictedName(idx.GetName()) specifiers = append(specifiers, zs) diff --git a/pkg/sql/show_create.go b/pkg/sql/show_create.go index 187c1705b560..f0e64bbe8708 100644 --- a/pkg/sql/show_create.go +++ b/pkg/sql/show_create.go @@ -155,7 +155,7 @@ func ShowCreateTable( // Build the PARTITION BY clause. var partitionBuf bytes.Buffer if err := ShowCreatePartitioning( - a, p.ExecCfg().Codec, desc, idx, &idx.IndexDesc().Partitioning, &partitionBuf, 1 /* indent */, 0, /* colOffset */ + a, p.ExecCfg().Codec, desc, idx, idx.GetPartitioning(), &partitionBuf, 1 /* indent */, 0, /* colOffset */ ); err != nil { return "", err } @@ -199,7 +199,7 @@ func ShowCreateTable( return "", err } if err := ShowCreatePartitioning( - a, p.ExecCfg().Codec, desc, desc.GetPrimaryIndex(), &desc.GetPrimaryIndex().IndexDesc().Partitioning, &f.Buffer, 0 /* indent */, 0, /* colOffset */ + a, p.ExecCfg().Codec, desc, desc.GetPrimaryIndex(), desc.GetPrimaryIndex().GetPartitioning(), &f.Buffer, 0 /* indent */, 0, /* colOffset */ ); err != nil { return "", err } diff --git a/pkg/sql/show_create_clauses.go b/pkg/sql/show_create_clauses.go index 4448f0167438..921a8cbfcb02 100644 --- a/pkg/sql/show_create_clauses.go +++ b/pkg/sql/show_create_clauses.go @@ -427,7 +427,7 @@ func ShowCreatePartitioning( codec keys.SQLCodec, tableDesc catalog.TableDescriptor, idx catalog.Index, - partDesc *descpb.PartitioningDescriptor, + part catalog.Partitioning, buf *bytes.Buffer, indent int, colOffset int, @@ -435,7 +435,7 @@ func ShowCreatePartitioning( isPrimaryKeyOfPartitionAllByTable := tableDesc.IsPartitionAllBy() && tableDesc.GetPrimaryIndexID() == idx.GetID() && colOffset == 0 - if partDesc.NumColumns == 0 && !isPrimaryKeyOfPartitionAllByTable { + if part.NumColumns() == 0 && !isPrimaryKeyOfPartitionAllByTable { return nil } // Do not print PARTITION BY clauses of non-primary indexes belonging to a table @@ -464,18 +464,18 @@ func ShowCreatePartitioning( buf.WriteString(`ALL `) } buf.WriteString(`BY `) - if len(partDesc.List) > 0 { + if part.NumLists() > 0 { buf.WriteString(`LIST`) - } else if len(partDesc.Range) > 0 { + } else if part.NumRanges() > 0 { buf.WriteString(`RANGE`) } else if isPrimaryKeyOfPartitionAllByTable { buf.WriteString(`NOTHING`) return nil } else { - return errors.Errorf(`invalid partition descriptor: %v`, partDesc) + return errors.Errorf(`invalid partition descriptor: %v`, part.PartitioningDesc()) } buf.WriteString(` (`) - for i := 0; i < int(partDesc.NumColumns); i++ { + for i := 0; i < part.NumColumns(); i++ { if i != 0 { buf.WriteString(", ") } @@ -483,58 +483,65 @@ func ShowCreatePartitioning( } buf.WriteString(`) (`) fmtCtx := tree.NewFmtCtx(tree.FmtSimple) - for i := range partDesc.List { - part := &partDesc.List[i] - if i != 0 { + isFirst := true + err := part.ForEachList(func(name string, values [][]byte, subPartitioning catalog.Partitioning) error { + if !isFirst { buf.WriteString(`, `) } + isFirst = false buf.WriteString("\n") buf.WriteString(indentStr) buf.WriteString("\tPARTITION ") - fmtCtx.FormatNameP(&part.Name) + fmtCtx.FormatNameP(&name) _, _ = fmtCtx.Buffer.WriteTo(buf) buf.WriteString(` VALUES IN (`) - for j, values := range part.Values { + for j, values := range values { if j != 0 { buf.WriteString(`, `) } tuple, _, err := rowenc.DecodePartitionTuple( - a, codec, tableDesc, idx, partDesc, values, fakePrefixDatums) + a, codec, tableDesc, idx, part, values, fakePrefixDatums) if err != nil { return err } buf.WriteString(tuple.String()) } buf.WriteString(`)`) - if err := ShowCreatePartitioning( - a, codec, tableDesc, idx, &part.Subpartitioning, buf, indent+1, - colOffset+int(partDesc.NumColumns), - ); err != nil { - return err - } + return ShowCreatePartitioning( + a, codec, tableDesc, idx, subPartitioning, buf, indent+1, colOffset+part.NumColumns(), + ) + }) + if err != nil { + return err } - for i, part := range partDesc.Range { - if i != 0 { + isFirst = true + err = part.ForEachRange(func(name string, from, to []byte) error { + if !isFirst { buf.WriteString(`, `) } + isFirst = false buf.WriteString("\n") buf.WriteString(indentStr) buf.WriteString("\tPARTITION ") - buf.WriteString(part.Name) + buf.WriteString(name) buf.WriteString(" VALUES FROM ") fromTuple, _, err := rowenc.DecodePartitionTuple( - a, codec, tableDesc, idx, partDesc, part.FromInclusive, fakePrefixDatums) + a, codec, tableDesc, idx, part, from, fakePrefixDatums) if err != nil { return err } buf.WriteString(fromTuple.String()) buf.WriteString(" TO ") toTuple, _, err := rowenc.DecodePartitionTuple( - a, codec, tableDesc, idx, partDesc, part.ToExclusive, fakePrefixDatums) + a, codec, tableDesc, idx, part, to, fakePrefixDatums) if err != nil { return err } buf.WriteString(toTuple.String()) + return nil + }) + if err != nil { + return err } buf.WriteString("\n") buf.WriteString(indentStr) diff --git a/pkg/sql/zone_config.go b/pkg/sql/zone_config.go index ee96a2bcb723..d8f804bd6d42 100644 --- a/pkg/sql/zone_config.go +++ b/pkg/sql/zone_config.go @@ -349,7 +349,7 @@ func resolveSubzone( partitionName := string(zs.Partition) if partitionName != "" { - if partitioning := tabledesc.FindIndexPartitionByName(index, partitionName); partitioning == nil { + if index.GetPartitioning().FindPartitionByName(partitionName) == nil { return nil, "", fmt.Errorf("partition %q does not exist on index %q", partitionName, indexName) } } @@ -362,20 +362,22 @@ func deleteRemovedPartitionZoneConfigs( txn *kv.Txn, tableDesc catalog.TableDescriptor, indexID descpb.IndexID, - oldPartDesc *descpb.PartitioningDescriptor, - newPartDesc *descpb.PartitioningDescriptor, + oldPart catalog.Partitioning, + newPart catalog.Partitioning, execCfg *ExecutorConfig, ) error { newNames := map[string]struct{}{} - for _, n := range newPartDesc.PartitionNames() { - newNames[n] = struct{}{} - } - removedNames := []string{} - for _, n := range oldPartDesc.PartitionNames() { - if _, exists := newNames[n]; !exists { - removedNames = append(removedNames, n) + _ = newPart.ForEachPartitionName(func(newName string) error { + newNames[newName] = struct{}{} + return nil + }) + removedNames := make([]string, 0, len(newNames)) + _ = oldPart.ForEachPartitionName(func(oldName string) error { + if _, exists := newNames[oldName]; !exists { + removedNames = append(removedNames, oldName) } - } + return nil + }) if len(removedNames) == 0 { return nil }