Convert dataframe to tuple to support CoGroupByKey

Viewed 79

I have a DataFrame with 3 columns. I want to convert this DataFrame to a python tuple cause CoGroupByKey only accept this type. I want the elements of first column be keys and the elements of other columns in same row be values. And can sink it to BigQuery cause BigQuery only accept json format schema

peakid bcdate smtdate
TKRG 3/6/10 3/12/10
AMPG 4/5/10
AMAD 9/5/10 4/21/10

And I expect result like this:

[('TKRG', 
    {'bcdate': ['3/6/10'],
     'smtdate': ['3/12/10']}),
 ('AMPG', 
    {'bcdate': ['4/5/10'],
     'smtdate': []}),
 ('AMAD', 
    {'bcdate': ['9/5/10'],
     'smtdate': ['4/21/10']})
]

Reference: https://beam.apache.org/documentation/programming-guide/#cogroupbykey

1 Answers

By Dataframe, you mean a pandas dataframe or a beam dataframe?

In case of beam dataframe, you can use to_pcollection to transform the dataframe to a pcollection and then use the following Map function:

  beam.Map(lambda e: (e['peakid'], {'bcdate': [e['bcdate']], 'smtdate':[e['smtdate']]}))

In case of your input is a Pandas dataframe, you can transform its records into a dictionary, and pass it to beam.Create(your_dict) and then process it with the same Map function

UPDATE:

You can read directly the CSV using beam.io.ReadFromText. This is an approach that you can use:

pipeline | 'Read input file' >> beam.io.ReadFromText(self.csv_path, skip_header_lines=1)
         | 'Parse file' >> beam.Map(parse_file))
  • parse_file is a function that you should write to parse the CSV content. each element would look like this: value1,value2,value3,etc, and you can split them and append them to you dictionary
  • skip_header_lines will help you to skip the header which in your case contains the column name

Else, you can read directly your CSV into a beam dataframe using read_csv() and then transform it into a pcollection using to_pcollection

In case of multiple columns, the lambda function mentioned in the solution can be replaced with a function. For example:

def process_my_elements(element):
    key = element['key']
    values = {}
    for k,v in element.items():
        # pick the values you want to append
    return (key, values)

beam.Map(process_my_elements)
Related