Extract constant value from a PCollection

Question

I have a connector which feeds me millions of entries. Some of those entries have a field date whose value is the same for all of them, but others don't. I need to get this value and set it to every entry on my pipeline.

If it was not a distributed pipepline, I would do something like this:

def set_date(entries):
    for entry in entries:
        mydate = entry.get('date')
        if mydate:
            break

    for entry in entries:
        entry['date'] = mydate

However, I have to treat each entry individually and not as a set, so I can not write the logic like this.

Moreover, we can not assure that every set of PCollection has at least one entry with the desired date field. It is unlikely, but it could still happen that all the entries on a PCollection does not have the date field, so the logic I have written above is not valid. I need to write some kind of PTransform which saves the date value to a class attribute, and on some type of "finally" clause, set the value on each entry.

What is the correct way of implementing this?

Answer 1

Great question! For Beam in a Batch pipeline you can do this with a side input. You would do something like this:

def get_date(entry):
  if 'date' in entry:
    yield entry['date']

def put_date(entry, date):
  entry['date'] = date[0]
  return entry

with beam.Pipeline(...) as p:
  entries = p | ReadEntries(....)

  # First obtain the date from the entires.
  # We do this by outputting the date from the entries with a date,
  # and then we 'sample' one (this is a simple combiner that lets us
  # pick one of the dates - since they're all the same).
  date_pc = beam.pvalue.AsList(
      entries
      | beam.FlatMap(get_date)
      | beam.Combine(Sample.FixedSizeGlobally(1))

  # We then use the date_pc side input to join it into entries.
  dated_entries = entries | beam.Map(put_date, date_pc)

Let me know if this helps.

If none of the entries has a date attribute, then we may need to do something else. Happy to iterate!

Answer 2

As I know the date value I want to set to my entries before running my pipeline, I have added it as a ValueProvier and fed it to my SetDate PTransform as a side input. If it is not your case, you can try the answer marked as correct instead.

This is my final implementation:

import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions

class CustomOptions(PipelineOptions):
    @classmethod
    def _add_argparse_args(cls, parser):
        parser.add_value_provider_argument(
            '--query',
            help='Query to retrieve from BigQuery acting as data source.')
        parser.add_value_provider_argument(
            '--mydate',
            help='Date with %Y-%m-%d format of the entries we are processing.')

options = PipelineOptions()
args = options.view_as(CustomOptions)

def run():
    with beam.Pipeline(options=options) as p:
        (
            p
            | "Read from BigQuery" >> beam.io.ReadFromBigQuery(
                                          query=args.query,
                                          use_standard_sql=True,
                                          flatten_results=False)
            | "Set date" >> beam.ParDo(SetDate(args.mydate))
        )

class SetDate(beam.DoFn):
    def __init__(self, mydate):
        self.mydate = mydate

    def process(self, entry):
        if 'date' not in entry:
            entry['date'] = self.mydate.get()
        yield entry