Source code for apache_beam.transforms.stats

#
# 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.
#

# cython: language_level=3

"""This module has all statistic related transforms.

This ApproximateUnique class will be deprecated [1]. PLease look into using
HLLCount in the zetasketch extension module [2].

[1] https://lists.apache.org/thread.html/501605df5027567099b81f18c080469661fb426
4a002615fa1510502%40%3Cdev.beam.apache.org%3E
[2] https://beam.apache.org/releases/javadoc/2.16.0/org/apache/beam/sdk/extensio
ns/zetasketch/HllCount.html
"""

# pytype: skip-file

import hashlib
import heapq
import itertools
import logging
import math
import typing
from typing import Any
from typing import Callable
from typing import List
from typing import Tuple

from apache_beam import coders
from apache_beam import typehints
from apache_beam.transforms.core import *
from apache_beam.transforms.display import DisplayDataItem
from apache_beam.transforms.ptransform import PTransform

__all__ = [
    'ApproximateQuantiles',
    'ApproximateUnique',
]

# Type variables
T = typing.TypeVar('T')
K = typing.TypeVar('K')
V = typing.TypeVar('V')

try:
  import mmh3  # pylint: disable=import-error

  def _mmh3_hash(value):
    # mmh3.hash64 returns two 64-bit unsigned integers
    return mmh3.hash64(value, seed=0, signed=False)[0]

  _default_hash_fn = _mmh3_hash
  _default_hash_fn_type = 'mmh3'
except ImportError:

  def _md5_hash(value):
    # md5 is a 128-bit hash, so we truncate the hexdigest (string of 32
    # hexadecimal digits) to 16 digits and convert to int to get the 64-bit
    # integer fingerprint.
    return int(hashlib.md5(value).hexdigest()[:16], 16)

  _default_hash_fn = _md5_hash
  _default_hash_fn_type = 'md5'


def _get_default_hash_fn():
  """Returns either murmurhash or md5 based on installation."""
  if _default_hash_fn_type == 'md5':
    logging.warning(
        'Couldn\'t find murmurhash. Install mmh3 for a faster implementation of'
        'ApproximateUnique.')
  return _default_hash_fn


