How to use CoGroupByKey sink to BigQuery in Apache Beam using Dataflow

Question

I have 2 csv files and need to join them with join key "peakid". I have already transformed them like this
expeditions- 2010s

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

peak

('ACHN', {'pkname': ['Aichyn'], 'heightm': ['6055']})
('AGLE', {'pkname': ['Agole East'], 'heightm': ['6675']})
('AMAD', {'pkname': ['Ama Dablam'], 'heightm': ['6814']})
('AMOT', {'pkname': ['Amotsang'], 'heightm': ['6393']})
('AMPG', {'pkname': ['Amphu Gyabjen'], 'heightm': ['5630']})
...

And when I used CoGroupByKey, the result look like

('ACHN', ([{'bcdate': [''], 'smtdate': ['9/25/15']}, {'bcdate': [''], 'smtdate': ['9/3/15']}, {'bcdate': [''], 'smtdate': ['']}], [{'pkname': ['Aichyn'], 'heightm': ['6055']}]))
('AGLE', ([], [{'pkname': ['Agole East'], 'heightm': ['6675']}]))
('AMAD', ([{'bcdate': ['4/5/10'], 'smtdate': ['4/21/10']}, {'bcdate': ['4/20/10'], 'smtdate': ['4/27/10']}, {'bcdate': ['4/4/10'], 'smtdate': ['4/10/10']},...

After that I write to BigQuery, it get an error BigQuery job beam_bq_job_LOAD_AUTOMATIC_JOB_NAME_LOAD_STEP_460_215864ba592a2e01f0c4e2157cc60c47_bc7734af2ebb4a53a0e268bbe6c40824 failed. Error Result: [while running 'Write To BigQuery/BigQueryBatchFileLoads/WaitForDestinationLoadJobs']

Below is code:

input_p1 = (
        p
         | 'Read From GCS input1' >> beam.Create([known_args.input1])
         | 'Parse csv file p1' >> beam.FlatMap(read_csv_file)
         | 'Tuple p1' >> beam.Map(lambda e: (e["peakid"], {'bcdate': [e["bcdate"]], 'smtdate':[e["smtdate"]]}))
    )

 input_p2 = (
        p
         | 'Read From GCS input2' >> beam.Create([known_args.input])
         | 'Parse csv file p2' >> beam.FlatMap(read_csv_file)
         | 'Tuple p2' >> beam.Map(lambda e: (e["peakid"], {'pkname': [e["pkname"]], 'heightm':[e["heightm"]]}))
    )

output = (
        (input_p1, input_p2)
        | 'Join' >> beam.CoGroupByKey()
        # | beam.Map(print)
        | 'Write To BigQuery' >> beam.io.gcp.bigquery.WriteToBigQuery(
           table='project-name.dataset.expeditions',
           schema='peakid:STRING,bcdate:DATE,pkname:STRING,heightm:INTEGER',
           method='FILE_LOADS',
           custom_gcs_temp_location='gs://dtnhu_test_dataflow_v1/input/temp',
           create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
           write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE)    
    )

Mazlum Tosun · Answer

I changed my answer after seeing the expected structure and more understanding the need :

I created a test with simulating your CSV files :

def test_pipeline(self):
        with TestPipeline() as p:
            csv1 = [
                ('TKRG', {'peakid': 'TKRG', 'bcdate': '3/6/10', 'smtdate': '3/12/10'}),
                ('AMPG', {'peakid': 'AMPG', 'bcdate': '4/5/10', 'smtdate': ''}),
                ('AMAD', {'peakid': 'AMAD', 'bcdate': '4/5/10', 'smtdate': '4/21/10'}),
                ('AMAD', {'peakid': 'AMAD', 'bcdate': '4/20/10', 'smtdate': '4/27/10'}),
                ('AMAD', {'peakid': 'AMAD', 'bcdate': '4/4/10', 'smtdate': '4/10/10'})
            ]

            csv2 = [
                ('ACHN', {'peakid': 'ACHN', 'pkname': 'Aichyn', 'heightm': '6055'}),
                ('AGLE', {'peakid': 'AGLE', 'pkname': 'Agole East', 'heightm': '6675'}),
                ('AMAD', {'peakid': 'AMAD', 'pkname': 'Ama Dablam', 'heightm': '6814'}),
                ('AMOT', {'peakid': 'AMOT', 'pkname': 'Amotsang', 'heightm': '6393'}),
                ('AMPG', {'peakid': 'AMPG', 'pkname': 'Amphu Gyabjen', 'heightm': '5630'})
            ]

            input1 = p | "CSV1" >> beam.Create(csv1)
            input2 = p | "CSV2" >> beam.Create(csv2)

            output = (
                    (input1, input2)
                    | 'Join' >> beam.CoGroupByKey()
                    | 'Final Map' >> beam.Map(lambda el: self.to_final_dict(el[1]))
                    | beam.Map(print)
            )

    def to_final_dict(self, list_tuple_of_tuple):
        result = {}
        for list_tuple in list_tuple_of_tuple:
            for el in list_tuple:
                result.update(el)

        return result

I have the following result :

{'peakid': 'AGLE', 'pkname': 'Agole East', 'heightm': '6675'}
{'peakid': 'AMAD', 'bcdate': '4/4/10', 'smtdate': '4/10/10', 'pkname': 'Ama Dablam', 'heightm': '6814'}
{'peakid': 'AMOT', 'pkname': 'Amotsang', 'heightm': '6393'}
{'peakid': 'AMPG', 'bcdate': '4/5/10', 'smtdate': '', 'pkname': 'Amphu Gyabjen', 'heightm': '5630'}
{'peakid': 'TKRG', 'bcdate': '3/6/10', 'smtdate': '3/12/10'}

In your code, you have to change a little things to have the same inputs as in my example, in the flatMap operation :

input_p1 = (
        p
         | 'Read From GCS input1' >> beam.Create([known_args.input1])
         | 'Parse csv file p1' >> beam.FlatMap(read_csv_file)
         | 'Tuple p1' >> beam.Map(lambda e: (e["peakid"], {'peakid': e["peakid"], 'bcdate': e["bcdate"], 'smtdate':e["smtdate"]}))
    )

input_p2 = (
        p
         | 'Read From GCS input2' >> beam.Create([known_args.input])
         | 'Parse csv file p2' >> beam.FlatMap(read_csv_file)
         | 'Tuple p2' >> beam.Map(lambda e: (e["peakid"], {'peakid': e["peakid"], 'pkname': e["pkname"], 'heightm':e["heightm"]}))
    )

I don't care about the naming of variables and function name in the to_final_dict method, please don't hesitate to improve them.

How to use CoGroupByKey sink to BigQuery in Apache Beam using Dataflow

1 Answers