Apache-Beam + Python: Writing JSON (or dictionaries) strings to output file

Question

I am trying to use a Beam pipeline in order to apply the SequenceMatcher function to a ton of words. I (hopefully) have figured everything out except the WriteToText part.

I have defined a custom ParDo (hereby called ProcessDataDoFn) that takes the main_input and the side_input, process them and output dictionaries like this one

{u'key': (u'string', float)}

My pipeline is quite simple

class ProcessDataDoFn(beam.DoFn):
    def process(self, element, side_input):

    ... Series of operations ...

    return output_dictionary

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

    # Main input
    main_input = p | 'ReadMainInput' >> beam.io.Read(
        beam.io.BigQuerySource(
            query=CUSTOM_SQL,
            use_standard_sql=True
        ))

    # Side input
    side_input = p | 'ReadSideInput' >> beam.io.Read(
        beam.io.BigQuerySource(
            project=PROJECT_ID,
            dataset=DATASET,
            table=TABLE
        ))

    output = (
        main_input
        | 'ProcessData' >> beam.ParDo(
            ProcessDataDoFn(),
            side_input=beam.pvalue.AsList(side_input))
        | 'WriteOutput' >> beam.io.WriteToText(GCS_BUCKET)
    )

Now the problem is that if I leave the pipeline like this it only output the key of the output_dictionary. If I change the return of ProcessDataDoFn to json.dumps(ouput_dictionary), The Json is written correctly but like this

{
'
k
e
y
'

:

[
'
s
t
r
i
n
g
'

,

f
l
o
a
t
]

How can I correctly output the results?

Answer 1

I actually partially solved the issue.

The ParDoFn that I wrote either return a dictionary or a JSON formatted string. In both cases, the problem arises when Beam tries to do something with said input. Beam seems to iterate over a given PCollection if said PCollection is a dictionary, it only gets its key, if said PCollection is a string, it iterates over all the character (that's why the JSON output is so strange). I find the solution to be rather simple: encapsulate either the dictionary or the string in a list. The JSON formatting part can either be done at the ParDoFn level or via a Transform like the one you showed.

Answer 2

It's unusual that your output looks like that. json.dumps should print json in a single line, and it should go out to files line-by-line.

Perhaps to have cleaner code, you can add an extra map operation that does your formatting however way you need. Something like so:

output = (
  main_input
  | 'ProcessData' >> beam.ParDo(
        ProcessDataDoFn(),
        side_input=beam.pvalue.AsList(side_input))
  | 'FormatOutput' >> beam.Map(json.dumps)
  | 'WriteOutput' >> beam.io.WriteToText(GCS_BUCKET)
)