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 <34928439+jrmccluskey@users.noreply.github.com>

Author: SreeramaYeshwanthGowd

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))

sdks/python/apache_beam/transforms/combiners.py

@@ -597,6 +597,35 @@ 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):
+      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 @@ 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,59 @@ 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_sample_any_negative_n(self):
+    with self.assertRaises(ValueError):
+      combine.Sample.Any(-1)
+
   def test_tuple_combine_fn(self):
     with TestPipeline() as p:
       result = (