combiners.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""A library of basic combiner PTransform subclasses."""

# pytype: skip-file

importcopy
importheapq
importitertools
importoperator
importrandom
fromcollections.abcimportIterable
fromtypingimportAny
fromtypingimportTypeVar
fromtypingimportUnion

importnumpyasnp

fromapache_beamimporttypehints
fromapache_beam.transformsimportcore
fromapache_beam.transformsimportcy_combiners
fromapache_beam.transformsimportptransform
fromapache_beam.transformsimportwindow
fromapache_beam.transforms.displayimportDisplayDataItem
fromapache_beam.typehintsimportwith_input_types
fromapache_beam.typehintsimportwith_output_types
fromapache_beam.utils.timestampimportDuration
fromapache_beam.utils.timestampimportTimestamp

__all__= [
'Count',
'Mean',
'Sample',
'Top',
'ToDict',
'ToList',
'ToSet',
'Latest',
'CountCombineFn',
'MeanCombineFn',
'SampleCombineFn',
'TopCombineFn',
'ToDictCombineFn',
'ToListCombineFn',
'ToSetCombineFn',
'LatestCombineFn',
]

# Type variables
T=TypeVar('T')
K=TypeVar('K')
V=TypeVar('V')
TimestampType=Union[int, float, Timestamp, Duration]


classCombinerWithoutDefaults(ptransform.PTransform):
"""Super class to inherit without_defaults to built-in Combiners."""
def__init__(self, has_defaults=True):
super().__init__()
self.has_defaults=has_defaults

defwith_defaults(self, has_defaults=True):
new=copy.copy(self)
new.has_defaults=has_defaults
returnnew

defwithout_defaults(self):
returnself.with_defaults(False)


classMean(object):
"""Combiners for computing arithmetic means of elements."""
classGlobally(CombinerWithoutDefaults):
"""combiners.Mean.Globally computes the arithmetic mean of the elements."""
defexpand(self, pcoll):
ifself.has_defaults:
returnpcoll|core.CombineGlobally(MeanCombineFn())
else:
returnpcoll|core.CombineGlobally(MeanCombineFn()).without_defaults()

classPerKey(ptransform.PTransform):
"""combiners.Mean.PerKey finds the means of the values for each key."""
defexpand(self, pcoll):
returnpcoll|core.CombinePerKey(MeanCombineFn())


# TODO(laolu): This type signature is overly restrictive. This should be
# more general.
@with_input_types(Union[float, int, np.int64, np.float64])
@with_output_types(float)
classMeanCombineFn(core.CombineFn):
"""CombineFn for computing an arithmetic mean."""
defcreate_accumulator(self):
return (0, 0)

defadd_input(self, sum_count, element):
 (sum_, count) =sum_count
returnsum_+element, count+1

defmerge_accumulators(self, accumulators):
sums, counts=zip(*accumulators)
returnsum(sums), sum(counts)

defextract_output(self, sum_count):
 (sum_, count) =sum_count
ifcount==0:
returnfloat('NaN')
returnsum_/float(count)

deffor_input_type(self, input_type):
ifinput_typeisint:
returncy_combiners.MeanInt64Fn()
elifinput_typeisfloat:
returncy_combiners.MeanFloatFn()
returnself


classCount(object):
"""Combiners for counting elements."""
@with_input_types(T)
@with_output_types(int)
classGlobally(CombinerWithoutDefaults):
"""combiners.Count.Globally counts the total number of elements."""
defexpand(self, pcoll):
ifself.has_defaults:
returnpcoll|core.CombineGlobally(CountCombineFn())
else:
returnpcoll|core.CombineGlobally(CountCombineFn()).without_defaults()

@with_input_types(tuple[K, V])
@with_output_types(tuple[K, int])
classPerKey(ptransform.PTransform):
"""combiners.Count.PerKey counts how many elements each unique key has."""
defexpand(self, pcoll):
returnpcoll|core.CombinePerKey(CountCombineFn())

@with_input_types(T)
@with_output_types(tuple[T, int])
classPerElement(ptransform.PTransform):
"""combiners.Count.PerElement counts how many times each element occurs."""
defexpand(self, pcoll):
paired_with_void_type=typehints.Tuple[pcoll.element_type, Any]
output_type=typehints.KV[pcoll.element_type, int]
return (
pcoll
| (
'%s:PairWithVoid'%self.label>>core.Map(
lambdax: (x, None)).with_output_types(paired_with_void_type))
|core.CombinePerKey(CountCombineFn()).with_output_types(output_type))


@with_input_types(Any)
@with_output_types(int)
classCountCombineFn(core.CombineFn):
"""CombineFn for computing PCollection size."""
defcreate_accumulator(self):
return0

defadd_input(self, accumulator, element):
returnaccumulator+1

defadd_inputs(self, accumulator, elements):
returnaccumulator+len(list(elements))

defmerge_accumulators(self, accumulators):
returnsum(accumulators)

defextract_output(self, accumulator):
returnaccumulator


classTop(object):
"""Combiners for obtaining extremal elements."""

# pylint: disable=no-self-argument
@with_input_types(T)
@with_output_types(list[T])
classOf(CombinerWithoutDefaults):
"""Returns the n greatest elements in the PCollection.

 This transform will retrieve the n greatest elements in the PCollection
 to which it is applied, where "greatest" is determined by a
 function supplied as the `key` or `reverse` arguments.
 """
def__init__(self, n, key=None, reverse=False):
"""Creates a global Top operation.

 The arguments 'key' and 'reverse' may be passed as keyword arguments,
 and have the same meaning as for Python's sort functions.

 Args:
 n: number of elements to extract from pcoll.
 key: (optional) a mapping of elements to a comparable key, similar to
 the key argument of Python's sorting methods.
 reverse: (optional) whether to order things smallest to largest, rather
 than largest to smallest
 """
super().__init__()
self._n=n
self._key=key
self._reverse=reverse

defdefault_label(self):
return'Top(%d)'%self._n

defexpand(self, pcoll):
ifpcoll.windowing.is_default():
# This is a more efficient global algorithm.
top_per_bundle=pcoll|core.ParDo(
_TopPerBundle(self._n, self._key, self._reverse))
# If pcoll is empty, we can't guarantee that top_per_bundle
# won't be empty, so inject at least one empty accumulator
# so that downstream is guaranteed to produce non-empty output.
empty_bundle= (
pcoll.pipeline|core.Create([(None, [])]).with_output_types(
top_per_bundle.element_type))
return ((top_per_bundle, empty_bundle) |core.Flatten()
|core.GroupByKey()
|core.ParDo(
_MergeTopPerBundle(self._n, self._key, self._reverse)))
else:
ifself.has_defaults:
returnpcoll|core.CombineGlobally(
TopCombineFn(self._n, self._key, self._reverse))
else:
returnpcoll|core.CombineGlobally(
TopCombineFn(self._n, self._key,
self._reverse)).without_defaults()

@with_input_types(tuple[K, V])
@with_output_types(tuple[K, list[V]])
classPerKey(ptransform.PTransform):
"""Identifies the N greatest elements associated with each key.

 This transform will produce a PCollection mapping unique keys in the input
 PCollection to the n greatest elements with which they are associated, where
 "greatest" is determined by a function supplied as the `key` or
 `reverse` arguments.
 """
def__init__(self, n, key=None, reverse=False):
"""Creates a per-key Top operation.

 The arguments 'key' and 'reverse' may be passed as keyword arguments,
 and have the same meaning as for Python's sort functions.

 Args:
 n: number of elements to extract from pcoll.
 key: (optional) a mapping of elements to a comparable key, similar to
 the key argument of Python's sorting methods.
 reverse: (optional) whether to order things smallest to largest, rather
 than largest to smallest
 """
self._n=n
self._key=key
self._reverse=reverse

defdefault_label(self):
return'TopPerKey(%d)'%self._n

defexpand(self, pcoll):
"""Expands the transform.

 Raises TypeCheckError: If the output type of the input PCollection is not
 compatible with tuple[A, B].

 Args:
 pcoll: PCollection to process

 Returns:
 the PCollection containing the result.
 """
returnpcoll|core.CombinePerKey(
TopCombineFn(self._n, self._key, self._reverse))

@staticmethod
@ptransform.ptransform_fn
defLargest(pcoll, n, has_defaults=True, key=None):
"""Obtain a list of the greatest N elements in a PCollection."""
ifhas_defaults:
returnpcoll|Top.Of(n, key)
else:
returnpcoll|Top.Of(n, key).without_defaults()

@staticmethod
@ptransform.ptransform_fn
defSmallest(pcoll, n, has_defaults=True, key=None):
"""Obtain a list of the least N elements in a PCollection."""
ifhas_defaults:
returnpcoll|Top.Of(n, key, reverse=True)
else:
returnpcoll|Top.Of(n, key, reverse=True).without_defaults()

@staticmethod
@ptransform.ptransform_fn
defLargestPerKey(pcoll, n, key=None):
"""Identifies the N greatest elements associated with each key."""
returnpcoll|Top.PerKey(n, key)

@staticmethod
@ptransform.ptransform_fn
defSmallestPerKey(pcoll, n, *, key=None):
"""Identifies the N least elements associated with each key."""
returnpcoll|Top.PerKey(n, key, reverse=True)


@with_input_types(T)
@with_output_types(tuple[None, list[T]])
class_TopPerBundle(core.DoFn):
def__init__(self, n, key, reverse):
self._n=n
self._compare=operator.gtifreverseelseNone
self._key=key

defstart_bundle(self):
self._heap= []

defprocess(self, element):
ifself._compareorself._key:
element=cy_combiners.ComparableValue(element, self._compare, self._key)
iflen(self._heap) <self._n:
heapq.heappush(self._heap, element)
else:
heapq.heappushpop(self._heap, element)

deffinish_bundle(self):
# Though sorting here results in more total work, this allows us to
# skip most elements in the reducer.
# Essentially, given s map bundles, we are trading about O(sn) compares in
# the (single) reducer for O(sn log n) compares across all mappers.
self._heap.sort()

# Unwrap to avoid serialization via pickle.
ifself._compareorself._key:
yieldwindow.GlobalWindows.windowed_value(
 (None, [wrapper.valueforwrapperinself._heap]))
else:
yieldwindow.GlobalWindows.windowed_value((None, self._heap))


@with_input_types(tuple[None, Iterable[list[T]]])
@with_output_types(list[T])
class_MergeTopPerBundle(core.DoFn):
def__init__(self, n, key, reverse):
self._n=n
self._compare=operator.gtifreverseelseNone
self._key=key

defprocess(self, key_and_bundles):
_, bundles=key_and_bundles

defpush(hp, e):
iflen(hp) <self._n:
heapq.heappush(hp, e)
returnFalse
elife<hp[0]:
# Because _TopPerBundle returns sorted lists, all other elements
# will also be smaller.
returnTrue
else:
heapq.heappushpop(hp, e)
returnFalse

ifself._compareorself._key:
heapc: list[cy_combiners.ComparableValue] = []
forbundleinbundles:
ifnotheapc:
heapc= [
cy_combiners.ComparableValue(element, self._compare, self._key)
forelementinbundle
 ]
continue
# TODO(https://github.com/apache/beam/issues/21205): Remove this
# workaround once legacy dataflow correctly handles coders with
# combiner packing and/or is deprecated.
ifnotisinstance(bundle, list):
bundle=list(bundle)
forelementinreversed(bundle):
ifpush(heapc,
cy_combiners.ComparableValue(element,
self._compare,
self._key)):
break
heapc.sort()
yield [wrapper.valueforwrapperinreversed(heapc)]

else:
heap= []
forbundleinbundles:
# TODO(https://github.com/apache/beam/issues/21205): Remove this
# workaround once legacy dataflow correctly handles coders with
# combiner packing and/or is deprecated.
ifnotisinstance(bundle, list):
bundle=list(bundle)
ifnotheap:
heap=bundle
continue
forelementinreversed(bundle):
ifpush(heap, element):
break
heap.sort()
yieldheap[::-1]


@with_input_types(T)
@with_output_types(list[T])
classTopCombineFn(core.CombineFn):
"""CombineFn doing the combining for all of the Top transforms.

 This CombineFn uses a `key` or `reverse` operator to rank the elements.

 Args:
 key: (optional) a mapping of elements to a comparable key, similar to
 the key argument of Python's sorting methods.
 reverse: (optional) whether to order things smallest to largest, rather
 than largest to smallest
 """
def__init__(self, n, key=None, reverse=False):
self._n=n
self._compare=operator.gtifreverseelseoperator.lt
self._key=key

def_hydrated_heap(self, heap):
ifheap:
first=heap[0]
ifisinstance(first, cy_combiners.ComparableValue):
iffirst.requires_hydration:
forcomparableinheap:
assertcomparable.requires_hydration
comparable.hydrate(self._compare, self._key)
assertnotcomparable.requires_hydration
returnheap
else:
returnheap
else:
return [
cy_combiners.ComparableValue(element, self._compare, self._key)
forelementinheap
 ]
else:
returnheap

defdisplay_data(self):
return {
'n': self._n,
'compare': DisplayDataItem(
self._compare.__name__ifhasattr(self._compare, '__name__') else
self._compare.__class__.__name__).drop_if_none()
 }

# The accumulator type is a tuple
# (bool, Union[list[T], list[ComparableValue[T]])
# where the boolean indicates whether the second slot contains a List of T
# (False) or List of ComparableValue[T] (True). In either case, the List
# maintains heap invariance. When the contents of the List are
# ComparableValue[T] they either all 'requires_hydration' or none do.
# This accumulator representation allows us to minimize the data encoding
# overheads. Creation of ComparableValues is elided for performance reasons
# when there is no need for complicated comparison functions.
defcreate_accumulator(self, *args, **kwargs):
return (False, [])

defadd_input(self, accumulator, element, *args, **kwargs):
# Caching to avoid paying the price of variadic expansion of args / kwargs
# when it's not needed (for the 'if' case below).
holds_comparables, heap=accumulator
ifself._compareisnotoperator.ltorself._key:
heap=self._hydrated_heap(heap)
holds_comparables=True
else:
assertnotholds_comparables

comparable= (
cy_combiners.ComparableValue(element, self._compare, self._key)
ifholds_comparableselseelement)

iflen(heap) <self._n:
heapq.heappush(heap, comparable)
else:
heapq.heappushpop(heap, comparable)
return (holds_comparables, heap)

defmerge_accumulators(self, accumulators, *args, **kwargs):
result_heap=None
holds_comparables=None
foraccumulatorinaccumulators:
holds_comparables, heap=accumulator
ifself._compareisnotoperator.ltorself._key:
heap=self._hydrated_heap(heap)
holds_comparables=True
else:
assertnotholds_comparables

ifresult_heapisNone:
result_heap=heap
else:
forcomparableinheap:
_, result_heap=self.add_input(
 (holds_comparables, result_heap),
comparable.valueifholds_comparableselsecomparable)

assertresult_heapisnotNoneandholds_comparablesisnotNone
return (holds_comparables, result_heap)

defcompact(self, accumulator, *args, **kwargs):
holds_comparables, heap=accumulator
# Unwrap to avoid serialization via pickle.
ifholds_comparables:
return (False, [comparable.valueforcomparableinheap])
else:
returnaccumulator

defextract_output(self, accumulator, *args, **kwargs):
holds_comparables, heap=accumulator
ifself._compareisnotoperator.ltorself._key:
ifnotholds_comparables:
heap=self._hydrated_heap(heap)
holds_comparables=True
else:
assertnotholds_comparables

assertlen(heap) <=self._n
heap.sort(reverse=True)
return [
comparable.valueifholds_comparableselsecomparable
forcomparableinheap
 ]


classLargest(TopCombineFn):
defdefault_label(self):
return'Largest(%s)'%self._n


classSmallest(TopCombineFn):
def__init__(self, n):
super().__init__(n, reverse=True)

defdefault_label(self):
return'Smallest(%s)'%self._n


classSample(object):
"""Combiners for sampling n elements without replacement."""

# pylint: disable=no-self-argument

@with_input_types(T)
@with_output_types(list[T])
classFixedSizeGlobally(CombinerWithoutDefaults):
"""Sample n elements from the input PCollection without replacement."""
def__init__(self, n):
super().__init__()
self._n=n

defexpand(self, pcoll):
ifself.has_defaults:
returnpcoll|core.CombineGlobally(SampleCombineFn(self._n))
else:
returnpcoll|core.CombineGlobally(SampleCombineFn(
self._n)).without_defaults()

defdisplay_data(self):
return {'n': self._n}

defdefault_label(self):
return'FixedSizeGlobally(%d)'%self._n

@with_input_types(tuple[K, V])
@with_output_types(tuple[K, list[V]])
classFixedSizePerKey(ptransform.PTransform):
"""Sample n elements associated with each key without replacement."""
def__init__(self, n):
self._n=n

defexpand(self, pcoll):
returnpcoll|core.CombinePerKey(SampleCombineFn(self._n))

defdisplay_data(self):
return {'n': self._n}

defdefault_label(self):
return'FixedSizePerKey(%d)'%self._n


@with_input_types(T)
@with_output_types(list[T])
classSampleCombineFn(core.CombineFn):
"""CombineFn for all Sample transforms."""
def__init__(self, n):
super().__init__()
# Most of this combiner's work is done by a TopCombineFn. We could just
# subclass TopCombineFn to make this class, but since sampling is not
# really a kind of Top operation, we use a TopCombineFn instance as a
# helper instead.
self._top_combiner=TopCombineFn(n)

defsetup(self):
self._top_combiner.setup()

defcreate_accumulator(self):
returnself._top_combiner.create_accumulator()

defadd_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.
returnself._top_combiner.add_input(heap, (random.random(), element))

defmerge_accumulators(self, heaps):
returnself._top_combiner.merge_accumulators(heaps)

defcompact(self, heap):
returnself._top_combiner.compact(heap)

defextract_output(self, heap):
# Here we strip off the random number keys we added in add_input.
return [efor_, einself._top_combiner.extract_output(heap)]

defteardown(self):
self._top_combiner.teardown()


class_TupleCombineFnBase(core.CombineFn):
def__init__(self, *combiners, merge_accumulators_batch_size=None):
self._combiners= [core.CombineFn.maybe_from_callable(c) forcincombiners]
self._named_combiners=combiners
# If the `merge_accumulators_batch_size` value is not specified, we chose a
# bounded default that is inversely proportional to the number of
# accumulators in merged tuples.
num_combiners=max(1, len(combiners))
self._merge_accumulators_batch_size= (
merge_accumulators_batch_sizeormax(10, 1000//num_combiners))

defdisplay_data(self):
combiners= [
c.__name__ifhasattr(c, '__name__') elsec.__class__.__name__
forcinself._named_combiners
 ]
return {
'combiners': str(combiners),
'merge_accumulators_batch_size': self._merge_accumulators_batch_size
 }

defsetup(self, *args, **kwargs):
forcinself._combiners:
c.setup(*args, **kwargs)

defcreate_accumulator(self, *args, **kwargs):
return [c.create_accumulator(*args, **kwargs) forcinself._combiners]

defmerge_accumulators(self, accumulators, *args, **kwargs):
# Make sure that `accumulators` is an iterator (so that the position is
# remembered).
accumulators=iter(accumulators)
result=next(accumulators)
whileTrue:
# Load accumulators into memory and merge in batches to decrease peak
# memory usage.
accumulators_batch= [result] +list(
itertools.islice(accumulators, self._merge_accumulators_batch_size))
iflen(accumulators_batch) ==1:
break
result= [
c.merge_accumulators(a, *args, **kwargs) forc,
ainzip(self._combiners, zip(*accumulators_batch))
 ]
returnresult

defcompact(self, accumulator, *args, **kwargs):
return [
c.compact(a, *args, **kwargs) forc,
ainzip(self._combiners, accumulator)
 ]

defextract_output(self, accumulator, *args, **kwargs):
returntuple(
c.extract_output(a, *args, **kwargs) forc,
ainzip(self._combiners, accumulator))

defteardown(self, *args, **kwargs):
forcinreversed(self._combiners):
c.teardown(*args, **kwargs)


classTupleCombineFn(_TupleCombineFnBase):
"""A combiner for combining tuples via a tuple of combiners.

 Takes as input a tuple of N CombineFns and combines N-tuples by
 combining the k-th element of each tuple with the k-th CombineFn,
 outputting a new N-tuple of combined values.
 """
defadd_input(self, accumulator, element, *args, **kwargs):
return [
c.add_input(a, e, *args, **kwargs) forc,
a,
einzip(self._combiners, accumulator, element)
 ]

defwith_common_input(self):
returnSingleInputTupleCombineFn(*self._combiners)


classSingleInputTupleCombineFn(_TupleCombineFnBase):
"""A combiner for combining a single value via a tuple of combiners.

 Takes as input a tuple of N CombineFns and combines elements by
 applying each CombineFn to each input, producing an N-tuple of
 the outputs corresponding to each of the N CombineFn's outputs.
 """
defadd_input(self, accumulator, element, *args, **kwargs):
return [
c.add_input(a, element, *args, **kwargs) forc,
ainzip(self._combiners, accumulator)
 ]


@with_input_types(T)
@with_output_types(list[T])
classToList(CombinerWithoutDefaults):
"""A global CombineFn that condenses a PCollection into a single list."""
defexpand(self, pcoll):
ifself.has_defaults:
returnpcoll|self.label>>core.CombineGlobally(ToListCombineFn())
else:
returnpcoll|self.label>>core.CombineGlobally(
ToListCombineFn()).without_defaults()


@with_input_types(T)
@with_output_types(list[T])
classToListCombineFn(core.CombineFn):
"""CombineFn for to_list."""
defcreate_accumulator(self):
return []

defadd_input(self, accumulator, element):
accumulator.append(element)
returnaccumulator

defmerge_accumulators(self, accumulators):
returnsum(accumulators, [])

defextract_output(self, accumulator):
returnaccumulator


@with_input_types(T)
@with_output_types(T)
classConcatListCombineFn(core.CombineFn):
"""CombineFn for concatenating lists together."""
defcreate_accumulator(self):
return []

defadd_input(self, accumulator, element):
returnaccumulator+element

defmerge_accumulators(self, accumulators):
returnsum(accumulators, [])

defextract_output(self, accumulator):
returnaccumulator


@with_input_types(tuple[K, V])
@with_output_types(dict[K, V])
classToDict(CombinerWithoutDefaults):
"""A global CombineFn that condenses a PCollection into a single dict.

 PCollections should consist of 2-tuples, notionally (key, value) pairs.
 If multiple values are associated with the same key, only one of the values
 will be present in the resulting dict.
 """
defexpand(self, pcoll):
ifself.has_defaults:
returnpcoll|self.label>>core.CombineGlobally(ToDictCombineFn())
else:
returnpcoll|self.label>>core.CombineGlobally(
ToDictCombineFn()).without_defaults()


@with_input_types(tuple[K, V])
@with_output_types(dict[K, V])
classToDictCombineFn(core.CombineFn):
"""CombineFn for to_dict."""
defcreate_accumulator(self):
return {}

defadd_input(self, accumulator, element):
key, value=element
accumulator[key] =value
returnaccumulator

defmerge_accumulators(self, accumulators):
result= {}
forainaccumulators:
result.update(a)
returnresult

defextract_output(self, accumulator):
returnaccumulator


@with_input_types(T)
@with_output_types(set[T])
classToSet(CombinerWithoutDefaults):
"""A global CombineFn that condenses a PCollection into a set."""
defexpand(self, pcoll):
ifself.has_defaults:
returnpcoll|self.label>>core.CombineGlobally(ToSetCombineFn())
else:
returnpcoll|self.label>>core.CombineGlobally(
ToSetCombineFn()).without_defaults()


@with_input_types(T)
@with_output_types(set[T])
classToSetCombineFn(core.CombineFn):
"""CombineFn for ToSet."""
defcreate_accumulator(self):
returnset()

defadd_input(self, accumulator, element):
accumulator.add(element)
returnaccumulator

defmerge_accumulators(self, accumulators):
returnset.union(*accumulators)

defextract_output(self, accumulator):
returnaccumulator


class_CurriedFn(core.CombineFn):
"""Wrapped CombineFn with extra arguments."""
def__init__(self, fn, args, kwargs):
self.fn=fn
self.args=args
self.kwargs=kwargs

defsetup(self):
self.fn.setup(*self.args, **self.kwargs)

defcreate_accumulator(self):
returnself.fn.create_accumulator(*self.args, **self.kwargs)

defadd_input(self, accumulator, element):
returnself.fn.add_input(accumulator, element, *self.args, **self.kwargs)

defmerge_accumulators(self, accumulators):
returnself.fn.merge_accumulators(accumulators, *self.args, **self.kwargs)

defcompact(self, accumulator):
returnself.fn.compact(accumulator, *self.args, **self.kwargs)

defextract_output(self, accumulator):
returnself.fn.extract_output(accumulator, *self.args, **self.kwargs)

defteardown(self):
self.fn.teardown(*self.args, **self.kwargs)

defapply(self, elements):
returnself.fn.apply(elements, *self.args, **self.kwargs)


defcurry_combine_fn(fn, args, kwargs):
ifnotargsandnotkwargs:
returnfn
else:
return_CurriedFn(fn, args, kwargs)


classPhasedCombineFnExecutor(object):
"""Executor for phases of combine operations."""
def__init__(self, phase, fn, args, kwargs):

self.combine_fn=curry_combine_fn(fn, args, kwargs)

ifphase=='all':
self.apply=self.full_combine
elifphase=='add':
self.apply=self.add_only
elifphase=='merge':
self.apply=self.merge_only
elifphase=='extract':
self.apply=self.extract_only
elifphase=='convert':
self.apply=self.convert_to_accumulator
else:
raiseValueError('Unexpected phase: %s'%phase)

deffull_combine(self, elements):
returnself.combine_fn.apply(elements)

defadd_only(self, elements):
returnself.combine_fn.add_inputs(
self.combine_fn.create_accumulator(), elements)

defmerge_only(self, accumulators):
returnself.combine_fn.merge_accumulators(accumulators)

defextract_only(self, accumulator):
returnself.combine_fn.extract_output(accumulator)

defconvert_to_accumulator(self, element):
returnself.combine_fn.add_input(
self.combine_fn.create_accumulator(), element)


classLatest(object):
"""Combiners for computing the latest element"""
@with_input_types(T)
@with_output_types(T)
classGlobally(CombinerWithoutDefaults):
"""Compute the element with the latest timestamp from a
 PCollection."""
@staticmethod
defadd_timestamp(element, timestamp=core.DoFn.TimestampParam):
return [(element, timestamp)]

defexpand(self, pcoll):
ifself.has_defaults:
return (
pcoll
|core.ParDo(self.add_timestamp).with_output_types(
tuple[T, TimestampType])
|core.CombineGlobally(LatestCombineFn()))
else:
return (
pcoll
|core.ParDo(self.add_timestamp).with_output_types(
tuple[T, TimestampType])
|core.CombineGlobally(LatestCombineFn()).without_defaults())

@with_input_types(tuple[K, V])
@with_output_types(tuple[K, V])
classPerKey(ptransform.PTransform):
"""Compute elements with the latest timestamp for each key
 from a keyed PCollection"""
@staticmethod
defadd_timestamp(element, timestamp=core.DoFn.TimestampParam):
key, value=element
return [(key, (value, timestamp))]

defexpand(self, pcoll):
return (
pcoll
|core.ParDo(self.add_timestamp).with_output_types(
tuple[K, tuple[T, TimestampType]])
|core.CombinePerKey(LatestCombineFn()))


@with_input_types(tuple[T, TimestampType])
@with_output_types(T)
classLatestCombineFn(core.CombineFn):
"""CombineFn to get the element with the latest timestamp
 from a PCollection."""
defcreate_accumulator(self):
return (None, window.MIN_TIMESTAMP)

defadd_input(self, accumulator, element):
ifaccumulator[1] >element[1]:
returnaccumulator
else:
returnelement

defmerge_accumulators(self, accumulators):
result=self.create_accumulator()
foraccumulatorinaccumulators:
result=self.add_input(result, accumulator)
returnresult

defextract_output(self, accumulator):
returnaccumulator[0]