Add Sample.Any to the Python SDK to match Java's Sample.any

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

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]

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 = (