Modify MinimalWordCount example to read from BigQuery

Question

I am trying to modify Apache Beam's MinimalWordCount python example to read from a BigQuery table. I have made the following modifications, and I appear to have the query working but the example.

Original Example Here:

 with beam.Pipeline(options=pipeline_options) as p:

    # Read the text file[pattern] into a PCollection.
    lines = p | ReadFromText(known_args.input)

    # Count the occurrences of each word.
    counts = (
        lines
        | 'Split' >> (beam.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x))
                      .with_output_types(unicode))
        | 'PairWithOne' >> beam.Map(lambda x: (x, 1))
        | 'GroupAndSum' >> beam.CombinePerKey(sum))

    # Format the counts into a PCollection of strings.
    output = counts | 'Format' >> beam.Map(lambda (w, c): '%s: %s' % (w, c))

    # Write the output using a "Write" transform that has side effects.
    # pylint: disable=expression-not-assigned
    output | WriteToText(known_args.output)

Rather than ReadFromText I am trying to adjust this to read from a column in a BigQuery table. To do this I have replaced lines = p | ReadFromText(known_args.input) with the following code:

query = 'SELECT text_column FROM `bigquery.table.goes.here` '
lines = p | 'ReadFromBigQuery' >> beam.io.Read(beam.io.BigQuerySource(query=query, use_standard_sql=True))

When I re-run the pipeline, I get the error: "WARNING:root:A task failed with exception. expected string or buffer [while running 'Split']"

I recognize that the 'Split' operation is expecting a string and it is clearly not getting a string. How can I modify 'ReadFromBigQuery' so that it is passing a string/buffer? Do I need to provide a table schema or something to convert the results of 'ReadFromBigQuery' into a buffer of strings?

user_1955346 · Answer

This is because BigQuerySource returns PCollection of dictionaries (dict), where every key in the dictionary represents a column. For your case the simplest thing to do will be just applying beam.Map after beam.io.Read(beam.io.BigQuerySource(query=query, use_standard_sql=True) like this:

lines = (p 
|"ReadFromBigQuery" >> beam.io.Read(beam.io.BigQuerySource(query=query, use_standard_sql=True)) 
| "Extract text column" >>  beam.Map(lambda row: row.get("text_column"))
         )

If you encounter problem with column name, try change it to u"text_column".

Alternatively you can modify your Split transform to extract the value of column there:

'Split' >> (beam.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x.get("text_column")))
                      .with_output_types(unicode))

Modify MinimalWordCount example to read from BigQuery

1 Answers