[Go SDK] Make pipelinex.computeRoots and xlangx.HandlerParams ordering deterministic

sdks/go/pkg/beam/core/runtime/pipelinex/replace.go

@@ -100,6 +100,7 @@ func computeRoots(xforms map[string]*pipepb.PTransform) []string {
 			roots = append(roots, id)
 		}
 	}
+	sort.Sort(idSorted(roots))
 	return TopologicalSort(xforms, roots)
 }
 

sdks/go/pkg/beam/core/runtime/pipelinex/replace_test.go

@@ -425,3 +425,20 @@ func buildEnvironment(t *testing.T, containerImg string) *pipepb.Environment {
 	env.Payload = plb
 	return env
 }
+
+func TestComputeRootsDeterminism(t *testing.T) {
+	xforms := map[string]*pipepb.PTransform{
+		"e10": {UniqueName: "root10"},
+		"e2":  {UniqueName: "root2", Inputs: map[string]string{"i0": "p1"}},
+		"e1":  {UniqueName: "root1", Outputs: map[string]string{"i0": "p1"}},
+		"s1":  {UniqueName: "comp1", Subtransforms: []string{"e3"}},
+		"e3":  {UniqueName: "child3"},
+	}
+	want := []string{"s1", "e10", "e1", "e2"}
+	for i := 0; i < 20; i++ {
+		got := computeRoots(xforms)
+		if diff := cmp.Diff(want, got); diff != "" {
+			t.Fatalf("iteration %d: computeRoots(%v) = %v, want %v (-want +got):\n%s", i, xforms, got, want, diff)
+		}
+	}
+}

sdks/go/pkg/beam/core/runtime/xlangx/registry.go

@@ -16,9 +16,11 @@
 package xlangx
 
 import (
+	"cmp"
 	"context"
 	"fmt"
 	"net/url"
+	"slices"
 	"strings"
 	"sync"
 
@@ -126,6 +128,13 @@ func makePCol(node *graph.Node, index int, local, namespace string) PCol {
 	}
 }
 
+func comparePCol(a, b PCol) int {
+	if c := cmp.Compare(a.Index, b.Index); c != 0 {
+		return c
+	}
+	return cmp.Compare(a.Local, b.Local)
+}
+
 // Outputs returns the provided output PCollections, if any, for expected outputs
 // for this expansion service request.
 //
@@ -136,6 +145,7 @@ func (p *HandlerParams) Outputs() []PCol {
 	for local, i := range p.ext.OutputsMap {
 		out = append(out, makePCol(p.edge.Output[i].To, i, local, p.Req.Namespace))
 	}
+	slices.SortFunc(out, comparePCol)
 	return out
 }
 
@@ -146,6 +156,7 @@ func (p *HandlerParams) Inputs() []PCol {
 	for local, i := range p.ext.InputsMap {
 		out = append(out, makePCol(p.edge.Input[i].From, i, local, p.Req.Namespace))
 	}
+	slices.SortFunc(out, comparePCol)
 	return out
 }
 

sdks/go/pkg/beam/core/runtime/xlangx/registry_test.go

@@ -20,6 +20,7 @@ import (
 	"reflect"
 	"testing"
 
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph"
 	jobpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/jobmanagement_v1"
 )
 
@@ -86,3 +87,65 @@ func TestUseAutomatedExpansionService(t *testing.T) {
 		t.Errorf("got %v, want %v", got, want)
 	}
 }
+
+func TestHandlerParamsInputsOutputsOrder(t *testing.T) {
+	in0 := &graph.Node{}
+	in1 := &graph.Node{}
+	in2 := &graph.Node{}
+	out0 := &graph.Node{}
+	out1 := &graph.Node{}
+
+	p := &HandlerParams{
+		Req: &jobpb.ExpansionRequest{Namespace: "ns"},
+		edge: &graph.MultiEdge{
+			Input: []*graph.Inbound{
+				{From: in0},
+				{From: in1},
+				{From: in2},
+			},
+			Output: []*graph.Outbound{
+				{To: out0},
+				{To: out1},
+			},
+		},
+		ext: &graph.ExternalTransform{
+			InputsMap: map[string]int{
+				"table_c": 2,
+				"table_a": 0,
+				"table_b": 1,
+			},
+			OutputsMap: map[string]int{
+				"out_b": 1,
+				"out_a": 0,
+			},
+		},
+	}
+
+	for iter := 0; iter < 20; iter++ {
+		ins := p.Inputs()
+		wantInLocals := []string{"table_a", "table_b", "table_c"}
+		for i, wantLocal := range wantInLocals {
+			if ins[i].Index != i || ins[i].Local != wantLocal {
+				t.Fatalf("iter %d: Inputs()[%d] = {Index: %d, Local: %q}, want {Index: %d, Local: %q}",
+					iter, i, ins[i].Index, ins[i].Local, i, wantLocal)
+			}
+		}
+
+		outs := p.Outputs()
+		wantOutLocals := []string{"out_a", "out_b"}
+		for i, wantLocal := range wantOutLocals {
+			if outs[i].Index != i || outs[i].Local != wantLocal {
+				t.Fatalf("iter %d: Outputs()[%d] = {Index: %d, Local: %q}, want {Index: %d, Local: %q}",
+					iter, i, outs[i].Index, outs[i].Local, i, wantLocal)
+			}
+		}
+	}
+
+	// Verify that a subset/sparse OutputsMap (such as when SetOutputCoder
+	// rewrites OutputsMap to a single entry at index > 0) does not panic.
+	p.ext.OutputsMap = map[string]int{graph.UnnamedOutputTag: 1}
+	outs := p.Outputs()
+	if len(outs) != 1 || outs[0].Index != 1 || outs[0].Local != graph.UnnamedOutputTag {
+		t.Fatalf("Outputs() with non-zero index = %v, want single entry at Index 1", outs)
+	}
+}