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) + } +}