[Python] Accept int tags on Partition results applied to a list

CHANGES.md

@@ -104,6 +104,7 @@
 * (Python) DaskRunner now preserves outputs emitted by `DoFn.finish_bundle()` instead of silently dropping them ([#40519](https://github.com/apache/beam/issues/40519)).
 * (Python) BigQuery streaming inserts now reset the retry budget and backoff for each batch instead of sharing them across a bundle ([#39995](https://github.com/apache/beam/pull/39995)).
 * (Python) State write failures reported by portable runners now fail the bundle instead of silently losing state updates ([#39992](https://github.com/apache/beam/pull/39992)).
+* (Python) The result of `beam.Partition` applied to a list, such as `[1, 2, 3] | beam.Partition(fn, 2)`, can now be indexed with ints (`partitions[0]`) as inside a pipeline, instead of raising `KeyError` ([#19154](https://github.com/apache/beam/issues/19154)).
 * (Go) Fixed a data race on the Prism runner's artifact cache map in JobServices ([#32656](https://github.com/apache/beam/issues/32656)).
 * (Go) Fixed the harness leaking Data/State gRPC streams after the worker stops, and a deadlock when Send returns EOF ([#40260](https://github.com/apache/beam/issues/40260)).
 * (Go) Wrap last dataWriter/timerWriter send errors on Close ([#40405](https://github.com/apache/beam/issues/40405)).

sdks/python/apache_beam/transforms/ptransform.py

@@ -212,6 +212,9 @@ def __init__(self, deferred, results_by_tag):
     self._results_by_tag = results_by_tag
 
   def __getitem__(self, tag):
+    # Accept int tags, like DoOutputsTuple does for Partition outputs.
+    if isinstance(tag, int):
+      tag = str(tag)
     if tag not in self._results_by_tag:
       raise KeyError(
           'Tag %r is not a defined output tag of %s.' % (tag, self._deferred))

sdks/python/apache_beam/transforms/ptransform_test.py

@@ -739,6 +739,13 @@ def test_partition_followed_by_flatten_and_groupbykey(self):
       grouped = flattened | 'D' >> beam.GroupByKey() | SortLists
       assert_that(grouped, equal_to([('aa', [1, 2]), ('bb', [2])]))
 
+  def test_partition_applied_to_list(self):
+    # https://github.com/apache/beam/issues/19154
+    partitions = [0, 1, 2, 3, 4] | beam.Partition(lambda e, n: e % n, 2)
+    self.assertCountEqual([0, 2, 4], partitions[0])
+    self.assertCountEqual([1, 3], partitions[1])
+    self.assertCountEqual([1, 3], partitions['1'])
+
   @pytest.mark.it_validatesrunner
   def test_flatten_pcollections(self):
     with TestPipeline() as pipeline: