[Python] Add PCollection.with_side_outputs for chainable side outputs#38268
[Python] Add PCollection.with_side_outputs for chainable side outputs#38268hjtran wants to merge 5 commits into
Conversation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enhances the Apache Beam Python SDK by enabling PCollections to carry named side outputs while remaining chainable. This solves a long-standing usability issue where transform authors had to choose between returning a chainable PCollection or a container that preserves side outputs. The changes introduce a new API for attaching side outputs and ensure these are correctly registered within the pipeline graph, maintaining compatibility with existing runner APIs and composite transform structures. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
Rename _FilterWithSideOutputs/_SplitOddDoFn helpers in pipeline_test.py
to RemoveEvens/_RemoveEvensDoFn for clarity, and flip the DoFn semantics
so the name matches the behavior (evens go to the 'dropped' side output).
Drop two tests that were not pulling their weight:
- test_pcollection_side_outputs_wraps_with_outputs: fully subsumed by
test_pcollection_side_outputs_end_to_end, which already asserts the
applied transform's outputs contain both tags.
- test_pcollection_side_outputs_allows_idempotent_tag_collision: only
reachable by reaching past the public API to set _side_outputs
manually; the defensive impl branch is still worth keeping but the
test is not.
Merge the two pvalue_test.py validation tests into a single
test_with_side_outputs_validation.
|
This pull request has been marked as stale due to 60 days of inactivity. It will be closed in 1 week if no further activity occurs. If you think that’s incorrect or this pull request requires a review, please simply write any comment. If closed, you can revive the PR at any time and @mention a reviewer or discuss it on the dev@beam.apache.org list. Thank you for your contributions. |
|
This pull request has been closed due to lack of activity. If you think that is incorrect, or the pull request requires review, you can revive the PR at any time. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for registering and accessing side outputs on PCollection objects in Apache Beam's Python SDK, adding a _SideOutputsContainer for attribute/index access, a with_side_outputs method, and updating pipeline construction and replacement logic. The review feedback focuses on improving robustness and usability: it suggests using getattr to safely access _side_outputs on reconstructed or uninitialized PCollection instances, expanding the underscore check in getattr to prevent infinite recursion, simplifying attribute assignment in _SideOutputsContainer, and implementing dir to support auto-completion in interactive environments.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if not result._side_outputs: | ||
| return | ||
|
|
||
| descendants = _descendant_applied_ptransforms(current) | ||
| for side_tag, side_pcoll in result._side_outputs.items(): |
There was a problem hiding this comment.
Accessing result._side_outputs directly can raise an AttributeError if the PCollection was reconstructed (e.g., via from_runner_api or other deserialization paths) where _side_outputs is not populated in the instance's __dict__. Using getattr(result, '_side_outputs', None) is safer and prevents potential crashes during pipeline replacement or other graph operations on reconstructed pipelines.
side_outputs = getattr(result, '_side_outputs', None)
if not side_outputs:
return
descendants = _descendant_applied_ptransforms(current)
for side_tag, side_pcoll in side_outputs.items():|
|
||
| @property | ||
| def side_outputs(self) -> _SideOutputsContainer: | ||
| return _SideOutputsContainer(self._side_outputs or {}) |
There was a problem hiding this comment.
Accessing self._side_outputs directly in the side_outputs property will raise an AttributeError if the PCollection was reconstructed (e.g., from the runner API) or created without calling __init__ (which is common in some serialization or mocking frameworks). Using getattr(self, '_side_outputs', None) ensures that this property safely returns an empty container instead of crashing.
| return _SideOutputsContainer(self._side_outputs or {}) | |
| return _SideOutputsContainer(getattr(self, '_side_outputs', None) or {}) |
| if tag.startswith('__'): | ||
| raise AttributeError(tag) |
There was a problem hiding this comment.
If _side_outputs is not yet initialized or is missing from __dict__ (for example, during unpickling or copying), accessing self._side_outputs inside __getattr__ will trigger __getattr__('_side_outputs'), leading to infinite recursion and a RecursionError. Restricting dynamic attribute lookup to exclude any attributes starting with an underscore (or specifically _side_outputs) prevents this infinite recursion.
| if tag.startswith('__'): | |
| raise AttributeError(tag) | |
| if tag.startswith('_'): | |
| raise AttributeError(tag) |
| def __init__(self, side_outputs: dict[str, 'PCollection']): | ||
| object.__setattr__(self, '_side_outputs', dict(side_outputs)) |
There was a problem hiding this comment.
Since _SideOutputsContainer does not override __setattr__, using object.__setattr__ is unnecessary and can be confusing. Standard attribute assignment self._side_outputs = dict(side_outputs) is cleaner and more idiomatic.
| def __init__(self, side_outputs: dict[str, 'PCollection']): | |
| object.__setattr__(self, '_side_outputs', dict(side_outputs)) | |
| def __init__(self, side_outputs: dict[str, 'PCollection']): | |
| self._side_outputs = dict(side_outputs) |
| def __contains__(self, tag): | ||
| return tag in self._side_outputs |
There was a problem hiding this comment.
Implementing __dir__ on _SideOutputsContainer allows interactive environments (like Jupyter notebooks or IPython) to auto-complete the available side output tags, greatly improving usability.
| def __contains__(self, tag): | |
| return tag in self._side_outputs | |
| def __contains__(self, tag): | |
| return tag in self._side_outputs | |
| def __dir__(self): | |
| return sorted(super().__dir__() + list(getattr(self, '_side_outputs', {}).keys())) |
|
Assigning reviewers: R: @jrmccluskey for label python. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
|
Whoops sorry I didnt mean to post this yet. No need to review yet @jrmccluskey |
1 participant
Summary
Adds
PCollection.with_side_outputs(**kwargs)to the Python SDK so transforms can return a single, chainablePCollectionthat also carries named side-output PCollections, accessible viaresult.side_outputs.<tag>.Today, transform authors face a tradeoff between returning a
PCollection(chainable, but loses side outputs) or adict/PCollectionTuple(preserves side outputs, but breaks chaining). This change lets them have both.Design
_SideOutputsContainerexposes attribute (container.dropped) and item (container["dropped"]) access plus iteration /len/in.PCollection.with_side_outputs(**kwargs)returns acopy.copy(self)with the side-output mapping attached. The original is unchanged. Validation: each value must be aPCollection(TypeError) on the same pipeline (ValueError).Pipeline._apply_internalandPipeline._replaceboth call a shared helper that registers the side outputs on the wrappingAppliedPTransform— only for the top-level returnedPCollection. This makes side outputs first-class in the pipeline graph (visible to runners, the proto, visualization, and YAML'sTransform.tagresolution).ValueError.PCollectionobject (idempotent); otherwiseValueError.__copy__is added toPCollectionsocopy.copy(self)doesn't dispatch to the existing__reduce_ex__anti-pickling hook.Chaining behavior
pc | A | Bis unchanged.BreceivesA's main output as a normal PCollection.A's side outputs are not carried forward — capture the intermediate result if you need them. This matches the behavior described in the proposal and requires no special code.Caveats
**kwargssyntax restricts tags to valid Python identifiers. Tags with hyphens/spaces are not supported here; users with arbitrary string tags should keep using the existing dict /DoOutputsTuplepatterns..side_outputsis a construction-time ergonomic, not a durable graph property. Afterfrom_runner_api, the reconstructed PCollection won't have_side_outputspopulated, but the side outputs themselves still exist as named outputs on the producer'sAppliedPTransform.DoOutputsTuple).Follow-ups (not in this PR)
ParDo.with_side_outputs(...)convenience method.get_pcollectionalready resolvesTransform.tagby walking the producer's outputs dict, so this should mostly "just work" once verified).Tests
pvalue_test.pycovers: copy semantics, attribute / index access, missing-tag errors, type/pipeline validation, empty container, double-call replace.pipeline_test.pycovers: end-to-end materialization withassert_that, the canonicalMyFilter-style composite wrappingwith_outputs, foreign-pcollection rejection, tag-collision rejection, idempotent same-object collision, replacement-path registration viaPipeline.replace_all, runner API round-trip, and the nested-return non-flattening case.Test plan
cd sdks/python && python -m pytest apache_beam/pvalue_test.py apache_beam/pipeline_test.py— 84 passed, 1 skippedruff check --config=sdks/python/ruff.tomlon the four changed files — cleanyapf --diffon the four changed files — clean