From e06cc8c58bb0ab041f875041f1e7632a0435401d Mon Sep 17 00:00:00 2001 From: Malo Denielou Date: Wed, 28 Oct 2026 20:30:00 +0000 Subject: [PATCH] [Go SDK] Make pipelinex.computeRoots and xlangx.HandlerParams ordering deterministic Fix two sources of map-iteration non-determinism in the Go SDK pipeline construction and cross-language expansion paths: 1. In pipelinex.computeRoots, sort the collected root transform IDs using sort.Sort(idSorted(roots)) before passing them to TopologicalSort. While TopologicalSort already sorts its internal frontier (ret.next) and output keys for deterministic traversal, the initial roots slice was populated by ranging over the xforms map in random order. 2. In xlangx.HandlerParams.Inputs() and Outputs(), sort the returned []PCol slice by PCol.Index (breaking ties by PCol.Local) after collecting entries from InputsMap and OutputsMap. Retaining append + slices.SortFunc (rather than indexing directly into a preallocated slice by index) also safely handles sparse index maps such as when xlang.SetOutputCoder rewrites OutputsMap to a single entry at index > 0. --- .../beam/core/runtime/pipelinex/replace.go | 1 + .../core/runtime/pipelinex/replace_test.go | 17 +++++ .../pkg/beam/core/runtime/xlangx/registry.go | 11 ++++ .../beam/core/runtime/xlangx/registry_test.go | 63 +++++++++++++++++++ 4 files changed, 92 insertions(+) diff --git a/sdks/go/pkg/beam/core/runtime/pipelinex/replace.go b/sdks/go/pkg/beam/core/runtime/pipelinex/replace.go index 9e527f2fd322..60be88b57a34 100644 --- a/sdks/go/pkg/beam/core/runtime/pipelinex/replace.go +++ b/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) } diff --git a/sdks/go/pkg/beam/core/runtime/pipelinex/replace_test.go b/sdks/go/pkg/beam/core/runtime/pipelinex/replace_test.go index 3024787e6163..b88da9fc3bca 100644 --- a/sdks/go/pkg/beam/core/runtime/pipelinex/replace_test.go +++ b/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) + } + } +} diff --git a/sdks/go/pkg/beam/core/runtime/xlangx/registry.go b/sdks/go/pkg/beam/core/runtime/xlangx/registry.go index 1957cc61bb95..67b348abca6f 100644 --- a/sdks/go/pkg/beam/core/runtime/xlangx/registry.go +++ b/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 } diff --git a/sdks/go/pkg/beam/core/runtime/xlangx/registry_test.go b/sdks/go/pkg/beam/core/runtime/xlangx/registry_test.go index 844a8cbd7b85..25b745ebce93 100644 --- a/sdks/go/pkg/beam/core/runtime/xlangx/registry_test.go +++ b/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) + } +}