diff --git a/CHANGES.md b/CHANGES.md index d06940d12cf4..6e0f205f195a 100644 --- a/CHANGES.md +++ b/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)). diff --git a/sdks/python/apache_beam/transforms/ptransform.py b/sdks/python/apache_beam/transforms/ptransform.py index 06ea822aa0be..f41a9e4136d9 100644 --- a/sdks/python/apache_beam/transforms/ptransform.py +++ b/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)) diff --git a/sdks/python/apache_beam/transforms/ptransform_test.py b/sdks/python/apache_beam/transforms/ptransform_test.py index 21691eebf3ae..54fea8d79e4e 100644 --- a/sdks/python/apache_beam/transforms/ptransform_test.py +++ b/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: