This is an automated email from the ASF dual-hosted git repository. tvalentyn pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push: new 9f582e564fe Add key param support for helper Combiners (#25895) 9f582e564fe is described below commit 9f582e564fededdd4506157a21bef66ada303529 Author: harrisonlimh <97203667+harrisonl...@users.noreply.github.com> AuthorDate: Wed Mar 29 14:08:31 2023 -0700 Add key param support for helper Combiners (#25895) Co-authored-by: tvalentyn <tvalen...@users.noreply.github.com> --- CHANGES.md | 1 + sdks/python/apache_beam/transforms/combiners.py | 20 ++++++++++---------- sdks/python/apache_beam/transforms/combiners_test.py | 14 ++++++++++++++ 3 files changed, 25 insertions(+), 10 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 0f54d233ad6..5e6c99f4adb 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -77,6 +77,7 @@ * If a main session fails to load, the pipeline will now fail at worker startup. ([#25401](https://github.com/apache/beam/issues/25401)). * Python pipeline options will now ignore unparsed command line flags prefixed with a single dash. ([#25943](https://github.com/apache/beam/issues/25943)). +* The SmallestPerKey combiner now requires keyword-only arguments for specifying optional parameters, such as `key` and `reverse`. ([#25888](https://github.com/apache/beam/issues/25888)). ## Deprecations diff --git a/sdks/python/apache_beam/transforms/combiners.py b/sdks/python/apache_beam/transforms/combiners.py index 7afa4e46e75..0e5f01196e6 100644 --- a/sdks/python/apache_beam/transforms/combiners.py +++ b/sdks/python/apache_beam/transforms/combiners.py @@ -293,33 +293,33 @@ class Top(object): @staticmethod @ptransform.ptransform_fn - def Largest(pcoll, n, has_defaults=True): + def Largest(pcoll, n, has_defaults=True, key=None): """Obtain a list of the greatest N elements in a PCollection.""" if has_defaults: - return pcoll | Top.Of(n) + return pcoll | Top.Of(n, key) else: - return pcoll | Top.Of(n).without_defaults() + return pcoll | Top.Of(n, key).without_defaults() @staticmethod @ptransform.ptransform_fn - def Smallest(pcoll, n, has_defaults=True): + def Smallest(pcoll, n, has_defaults=True, key=None): """Obtain a list of the least N elements in a PCollection.""" if has_defaults: - return pcoll | Top.Of(n, reverse=True) + return pcoll | Top.Of(n, key, reverse=True) else: - return pcoll | Top.Of(n, reverse=True).without_defaults() + return pcoll | Top.Of(n, key, reverse=True).without_defaults() @staticmethod @ptransform.ptransform_fn - def LargestPerKey(pcoll, n): + def LargestPerKey(pcoll, n, key=None): """Identifies the N greatest elements associated with each key.""" - return pcoll | Top.PerKey(n) + return pcoll | Top.PerKey(n, key) @staticmethod @ptransform.ptransform_fn - def SmallestPerKey(pcoll, n, reverse=True): + def SmallestPerKey(pcoll, n, *, key=None, reverse=None): """Identifies the N least elements associated with each key.""" - return pcoll | Top.PerKey(n, reverse=True) + return pcoll | Top.PerKey(n, key, reverse=True) @with_input_types(T) diff --git a/sdks/python/apache_beam/transforms/combiners_test.py b/sdks/python/apache_beam/transforms/combiners_test.py index 34fb463d00a..385b3332e0c 100644 --- a/sdks/python/apache_beam/transforms/combiners_test.py +++ b/sdks/python/apache_beam/transforms/combiners_test.py @@ -191,6 +191,20 @@ class CombineTest(unittest.TestCase): | combine.Top.Of(3, key=len, reverse=True), [['c', 'aa', 'bbb']]) + self.assertEqual(['xc', 'zb', 'yd', 'wa'] + | combine.Top.Largest(3, key=lambda x: x[-1]), + [['yd', 'xc', 'zb']]) + self.assertEqual(['xc', 'zb', 'yd', 'wa'] + | combine.Top.Smallest(3, key=lambda x: x[-1]), + [['wa', 'zb', 'xc']]) + + self.assertEqual([('a', x) for x in [1, 2, 3, 4, 1, 1]] + | combine.Top.LargestPerKey(3, key=lambda x: -x), + [('a', [1, 1, 1])]) + self.assertEqual([('a', x) for x in [1, 2, 3, 4, 1, 1]] + | combine.Top.SmallestPerKey(3, key=lambda x: -x), + [('a', [4, 3, 2])]) + def test_sharded_top_combine_fn(self): def test_combine_fn(combine_fn, shards, expected): accumulators = [