From a778336059c315c2fc0bdf70a56fd8f795b15ea4 Mon Sep 17 00:00:00 2001 From: Pablo Rodriguez Defino Date: Fri, 16 Oct 2026 14:17:10 -0700 Subject: [PATCH 1/2] [Prism] Finalize bundles of DoFns fused after another transform Problem: In finalizeStage (sdks/go/pkg/beam/runners/prism/internal/preprocess.go), Prism sets stage.finalize only while it inspects the consumers of the stage's input PCollections. The loop skips every PCollection that a transform of the same stage produces: for pid, plinks := range stageFacts.PcolConsumers { if _, ok := stageFacts.PcolProducers[pid]; ok { continue } ... if pardo.GetRequestsFinalization() { stg.finalize = true } So a ParDo that sets requests_finalization but is fused after another transform (a Map after Create, the expand step of a Reshuffle, any element-wise chain) never marks its stage. The SDK harness registers the callbacks and returns requires_finalization, but Prism never sends FinalizeBundleRequest, and the callbacks never run. This affects every SDK. The existing TestParDoBundleFinalizer cases apply the finalizing DoFn directly to Impulse, so the DoFn always consumes the stage input and the tests pass. Fix: Set stage.finalize from every ParDo-like transform of the stage (ParDo and the splittable DoFn stage URNs), not only from the consumers of the stage inputs. Statefulness, the on-window-expiration timer and side inputs still come from the input consumers. Add a TestParDoBundleFinalizer case with the finalizing DoFn fused after another ParDo. Reproduction (Python SDK on PrismRunner): ```python import os import tempfile import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions class MarkFinalization(beam.DoFn): """Registers a bundle finalization callback that creates the file `path`.""" def __init__(self, path): self._path = path def process(self, element, bundle_finalizer=beam.DoFn.BundleFinalizerParam): path = self._path bundle_finalizer.register(lambda: open(path, "w").close()) yield element if __name__ == "__main__": markers = tempfile.mkdtemp() at_input = os.path.join(markers, "at_stage_input") fused = os.path.join(markers, "after_fused_map") options = PipelineOptions(["--runner=PrismRunner", "--environment_type=LOOPBACK"]) with beam.Pipeline(options=options) as p: impulse = p | beam.Impulse() # Consumes the input of its stage: Prism finalizes the bundle. _ = impulse | "AtStageInput" >> beam.ParDo(MarkFinalization(at_input)) # Fused after a Map in the same stage: Prism never finalizes the bundle. _ = ( impulse | "Map" >> beam.Map(lambda x: x) | "AfterFusedMap" >> beam.ParDo(MarkFinalization(fused)) ) print("finalized:", sorted(os.listdir(markers))) assert os.path.exists(at_input), "the DoFn at the stage input was not finalized" assert os.path.exists(fused), "the DoFn fused after Map was not finalized" ``` Before the fix: `finalized: ['at_stage_input']`, and the second assertion fails. After the fix: `finalized: ['after_fused_map', 'at_stage_input']`. --- .../beam/runners/prism/internal/preprocess.go | 19 ++++++++++++++++--- sdks/go/test/integration/primitives/pardo.go | 15 +++++++++++++++ .../test/integration/primitives/pardo_test.go | 5 +++++ 3 files changed, 36 insertions(+), 3 deletions(-) diff --git a/sdks/go/pkg/beam/runners/prism/internal/preprocess.go b/sdks/go/pkg/beam/runners/prism/internal/preprocess.go index 3311bcced9f4..0abb680dd0bd 100644 --- a/sdks/go/pkg/beam/runners/prism/internal/preprocess.go +++ b/sdks/go/pkg/beam/runners/prism/internal/preprocess.go @@ -446,6 +446,22 @@ 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 { + t := comps.GetTransforms()[tid] + switch t.GetSpec().GetUrn() { + case urns.TransformParDo, urns.TransformProcessSizedElements, urns.TransformPairWithRestriction, urns.TransformSplitAndSize, urns.TransformTruncate: + pardo := &pipepb.ParDoPayload{} + if err := (proto.UnmarshalOptions{}).Unmarshal(t.GetSpec().GetPayload(), pardo); err != nil { + return fmt.Errorf("unable to decode ParDoPayload for %v", tid) + } + if pardo.GetRequestsFinalization() { + stg.finalize = true + } + } + } + // Now we can see which consumers (inputs) aren't covered by the producers (outputs). mainInputs := map[string]string{} var sideInputs []engine.LinkID @@ -468,9 +484,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 } 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, From e9df3d674a85f64034870bc574a4a8ed9ee7a6c2 Mon Sep 17 00:00:00 2001 From: Pablo Rodriguez Defino Date: Fri, 16 Oct 2026 14:47:47 -0700 Subject: [PATCH 2/2] [Prism] Test the finalization flag of a stage, and drop an unreachable error Move the check of requests_finalization into requestsFinalization. finalizeStage runs computeFacts first, which already rejects a ParDo payload that does not decode, so the error return of the earlier loop could not run. Add TestFinalizeStage_RequestsFinalization: the flag stays false when no transform asks for it, and is set when the transform at the stage input, a ParDo fused after another transform, or a fused splittable DoFn process stage asks for it. --- .../beam/runners/prism/internal/preprocess.go | 27 +++++--- .../runners/prism/internal/preprocess_test.go | 65 +++++++++++++++++++ 2 files changed, 82 insertions(+), 10 deletions(-) diff --git a/sdks/go/pkg/beam/runners/prism/internal/preprocess.go b/sdks/go/pkg/beam/runners/prism/internal/preprocess.go index 0abb680dd0bd..8f6000276a7f 100644 --- a/sdks/go/pkg/beam/runners/prism/internal/preprocess.go +++ b/sdks/go/pkg/beam/runners/prism/internal/preprocess.go @@ -449,16 +449,9 @@ func finalizeStage(stg *stage, comps *pipepb.Components, pipelineFacts *fusionFa // 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 { - t := comps.GetTransforms()[tid] - switch t.GetSpec().GetUrn() { - case urns.TransformParDo, urns.TransformProcessSizedElements, urns.TransformPairWithRestriction, urns.TransformSplitAndSize, urns.TransformTruncate: - pardo := &pipepb.ParDoPayload{} - if err := (proto.UnmarshalOptions{}).Unmarshal(t.GetSpec().GetPayload(), pardo); err != nil { - return fmt.Errorf("unable to decode ParDoPayload for %v", tid) - } - if pardo.GetRequestsFinalization() { - stg.finalize = true - } + if requestsFinalization(comps.GetTransforms()[tid]) { + stg.finalize = true + break } } @@ -548,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) + } + }) + } +}