diff --git a/CHANGES.md b/CHANGES.md index 4ea769676172..65cc4800ca28 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -93,6 +93,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 8d35405f3fff..c45ba4e89b9a 100644 --- a/sdks/python/apache_beam/transforms/combiners.py +++ b/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] diff --git a/sdks/python/apache_beam/transforms/combiners_test.py b/sdks/python/apache_beam/transforms/combiners_test.py index a7f357719617..14348bb8ce78 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,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 = (