[docs]class ApproximateUnique(object): """ Hashes input elements and uses those to extrapolate the size of the entire set of hash values by assuming the rest of the hash values are as densely distributed as the sample space. """ _NO_VALUE_ERR_MSG = 'Either size or error should be set. Received {}.' _MULTI_VALUE_ERR_MSG = 'Either size or error should be set. ' \ 'Received {size = %s, error = %s}.' _INPUT_SIZE_ERR_MSG = 'ApproximateUnique needs a size >= 16 for an error ' \ '<= 0.50. In general, the estimation error is about ' \ '2 / sqrt(sample_size). Received {size = %s}.' _INPUT_ERROR_ERR_MSG = 'ApproximateUnique needs an estimation error ' \ 'between 0.01 and 0.50. Received {error = %s}.'
[docs] @staticmethod def parse_input_params(size=None, error=None): """ Check if input params are valid and return sample size. :param size: an int not smaller than 16, which we would use to estimate number of unique values. :param error: max estimation error, which is a float between 0.01 and 0.50. If error is given, sample size will be calculated from error with _get_sample_size_from_est_error function. :return: sample size :raises: ValueError: If both size and error are given, or neither is given, or values are out of range. """ if None not in (size, error): raise ValueError(ApproximateUnique._MULTI_VALUE_ERR_MSG % (size, error)) elif size is None and error is None: raise ValueError(ApproximateUnique._NO_VALUE_ERR_MSG) elif size is not None: if not isinstance(size, int) or size < 16: raise ValueError(ApproximateUnique._INPUT_SIZE_ERR_MSG % (size)) else: return size else: if error < 0.01 or error > 0.5: raise ValueError(ApproximateUnique._INPUT_ERROR_ERR_MSG % (error)) else: return ApproximateUnique._get_sample_size_from_est_error(error)
@staticmethod def _get_sample_size_from_est_error(est_err): """ :return: sample size Calculate sample size from estimation error """ return math.ceil(4.0 / math.pow(est_err, 2.0))
[docs] @typehints.with_input_types(T) @typehints.with_output_types(int) class Globally(PTransform): """ Approximate.Globally approximate number of unique values""" def __init__(self, size=None, error=None): self._sample_size = ApproximateUnique.parse_input_params(size, error)
[docs] def expand(self, pcoll): coder = coders.registry.get_coder(pcoll) return pcoll \ | 'CountGlobalUniqueValues' \ >> (CombineGlobally(ApproximateUniqueCombineFn(self._sample_size, coder)))
[docs] @typehints.with_input_types(typing.Tuple[K, V]) @typehints.with_output_types(typing.Tuple[K, int]) class PerKey(PTransform): """ Approximate.PerKey approximate number of unique values per key""" def __init__(self, size=None, error=None): self._sample_size = ApproximateUnique.parse_input_params(size, error)
[docs] def expand(self, pcoll): coder = coders.registry.get_coder(pcoll) return pcoll \ | 'CountPerKeyUniqueValues' \ >> (CombinePerKey(ApproximateUniqueCombineFn(self._sample_size, coder)))
class _LargestUnique(object): """ An object to keep samples and calculate sample hash space. It is an accumulator of a combine function. """ # We use unsigned 64-bit integer hashes. _HASH_SPACE_SIZE = 2.0**64 def __init__(self, sample_size): self._sample_size = sample_size self._min_hash = 2.0**64 self._sample_heap = [] self._sample_set = set() def add(self, element): """ :param an element from pcoll. :return: boolean type whether the value is in the heap Adds a value to the heap, returning whether the value is (large enough to be) in the heap. """ if len(self._sample_heap) >= self._sample_size and element < self._min_hash: return False if element not in self._sample_set: self._sample_set.add(element) heapq.heappush(self._sample_heap, element) if len(self._sample_heap) > self._sample_size: temp = heapq.heappop(self._sample_heap) self._sample_set.remove(temp) self._min_hash = self._sample_heap[0] elif element < self._min_hash: self._min_hash = element return True def get_estimate(self): """ :return: estimation count of unique values If heap size is smaller than sample size, just return heap size. Otherwise, takes into account the possibility of hash collisions, which become more likely than not for 2^32 distinct elements. Note that log(1+x) ~ x for small x, so for sampleSize << maxHash log(1 - sample_size/sample_space) / log(1 - 1/sample_space) ~ sample_size and hence estimate ~ sample_size * hash_space / sample_space as one would expect. Given sample_size / sample_space = est / hash_space est = sample_size * hash_space / sample_space Given above sample_size approximate, est = log1p(-sample_size/sample_space) / log1p(-1/sample_space) * hash_space / sample_space """ if len(self._sample_heap) < self._sample_size: return len(self._sample_heap) else: sample_space_size = self._HASH_SPACE_SIZE - 1.0 * self._min_hash est = ( math.log1p(-self._sample_size / sample_space_size) / math.log1p(-1 / sample_space_size) * self._HASH_SPACE_SIZE / sample_space_size) return round(est) class ApproximateUniqueCombineFn(CombineFn): """ ApproximateUniqueCombineFn computes an estimate of the number of unique values that were combined. """ def __init__(self, sample_size, coder): self._sample_size = sample_size coder = coders.typecoders.registry.verify_deterministic( coder, 'ApproximateUniqueCombineFn') self._coder = coder self._hash_fn = _get_default_hash_fn() def create_accumulator(self, *args, **kwargs): return _LargestUnique(self._sample_size) def add_input(self, accumulator, element, *args, **kwargs): try: hashed_value = self._hash_fn(self._coder.encode(element)) accumulator.add(hashed_value) return accumulator except Exception as e: raise RuntimeError("Runtime exception: %s" % e) # created an issue https://issues.apache.org/jira/browse/BEAM-7285 to speed up # merge process. def merge_accumulators(self, accumulators, *args, **kwargs): merged_accumulator = self.create_accumulator() for accumulator in accumulators: for i in accumulator._sample_heap: merged_accumulator.add(i) return merged_accumulator @staticmethod def extract_output(accumulator): return accumulator.get_estimate() def display_data(self): return {'sample_size': self._sample_size}
[docs]class ApproximateQuantiles(object): """ PTransform for getting the idea of data distribution using approximate N-tile (e.g. quartiles, percentiles etc.) either globally or per-key. Examples: in: list(range(101)), num_quantiles=5 out: [0, 25, 50, 75, 100] in: [(i, 1 if i<10 else 1e-5) for i in range(101)], num_quantiles=5, weighted=True out: [0, 2, 5, 7, 100] in: [list(range(10)), ..., list(range(90, 101))], num_quantiles=5, input_batched=True out: [0, 25, 50, 75, 100] in: [(list(range(10)), [1]*10), (list(range(10)), [0]*10), ..., (list(range(90, 101)), [0]*11)], num_quantiles=5, input_batched=True, weighted=True out: [0, 2, 5, 7, 100] """ @staticmethod def _display_data(num_quantiles, key, reverse, weighted, input_batched): return { 'num_quantiles': DisplayDataItem(num_quantiles, label='Quantile Count'), 'key': DisplayDataItem( key.__name__ if hasattr(key, '__name__') else key.__class__.__name__, label='Record Comparer Key'), 'reverse': DisplayDataItem(str(reverse), label='Is Reversed'), 'weighted': DisplayDataItem(str(weighted), label='Is Weighted'), 'input_batched': DisplayDataItem( str(input_batched), label='Is Input Batched'), }
[docs] @typehints.with_input_types( typehints.Union[typing.Sequence[T], typing.Tuple[T, float]]) @typehints.with_output_types(typing.List[T]) class Globally(PTransform): """ PTransform takes PCollection and returns a list whose single value is approximate N-tiles of the input collection globally. Args: num_quantiles: number of elements in the resulting quantiles values list. key: (optional) Key is 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. weighted: (optional) if set to True, the transform returns weighted quantiles. The input PCollection is then expected to contain tuples of input values with the corresponding weight. input_batched: (optional) if set to True, the transform expects each element of input PCollection to be a batch, which is a list of elements for non-weighted case and a tuple of lists of elements and weights for weighted. Provides a way to accumulate multiple elements at a time more efficiently. """ def __init__( self, num_quantiles, key=None, reverse=False, weighted=False, input_batched=False): self._num_quantiles = num_quantiles self._key = key self._reverse = reverse self._weighted = weighted self._input_batched = input_batched
[docs] def expand(self, pcoll): return pcoll | CombineGlobally( ApproximateQuantilesCombineFn.create( num_quantiles=self._num_quantiles, key=self._key, reverse=self._reverse, weighted=self._weighted, input_batched=self._input_batched))
[docs] def display_data(self): return ApproximateQuantiles._display_data( num_quantiles=self._num_quantiles, key=self._key, reverse=self._reverse, weighted=self._weighted, input_batched=self._input_batched)
[docs] @typehints.with_input_types( typehints.Union[typing.Tuple[K, V], typing.Tuple[K, typing.Tuple[V, float]]]) @typehints.with_output_types(typing.Tuple[K, typing.List[V]]) class PerKey(PTransform): """ PTransform takes PCollection of KV and returns a list based on each key whose single value is list of approximate N-tiles of the input element of the key. Args: num_quantiles: number of elements in the resulting quantiles values list. key: (optional) Key is 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. weighted: (optional) if set to True, the transform returns weighted quantiles. The input PCollection is then expected to contain tuples of input values with the corresponding weight. input_batched: (optional) if set to True, the transform expects each element of input PCollection to be a batch, which is a list of elements for non-weighted case and a tuple of lists of elements and weights for weighted. Provides a way to accumulate multiple elements at a time more efficiently. """ def __init__( self, num_quantiles, key=None, reverse=False, weighted=False, input_batched=False): self._num_quantiles = num_quantiles self._key = key self._reverse = reverse self._weighted = weighted self._input_batched = input_batched
[docs] def expand(self, pcoll): return pcoll | CombinePerKey( ApproximateQuantilesCombineFn.create( num_quantiles=self._num_quantiles, key=self._key, reverse=self._reverse, weighted=self._weighted, input_batched=self._input_batched))
[docs] def display_data(self): return ApproximateQuantiles._display_data( num_quantiles=self._num_quantiles, key=self._key, reverse=self._reverse, weighted=self._weighted, input_batched=self._input_batched)
class _QuantileSpec(object): """Quantiles computation specifications.""" def __init__(self, buffer_size, num_buffers, weighted, key, reverse): # type: (int, int, bool, Any, bool) -> None self.buffer_size = buffer_size self.num_buffers = num_buffers self.weighted = weighted self.key = key self.reverse = reverse # Used to sort tuples of values and weights. self.weighted_key = None if key is None else (lambda x: key(x[0])) # Used to compare values. if reverse and key is None: self.less_than = lambda a, b: a > b elif reverse: self.less_than = lambda a, b: key(a) > key(b) elif key is None: self.less_than = lambda a, b: a < b else: self.less_than = lambda a, b: key(a) < key(b) def get_argsort_key(self, elements): # type: (List) -> Callable[[int], Any] """Returns a key for sorting indices of elements by element's value.""" if self.key is None: return elements.__getitem__ else: return lambda idx: self.key(elements[idx]) def __reduce__(self): return ( self.__class__, ( self.buffer_size, self.num_buffers, self.weighted, self.key, self.reverse)) class _QuantileBuffer(object): """A single buffer in the sense of the referenced algorithm. (see http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.6.6513&rep=rep1 &type=pdf and ApproximateQuantilesCombineFn for further information)""" def __init__( self, elements, weights, weighted, level=0, min_val=None, max_val=None): # type: (List, List, bool, int, Any, Any) -> None self.elements = elements # In non-weighted case weights contains a single element representing weight # of the buffer in the sense of the original algorithm. In weighted case, # it stores weights of individual elements. self.weights = weights self.weighted = weighted self.level = level if min_val is None or max_val is None: # Buffer is always initialized with sorted elements. self.min_val = elements[0] self.max_val = elements[-1] else: # Note that collapsed buffer may not contain min and max in the list of # elements. self.min_val = min_val self.max_val = max_val def __iter__(self): return zip( self.elements, self.weights if self.weighted else itertools.repeat(self.weights[0])) def __lt__(self, other): return self.level < other.level class _QuantileState(object): """ Compact summarization of a collection on which quantiles can be estimated. """ def __init__(self, unbuffered_elements, unbuffered_weights, buffers, spec): # type: (List, List, List[_QuantileBuffer], _QuantileSpec) -> None self.buffers = buffers self.spec = spec if spec.weighted: self.add_unbuffered = self._add_unbuffered_weighted else: self.add_unbuffered = self._add_unbuffered # The algorithm requires that the manipulated buffers always be filled to # capacity to perform the collapse operation. This operation can be extended # to buffers of varying sizes by introducing the notion of fractional # weights, but it's easier to simply combine the remainders from all shards # into new, full buffers and then take them into account when computing the # final output. self.unbuffered_elements = unbuffered_elements self.unbuffered_weights = unbuffered_weights # This is needed for pickling to work when Cythonization is enabled. def __reduce__(self): return ( self.__class__, ( self.unbuffered_elements, self.unbuffered_weights, self.buffers, self.spec)) def is_empty(self): # type: () -> bool """Check if the buffered & unbuffered elements are empty or not.""" return not self.unbuffered_elements and not self.buffers def _add_unbuffered(self, elements, offset_fn): # type: (List, Any) -> None """ Add elements to the unbuffered list, creating new buffers and collapsing if needed. """ self.unbuffered_elements.extend(elements) num_new_buffers = len(self.unbuffered_elements) // self.spec.buffer_size for idx in range(num_new_buffers): to_buffer = sorted( self.unbuffered_elements[idx * self.spec.buffer_size:(idx + 1) * self.spec.buffer_size], key=self.spec.key, reverse=self.spec.reverse) heapq.heappush( self.buffers, _QuantileBuffer(elements=to_buffer, weights=[1], weighted=False)) if num_new_buffers > 0: self.unbuffered_elements = self.unbuffered_elements[num_new_buffers * self.spec. buffer_size:] self.collapse_if_needed(offset_fn) def _add_unbuffered_weighted(self, elements, offset_fn): # type: (List, Any) -> None """ Add elements with weights to the unbuffered list, creating new buffers and collapsing if needed. """ if len(elements) == 1: self.unbuffered_elements.append(elements[0][0]) self.unbuffered_weights.append(elements[0][1]) else: self.unbuffered_elements.extend(elements[0]) self.unbuffered_weights.extend(elements[1]) num_new_buffers = len(self.unbuffered_elements) // self.spec.buffer_size argsort_key = self.spec.get_argsort_key(self.unbuffered_elements) for idx in range(num_new_buffers): argsort = sorted( range(idx * self.spec.buffer_size, (idx + 1) * self.spec.buffer_size), key=argsort_key, reverse=self.spec.reverse) elements_to_buffer = [self.unbuffered_elements[idx] for idx in argsort] weights_to_buffer = [self.unbuffered_weights[idx] for idx in argsort] heapq.heappush( self.buffers, _QuantileBuffer( elements=elements_to_buffer, weights=weights_to_buffer, weighted=True)) if num_new_buffers > 0: self.unbuffered_elements = self.unbuffered_elements[num_new_buffers * self.spec. buffer_size:] self.unbuffered_weights = self.unbuffered_weights[num_new_buffers * self.spec.buffer_size:] self.collapse_if_needed(offset_fn) def finalize(self): # type: () -> None """ Creates a new buffer using all unbuffered elements. Called before extracting an output. Note that the buffer doesn't have to be put in a proper position since _collapse is not going to be called after. """ if self.unbuffered_elements and self.spec.weighted: argsort_key = self.spec.get_argsort_key(self.unbuffered_elements) argsort = sorted( range(len(self.unbuffered_elements)), key=argsort_key, reverse=self.spec.reverse) self.unbuffered_elements = [ self.unbuffered_elements[idx] for idx in argsort ] self.unbuffered_weights = [ self.unbuffered_weights[idx] for idx in argsort ] self.buffers.append( _QuantileBuffer( self.unbuffered_elements, self.unbuffered_weights, weighted=True)) self.unbuffered_weights = [] elif self.unbuffered_elements: self.unbuffered_elements.sort( key=self.spec.key, reverse=self.spec.reverse) self.buffers.append( _QuantileBuffer( self.unbuffered_elements, weights=[1], weighted=False)) self.unbuffered_elements = [] def collapse_if_needed(self, offset_fn): # type: (Any) -> None """ Checks if summary has too many buffers and collapses some of them until the limit is restored. """ while len(self.buffers) > self.spec.num_buffers: to_collapse = [heapq.heappop(self.buffers), heapq.heappop(self.buffers)] min_level = to_collapse[1].level while self.buffers and self.buffers[0].level <= min_level: to_collapse.append(heapq.heappop(self.buffers)) heapq.heappush(self.buffers, _collapse(to_collapse, offset_fn, self.spec)) def _collapse(buffers, offset_fn, spec): # type: (List[_QuantileBuffer], Any, _QuantileSpec) -> _QuantileBuffer """ Approximates elements from multiple buffers and produces a single buffer. """ new_level = 0 new_weight = 0 for buffer in buffers: # As presented in the paper, there should always be at least two # buffers of the same (minimal) level to collapse, but it is possible # to violate this condition when combining buffers from independently # computed shards. If they differ we take the max. new_level = max([new_level, buffer.level + 1]) new_weight = new_weight + sum(buffer.weights) if spec.weighted: step = new_weight / (spec.buffer_size - 1) offset = new_weight / (2 * spec.buffer_size) else: step = new_weight offset = offset_fn(new_weight) new_elements, new_weights, min_val, max_val = \ _interpolate(buffers, spec.buffer_size, step, offset, spec) if not spec.weighted: new_weights = [new_weight] return _QuantileBuffer( new_elements, new_weights, spec.weighted, new_level, min_val, max_val) def _interpolate(buffers, count, step, offset, spec): # type: (List[_QuantileBuffer], int, float, float, _QuantileSpec) -> Tuple[List, List, Any, Any] """ Emulates taking the ordered union of all elements in buffers, repeated according to their weight, and picking out the (k * step + offset)-th elements of this list for `0 <= k < count`. """ buffer_iterators = [] min_val = buffers[0].min_val max_val = buffers[0].max_val for buffer in buffers: # Calculate extreme values for the union of buffers. min_val = buffer.min_val if spec.less_than( buffer.min_val, min_val) else min_val max_val = buffer.max_val if spec.less_than( max_val, buffer.max_val) else max_val buffer_iterators.append(iter(buffer)) # Note that `heapq.merge` can also be used here since the buffers are sorted. # In practice, however, `sorted` uses natural order in the union and # significantly outperforms `heapq.merge`. sorted_elements = sorted( itertools.chain.from_iterable(buffer_iterators), key=spec.weighted_key, reverse=spec.reverse) if not spec.weighted: # If all buffers have the same weight, then quantiles' indices are evenly # distributed over a range [0, len(sorted_elements)]. buffers_have_same_weight = True weight = buffers[0].weights[0] for buffer in buffers: if buffer.weights[0] != weight: buffers_have_same_weight = False break if buffers_have_same_weight: offset = offset / weight step = step / weight max_idx = len(sorted_elements) - 1 result = [ sorted_elements[min(int(j * step + offset), max_idx)][0] for j in range(count) ] return result, [], min_val, max_val sorted_elements_iter = iter(sorted_elements) weighted_element = next(sorted_elements_iter) new_elements = [] new_weights = [] j = 0 current_weight = weighted_element[1] previous_weight = 0 while j < count: target_weight = j * step + offset j += 1 try: while current_weight <= target_weight: weighted_element = next(sorted_elements_iter) current_weight += weighted_element[1] except StopIteration: pass new_elements.append(weighted_element[0]) if spec.weighted: new_weights.append(current_weight - previous_weight) previous_weight = current_weight return new_elements, new_weights, min_val, max_val class ApproximateQuantilesCombineFn(CombineFn): """ This combiner gives an idea of the distribution of a collection of values using approximate N-tiles. The output of this combiner is the list of size of the number of quantiles (num_quantiles), containing the input values of the minimum value item of the list, the intermediate values (n-tiles) and the maximum value item of the list, in the sort order provided via key (similar to the key argument of Python's sorting methods). If there are fewer values to combine than the number of quantile (num_quantiles), then the resulting list will contain all the values being combined, in sorted order. If no `key` is provided, then the results are sorted in the natural order. To evaluate the quantiles, we use the "New Algorithm" described here: [MRL98] Manku, Rajagopalan & Lindsay, "Approximate Medians and other Quantiles in One Pass and with Limited Memory", Proc. 1998 ACM SIGMOD, Vol 27, No 2, p 426-435, June 1998. http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.6.6513&rep=rep1 &type=pdf Note that the weighted quantiles are evaluated using a generalized version of the algorithm referenced in the paper. The default error bound is (1 / num_quantiles) for uniformly distributed data and min(1e-2, 1 / num_quantiles) for weighted case, though in practice the accuracy tends to be much better. Args: num_quantiles: Number of quantiles to produce. It is the size of the final output list, including the mininum and maximum value items. buffer_size: The size of the buffers, corresponding to k in the referenced paper. num_buffers: The number of buffers, corresponding to b in the referenced paper. key: (optional) Key is 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. weighted: (optional) if set to True, the combiner produces weighted quantiles. The input elements are then expected to be tuples of input values with the corresponding weight. input_batched: (optional) if set to True, inputs are expected to be batches of elements. """ # For alternating between biasing up and down in the above even weight # collapse operation. _offset_jitter = 0 # The cost (in time and space) to compute quantiles to a given accuracy is a # function of the total number of elements in the data set. If an estimate is # not known or specified, we use this as an upper bound. If this is too low, # errors may exceed the requested tolerance; if too high, efficiency may be # non-optimal. The impact is logarithmic with respect to this value, so this # default should be fine for most uses. _MAX_NUM_ELEMENTS = 1e9 _qs = None # type: _QuantileState def __init__( self, num_quantiles, # type: int buffer_size, # type: int num_buffers, # type: int key=None, reverse=False, weighted=False, input_batched=False): self._num_quantiles = num_quantiles self._spec = _QuantileSpec(buffer_size, num_buffers, weighted, key, reverse) self._input_batched = input_batched if self._input_batched: setattr(self, 'add_input', self._add_inputs) def __reduce__(self): return ( self.__class__, ( self._num_quantiles, self._spec.buffer_size, self._spec.num_buffers, self._spec.key, self._spec.reverse, self._spec.weighted, self._input_batched)) @classmethod def create( cls, num_quantiles, # type: int epsilon=None, max_num_elements=None, key=None, reverse=False, weighted=False, input_batched=False): # type: (...) -> ApproximateQuantilesCombineFn """ Creates an approximate quantiles combiner with the given key and desired number of quantiles. Args: num_quantiles: Number of quantiles to produce. It is the size of the final output list, including the mininum and maximum value items. epsilon: (optional) The default error bound is `epsilon`, which holds as long as the number of elements is less than `_MAX_NUM_ELEMENTS`. Specifically, if one considers the input as a sorted list x_1, ..., x_N, then the distance between each exact quantile x_c and its approximation x_c' is bounded by `|c - c'| < epsilon * N`. Note that these errors are worst-case scenarios. In practice the accuracy tends to be much better. max_num_elements: (optional) The cost (in time and space) to compute quantiles to a given accuracy is a function of the total number of elements in the data set. key: (optional) Key is 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. weighted: (optional) if set to True, the combiner produces weighted quantiles. The input elements are then expected to be tuples of values with the corresponding weight. input_batched: (optional) if set to True, inputs are expected to be batches of elements. """ max_num_elements = max_num_elements or cls._MAX_NUM_ELEMENTS if not epsilon: epsilon = min(1e-2, 1.0 / num_quantiles) \ if weighted else (1.0 / num_quantiles) # Note that calculation of the buffer size and the number of buffers here # is based on technique used in the Munro-Paterson algorithm. Switching to # the logic used in the "New Algorithm" may result in memory savings since # it results in lower values for b and k in practice. b = 2 while (b - 2) * (1 << (b - 2)) < epsilon * max_num_elements: b = b + 1 b = b - 1 k = max(2, int(math.ceil(max_num_elements / float(1 << (b - 1))))) return cls( num_quantiles=num_quantiles, buffer_size=k, num_buffers=b, key=key, reverse=reverse, weighted=weighted, input_batched=input_batched) def _offset(self, new_weight): # type: (int) -> float """ If the weight is even, we must round up or down. Alternate between these two options to avoid a bias. """ if new_weight % 2 == 1: return (new_weight + 1) / 2 else: self._offset_jitter = 2 - self._offset_jitter return (new_weight + self._offset_jitter) / 2 # TODO(BEAM-7746): Signature incompatible with supertype def create_accumulator(self): # type: ignore[override] # type: () -> _QuantileState self._qs = _QuantileState( unbuffered_elements=[], unbuffered_weights=[], buffers=[], spec=self._spec) return self._qs def add_input(self, quantile_state, element): """ Add a new element to the collection being summarized by quantile state. """ quantile_state.add_unbuffered([element], self._offset) return quantile_state def _add_inputs(self, quantile_state, elements): # type: (_QuantileState, List) -> _QuantileState """ Add a batch of elements to the collection being summarized by quantile state. """ if len(elements) == 0: return quantile_state quantile_state.add_unbuffered(elements, self._offset) return quantile_state def merge_accumulators(self, accumulators): """Merges all the accumulators (quantile state) as one.""" qs = self.create_accumulator() for accumulator in accumulators: if accumulator.is_empty(): continue if self._spec.weighted: qs.add_unbuffered( [accumulator.unbuffered_elements, accumulator.unbuffered_weights], self._offset) else: qs.add_unbuffered(accumulator.unbuffered_elements, self._offset) qs.buffers.extend(accumulator.buffers) heapq.heapify(qs.buffers) qs.collapse_if_needed(self._offset) return qs def extract_output(self, accumulator): """ Outputs num_quantiles elements consisting of the minimum, maximum and num_quantiles - 2 evenly spaced intermediate elements. Returns the empty list if no elements have been added. """ if accumulator.is_empty(): return [] accumulator.finalize() all_elems = accumulator.buffers total_weight = 0 if self._spec.weighted: for buffer_elem in all_elems: total_weight += sum(buffer_elem.weights) else: for buffer_elem in all_elems: total_weight += len(buffer_elem.elements) * buffer_elem.weights[0] step = total_weight / (self._num_quantiles - 1) offset = (total_weight - 1) / (self._num_quantiles - 1) quantiles, _, min_val, max_val = \ _interpolate(all_elems, self._num_quantiles - 2, step, offset, self._spec) return [min_val] + quantiles + [max_val]