[Prism] Finalize bundles of DoFns fused after another transform

[prodriguezdefino]

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

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'].

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 56.22%. Comparing base (76f6e6d) to head (e9df3d6).
⚠️ Report is 215 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #40640      +/-   ##
============================================
+ Coverage     56.16%   56.22%   +0.05%     
============================================
  Files          1124      868     -256     
  Lines        178336   153546   -24790     
  Branches       1489        0    -1489     
============================================
- Hits         100170    86327   -13843     
+ Misses        75645    65382   -10263     
+ Partials       2521     1837     -684     

Flag	Coverage Δ
go	28.64% <100.00%> (-0.50%)	⬇️
java	?
python	79.93% <ø> (+0.37%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[github-actions]

📋 Review Summary

This PR resolves a bug in Prism where bundle finalization callbacks were skipped for ParDo transforms fused downstream of another transform within the same stage. The fix properly inspects all transforms in the stage for requests_finalization rather than only the transforms consuming the stage's external inputs, and includes both unit and integration tests.

🔍 General Feedback

• The root cause analysis and fix in finalizeStage are clean, accurate, and properly account for both standard ParDo and expanded Splittable DoFn transform URNs.
• Unit test coverage in preprocess_test.go and integration test coverage in primitives/pardo_test.go thoroughly verify the fix.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

[github-actions]

Assigning reviewers:

R: @lostluck for label go.
Backup reviewer: shunping (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

Note: If you would like to opt out of this review, comment assign to next reviewer.

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

[lostluck]

Good find! Looking presently.

[lostluck]

LGTM