Executing a sql query inside a dataflow using python on PCollection

Question

I am trying to implement one sql query as a transform in dataflow. I loaded a table from bigquery as PCollection. I want to aggregate my data like below query.

SELECT name, user_id, place, SUM(amount) as some_amount , SUM(cost) as sum_cost FROM [project:test.day_0_test] GROUP BY 1,2,3 How I can implement it easily. I heard that Data flow with Java support running sql kind query on P Collection, but correctly python is not supporting. Can any one help me to solve this

Note:

I want to implement this query on a P Collection .. Not to read from bigquery directly

You need to use a GroupByKey. See here: cloud.google.com/dataflow/model/group-by-key#groupbykey — Graham Polley

Willian Fuks Willian Fuks · Accepted Answer · 2018-01-02T23:07:58

(I edited my answer as you commented on not wanting to run a SQL query directly in BigQuery)

I simulated a file input.csv that contains:

#input.csv
name1,1,place1,2.,1.5
name1,1,place1,3.,0.5
name1,1,place2,1.,1
name1,2,place3,2.,1.5
name2,2,place3,3.,0.5

This is the data that seems you are retrieving from BQ. Your SQL query could be implemented in Beam like:

def sum_l(l):                       
    s0, s1 = 0, 0                                         
    for i in range(len(l)):                                        
        s0 += l[i][0]                                                      
        s1 += l[i][1]                
    return [s0, s1] 

with beam.Pipeline(options=po) as p:
     (p | 'Read Input' >> beam.io.ReadFromText("input.csv")
        | 'Split Commas' >> beam.Map(lambda x: x.strip().split(','))
        | 'Prepare Keys' >> beam.Map(lambda x: (x[:-2], map(float, x[-2:])))
        | 'Group Each Key' >> beam.GroupByKey()
        | 'Make Summation' >> beam.Map(lambda x: [x[0], sum_l([e for e in x[1]])])
        | 'Write Results' >> beam.io.WriteToText('results.csv'))

Results are:

#results.csv-00000-of-00001
[[u'name1', u'1', u'place2'], [1.0, 1.0]]
[[u'name1', u'2', u'place3'], [2.0, 1.5]]
[[u'name1', u'1', u'place1'], [5.0, 2.0]]
[[u'name2', u'2', u'place3'], [3.0, 0.5]]

It's basically the straightforward MapReduce implementation of your query: a key is built for each row, they are grouped together and the final summation happens in the Map operation using the function sum_l.

I'm not sure why you want to run the query operations in Beam instead of BigQuery though. I recommend trying both approaches as probably it's not possible to be as efficient in Beam as you can be in BigQuery in this case.

Executing a sql query inside a dataflow using python on PCollection

1 Answers