Select some columns from PCollection (Apache Beam, Python)

Question

I have the following PCollection:

And I want to select only 2 columns from that PCollection. I tried to do:

def cut_data(data):
    return data[["WebSpeedRef", "WebSpeedAct"]]

data_min = data_json | 'min' >> beam.Map(cut_data)

but got an error. What is the simplest way to accomplish this.

Answer 1

You could do this:

disired_columns = (
    dataset
    | beam.Map(lambda x: [x["column1"], x["column2"]])
)