Skip to content

[Python] Optimize Count, Top, Sample, ToList, and ConcatList CombineFns #41062

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
damccorm wants to merge 1 commit into
base: master
Choose a base branch
from
+27 −11
Draft

[Python] Optimize Count, Top, Sample, ToList, and ConcatList CombineFns #41062

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 27 additions & 11 deletions sdks/python/apache_beam/transforms/combiners.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import operator
import random
from collections.abc import Iterable
from collections.abc import Sized
from typing import Any
from typing import TypeVar
from typing import Union
Expand Down Expand Up @@ -176,7 +177,9 @@ def add_input(self, accumulator, element):
return accumulator + 1

def add_inputs(self, accumulator, elements):
return accumulator + len(list(elements))
if isinstance(elements, Sized):
return accumulator + len(elements)
return accumulator + sum(1 for _ in elements)

def merge_accumulators(self, accumulators):
return sum(accumulators)
Expand Down Expand Up @@ -498,21 +501,24 @@ def add_input(self, accumulator, element, *args, **kwargs):
def merge_accumulators(self, accumulators, *args, **kwargs):
result_heap = None
holds_comparables = None
use_comparables = self._compare is not operator.lt or bool(self._key)
n = self._n
for accumulator in accumulators:
holds_comparables, heap = accumulator
if self._compare is not operator.lt or self._key:
if use_comparables:
heap = self._hydrated_heap(heap)
holds_comparables = True
else:
assert not holds_comparables

if result_heap is None:
result_heap = heap
else:
result_heap = list(heap)
elif n > 0:

@github-actions github-actions Bot Nov 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟢 Nit: If n <= 0 (e.g. TopCombineFn(0)), add_input will still push up to 1 element into heap because len(heap) < self._n is False and else: heapq.heappushpop(heap, comparable) is executed on an empty list (which pushes the element and pops it, leaving heap empty—wait: heappushpop([], x) actually returns x and leaves the list empty!). However, if the first accumulator in merge_accumulators somehow had elements or if n <= 0 is guarded in add_input as well (like in SampleCombineFn.add_input), guarding if self._n <= 0: return accumulator in TopCombineFn.add_input would make TopCombineFn(0) consistent and avoid unnecessary ComparableValue allocations.

for comparable in heap:
_, result_heap = self.add_input(
(holds_comparables, result_heap),
comparable.value if holds_comparables else comparable)
if len(result_heap) < n:
heapq.heappush(result_heap, comparable)
else:
heapq.heappushpop(result_heap, comparable)

assert result_heap is not None and holds_comparables is not None
return (holds_comparables, result_heap)
Expand Down Expand Up @@ -649,7 +655,16 @@ def add_input(self, heap, element):
# Before passing elements to the Top combiner, we pair them with random
# numbers. The elements with the n largest random number "keys" will be
# selected for the output.
return self._top_combiner.add_input(heap, (random.random(), element))
n = self._top_combiner._n
if n <= 0:
return heap
inner_heap = heap[1]
r = random.random()
if len(inner_heap) < n:
heapq.heappush(inner_heap, (r, element))
elif r > inner_heap[0][0]:
heapq.heapreplace(inner_heap, (r, element))
return heap

def merge_accumulators(self, heaps):
return self._top_combiner.merge_accumulators(heaps)
Expand Down Expand Up @@ -811,7 +826,7 @@ def add_input(self, accumulator, element):
return accumulator

def merge_accumulators(self, accumulators):
return sum(accumulators, [])
return list(itertools.chain.from_iterable(accumulators))

def extract_output(self, accumulator):
return accumulator
Expand All @@ -825,10 +840,11 @@ def create_accumulator(self):
return []

def add_input(self, accumulator, element):
return accumulator + element
accumulator.extend(element)
return accumulator

def merge_accumulators(self, accumulators):
return sum(accumulators, [])
return list(itertools.chain.from_iterable(accumulators))

def extract_output(self, accumulator):
return accumulator
Expand Down
Loading