This is an automated email from the ASF dual-hosted git repository.
claudevdm 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 d97899b7ab8 Add Sample.Any to the Python SDK to match Java's
Sample.any (#39442)
d97899b7ab8 is described below
commit d97899b7ab8998499d5ee56f4052faa371facd47
Author: SreeramaYeshwanthGowd <[email protected]>
AuthorDate: Mon Aug 10 17:09:31 2026 +0530
Add Sample.Any to the Python SDK to match Java's Sample.any (#39442)
* 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
* Reject negative n in Sample.Any
Match Java's Sample.any, which rejects a negative limit at construction
time. Add a regression test.
* Add explicit type hints to the FlatMap in Sample.Any
Address review feedback: declare with_input_types(list[T]) and
with_output_types(T) on the FlatMap that flattens the combiner's output.
---------
Co-authored-by: Jack McCluskey
<[email protected]>
---
CHANGES.md | 1 +
sdks/python/apache_beam/transforms/combiners.py | 58 ++++++++++++++++++++++
.../apache_beam/transforms/combiners_test.py | 54 ++++++++++++++++++++
3 files changed, 113 insertions(+)
diff --git a/CHANGES.md b/CHANGES.md
index 6f61ef415b4..e2c3f262536 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -123,6 +123,7 @@
* (Python) Added `Watch`, a transform that polls a growing set of outputs for
each input element, deduplicates outputs across poll rounds, and stops per a
user-supplied termination condition
([#21521](https://github.com/apache/beam/issues/21521)).
* (Python) Added support to analyze core dumps created after python worker
segmentation faults with `pystack` (or `gdb` if installed) using the
`--profiler_agent=coredump` pipeline option.
([#39484](https://github.com/apache/beam/issues/39484)).
+* (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)).
* (Java) Added per-element OpenTelemetry trace propagation across stages in
the Dataflow Streaming Runner. Enable it with
`--experiments=enable_otel_defaults,element_metadata_supported,disable_portable_worker`.
Cloud Trace incurs additional cost.
([#33176](https://github.com/apache/beam/issues/33176))
* (Java) Added OpenTelemetry header propagation support for both reads and
writes in KafkaIO and PubSubIO.
([#33176](https://github.com/apache/beam/issues/33176))
* (Java) Added OpenTelemetry tracing support for SpannerIO change streams
([#33176](https://github.com/apache/beam/issues/33176))
diff --git a/sdks/python/apache_beam/transforms/combiners.py
b/sdks/python/apache_beam/transforms/combiners.py
index 8d35405f3ff..c45ba4e89b9 100644
--- a/sdks/python/apache_beam/transforms/combiners.py
+++ b/sdks/python/apache_beam/transforms/combiners.py
@@ -597,6 +597,35 @@ class Sample(object):
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):
+ if n < 0:
+ raise ValueError('Expected non-negative n, received %s.' % n)
+ self._n = n
+
+ def expand(self, pcoll):
+ return (
+ pcoll
+ | core.CombineGlobally(_SampleAnyCombineFn(
+ self._n)).without_defaults()
+ | core.FlatMap(lambda elements: elements).with_input_types(
+ list[T]).with_output_types(T))
+
+ 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 +665,35 @@ class SampleCombineFn(core.CombineFn):
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 a7f35771961..14348bb8ce7 100644
--- a/sdks/python/apache_beam/transforms/combiners_test.py
+++ b/sdks/python/apache_beam/transforms/combiners_test.py
@@ -253,6 +253,7 @@ class CombineTest(unittest.TestCase):
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,59 @@ class CombineTest(unittest.TestCase):
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_sample_any_negative_n(self):
+ with self.assertRaises(ValueError):
+ combine.Sample.Any(-1)
+
def test_tuple_combine_fn(self):
with TestPipeline() as p:
result = (
