diff --git a/sdks/go/pkg/beam/runners/prism/internal/preprocess.go b/sdks/go/pkg/beam/runners/prism/internal/preprocess.go index 3311bcced9f4..8f6000276a7f 100644 --- a/sdks/go/pkg/beam/runners/prism/internal/preprocess.go +++ b/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. 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 8baaa8538690..f4064e814b39 100644 --- a/sdks/go/pkg/beam/runners/prism/internal/preprocess_test.go +++ b/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) + } + }) + } +} diff --git a/sdks/go/test/integration/primitives/pardo.go b/sdks/go/test/integration/primitives/pardo.go index dc59d8f67b80..a799f142745a 100644 --- a/sdks/go/test/integration/primitives/pardo.go +++ b/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 { } diff --git a/sdks/go/test/integration/primitives/pardo_test.go b/sdks/go/test/integration/primitives/pardo_test.go index aa6cb3de2008..1c169272b049 100644 --- a/sdks/go/test/integration/primitives/pardo_test.go +++ b/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,