Branching Apache Beam pipelines with loops

user_7807920 0
Viewed 2108

I am trying to perform a de-noramlization operation, where I need to reorganize a table with the following logic:

| itemid | class | value |
+--------+-------+-------+
| 1      | A     | 0.2   |       | itemid | value A | value B | value C |
| 1      | B     | 10.3  |  ==>  +--------+---------+---------+---------+
| 2      | A     | 3.0   |  ==>  | 1      |   0.2   |  10.3   |         |
| 2      | B     | 0.2   |  ==>  | 2      |   3.0   |   0.2   |         |
| 3      | A     | 0.0   |       | 3      |   0.0   |   1.2   |   5.4   | 
| 3      | B     | 1.2   |  
| 3      | C     | 5.4   |

My approach has been to execute a for loop in order to filter by class, given that I know the list of classes a-priori and then to join the resulting pcollections.

The high-level code:

CLASSES = ["A", "B", "C"]

tables = [  
    (
        data
        | "Filter by Language" >> beam.Filter(lambda elem: elem["class"]==c)
        | "Add id as key" >> beam.Map(lambda elem: (elem["itemid"], elem))
    )
    for cin CLASSES
]

and the join:

_ = ( 
    tables
    | "Flatten" >> beam.Flatten()
    | "Join Collections" >> beam.GroupByKey()
    | "Remove key" >> beam.MapTuple(lambda _, val: val)
    | "Merge dicts" >> beam.ParDo(mergeDicts())
    | "Write to GCS" >> beam.io.WriteToText(output_file)
)

with (edit following Peter Kim's suggestion):

class mergeDicts(beam.DoFn):
    process(self, elements):
        result = {}
        for dictionary in elements:
            if len(dictionary)>0:
                result["itemid"] = dictionary["itemid"]
                result["value {}".format(dictionary["class"])] = dictionary["value"]
        yield result

My issue here is that when the pipeline is executed within the Apache Beam compute engine, I obtain identical pcollections filtered by the last element of the list, which in this case is C.

[ADDED] It looks like the Apache Beam engine takes the iterating variable in its final state, which means the last element of the iterated list, for all the invoked branches.

I am obviously following the wrong approach, but which should be the best way to perform this operation?

apache-beam google-cloud-dataflow python
3 Answers

What you're running into is the surprising gotcha with respect to closures, loops, and Python scoping. You can get around this by assigning the variable rather than pulling it from the closure. E.g.

tables = [  
    (
        data
        # Pass it as a side input to Filter.
        | "Filter by Language" >> beam.Filter(lambda elem, cls: elem["class"], c)
        | "Add id as key" >> beam.Map(lambda elem: (elem["itemid"], elem))
    )
    for c in CLASSES
]

or

tables = [  
    (
        data
        # Explicitly capture it as a default value in the lambda.
        | "Filter by Language" >> beam.Filter(lambda elem, cls=c: elem["class"])
        | "Add id as key" >> beam.Map(lambda elem: (elem["itemid"], elem))
    )
    for c in CLASSES
]

Partition works well here too, both in avoiding this pitfall and in expressing your intent.

user_5823330

Based on the resulting table you showed, I assume you want your output to look like this:

{'itemid': '1', 'value B': 10.3, 'value A': 0.2}
{'itemid': '2', 'value B': 0.2, 'value A': 3.0}
{'itemid': '3', 'value B': 1.2, 'value A': 0.0, 'value C': 5.4}

Your mergeDicts was overwriting values because dictionaries can only hold one value per key. Update mergeDicts to something like this to specify the keys:

class mergeDicts(beam.DoFn):
    process(self, elements):
        result = {}
        for dictionary in elements:
            if len(dictionary)>0:
                result["itemid"] = dictionary["itemid"]
                result["value {}".format(dictionary["class"])] = dictionary["value"]
        yield result
user_47562790

I am posting here a solution that I found by myself, but I am not checking it as the correct answer, because I would like to understand better the execution logic of the Beam engine.

Instead of filtering the table by the items in a loop, in order to obtain separate pcollections on a condition, I used the beam.Partition class. By applying the code sample in the documentation directly, I partitioned the pcollection into multiple tables, ready to be joined.

This way the issue is avoided, however it is not clear to me why the for loop doesn't work as I expected.

user_78079200
Related