schemachanger: refactor subzone config elements

TODO

pkg/ccl/schemachangerccl/testdata/decomp/multiregion

@@ -1373,7 +1373,10 @@ ElementState:
           numReplicas: 2
       indexId: 1
       partitionName: us-east1
-    subzoneSpans: []
+    subzoneSpan:
+      endKey: null
+      key: null
+      subzoneIndex: 0
     tableId: 108
   Status: PUBLIC
 - IndexZoneConfig:
@@ -1406,7 +1409,10 @@ ElementState:
           numReplicas: 2
       indexId: 1
       partitionName: us-east2
-    subzoneSpans: []
+    subzoneSpan:
+      endKey: null
+      key: null
+      subzoneIndex: 0
     tableId: 108
   Status: PUBLIC
 - IndexZoneConfig:
@@ -1439,7 +1445,10 @@ ElementState:
           numReplicas: 2
       indexId: 1
       partitionName: us-east3
-    subzoneSpans: []
+    subzoneSpan:
+      endKey: null
+      key: null
+      subzoneIndex: 0
     tableId: 108
   Status: PUBLIC
 - Namespace:
@@ -1483,7 +1492,10 @@ ElementState:
           numReplicas: 2
       indexId: 1
       partitionName: us-east1
-    subzoneSpans: []
+    subzoneSpan:
+      endKey: null
+      key: null
+      subzoneIndex: 0
     tableId: 108
   Status: PUBLIC
 - PartitionZoneConfig:
@@ -1517,7 +1529,10 @@ ElementState:
           numReplicas: 2
       indexId: 1
       partitionName: us-east2
-    subzoneSpans: []
+    subzoneSpan:
+      endKey: null
+      key: null
+      subzoneIndex: 0
     tableId: 108
   Status: PUBLIC
 - PartitionZoneConfig:
@@ -1551,7 +1566,10 @@ ElementState:
           numReplicas: 2
       indexId: 1
       partitionName: us-east3
-    subzoneSpans: []
+    subzoneSpan:
+      endKey: null
+      key: null
+      subzoneIndex: 0
     tableId: 108
   Status: PUBLIC
 - PrimaryIndex:

pkg/config/zonepb/zone.go

@@ -10,6 +10,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"slices"
 	"strconv"
 	"strings"
 
@@ -1153,6 +1154,23 @@ func (z *ZoneConfig) DeleteSubzone(indexID uint32, partition string) bool {
 	return false
 }
 
+// SetSubzoneSpan installs subzoneSpan into the ZoneConfig, overwriting any
+// existing subzone span with the same SubzoneIndex. If we end up adding a new
+// span to the list, we sort the list in the end.
+func (z *ZoneConfig) SetSubzoneSpan(subzoneSpan SubzoneSpan) {
+	for i, s := range z.SubzoneSpans {
+		if s.SubzoneIndex == subzoneSpan.SubzoneIndex {
+			z.SubzoneSpans[i] = subzoneSpan
+			return
+		}
+	}
+
+	z.SubzoneSpans = append(z.SubzoneSpans, subzoneSpan)
+	slices.SortFunc(z.SubzoneSpans, func(a, b SubzoneSpan) int {
+		return a.Key.Compare(b.Key)
+	})
+}
+
 // DeleteIndexSubzones deletes all subzones that refer to the index with the
 // specified ID. This includes subzones for partitions of the index as well as
 // the index subzone itself.

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/index_zone_config.go

@@ -50,11 +50,26 @@ func (izo *indexZoneConfigObj) getZoneConfigElem(b BuildCtx) scpb.Element {
 		panic(err)
 	}
 
+	// Find the index of our subzone in our subzones list, along with the
+	// corresponding span.
+	var subzoneIdx int
+	for i, sub := range subzones {
+		if sub.IndexID == uint32(izo.indexID) && sub.PartitionName == "" {
+			subzoneIdx = i
+		}
+	}
+	var subzoneSpans []zonepb.SubzoneSpan
+	for _, s := range ss {
+		if s.SubzoneIndex == int32(subzoneIdx) {
+			subzoneSpans = append(subzoneSpans, s)
+		}
+	}
+
 	elem := &scpb.IndexZoneConfig{
 		TableID:      izo.tableID,
 		IndexID:      izo.indexID,
 		Subzone:      *izo.indexSubzone,
-		SubzoneSpans: ss,
+		SubzoneSpans: subzoneSpans,
 		SeqNum:       izo.seqNum,
 	}
 	return elem

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/partition_zone_config.go

@@ -54,12 +54,27 @@ func (pzo *partitionZoneConfigObj) getZoneConfigElem(b BuildCtx) scpb.Element {
 		panic(err)
 	}
 
+	// Find the index of our subzone in our subzones list, along with the
+	// corresponding span.
+	var subzoneIdx int
+	for i, sub := range subzones {
+		if sub.IndexID == uint32(pzo.indexID) && sub.PartitionName == pzo.partitionName {
+			subzoneIdx = i
+		}
+	}
+	var subzoneSpans []zonepb.SubzoneSpan
+	for _, s := range ss {
+		if s.SubzoneIndex == int32(subzoneIdx) {
+			subzoneSpans = append(subzoneSpans, s)
+		}
+	}
+
 	elem := &scpb.PartitionZoneConfig{
 		TableID:       pzo.tableID,
 		IndexID:       pzo.indexID,
 		PartitionName: pzo.partitionName,
 		Subzone:       *pzo.partitionSubzone,
-		SubzoneSpans:  ss,
+		SubzoneSpans:  subzoneSpans,
 		SeqNum:        pzo.seqNum,
 	}
 	return elem
