Split each row of a pcollection into multiple pcollections?

user_15316893 0
Viewed 908

After I do some processing and a group by key, I have a dataset like the one below. I now need to do some processing on each row of the data to get the output below. I have tried flatmap it is really slow because the length of the "value" list can be arbitrarily long. I figured I can split each row into separate pcollections, process in parallel and then flatten them together. How can I split each row into a different pcollection? If that isn't feasible, is there another way I can speed up computation?

Input

key, value
1    [A, B, B, B]
2    [A, B, B, B]
3    [A, B, B, B]
4    [A, B, B, B]
5    [A, B, B, B]

Output:

key, value
1    (A, 0)
1    (B, 1)
1    (B, 2)
1    (B, 3)
2    (A, 0)
2    (B, 1)
2    (B, 2)
2    (B, 3)
...
apache-beam python
1 Answers

When using the Apache Beam model, it is a common misconception that the parallelization scheme is defined by the PCollection (understandable since this is short for Parallel Collection). In actuality, the parallelization is defined per key in each PCollection[1]. In other words, the Beam model processes keys in parallel but values in a single key sequentially.

The problem you are coming up against is commonly referred to as a hot key. This happens when too many values are paired to a single key, thus limiting parallelism.

To manipulate the data to the expected output you will have to edit your existing pipeline to emit the values in such a way that not all elements go to a single key. This is a little tough because it looks like in your example you wish to output the index with the element. If this is the case, then no matter how you cut it, you will have to merge all the values somewhere to a key in memory to get the correct index.

If you don't care about getting the specific index like you have in the above example then take a look at the following code. This code assigns each element to a random partition within each key. This helps to break up the number of elements per key into something manageable.

data = [
  (k, c) for k in range(1, 6) for c in ('A', 'B', 'B', 'B')
]

p = beam.Pipeline()
elems = p | beam.Create(data)

num_buckets = 4

class Preprocess(beam.DoFn):
  def process(self, el):
    key = str(el[0])
    partition = random.randint(0, num_buckets)
    yield (key, partition), el

class Postprocess(beam.DoFn):
  def process(self, el):
    (key, partition), values = el
    index = 0
    for el in values:
      yield key, (el[1], partition, index)
      index += 1
    
out = (elems | beam.ParDo(Preprocess())
             | beam.GroupByKey()
             | beam.ParDo(Postprocess()))

Input

key,value
1   A
1   B
1   B
1   B
2   A
2   B
2   B
2   B
3   A
3   B
...

Potential Output

key,(value,partition,index)
1   ('A', 1, 0)
1   ('B', 1, 1)
1   ('B', 2, 0)
1   ('B', 3, 0)
2   ('A', 3, 0)
2   ('B', 3, 1)
2   ('B', 1, 0)
2   ('B', 1, 1)
3   ('A', 3, 0)
3   ('B', 2, 0)
...

[1] When using streaming it is defined per-key-per-window

user_37053630
Related