[Prism] Finalize bundles of DoFns fused after another transform

sdks/go/pkg/beam/runners/prism/internal/preprocess.go

@@ -446,6 +446,15 @@ func finalizeStage(stg *stage, comps *pipepb.Components, pipelineFacts *fusionFa
 		transformSet[tid] = true
 	}
 
+	// Any ParDo in the stage can register bundle finalization callbacks, not only the
+	// transforms that consume the stage's inputs, so every one of them must be checked.
+	for _, tid := range stg.transforms {
+		if requestsFinalization(comps.GetTransforms()[tid]) {
+			stg.finalize = true
+			break
+		}
+	}
+
 	// Now we can see which consumers (inputs) aren't covered by the producers (outputs).
 	mainInputs := map[string]string{}
 	var sideInputs []engine.LinkID
@@ -468,9 +477,6 @@ func finalizeStage(stg *stage, comps *pipepb.Components, pipelineFacts *fusionFa
 				if err := (proto.UnmarshalOptions{}).Unmarshal(t.GetSpec().GetPayload(), pardo); err != nil {
 					return fmt.Errorf("unable to decode ParDoPayload for %v", link.Transform)
 				}
-				if pardo.GetRequestsFinalization() {
-					stg.finalize = true
-				}
 				if len(pardo.GetTimerFamilySpecs())+len(pardo.GetStateSpecs())+len(pardo.GetOnWindowExpirationTimerFamilySpec()) > 0 {
 					stg.stateful = true
 				}
@@ -535,6 +541,20 @@ func finalizeStage(stg *stage, comps *pipepb.Components, pipelineFacts *fusionFa
 	return nil
 }
 
+// requestsFinalization reports whether t is a ParDo, or a stage of an expanded splittable
+// DoFn, whose payload requests bundle finalization. computeFacts has already rejected
+// undecodable ParDo payloads, so a payload that does not decode reports false.
+func requestsFinalization(t *pipepb.PTransform) bool {
+	switch t.GetSpec().GetUrn() {
+	case urns.TransformParDo, urns.TransformProcessSizedElements, urns.TransformPairWithRestriction, urns.TransformSplitAndSize, urns.TransformTruncate:
+		pardo := &pipepb.ParDoPayload{}
+		err := (proto.UnmarshalOptions{}).Unmarshal(t.GetSpec().GetPayload(), pardo)
+		return err == nil && pardo.GetRequestsFinalization()
+	default:
+		return false
+	}
+}
+
 // greedyFusion produces a pipeline as tightly fused as possible.
 //
 // Fusion is a critical optimization for performance of pipeline execution.

sdks/go/pkg/beam/runners/prism/internal/preprocess_test.go

@@ -789,3 +789,68 @@ func TestComputeFacts(t *testing.T) {
 		})
 	}
 }
+
+func TestFinalizeStage_RequestsFinalization(t *testing.T) {
+	pardo := func(urn string, finalize bool) *pipepb.FunctionSpec {
+		return &pipepb.FunctionSpec{
+			Urn:     urn,
+			Payload: protox.MustEncode(&pipepb.ParDoPayload{RequestsFinalization: finalize}),
+		}
+	}
+	// first consumes the stage input n0; second is fused after it and consumes n1.
+	stageOf := func(first, second *pipepb.FunctionSpec) *pipepb.Components {
+		return &pipepb.Components{
+			Transforms: map[string]*pipepb.PTransform{
+				"first": {
+					Spec:    first,
+					Inputs:  map[string]string{"i0": "n0"},
+					Outputs: map[string]string{"o0": "n1"},
+				},
+				"second": {
+					Spec:    second,
+					Inputs:  map[string]string{"i0": "n1"},
+					Outputs: map[string]string{"o0": "n2"},
+				},
+			},
+		}
+	}
+
+	tests := []struct {
+		name         string
+		comps        *pipepb.Components
+		wantFinalize bool
+	}{
+		{
+			name:  "no_transform_requests_finalization",
+			comps: stageOf(&pipepb.FunctionSpec{Urn: urns.TransformAssignWindows}, pardo(urns.TransformParDo, false)),
+		}, {
+			name:         "transform_at_stage_input",
+			comps:        stageOf(pardo(urns.TransformParDo, true), pardo(urns.TransformParDo, false)),
+			wantFinalize: true,
+		}, {
+			name:         "pardo_fused_after_another_transform",
+			comps:        stageOf(pardo(urns.TransformParDo, false), pardo(urns.TransformParDo, true)),
+			wantFinalize: true,
+		}, {
+			name:         "sdf_process_fused_after_another_transform",
+			comps:        stageOf(pardo(urns.TransformParDo, false), pardo(urns.TransformProcessSizedElements, true)),
+			wantFinalize: true,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			topological := []string{"first", "second"}
+			facts, err := computeFacts(topological, test.comps)
+			if err != nil {
+				t.Fatalf("computeFacts() = %v", err)
+			}
+			stg := &stage{transforms: topological}
+			if err := finalizeStage(stg, test.comps, facts); err != nil {
+				t.Fatalf("finalizeStage() = %v", err)
+			}
+			if stg.finalize != test.wantFinalize {
+				t.Errorf("stage.finalize = %v, want %v", stg.finalize, test.wantFinalize)
+			}
+		})
+	}
+}

sdks/go/test/integration/primitives/pardo.go

@@ -34,6 +34,7 @@ func init() {
 	register.Function3x2(asymJoinFn)
 	register.Function5x0(splitByName)
 	register.Function2x0(emitPipelineOptions)
+	register.Function1x1(passBytesFn)
 	register.DoFn2x0[beam.BundleFinalization, []byte]((*processElemBundleFinalizer)(nil))
 	register.DoFn2x0[beam.BundleFinalization, []byte]((*finalizerInFinishBundle)(nil))
 	register.DoFn2x0[beam.BundleFinalization, []byte]((*finalizerInAll)(nil))
@@ -213,6 +214,20 @@ func ParDoProcessElementBundleFinalizer(s beam.Scope) {
 	beam.ParDo0(s, &processElemBundleFinalizer{}, imp)
 }
 
+// ParDoBundleFinalizerAfterFusedParDo creates a beam.Pipeline with a beam.ParDo0 that processes
+// a DoFn with a beam.BundleFinalization in its ProcessElement method, after another ParDo.
+// The runner fuses both ParDos into one stage, so the DoFn that requests finalization does not
+// consume the input of the stage.
+func ParDoBundleFinalizerAfterFusedParDo(s beam.Scope) {
+	imp := beam.Impulse(s)
+	passed := beam.ParDo(s, passBytesFn, imp)
+	beam.ParDo0(s, &processElemBundleFinalizer{}, passed)
+}
+
+func passBytesFn(b []byte) []byte {
+	return b
+}
+
 type processElemBundleFinalizer struct {
 }
 

sdks/go/test/integration/primitives/pardo_test.go

@@ -64,6 +64,11 @@ func TestParDoBundleFinalizer(t *testing.T) {
 			pipelineFn: ParDoProcessElementBundleFinalizer,
 			want:       BundleFinalizerProcess,
 		},
+		{
+			name:       "InProcessElementAfterFusedParDo",
+			pipelineFn: ParDoBundleFinalizerAfterFusedParDo,
+			want:       BundleFinalizerProcess,
+		},
 		{
 			name:       "InFinishBundle",
 			pipelineFn: ParDoFinishBundleFinalizer,