Repository: apache/beam
Pull request #34566: Fix Prism on handling empty and identity transforms
Author: login: shunping
Labels: name: go
name: runners
name: Next Action: Reviewers
name: prism
Created: 2025-04-07T18:16:22Z | Merged: 2025-04-08T17:27:22Z

addresses #34564


Diff:
diff --git a/sdks/go/pkg/beam/runners/prism/internal/jobservices/management.go b/sdks/go/pkg/beam/runners/prism/internal/jobservices/management.go
index db28ddc4ae58..f00838152111 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/jobservices/management.go
+++ b/sdks/go/pkg/beam/runners/prism/internal/jobservices/management.go
@@ -72,6 +72,13 @@ func (e *joinError) Error() string {
 	return string(b)
 }
 
+func getOnlyValue[K comparable, V any](in map[K]V) V {
+	for _, v := range in {
+		return v
+	}
+	panic("unreachable")
+}
+
 func (s *Server) Prepare(ctx context.Context, req *jobpb.PrepareJobRequest) (_ *jobpb.PrepareJobResponse, err error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -248,6 +255,16 @@ func (s *Server) Prepare(ctx context.Context, req *jobpb.PrepareJobRequest) (_ *
 			if len(t.GetSpec().GetPayload()) == 0 {
 				continue
 			}
+			// Another type of "empty" composite transforms without subtransforms but with
+			// a non-empty payload and identical input/output pcollections
+			if len(t.GetInputs()) == 1 && len(t.GetOutputs()) == 1 {
+				inputID := getOnlyValue(t.GetInputs())
+				outputID := getOnlyValue(t.GetOutputs())
+				if inputID == outputID {
+					slog.Warn("empty transform, with payload and identical input and output pcollection", "urn", urn, "name", t.GetUniqueName(), "pcoll", inputID)
+					continue
+				}
+			}
 			// Otherwise fail.
 			slog.Warn("unknown transform, with payload", "urn", urn, "name", t.GetUniqueName(), "payload", t.GetSpec().GetPayload())
 			check("PTransform.Spec.Urn", urn+" "+t.GetUniqueName(), "<doesn't exist>")
diff --git a/sdks/go/pkg/beam/runners/prism/internal/preprocess.go b/sdks/go/pkg/beam/runners/prism/internal/preprocess.go
index 2048d32e4ad4..5c64b08f1dc4 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/preprocess.go
+++ b/sdks/go/pkg/beam/runners/prism/internal/preprocess.go
@@ -106,6 +106,15 @@ func (p *preprocessor) preProcessGraph(comps *pipepb.Components, j *jobservices.
 
 			// If there's an unknown urn, and it's not composite, simply add it to the leaves.
 			if len(t.GetSubtransforms()) == 0 {
+				// However, if it is an empty transform with identical input/output pcollections,
+				// it will be discarded.
+				if len(t.GetInputs()) == 1 && len(t.GetOutputs()) == 1 {
+					inputID := getOnlyValue(t.GetInputs())
+					outputID := getOnlyValue(t.GetOutputs())
+					if inputID == outputID {
+						continue
+					}
+				}
 				leaves[tid] = struct{}{}
 			}
 			continue
diff --git a/sdks/go/pkg/beam/runners/prism/internal/preprocess_test.go b/sdks/go/pkg/beam/runners/prism/internal/preprocess_test.go
index 56879a3455f2..8baaa8538690 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/preprocess_test.go
+++ b/sdks/go/pkg/beam/runners/prism/internal/preprocess_test.go
@@ -125,6 +125,98 @@ func Test_preprocessor_preProcessGraph(t *testing.T) {
 					"env1": {Urn: "env1"},
 				},
 			},
+		}, {
+			name: "ignoreEmptyAndIdentityTransform",
+			input: &pipepb.Components{
+				Transforms: map[string]*pipepb.PTransform{
+					// e1 is an empty transform with identical input/output pcollections
+					"e1": {
+						UniqueName: "e1",
+						Spec: &pipepb.FunctionSpec{
+							Urn: "defaultUrn",
+						},
+						Inputs: map[string]string{
+							"in": "coll1",
+						},
+						Outputs: map[string]string{
+							"out": "coll1",
+						},
+					},
+					"e2": {
+						UniqueName: "e2",
+						Spec: &pipepb.FunctionSpec{
+							Urn: "defaultUrn",
+						},
+						Inputs: map[string]string{
+							"in": "coll1",
+						},
+						Outputs: map[string]string{
+							"out": "coll2",
+						},
+					},
+					"e3": {
+						UniqueName: "e3",
+						Spec: &pipepb.FunctionSpec{
+							Urn: "defaultUrn",
+						},
+						Inputs: map[string]string{
+							"in1": "coll1",
+							"in2": "coll2",
+						},
+						Outputs: map[string]string{
+							"out": "coll3",
+						},
+					},
+				},
+			},
+			wantStages: []*stage{
+				{
+					transforms:   []string{"e2", "e3"}, // e1 is ignored
+					primaryInput: "coll1",
+					internalCols: []string{"coll2", "coll3"},
+				},
+			},
+			wantComponents: &pipepb.Components{
+				Transforms: map[string]*pipepb.PTransform{
+					"e1": {
+						UniqueName: "e1",
+						Spec: &pipepb.FunctionSpec{
+							Urn: "defaultUrn",
+						},
+						Inputs: map[string]string{
+							"in": "coll1",
+						},
+						Outputs: map[string]string{
+							"out": "coll1",
+						},
+					},
+					"e2": {
+						UniqueName: "e2",
+						Spec: &pipepb.FunctionSpec{
+							Urn: "defaultUrn",
+						},
+						Inputs: map[string]string{
+							"in": "coll1",
+						},
+						Outputs: map[string]string{
+							"out": "coll2",
+						},
+					},
+					"e3": {
+						UniqueName: "e3",
+						Spec: &pipepb.FunctionSpec{
+							Urn: "defaultUrn",
+						},
+						Inputs: map[string]string{
+							"in1": "coll1",
+							"in2": "coll2",
+						},
+						Outputs: map[string]string{
+							"out": "coll3",
+						},
+					},
+				},
+			},
 		},
 	}
 	for _, test := range tests {
diff --git a/sdks/python/apache_beam/runners/portability/prism_runner_test.py b/sdks/python/apache_beam/runners/portability/prism_runner_test.py
index 1e8880a60a35..57bd78be3068 100644
--- a/sdks/python/apache_beam/runners/portability/prism_runner_test.py
+++ b/sdks/python/apache_beam/runners/portability/prism_runner_test.py
@@ -34,6 +34,7 @@
 from apache_beam.options.pipeline_options import PortableOptions
 from apache_beam.runners.portability import portable_runner_test
 from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
 
 # Run as
 #
@@ -178,6 +179,19 @@ def test_read(self):
       lines = p | beam.io.ReadFromText('/etc/profile')
       assert_that(lines, lambda lines: len(lines) > 0)
 
+  def test_create_transform(self):
+    with self.create_pipeline() as p:
+      out = (p | beam.Create([1]))
+      assert_that(out, equal_to([1]))
+
+    with self.create_pipeline() as p:
+      out = (p | beam.Create([1, 2], reshuffle=False))
+      assert_that(out, equal_to([1, 2]))
+
+    with self.create_pipeline() as p:
+      out = (p | beam.Create([1, 2], reshuffle=True))
+      assert_that(out, equal_to([1, 2]))
+
   def test_external_transform(self):
     raise unittest.SkipTest("Requires an expansion service to execute.")
 


Reviews:
user: login: lostluck
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-07T18:25:36Z
user: login: shunping
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-07T18:35:49Z
user: login: lostluck
body: LGTM.

Weird, but follows from it being an "empty" composite. This will probably make a bunch of similar failures pass too. Thanks!
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-08T17:18:53Z

Review comments:
diff_hunk: @@ -106,7 +106,21 @@ func (p *preprocessor) preProcessGraph(comps *pipepb.Components, j *jobservices.
 
 			// If there's an unknown urn, and it's not composite, simply add it to the leaves.
 			if len(t.GetSubtransforms()) == 0 {
-				leaves[tid] = struct{}{}
+				shouldAdd := true
+
+				// However, if it is an empty transform with identical input/output pcollections,
+				// it will be discarded.
+				if len(t.GetInputs()) == 1 && len(t.GetOutputs()) == 1 {
+					inputID := getOnlyValue(t.GetInputs())
+					outputID := getOnlyValue(t.GetOutputs())
+					if inputID == outputID {
+						shouldAdd = false
path: sdks/go/pkg/beam/runners/prism/internal/preprocess.go
user: login: lostluck
body: I'd use `continue` here instead of adding a boolean to set and check. Reduces the cyclomatic complexity of the surrounding code.
created_at: 2025-04-07T18:25:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -106,7 +106,21 @@ func (p *preprocessor) preProcessGraph(comps *pipepb.Components, j *jobservices.
 
 			// If there's an unknown urn, and it's not composite, simply add it to the leaves.
 			if len(t.GetSubtransforms()) == 0 {
-				leaves[tid] = struct{}{}
+				shouldAdd := true
+
+				// However, if it is an empty transform with identical input/output pcollections,
+				// it will be discarded.
+				if len(t.GetInputs()) == 1 && len(t.GetOutputs()) == 1 {
+					inputID := getOnlyValue(t.GetInputs())
+					outputID := getOnlyValue(t.GetOutputs())
+					if inputID == outputID {
+						shouldAdd = false
path: sdks/go/pkg/beam/runners/prism/internal/preprocess.go
user: login: shunping
body: Fixed.
created_at: 2025-04-07T18:35:48Z
author_association: COLLABORATOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-07T19:06:02Z
body: Assigning reviewers. If you would like to opt out of this review, comment `assign to next reviewer`:

R: @lostluck for label go.

Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).
author_association: CONTRIBUTOR