You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by ro...@apache.org on 2016/07/28 18:08:46 UTC
[1/3] incubator-beam git commit: Better top implementation.
Repository: incubator-beam
Updated Branches:
refs/heads/python-sdk 53ab635c7 -> b4716d9dc
Better top implementation.
When selecting the top k of n, it is common that k << n.
Using a heap is O(n log k) while select algorithms can
achieve O(n + k log k).
This also avoids the ugliness that heapq does not take the
comparator as an argument, resulting in _HeapItem classes that
were cumbersome and expensive to serialize.
Project: http://git-wip-us.apache.org/repos/asf/incubator-beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-beam/commit/adb3ed93
Tree: http://git-wip-us.apache.org/repos/asf/incubator-beam/tree/adb3ed93
Diff: http://git-wip-us.apache.org/repos/asf/incubator-beam/diff/adb3ed93
Branch: refs/heads/python-sdk
Commit: adb3ed93053c83b4e28e7baa879e9aee82f02785
Parents: 53ab635
Author: Robert Bradshaw <ro...@gmail.com>
Authored: Wed Jul 27 10:09:49 2016 -0700
Committer: Robert Bradshaw <ro...@gmail.com>
Committed: Thu Jul 28 11:08:05 2016 -0700
----------------------------------------------------------------------
sdks/python/apache_beam/transforms/combiners.py | 111 +++++++++----------
1 file changed, 51 insertions(+), 60 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/adb3ed93/sdks/python/apache_beam/transforms/combiners.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/combiners.py b/sdks/python/apache_beam/transforms/combiners.py
index 8c56e5a..453c0f8 100644
--- a/sdks/python/apache_beam/transforms/combiners.py
+++ b/sdks/python/apache_beam/transforms/combiners.py
@@ -228,76 +228,67 @@ class TopCombineFn(core.CombineFn):
apply call become additional arguments to the comparator.
"""
- # Actually pickling the comparison operators (including, often, their
- # entire globals) can be very expensive. Instead refer to them by index
- # in this dictionary, which is populated on construction (including
- # unpickling).
- compare_by_id = {}
-
- def __init__(self, n, compare, _compare_id=None): # pylint: disable=invalid-name
+ # TODO(robertwb): Allow taking a key rather than a compare.
+ def __init__(self, n, compare):
self._n = n
+ self._buffer_size = min(2 * n, n + 1000)
self._compare = compare
- self._compare_id = _compare_id or id(compare)
- TopCombineFn.compare_by_id[self._compare_id] = self._compare
-
- def __reduce_ex__(self, _):
- return TopCombineFn, (self._n, self._compare, self._compare_id)
- class _HeapItem(object):
- """A wrapper for values supporting arbitrary comparisons.
-
- The heap implementation supplied by Python is a min heap that always uses
- the __lt__ operator if one is available. This wrapper overloads __lt__,
- letting us specify arbitrary precedence for elements in the PCollection.
- """
+ def create_accumulator(self, *args, **kwargs):
+ return None, []
- def __init__(self, item, compare_id, *args, **kwargs):
- # item: wrapped item.
- # compare: an implementation of the pairwise < operator.
- # args, kwargs: extra arguments supplied to the compare function.
- self.item = item
- self.compare_id = compare_id
- self.args = args
- self.kwargs = kwargs
+ def add_input(self, accumulator, element, *args, **kwargs):
+ if args or kwargs:
+ lt = lambda a, b: self._compare(a, b, *args, **kwargs)
+ else:
+ lt = self._compare
- def __lt__(self, other):
- return TopCombineFn.compare_by_id[self.compare_id](
- self.item, other.item, *self.args, **self.kwargs)
+ threshold, buffer = accumulator
+ if len(buffer) < self._n:
+ if not buffer:
+ return element, [element]
+ else:
+ buffer.append(element)
+ if lt(element, threshold): # element < threshold
+ return element, buffer
+ else:
+ return accumulator # with mutated buffer
+ elif lt(threshold, element): # threshold < element
+ buffer.append(element)
+ if len(buffer) < self._buffer_size:
+ return accumulator
+ else:
+ buffer.sort(cmp=lambda a, b: (not lt(a, b)) - (not lt(b, a)))
+ return buffer[-self._n], buffer[-self._n:]
+ else:
+ return accumulator
- def create_accumulator(self, *args, **kwargs):
- return [] # Empty heap.
-
- def add_input(self, heap, element, *args, **kwargs):
- # Note that because heap is a min heap, heappushpop will discard incoming
- # elements that are lesser (according to compare) than those in the heap
- # (since that's what you would get if you pushed a small element on and
- # popped the smallest element off). So, filtering a collection with a
- # min-heap gives you the largest elements in the collection.
- item = self._HeapItem(element, self._compare_id, *args, **kwargs)
- if len(heap) < self._n:
- heapq.heappush(heap, item)
+ def merge_accumulators(self, accumulators, *args, **kwargs):
+ accumulators = list(accumulators)
+ if args or kwargs:
+ add_input = lambda accumulator, element: self.add_input(
+ accumulator, element, *args, **kwargs)
else:
- heapq.heappushpop(heap, item)
- return heap
-
- def merge_accumulators(self, heaps, *args, **kwargs):
- heap = []
- for e in itertools.chain(*heaps):
- if len(heap) < self._n:
- heapq.heappush(heap, e)
- else:
- heapq.heappushpop(heap, e)
- return heap
+ add_input = self.add_input
- def extract_output(self, heap, *args, **kwargs):
- # Items in the heap are heap-ordered. We put them in sorted order, but we
- # have to use the reverse order because the result is expected to go
- # from greatest to least (as defined by the supplied comparison function).
- return [e.item for e in sorted(heap, reverse=True)]
+ total_accumulator = None
+ for accumulator in accumulators:
+ if total_accumulator is None:
+ total_accumulator = accumulator
+ else:
+ for element in accumulator[1]:
+ total_accumulator = add_input(total_accumulator, element)
+ return total_accumulator
+ def extract_output(self, accumulator, *args, **kwargs):
+ if args or kwargs:
+ lt = lambda a, b: self._compare(a, b, *args, **kwargs)
+ else:
+ lt = self._compare
-# Python's pickling is broken for nested classes.
-_HeapItem = TopCombineFn._HeapItem # pylint: disable=protected-access
+ _, buffer = accumulator
+ buffer.sort(cmp=lambda a, b: (not lt(a, b)) - (not lt(b, a)))
+ return buffer[:-self._n-1:-1]
class Largest(TopCombineFn):
[3/3] incubator-beam git commit: Closes #741
Posted by ro...@apache.org.
Closes #741
Project: http://git-wip-us.apache.org/repos/asf/incubator-beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-beam/commit/b4716d9d
Tree: http://git-wip-us.apache.org/repos/asf/incubator-beam/tree/b4716d9d
Diff: http://git-wip-us.apache.org/repos/asf/incubator-beam/diff/b4716d9d
Branch: refs/heads/python-sdk
Commit: b4716d9dcf9c3f198fc72e181adc0bae8b6fa818
Parents: 53ab635 866a09d
Author: Robert Bradshaw <ro...@gmail.com>
Authored: Thu Jul 28 11:08:06 2016 -0700
Committer: Robert Bradshaw <ro...@gmail.com>
Committed: Thu Jul 28 11:08:06 2016 -0700
----------------------------------------------------------------------
sdks/python/apache_beam/transforms/combiners.py | 201 ++++++++++++-------
.../apache_beam/transforms/combiners_test.py | 42 +++-
2 files changed, 164 insertions(+), 79 deletions(-)
----------------------------------------------------------------------
[2/3] incubator-beam git commit: Allow Top operations to take key
argument rather than compare.
Posted by ro...@apache.org.
Allow Top operations to take key argument rather than compare.
Also support ordering by the natural ordering if neither is specified,
and a reverse option. This is both an API and performance improvement.
Project: http://git-wip-us.apache.org/repos/asf/incubator-beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-beam/commit/866a09d6
Tree: http://git-wip-us.apache.org/repos/asf/incubator-beam/tree/866a09d6
Diff: http://git-wip-us.apache.org/repos/asf/incubator-beam/diff/866a09d6
Branch: refs/heads/python-sdk
Commit: 866a09d6c5c4c1c3c39d3a1f2e7b4cb7ec319f6f
Parents: adb3ed9
Author: Robert Bradshaw <ro...@gmail.com>
Authored: Wed Jul 27 12:23:02 2016 -0700
Committer: Robert Bradshaw <ro...@gmail.com>
Committed: Thu Jul 28 11:08:06 2016 -0700
----------------------------------------------------------------------
sdks/python/apache_beam/transforms/combiners.py | 116 ++++++++++++++-----
.../apache_beam/transforms/combiners_test.py | 42 ++++++-
2 files changed, 126 insertions(+), 32 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/866a09d6/sdks/python/apache_beam/transforms/combiners.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/combiners.py b/sdks/python/apache_beam/transforms/combiners.py
index 453c0f8..155dcc6 100644
--- a/sdks/python/apache_beam/transforms/combiners.py
+++ b/sdks/python/apache_beam/transforms/combiners.py
@@ -19,8 +19,7 @@
from __future__ import absolute_import
-import heapq
-import itertools
+import operator
import random
from apache_beam.transforms import core
@@ -148,7 +147,7 @@ class Top(object):
# pylint: disable=no-self-argument
@ptransform.ptransform_fn
- def Of(pcoll, n, compare, *args, **kwargs):
+ def Of(pcoll, n, compare=None, *args, **kwargs):
"""Obtain a list of the compare-most N elements in a PCollection.
This transform will retrieve the n greatest elements in the PCollection
@@ -157,7 +156,11 @@ class Top(object):
compare should be an implementation of "a < b" taking at least two arguments
(a and b). Additional arguments and side inputs specified in the apply call
- become additional arguments to the comparator.
+ become additional arguments to the comparator. Defaults to the natural
+ ordering of the elements.
+
+ The arguments 'key' and 'reverse' may instead be passed as keyword
+ arguments, and have the same meaning as for Python's sort functions.
Args:
pcoll: PCollection to process.
@@ -166,11 +169,13 @@ class Top(object):
*args: as described above.
**kwargs: as described above.
"""
+ key = kwargs.pop('key', None)
+ reverse = kwargs.pop('reverse', False)
return pcoll | core.CombineGlobally(
- TopCombineFn(n, compare), *args, **kwargs)
+ TopCombineFn(n, compare, key, reverse), *args, **kwargs)
@ptransform.ptransform_fn
- def PerKey(pcoll, n, compare, *args, **kwargs):
+ def PerKey(pcoll, n, compare=None, *args, **kwargs):
"""Identifies the compare-most N elements associated with each key.
This transform will produce a PCollection mapping unique keys in the input
@@ -180,7 +185,11 @@ class Top(object):
compare should be an implementation of "a < b" taking at least two arguments
(a and b). Additional arguments and side inputs specified in the apply call
- become additional arguments to the comparator.
+ become additional arguments to the comparator. Defaults to the natural
+ ordering of the elements.
+
+ The arguments 'key' and 'reverse' may instead be passed as keyword
+ arguments, and have the same meaning as for Python's sort functions.
Args:
pcoll: PCollection to process.
@@ -193,28 +202,30 @@ class Top(object):
TypeCheckError: If the output type of the input PCollection is not
compatible with KV[A, B].
"""
+ key = kwargs.pop('key', None)
+ reverse = kwargs.pop('reverse', False)
return pcoll | core.CombinePerKey(
- TopCombineFn(n, compare), *args, **kwargs)
+ TopCombineFn(n, compare, key, reverse), *args, **kwargs)
@ptransform.ptransform_fn
def Largest(pcoll, n):
"""Obtain a list of the greatest N elements in a PCollection."""
- return pcoll | Top.Of(n, lambda a, b: a < b)
+ return pcoll | Top.Of(n)
@ptransform.ptransform_fn
def Smallest(pcoll, n):
"""Obtain a list of the least N elements in a PCollection."""
- return pcoll | Top.Of(n, lambda a, b: b < a)
+ return pcoll | Top.Of(n, reverse=True)
@ptransform.ptransform_fn
def LargestPerKey(pcoll, n):
"""Identifies the N greatest elements associated with each key."""
- return pcoll | Top.PerKey(n, lambda a, b: a < b)
+ return pcoll | Top.PerKey(n)
@ptransform.ptransform_fn
- def SmallestPerKey(pcoll, n):
+ def SmallestPerKey(pcoll, n, reverse=True):
"""Identifies the N least elements associated with each key."""
- return pcoll | Top.PerKey(n, lambda a, b: b < a)
+ return pcoll | Top.PerKey(n, reverse=True)
@with_input_types(T)
@@ -222,17 +233,56 @@ class Top(object):
class TopCombineFn(core.CombineFn):
"""CombineFn doing the combining for all of the Top transforms.
- The comparator function supplied as an argument to the apply call invoking
- TopCombineFn should be an implementation of "a < b" taking at least two
- arguments (a and b). Additional arguments and side inputs specified in the
- apply call become additional arguments to the comparator.
+ This CombineFn uses a key or comparison operator to rank the elements.
+
+ Args:
+ compare: (optional) an implementation of "a < b" taking at least two
+ arguments (a and b). Additional arguments and side inputs specified
+ in the apply call become additional arguments to the comparator.
+ 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
"""
+ _MIN_BUFFER_OVERSIZE = 100
+ _MAX_BUFFER_OVERSIZE = 1000
+
# TODO(robertwb): Allow taking a key rather than a compare.
- def __init__(self, n, compare):
+ def __init__(self, n, compare=None, key=None, reverse=False):
self._n = n
- self._buffer_size = min(2 * n, n + 1000)
- self._compare = compare
+ self._buffer_size = max(
+ min(2 * n, n + TopCombineFn._MAX_BUFFER_OVERSIZE),
+ n + TopCombineFn._MIN_BUFFER_OVERSIZE)
+
+ if compare is operator.lt:
+ compare = None
+ elif compare is operator.gt:
+ compare = None
+ reverse = not reverse
+
+ if compare:
+ self._compare = (
+ (lambda a, b, *args, **kwargs: not compare(a, b, *args, **kwargs))
+ if reverse
+ else compare)
+ else:
+ self._compare = operator.gt if reverse else operator.lt
+
+ self._key_fn = key
+ self._reverse = reverse
+
+ def _sort_buffer(self, buffer, lt):
+ if lt in (operator.gt, operator.lt):
+ buffer.sort(key=self._key_fn, reverse=self._reverse)
+ else:
+ buffer.sort(cmp=lambda a, b: (not lt(a, b)) - (not lt(b, a)),
+ key=self._key_fn)
+
+ # The accumulator type is a tuple (threshold, buffer), where threshold
+ # is the smallest element [key] that could possibly be in the top n based
+ # on the elements observed so far, and buffer is a (periodically sorted)
+ # list of candidates of bounded size.
def create_accumulator(self, *args, **kwargs):
return None, []
@@ -244,22 +294,26 @@ class TopCombineFn(core.CombineFn):
lt = self._compare
threshold, buffer = accumulator
+ element_key = self._key_fn(element) if self._key_fn else element
+
if len(buffer) < self._n:
if not buffer:
- return element, [element]
+ return element_key, [element]
else:
buffer.append(element)
- if lt(element, threshold): # element < threshold
- return element, buffer
+ if lt(element_key, threshold): # element_key < threshold
+ return element_key, buffer
else:
return accumulator # with mutated buffer
- elif lt(threshold, element): # threshold < element
+ elif lt(threshold, element_key): # threshold < element_key
buffer.append(element)
if len(buffer) < self._buffer_size:
return accumulator
else:
- buffer.sort(cmp=lambda a, b: (not lt(a, b)) - (not lt(b, a)))
- return buffer[-self._n], buffer[-self._n:]
+ self._sort_buffer(buffer, lt)
+ min_element = buffer[-self._n]
+ threshold = self._key_fn(min_element) if self._key_fn else min_element
+ return threshold, buffer[-self._n:]
else:
return accumulator
@@ -287,14 +341,14 @@ class TopCombineFn(core.CombineFn):
lt = self._compare
_, buffer = accumulator
- buffer.sort(cmp=lambda a, b: (not lt(a, b)) - (not lt(b, a)))
- return buffer[:-self._n-1:-1]
+ self._sort_buffer(buffer, lt)
+ return buffer[:-self._n-1:-1] # tail, reversed
class Largest(TopCombineFn):
def __init__(self, n):
- super(Largest, self).__init__(n, lambda a, b: a < b)
+ super(Largest, self).__init__(n)
def default_label(self):
return 'Largest(%s)' % self._n
@@ -303,7 +357,7 @@ class Largest(TopCombineFn):
class Smallest(TopCombineFn):
def __init__(self, n):
- super(Smallest, self).__init__(n, lambda a, b: b < a)
+ super(Smallest, self).__init__(n, reverse=True)
def default_label(self):
return 'Smallest(%s)' % self._n
@@ -333,7 +387,7 @@ class SampleCombineFn(core.CombineFn):
# 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, lambda a, b: a < b)
+ self._top_combiner = TopCombineFn(n)
def create_accumulator(self):
return self._top_combiner.create_accumulator()
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/866a09d6/sdks/python/apache_beam/transforms/combiners_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/combiners_test.py b/sdks/python/apache_beam/transforms/combiners_test.py
index bfe168d..db3b85b 100644
--- a/sdks/python/apache_beam/transforms/combiners_test.py
+++ b/sdks/python/apache_beam/transforms/combiners_test.py
@@ -69,14 +69,21 @@ class CombineTest(unittest.TestCase):
pcoll = pipeline | 'start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6])
result_top = pcoll | 'top' >> combine.Top.Largest(5)
result_bot = pcoll | 'bot' >> combine.Top.Smallest(4)
- result_cmp = pcoll | combine.Top.Of(
+ result_cmp = pcoll | 'cmp' >> combine.Top.Of(
'cmp',
6,
lambda a, b, names: len(names[a]) < len(names[b]),
names) # Note parameter passed to comparator.
+ result_cmp_rev = pcoll | 'cmp_rev' >> combine.Top.Of(
+ 'cmp',
+ 3,
+ lambda a, b, names: len(names[a]) < len(names[b]),
+ names, # Note parameter passed to comparator.
+ reverse=True)
assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='assert:top')
assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='assert:bot')
assert_that(result_cmp, equal_to([[9, 6, 6, 5, 3, 2]]), label='assert:cmp')
+ assert_that(result_cmp_rev, equal_to([[0, 1, 1]]), label='assert:cmp_rev')
# Again for per-key combines.
pcoll = pipeline | 'start-perkye' >> Create(
@@ -95,6 +102,37 @@ class CombineTest(unittest.TestCase):
label='key:cmp')
pipeline.run()
+ def test_top_key(self):
+ self.assertEqual(
+ ['aa', 'bbb', 'c', 'dddd'] | combine.Top.Of(3, key=len),
+ [['dddd', 'bbb', 'aa']])
+ self.assertEqual(
+ ['aa', 'bbb', 'c', 'dddd'] | combine.Top.Of(3, key=len, reverse=True),
+ [['c', 'aa', 'bbb']])
+
+ # The largest elements compared by their length mod 5.
+ self.assertEqual(
+ ['aa', 'bbbb', 'c', 'ddddd', 'eee', 'ffffff'] | combine.Top.Of(
+ 3,
+ compare=lambda len_a, len_b, m: len_a % m > len_b % m,
+ key=len,
+ reverse=True,
+ m=5),
+ [['bbbb', 'eee', 'aa']])
+
+ def test_sharded_top_combine_fn(self):
+ def test_combine_fn(combine_fn, shards, expected):
+ accumulators = [
+ combine_fn.add_inputs(combine_fn.create_accumulator(), shard)
+ for shard in shards]
+ final_accumulator = combine_fn.merge_accumulators(accumulators)
+ self.assertEqual(combine_fn.extract_output(final_accumulator), expected)
+
+ test_combine_fn(combine.TopCombineFn(3), [range(10), range(10)], [9, 9, 8])
+ test_combine_fn(combine.TopCombineFn(5),
+ [range(1000), range(100), range(1001)],
+ [1000, 999, 999, 998, 998])
+
def test_top_shorthands(self):
pipeline = Pipeline('DirectPipelineRunner')
@@ -227,4 +265,6 @@ class CombineTest(unittest.TestCase):
if __name__ == '__main__':
+ # Sort more often for more rigorous testing on small data sets.
+ combiners.TopCombineFn._MIN_BUFFER_OVERSIZE = 1
unittest.main()