You are viewing a plain text version of this content. The canonical link for it is here.

Posted to github@beam.apache.org by "hjtran (via GitHub)" <gi...@apache.org> on 2026/07/14 20:43:21 UTC

[PR] [Python] Add PCollection.with_side_outputs for chainable side outputs [beam]

hjtran opened a new pull request, #38268:
URL: https://github.com/apache/beam/pull/38268

   ## Summary
   
   Adds `PCollection.with_side_outputs(**kwargs)` to the Python SDK so transforms can return a single, chainable `PCollection` that *also* carries named side-output PCollections, accessible via `result.side_outputs.<tag>`.
   
   ```python
   class MyFilter(beam.PTransform):
       def expand(self, pcoll):
           results = pcoll | beam.ParDo(FilterDoFn()).with_outputs('dropped', main='main')
           return results['main'].with_side_outputs(dropped=results['dropped'])
   
   # Most users — just chain
   inputs | MyFilter() | NextTransform()
   
   # When you need the side output
   result = inputs | MyFilter()
   result.side_outputs.dropped
   ```
   
   Today, transform authors face a tradeoff between returning a `PCollection` (chainable, but loses side outputs) or a `dict`/`PCollectionTuple` (preserves side outputs, but breaks chaining). This change lets them have both.
   
   ## Design
   
   - New `_SideOutputsContainer` exposes attribute (`container.dropped`) and item (`container["dropped"]`) access plus iteration / `len` / `in`.
   - `PCollection.with_side_outputs(**kwargs)` returns a `copy.copy(self)` with the side-output mapping attached. The original is unchanged. Validation: each value must be a `PCollection` (TypeError) on the same pipeline (ValueError).
   - `Pipeline._apply_internal` and `Pipeline._replace` both call a shared helper that registers the side outputs on the wrapping `AppliedPTransform` — *only* for the top-level returned `PCollection`. This makes side outputs first-class in the pipeline graph (visible to runners, the proto, visualization, and YAML's `Transform.tag` resolution).
   - The helper enforces two apply-time invariants:
     - **Provenance:** each side output's producer must be the wrapping transform itself or one of its descendants. Foreign PCollections are rejected with `ValueError`.
     - **Tag collision:** if a side-output tag already exists in the wrapping transform's outputs, the value must be the same `PCollection` object (idempotent); otherwise `ValueError`.
   - `__copy__` is added to `PCollection` so `copy.copy(self)` doesn't dispatch to the existing `__reduce_ex__` anti-pickling hook.
   
   ## Chaining behavior
   
   `pc | A | B` is unchanged. `B` receives `A`'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
   
   - `**kwargs` syntax 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 / `DoOutputsTuple` patterns.
   - `.side_outputs` is a construction-time ergonomic, not a durable graph property. After `from_runner_api`, the reconstructed PCollection won't have `_side_outputs` populated, but the side outputs themselves still exist as named outputs on the producer's `AppliedPTransform`.
   - Side outputs do not participate in composite-boundary type checking (matches the existing behavior of side outputs accessed via `DoOutputsTuple`).
   
   ## Follow-ups (not in this PR)
   
   - `ParDo.with_side_outputs(...)` convenience method.
   - Beam YAML dot-modifier integration (the existing `get_pcollection` already resolves `Transform.tag` by walking the producer's outputs dict, so this should mostly "just work" once verified).
   
   ## Tests
   
   `pvalue_test.py` covers: copy semantics, attribute / index access, missing-tag errors, type/pipeline validation, empty container, double-call replace.
   
   `pipeline_test.py` covers: end-to-end materialization with `assert_that`, the canonical `MyFilter`-style composite wrapping `with_outputs`, foreign-pcollection rejection, tag-collision rejection, idempotent same-object collision, replacement-path registration via `Pipeline.replace_all`, runner API round-trip, and the nested-return non-flattening case.
   
   ## Test plan
   
   - [x] `cd sdks/python && python -m pytest apache_beam/pvalue_test.py apache_beam/pipeline_test.py` — 84 passed, 1 skipped
   - [x] `ruff check --config=sdks/python/ruff.toml` on the four changed files — clean
   - [x] `yapf --diff` on the four changed files — clean


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org