flatten list of lists in (key, List[List]) PCollection, retaining the key

user_1225744 0
Viewed 365

I have a PCollection of (key, value) pairs where the value is itself a list

{'apple': ['1','2'],
 'watermelon': ['4','5']
 'apple': ['9','10']
 }

I would like to obtain another collection with the same key, but where the list is a GroupByKey on the original collection, but where the resulting value is a flattened.

{'apple': ['1','2','9','10'],
 'watermelon': ['4','5']
 }

instead of

{'apple': [['1','2], ['9','10']],
 'watermelon': [['4','5']]
 }

I have tried a few things but it's not obvious to me how to go about it.

I guess one can always write a ParDo function that flattens the list on a per-element basis but it feels there must be an easier solution.

apache-beam
2 Answers

The logical operation you are looking for is "combine per key" where the combiner is list concatenation (known as + in Python).

Depending on the characteristics of your data, you may wish to implement this operation using a new CombineFn that mutates a list as its accumulator. This is nearly equivalent to performing a GroupByKey and then flattening the lists afterwards in a ParDo.

The difference in execution is that a Combine operation can be performed before shuffling your data, because the Beam runner knows that it is an associative and commutative operation. In this case, it will not reduce the amount of shuffled data so it is not so important.

user_9493270

Are you looking for something like this

import apache_beam as beam


class Combiner(beam.CombineFn):

    def create_accumulator(self, *args, **kwargs):
        return {}
    
    def add_input(self, acc, element):
        key = element[0]
        value = element[1]
        if key in acc:
            value.extend(acc[key])
        acc[key]=value
        return acc
    
    def merge_accumulators(self, accumulators):
        return accumulators

    def extract_output(self, accumulator):
        return accumulator

with beam.Pipeline() as p:
    pipe =(
        p
          |'create'>> beam.Create(
                
                [
                ('apple', ['1','2']),
                ('watermelon', ['4','5']),
                ('apple', ['9','10']),
                ('watermelon', ['49','50'])
                ]
                
 )
          |'combiner'>>beam.CombineGlobally(Combiner())
          |'print'>>beam.Map(print)
             
    )

#o/p - [{'apple': ['9', '10', '1', '2'], 'watermelon': ['49', '50', '4', '5']}]
user_60622920
Related