@@ -101,10 +116,6 @@ func (pzo *partitionZoneConfigObj) retrieveCompleteZoneConfig(
 	if getInheritedDefault {
 		zc, err = pzo.getInheritedDefaultZoneConfig(b)
 	} else {
-		//zc, err = pzo.tableZoneConfigObj.getZoneConfig(b, false /* inheritDefaultRange */)
-		//if err != nil {
-		//	return nil, nil, err
-		//}
 		placeholder, err = pzo.getZoneConfig(b, false /* inheritDefaultRange */)
 	}
 	if err != nil {

pkg/sql/schemachanger/scdeps/exec_deps.go

@@ -239,11 +239,12 @@ func (d *txnDeps) UpdateZoneConfig(ctx context.Context, id descpb.ID, zc *zonepb
 }
 
 // UpdateSubzoneConfig implements the scexec.Catalog interface.
+//
+// N.B. UpdateSubzoneConfig assumes that a subzone at index i in the subzones
+// list has the corresponding subzone span in the subzone span list at the same
+// index.
 func (d *txnDeps) UpdateSubzoneConfig(
-	ctx context.Context,
-	tableID descpb.ID,
-	subzones []zonepb.Subzone,
-	subzoneSpans []zonepb.SubzoneSpan,
+	ctx context.Context, tableID descpb.ID, subzone zonepb.Subzone, subzoneSpans []zonepb.SubzoneSpan,
 ) error {
 	var newZc catalog.ZoneConfig
 	oldZc, err := d.descsCollection.GetZoneConfig(ctx, d.txn.KV(), tableID)
@@ -260,16 +261,25 @@ func (d *txnDeps) UpdateSubzoneConfig(
 		rawBytes = oldZc.GetRawBytesInStorage()
 		zc = oldZc.ZoneConfigProto()
 	} else {
-		// If no zone config exists, create a new one that is a subzone placeholder.
+		// If no zone config exists, create a new one that is a s placeholder.
 		zc = zonepb.NewZoneConfig()
 		zc.DeleteTableConfig()
 	}
 
-	// Update the subzones in the zone config.
-	for _, s := range subzones {
-		zc.SetSubzone(s)
+	// Update the subzone in the zone config.
+	zc.SetSubzone(subzone)
+	// Update the subzone spans.
+	subzoneSpansToWrite := subzoneSpans
+	// If there are subzone spans that currently exist, merge those with the new
+	// spans we are updating. Otherwise, the zone config's set of subzone spans
+	// will be our input subzoneSpans.
+	if len(zc.SubzoneSpans) != 0 {
+		for _, s := range subzoneSpansToWrite {
+			zc.SetSubzoneSpan(s)
+		}
+		subzoneSpansToWrite = zc.SubzoneSpans
 	}
-	zc.SubzoneSpans = subzoneSpans
+	zc.SubzoneSpans = subzoneSpansToWrite
 
 	newZc = zone.NewZoneConfigWithRawBytes(zc, rawBytes)
 	return d.descsCollection.WriteZoneConfigToBatch(ctx, d.kvTrace, d.getOrCreateBatch(),

pkg/sql/schemachanger/scdeps/sctestdeps/test_deps.go

@@ -820,24 +820,21 @@ func (s *TestState) UpdateZoneConfig(
 
 // UpdateSubzoneConfig implements the scexec.Catalog interface.
 func (s *TestState) UpdateSubzoneConfig(
-	ctx context.Context,
-	tableID descpb.ID,
-	subzones []zonepb.Subzone,
-	subzoneSpans []zonepb.SubzoneSpan,
+	ctx context.Context, tableID descpb.ID, subzone zonepb.Subzone, subzoneSpans []zonepb.SubzoneSpan,
 ) error {
 	if s.catalogChanges.zoneConfigsToUpdate == nil {
 		s.catalogChanges.zoneConfigsToUpdate = make(map[descpb.ID]*zonepb.ZoneConfig)
 	}
 	var zc *zonepb.ZoneConfig
 	if czc, ok := s.catalogChanges.zoneConfigsToUpdate[tableID]; ok {
-		czc.Subzones = subzones
-		czc.SubzoneSpans = subzoneSpans
 		zc = czc
 	} else {
 		zc = zonepb.NewZoneConfig()
 		zc.DeleteTableConfig()
-		zc.Subzones = subzones
-		zc.SubzoneSpans = subzoneSpans
+	}
+	zc.SetSubzone(subzone)
+	for _, ss := range subzoneSpans {
+		zc.SetSubzoneSpan(ss)
 	}
 	s.catalogChanges.zoneConfigsToUpdate[tableID] = zc
 	return nil

pkg/sql/schemachanger/scexec/dependencies.go

@@ -75,7 +75,7 @@ type Catalog interface {
 	UpdateSubzoneConfig(
 		ctx context.Context,
 		tableID descpb.ID,
-		subzones []zonepb.Subzone,
+		subzone zonepb.Subzone,
 		subzoneSpans []zonepb.SubzoneSpan,
 	) error
 

pkg/sql/schemachanger/scexec/exec_immediate_mutation.go

@@ -236,7 +236,9 @@ func (s *immediateState) exec(ctx context.Context, c Catalog) error {
 
 	for _, zcToUpdate := range s.modifiedZoneConfigs {
 		if zcToUpdate.isSubzoneConfig {
-			if err := c.UpdateSubzoneConfig(ctx, zcToUpdate.id, zcToUpdate.zc.Subzones,
+			// TODO(before merge): unsatisfied with this; would it be better to
+			//  separate subzone configs from the modifiedZoneConfigs list?
+			if err := c.UpdateSubzoneConfig(ctx, zcToUpdate.id, zcToUpdate.zc.Subzones[0],
 				zcToUpdate.zc.SubzoneSpans); err != nil {
 				return err
 			}