This is an automated email from the ASF dual-hosted git repository.
lostluck pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new c31f2327041 Fix Reshuffle handling in Prism to recursively remove
nested sub-transforms (#38735)
c31f2327041 is described below
commit c31f23270416ca446dd4bfcdb896004106149916
Author: Shunping Huang <[email protected]>
AuthorDate: Fri May 29 17:56:06 2026 -0400
Fix Reshuffle handling in Prism to recursively remove nested sub-transforms
(#38735)
Updates `handleReshuffle` in the Go Prism runner to recursively clean up all
nested sub-transforms of a Reshuffle composite transform using
`removeSubTransforms`. This prevents nested sub-transforms from being
processed in subsequent graph optimization stages.
---
.../beam/runners/prism/internal/handlerunner.go | 4 +-
.../runners/prism/internal/handlerunner_test.go | 78 ++++++++++++++++++++++
2 files changed, 80 insertions(+), 2 deletions(-)
diff --git a/sdks/go/pkg/beam/runners/prism/internal/handlerunner.go
b/sdks/go/pkg/beam/runners/prism/internal/handlerunner.go
index 3ac0d98850d..ab7d505f8c6 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/handlerunner.go
+++ b/sdks/go/pkg/beam/runners/prism/internal/handlerunner.go
@@ -209,8 +209,8 @@ func (h *runner) handleReshuffle(tid string, t
*pipepb.PTransform, comps *pipepb
}
}
- // And all the sub transforms.
- toRemove = append(toRemove, t.GetSubtransforms()...)
+ // Also recursively remove all sub-transforms.
+ toRemove = append(toRemove, removeSubTransforms(comps,
t.GetSubtransforms())...)
// Return the new components which is the transforms consumer
return prepareResult{
diff --git a/sdks/go/pkg/beam/runners/prism/internal/handlerunner_test.go
b/sdks/go/pkg/beam/runners/prism/internal/handlerunner_test.go
new file mode 100644
index 00000000000..da3ac9d14f0
--- /dev/null
+++ b/sdks/go/pkg/beam/runners/prism/internal/handlerunner_test.go
@@ -0,0 +1,78 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package internal
+
+import (
+ "testing"
+
+ pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
+ "github.com/google/go-cmp/cmp"
+ "google.golang.org/protobuf/testing/protocmp"
+)
+
+func TestHandleReshuffle(t *testing.T) {
+ h := &runner{
+ config: RunnerCharacteristic{
+ SDKReshuffle: false,
+ },
+ }
+
+ comps := &pipepb.Components{
+ Transforms: map[string]*pipepb.PTransform{
+ "reshuffle": {
+ UniqueName: "reshuffle",
+ Inputs: map[string]string{
+ "in": "pcol_in",
+ },
+ Outputs: map[string]string{
+ "out": "pcol_out",
+ },
+ Subtransforms: []string{"sub_1"},
+ },
+ "sub_1": {
+ UniqueName: "sub_1",
+ Subtransforms: []string{"sub_2"},
+ },
+ "sub_2": {
+ UniqueName: "sub_2",
+ },
+ "consumer": {
+ UniqueName: "consumer",
+ Inputs: map[string]string{
+ "in": "pcol_out",
+ },
+ },
+ },
+ }
+
+ got := h.handleReshuffle("reshuffle",
comps.GetTransforms()["reshuffle"], comps)
+
+ want := prepareResult{
+ SubbedComps: nil,
+ RemovedLeaves: []string{"sub_1", "sub_2"},
+ ForcedRoots: []string{"consumer"},
+ }
+
+ if d := cmp.Diff(want, got, protocmp.Transform()); d != "" {
+ t.Errorf("handleReshuffle() diff (-want, +got):\n%v", d)
+ }
+
+ // Verify that the consumer's input has been remapped to the input of
the reshuffle
+ gotInput := comps.GetTransforms()["consumer"].GetInputs()["in"]
+ if gotInput != "pcol_in" {
+ t.Errorf("consumer input got %q, want %q", gotInput, "pcol_in")
+ }
+}
