How to add commas to Python Dictionaries output from beam's ReadFromBigQuery in Python? - google-cloud-dataflow

Question

I am trying to run a Dataflow pipeline in Python (3.9) via a FlexTemplate that runs a query in BigQuery (non-nested/non-repeated records) and writes the data to some database with beam_nuggets. My test database is Postgres. The data from BigQuery yields Python dictionaries as expected, but when I try to upload this to the database the pipeline fails.
Test data written to GCS from BigQuery.
{'order_id': 'CM-2011-110', 'order_date': '4/10/2011', 'ship_date': '10/10/2011', 'ship_mode': 'Standard Class', 'customer_name': 'Alejandro Grove', 'segment': 'Consumer', 'state': 'Est', 'country': 'Cameroon', 'market': 'Africa', 'region': 'Africa', 'product_id': 'OFF-CAR-10002031', 'category': 'Office Supplies', 'sub_category': 'Binders', 'product_name': 'Cardinal 3-Hole Punch, Durable', 'sales': 30, 'quantity': 1, 'discount': 0.0, 'profit': 13.92, 'shipping_cost': 2.57, 'order_priority': 'Medium', 'year': 2011}
{'order_id': 'CM-2011-110', 'order_date': '4/10/2011', 'ship_date': '10/10/2011', 'ship_mode': 'Standard Class', 'customer_name': 'Alejandro Grove', 'segment': 'Consumer', 'state': 'Est', 'country': 'Cameroon', 'market': 'Africa', 'region': 'Africa', 'product_id': 'TEC-CAN-10002879', 'category': 'Technology', 'sub_category': 'Copiers', 'product_name': 'Canon Copy Machine, High-Speed', 'sales': 521, 'quantity': 2, 'discount': 0.0, 'profit': 93.72, 'shipping_cost': 30.83, 'order_priority': 'Medium', 'year': 2011}
Current Code:
def run(save_main_session=True):
beam_options = PipelineOptions()
args = beam_options.view_as(MyOptions)
with beam.Pipeline(options=beam_options) as p:
db_source = db_reader(args.bqprojectid, args.dataset, args.table, args.limit)
db = db_writer(args.destinationip, args.port, args.destinationusername, args.destinationpassword, args.destinationtable, args.database_name)
result = (
p | beam.io.ReadFromBigQuery(use_standard_sql=True, query=db_source.sql_query()
)
# I would like to remove this step, but it helps with debugging.
|'Write to GCS' >> WriteToText('SOMEURI')
|'Write to Database' >> relational_db.Write(
source_config = (db.sink_config()),
table_config = (db.table_config())
))
I went through this answer: Write BigQuery results to GCS in CSV format using Apache Beam but was wondering if there was an easier way to transform the output from BigQuery to mirror this structure:
{'name': 'Jan', 'num': 1},
{'name': 'Feb', 'num': 2},
I was considering regex, but I do not thing that would be a best practice.
Please let me know if there is anything I can clarify.
Have a great day!
Error Message
File "/usr/local/lib/python3.9/site-packages/beam_nuggets/io/relational_db.py", line 181, in process
assert isinstance(element, dict)
I checked the credentials and realized that the structure of the data is the issue by creating a static p-collection via the beam nuggets docs. This uploaded successfully (months).
months = p | "Reading month records" >> beam.Create([
{'name': 'Jan', 'num': 1},
{'name': 'Feb', 'num': 2},
])
Resolution via answer provided below.
def map_to_beam_nuggets_data(element):
return {
'order_id': element['order_id'],
'order_date': element['order_date'],
'ship_date': element['ship_date'],
'ship_mode': element['ship_mode'],
'customer_name': element['customer_name'],
'segment': element['segment'],
'state': element['state'],
'country': element['country'],
'market': element['market'],
'region': element['region'],
'product_id': element['product_id'],
'category': element['category'],
'sub_category': element['sub_category'],
'product_name': element['product_name'],
'sales': element['sales'],
'quantity': element['quantity'],
'discount': element['discount'],
'profit': element['profit'],
'shipping_cost': element['shipping_cost'],
'order_priority': element['order_priority'],
}
# this is the controller function
def run(save_main_session=True):
beam_options = PipelineOptions()
args = beam_options.view_as(MyOptions)
with beam.Pipeline(options=beam_options) as p:
db_source = db_reader(args.bqprojectid, args.dataset, args.table, args.limit)
db = db_writer(args.destinationip, args.port, args.destinationusername, args.destinationpassword, args.destinationtable, args.database_name)
# Result from BigQuery.
result = (p | beam.io.ReadFromBigQuery(use_standard_sql=True, query=db_source.sql_query()))
# Sink result to GCS.
result | 'Write to GCS' >> WriteToText('SOMEURI')
# Map to result to Beam Nuggets data and Sink result to the database.
(result
| 'Map to beam nuggets data' >> beam.Map(map_to_beam_nuggets_data)
| 'Write to Database' >> relational_db.Write(
source_config=(db.sink_config()),
table_config=(db.table_config())
))
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
run()

Answer 1

I hope it can help.
In this case, you can use multi sink output :
Result PCollection from BigQuery
Sink the result to GCS
Transform the result to the expected output for beam nuggets data and sink it to the database
def run_pipeline(save_main_session=True):
beam_options = PipelineOptions()
args = beam_options.view_as(MyOptions)
with Pipeline(options=beam_options) as p:
db_source = db_reader(args.bqprojectid, args.dataset, args.table, args.limit)
db = db_writer(args.destinationip, args.port, args.destinationusername, args.destinationpassword,
args.destinationtable, args.database_name)
# Result from BigQuery.
result = (p | beam.io.ReadFromBigQuery(use_standard_sql=True, query=db_source.sql_query()))
# Sink result to GCS.
result | 'Write to GCS' >> WriteToText('SOMEURI')
# Map to result to Beam Nuggets data and Sink result to the database.
(result
| 'Map to beam nuggets data' >> beam.Map(map_to_beam_nuggets_data)
| 'Write to Database' >> relational_db.Write(
source_config=(db.sink_config()),
table_config=(db.table_config())
)
)
def map_to_beam_nuggets_data(element: Dict) -> Dict:
return {
'name': element['name'],
'num': element['num'],
}
To simulate your input data from BigQuery, I created an unit test with a mock and pytest. The ReadFromBigQuery source should returns a PCollection of Dict :
import apache_beam as beam
import pytest
from apache_beam.testing.test_pipeline import TestPipeline
def test_pipeline(self):
mock = [
{
'name': 'toto',
'num': 'num',
'other_field': 'other'
}
]
with TestPipeline() as p:
result = (
p
| beam.Create(mock)
| 'Map to Beam nuggets data' >> beam.Map(self.map_to_beam_nuggets_data)
)
result | "Print outputs" >> beam.Map(print)
def map_to_beam_nuggets_data(self, element: Dict) -> Dict:
return {
'name': element['name'],
'num': element['num'],
}
The result is :
{'name': 'toto', 'num': 'num'}