[Python] Speed up ApproximateUnique accumulator merging

CHANGES.md

@@ -71,6 +71,7 @@
 ## New Features / Improvements
 
 * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
+* (Python) Reduced the cost of merging `ApproximateUnique` accumulators by avoiding an unnecessary copy and heap rebuild ([#19459](https://github.com/apache/beam/issues/19459)).
 * (Java/Python) `Watch` can bound its deduplication state by event time, retiring an output key once the greatest emitted timestamp has moved more than the allowed lateness past it. Java adds `Watch.growthOf(...).withTimestampCursor()`. Python adds `allowed_lateness` for the existing `timestamp_cursor` option ([#18459](https://github.com/apache/beam/issues/18459)).
 * (Java) Spark Structured Streaming runner: stateful ParDo with state, timers, `@RequiresTimeSortedInput` and tagged outputs is now supported in batch mode ([#39779](https://github.com/apache/beam/issues/39779)).
 * (Python) Added support for Vertex AI Model Monitoring V2 in RunInference ([#39738](https://github.com/apache/beam/issues/39738)).

sdks/python/apache_beam/transforms/stats.py

@@ -262,11 +262,14 @@ def add_input(self, accumulator, element, *args, **kwargs):
     except Exception as e:
       raise RuntimeError("Runtime exception: %s" % e)
 
-  # created an issue https://github.com/apache/beam/issues/19459 to speed up
-  # merge process.
   def merge_accumulators(self, accumulators, *args, **kwargs):
-    merged_accumulator = self.create_accumulator()
-    for accumulator in accumulators:
+    accumulator_iter = iter(accumulators)
+    try:
+      merged_accumulator = next(accumulator_iter)
+    except StopIteration:
+      return self.create_accumulator()
+
+    for accumulator in accumulator_iter:
       for i in accumulator._sample_heap:
         merged_accumulator.add(i)
 

sdks/python/apache_beam/transforms/stats_test.py

@@ -219,6 +219,40 @@ def test_approximate_unique_combine_fn_requires_compatible_coder(self):
 
     self.assertRegex(e.exception.args[0], 'Runtime exception')
 
+  def test_approximate_unique_merge_accumulators_reuses_first(self):
+    sample_size = 16
+    combine_fn = ApproximateUniqueCombineFn(sample_size, coders.VarIntCoder())
+    accumulators = [combine_fn.create_accumulator() for _ in range(3)]
+    for accumulator, values in zip(
+        accumulators, [range(16), range(8, 24), range(24, 40)]):
+      for value in values:
+        accumulator.add(value)
+
+    later_accumulator_states = [(
+        list(accumulator._sample_heap),
+        set(accumulator._sample_set),
+        accumulator._min_hash) for accumulator in accumulators[1:]]
+
+    merged_accumulator = combine_fn.merge_accumulators(iter(accumulators))
+
+    self.assertIs(merged_accumulator, accumulators[0])
+    self.assertEqual(set(range(24, 40)), merged_accumulator._sample_set)
+    self.assertEqual(24, merged_accumulator._min_hash)
+    self.assertEqual(
+        later_accumulator_states,
+        [(
+            list(accumulator._sample_heap),
+            set(accumulator._sample_set),
+            accumulator._min_hash) for accumulator in accumulators[1:]])
+
+  def test_approximate_unique_merge_accumulators_empty(self):
+    combine_fn = ApproximateUniqueCombineFn(16, coders.VarIntCoder())
+
+    merged_accumulator = combine_fn.merge_accumulators(iter(()))
+
+    self.assertEqual([], merged_accumulator._sample_heap)
+    self.assertEqual(set(), merged_accumulator._sample_set)
+
   def test_get_sample_size_from_est_error(self):
     # test if get correct sample size from input error.
     assert beam.ApproximateUnique._get_sample_size_from_est_error(0.5) == 16