This is an automated email from the ASF dual-hosted git repository. udim pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push: new a4abc91 [BEAM-12434] Implement side_input for num_shards in iobase (#14916) a4abc91 is described below commit a4abc91436eadf78abda304664a3e0f64019ea61 Author: hoshimura <johan.ster...@gmail.com> AuthorDate: Fri Jun 4 19:19:20 2021 +0200 [BEAM-12434] Implement side_input for num_shards in iobase (#14916) Co-authored-by: Udi Meiri <u...@users.noreply.github.com> Co-authored-by: Johan Sternby <joha...@axis.com> --- sdks/python/apache_beam/io/iobase.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/sdks/python/apache_beam/io/iobase.py b/sdks/python/apache_beam/io/iobase.py index 71d8037..5d8e5df 100644 --- a/sdks/python/apache_beam/io/iobase.py +++ b/sdks/python/apache_beam/io/iobase.py @@ -1125,7 +1125,7 @@ class WriteImpl(ptransform.PTransform): if min_shards == 1: keyed_pcoll = pcoll | core.Map(lambda x: (None, x)) else: - keyed_pcoll = pcoll | core.ParDo(_RoundRobinKeyFn(min_shards)) + keyed_pcoll = pcoll | core.ParDo(_RoundRobinKeyFn(), count=min_shards) write_result_coll = ( keyed_pcoll | core.WindowInto(window.GlobalWindows()) @@ -1226,17 +1226,13 @@ def _finalize_write( class _RoundRobinKeyFn(core.DoFn): - def __init__(self, count): - # type: (int) -> None - self.count = count - def start_bundle(self): - self.counter = random.randint(0, self.count - 1) + self.counter = None - def process(self, element): - self.counter += 1 - if self.counter >= self.count: - self.counter -= self.count + def process(self, element, count): + if self.counter is None: + self.counter = random.randrange(0, count) + self.counter = (1 + self.counter) % count yield self.counter, element