[#34564] Fix Prism on handling empty and identity transforms (#34566)

* Fix an edge case of handling empty composite transform from Python SDK.

* A more general fix to cover empty transform with identical input and output pcoll

* Simplify the code with continue.

* Add unit tests.

Author: shunping

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>")

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

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 {

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.")