From 76e4c46aae2736bfd171bc40b95bb9232e543f99 Mon Sep 17 00:00:00 2001 From: SreeramaYeshwanthGowd Date: Wed, 22 Jul 2026 23:24:24 +0530 Subject: [PATCH] Add Sample.Any to the Python SDK to match Java's Sample.any Python only had Sample.FixedSizeGlobally, which runs a uniform reservoir sample and returns a single list. Add Sample.Any, the equivalent of Java's Sample.any, which returns up to n arbitrary elements as a PCollection without the random sampling cost. If the input has fewer than n elements, all are returned. Includes unit tests on the DirectRunner and a CHANGES.md entry. Fixes #18552 --- CHANGES.md | 1 + .../apache_beam/transforms/combiners.py | 55 +++++++++++++++++++ .../apache_beam/transforms/combiners_test.py | 50 +++++++++++++++++ 3 files changed, 106 insertions(+) diff --git a/CHANGES.md b/CHANGES.md index 59c7ac7b24ba..666716d3622a 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -70,6 +70,7 @@ * (Python) Removed the `envoy-data-plane` (and transitive `betterproto`) dependency; `EnvoyRateLimiter` now uses a small vendored protobuf definition instead, resolving dependency conflicts for downstream projects ([#37854](https://github.com/apache/beam/issues/37854)). * (Java) Supported acknowledge mode for JmsIO ([#39253](https://github.com/apache/beam/issues/39253)). +* (Python) Added `Sample.Any`, the Python equivalent of Java's `Sample.any`, which returns up to n arbitrary elements from a PCollection ([#18552](https://github.com/apache/beam/issues/18552)). * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Breaking Changes diff --git a/sdks/python/apache_beam/transforms/combiners.py b/sdks/python/apache_beam/transforms/combiners.py index 8d35405f3fff..37a5b97b1cfb 100644 --- a/sdks/python/apache_beam/transforms/combiners.py +++ b/sdks/python/apache_beam/transforms/combiners.py @@ -597,6 +597,32 @@ def display_data(self): def default_label(self): return 'FixedSizePerKey(%d)' % self._n + @with_input_types(T) + @with_output_types(T) + class Any(ptransform.PTransform): + """Returns up to n arbitrary elements from the input PCollection. + + This is the Python equivalent of Java's ``Sample.any``. Unlike + ``FixedSizeGlobally`` it does not sample uniformly at random, and it returns + the selected elements rather than a single list. If the input has fewer than + n elements, all of them are returned. + """ + def __init__(self, n): + self._n = n + + def expand(self, pcoll): + return ( + pcoll + | core.CombineGlobally(_SampleAnyCombineFn( + self._n)).without_defaults() + | core.FlatMap(lambda elements: elements)) + + def display_data(self): + return {'n': self._n} + + def default_label(self): + return 'Any(%d)' % self._n + @with_input_types(T) @with_output_types(list[T]) @@ -636,6 +662,35 @@ def teardown(self): self._top_combiner.teardown() +@with_input_types(T) +@with_output_types(list[T]) +class _SampleAnyCombineFn(core.CombineFn): + """CombineFn that keeps up to n arbitrary elements (no random sampling).""" + def __init__(self, n): + super().__init__() + self._n = n + + def create_accumulator(self): + return [] + + def add_input(self, accumulator, element): + if len(accumulator) < self._n: + accumulator.append(element) + return accumulator + + def merge_accumulators(self, accumulators): + result = [] + for accumulator in accumulators: + for element in accumulator: + if len(result) >= self._n: + return result + result.append(element) + return result + + def extract_output(self, accumulator): + return accumulator + + class _TupleCombineFnBase(core.CombineFn): def __init__(self, *combiners, merge_accumulators_batch_size=None): self._combiners = [core.CombineFn.maybe_from_callable(c) for c in combiners] diff --git a/sdks/python/apache_beam/transforms/combiners_test.py b/sdks/python/apache_beam/transforms/combiners_test.py index a7f357719617..5322ba7d24a2 100644 --- a/sdks/python/apache_beam/transforms/combiners_test.py +++ b/sdks/python/apache_beam/transforms/combiners_test.py @@ -253,6 +253,7 @@ def individual_test_per_key_dd(sampleFn, n): individual_test_per_key_dd(combine.Sample.FixedSizePerKey, 5) individual_test_per_key_dd(combine.Sample.FixedSizeGlobally, 5) + individual_test_per_key_dd(combine.Sample.Any, 5) def test_combine_globally_display_data(self): transform = beam.CombineGlobally(combine.Smallest(5)) @@ -359,6 +360,55 @@ def match(actual): assert_that(result, matcher()) + def test_sample_any(self): + with TestPipeline() as pipeline: + pcoll = pipeline | 'start' >> Create([1, 2, 3, 4, 5]) + result = pcoll | 'sample-any' >> combine.Sample.Any(3) + + def check(actual): + assert len(actual) == 3, actual + for element in actual: + assert element in [1, 2, 3, 4, 5], element + + assert_that(result, check) + + def test_sample_any_at_most_input_size(self): + with TestPipeline() as pipeline: + pcoll = pipeline | 'start' >> Create([1, 2]) + result = pcoll | 'sample-any' >> combine.Sample.Any(5) + assert_that(result, equal_to([1, 2])) + + def test_sample_any_windowed(self): + with TestPipeline() as pipeline: + pcoll = ( + pipeline + | 'start' >> Create([1, 2, 3, 4]) + | 'timestamp' >> Map(lambda x: TimestampedValue(x, x * 10)) + | 'window' >> WindowInto(FixedWindows(15))) + result = pcoll | 'sample-any' >> combine.Sample.Any(1) + + def check(actual): + # Timestamps 10, 20, 30, 40 fall into fixed windows [0, 15), [15, 30) + # and [30, 45), holding {1}, {2} and {3, 4}. One element is sampled from + # each window that has elements. + assert len(actual) == 3, actual + for element in actual: + assert element in [1, 2, 3, 4], element + + assert_that(result, check) + + def test_sample_any_empty(self): + with TestPipeline() as pipeline: + pcoll = pipeline | 'start' >> Create([]) + result = pcoll | 'sample-any' >> combine.Sample.Any(3) + assert_that(result, equal_to([])) + + def test_sample_any_zero(self): + with TestPipeline() as pipeline: + pcoll = pipeline | 'start' >> Create([1, 2, 3]) + result = pcoll | 'sample-any' >> combine.Sample.Any(0) + assert_that(result, equal_to([])) + def test_tuple_combine_fn(self): with TestPipeline() as p: result = (