backup: deduplicate routing logic with physicalplan package

Previously, the backup package contained duplicate implementations of
routingDatumsForSQLInstance and routingSpanForSQLInstance that were
identical to functions recently added to pkg/sql/physicalplan. This
commit removes these duplicates and uses the shared implementations:
RoutingDatumsForSQLInstance and MakeInstanceRouter.

This reduces code duplication and ensures consistent routing behavior
across DistSQL physical planning for both backup/restore and bulk merge
operations.

Informs #156580
Release note: None

Author: mw5h

pkg/backup/generative_split_and_scatter_processor.go

@@ -7,7 +7,6 @@ package backup
 
 import (
 	"context"
-	"fmt"
 	"hash/fnv"
 	"math/rand"
 	"strings"
@@ -22,9 +21,9 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql"
-	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catenumpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
+	"github.com/cockroachdb/cockroach/pkg/sql/physicalplan"
 	"github.com/cockroachdb/cockroach/pkg/sql/rowenc"
 	"github.com/cockroachdb/cockroach/pkg/sql/rowexec"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
@@ -249,14 +248,6 @@ func (s dbSplitAndScatterer) findDestination(res *kvpb.AdminScatterResponse) roa
 	return roachpb.NodeID(0)
 }
 
-func routingDatumsForSQLInstance(
-	sqlInstanceID base.SQLInstanceID,
-) (rowenc.EncDatum, rowenc.EncDatum) {
-	routingBytes := roachpb.Key(fmt.Sprintf("node%d", sqlInstanceID))
-	startDatum := rowenc.DatumToEncDatumUnsafe(types.Bytes, tree.NewDBytes(tree.DBytes(routingBytes)))
-	endDatum := rowenc.DatumToEncDatumUnsafe(types.Bytes, tree.NewDBytes(tree.DBytes(routingBytes.Next())))
-	return startDatum, endDatum
-}
 
 type entryNode struct {
 	entry execinfrapb.RestoreSpanEntry
@@ -393,7 +384,7 @@ func (gssp *generativeSplitAndScatterProcessor) Next() (
 		// The routing datums informs the router which output stream should be used.
 		routingDatum, ok := gssp.routingDatumCache.getRoutingDatum(scatteredEntry.node)
 		if !ok {
-			routingDatum, _ = routingDatumsForSQLInstance(base.SQLInstanceID(scatteredEntry.node))
+			routingDatum, _ = physicalplan.RoutingDatumsForSQLInstance(base.SQLInstanceID(scatteredEntry.node))
 			gssp.routingDatumCache.putRoutingDatum(scatteredEntry.node, routingDatum)
 		}
 
@@ -745,23 +736,6 @@ var splitAndScatterOutputTypes = []*types.T{
 	types.Bytes, // RestoreDataEntry bytes
 }
 
-// routingSpanForSQLInstance provides the mapping to be used during distsql planning
-// when setting up the output router.
-func routingSpanForSQLInstance(sqlInstanceID base.SQLInstanceID) ([]byte, []byte, error) {
-	var alloc tree.DatumAlloc
-	startDatum, endDatum := routingDatumsForSQLInstance(sqlInstanceID)
-
-	startBytes, endBytes := make([]byte, 0), make([]byte, 0)
-	startBytes, err := startDatum.Encode(splitAndScatterOutputTypes[0], &alloc, catenumpb.DatumEncoding_ASCENDING_KEY, startBytes)
-	if err != nil {
-		return nil, nil, err
-	}
-	endBytes, err = endDatum.Encode(splitAndScatterOutputTypes[0], &alloc, catenumpb.DatumEncoding_ASCENDING_KEY, endBytes)
-	if err != nil {
-		return nil, nil, err
-	}
-	return startBytes, endBytes, nil
-}
 
 func init() {
 	rowexec.NewGenerativeSplitAndScatterProcessor = newGenerativeSplitAndScatterProcessor

pkg/backup/restore_processor_planning.go

@@ -6,10 +6,8 @@
 package backup
 
 import (
-	"bytes"
 	"context"
 	"math"
-	"slices"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/cloud"
@@ -19,7 +17,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql"
-	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catenumpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/physicalplan"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
@@ -123,34 +120,10 @@ func distRestore(
 		// Plan SplitAndScatter on the coordinator node.
 		splitAndScatterStageID := p.NewStageOnNodes(sqlInstanceIDs)
 
-		defaultStream := int32(0)
-		rangeRouterSpec := execinfrapb.OutputRouterSpec_RangeRouterSpec{
-			Spans:       nil,
-			DefaultDest: &defaultStream,
-			Encodings: []execinfrapb.OutputRouterSpec_RangeRouterSpec_ColumnEncoding{
-				{
-					Column:   0,
-					Encoding: catenumpb.DatumEncoding_ASCENDING_KEY,
-				},
-			},
-		}
-		for stream, sqlInstanceID := range sqlInstanceIDs {
-			startBytes, endBytes, err := routingSpanForSQLInstance(sqlInstanceID)
-			if err != nil {
-				return nil, nil, err
-			}
-
-			span := execinfrapb.OutputRouterSpec_RangeRouterSpec_Span{
-				Start:  startBytes,
-				End:    endBytes,
-				Stream: int32(stream),
-			}
-			rangeRouterSpec.Spans = append(rangeRouterSpec.Spans, span)
+		rangeRouterSpec, err := physicalplan.MakeInstanceRouter(sqlInstanceIDs)
+		if err != nil {
+			return nil, nil, err
 		}
-		// The router expects the spans to be sorted.
-		slices.SortFunc(rangeRouterSpec.Spans, func(a, b execinfrapb.OutputRouterSpec_RangeRouterSpec_Span) int {
-			return bytes.Compare(a.Start, b.Start)
-		})
 
 		// TODO(pbardea): This not super principled. I just wanted something that
 		// wasn't a constant and grew slower than linear with the length of