Apache Beam Tutorial

We assume that the data is stored in a text file, which we will use as the input. Since we are going to process a static file, it will be an example of batch processing. Still, the same rules apply for the stream processing as well, where instead of reading from the static file, we stream the data from a stream source such as Apache Kafka, Google Pub/Sub, or some database.

The first question the hospital needs to answer is how many times a patient visits the hospital, so we need to aggregate the data. The following code shows a complete example of counting visits for the cardio department.

import apache_beam as beam

p1 = beam.Pipeline()
visit_count = (
 
  p1
  |beam.io.ReadFromText('data.txt')

  |beam.Map(lambda record: record.split(','))
  |beam.Filter(lambda record: record[3] == 'cardio')
  |beam.Map(lambda record: (record[1], 1))
  |beam.CombinePerKey(sum)
 
  |beam.io.WriteToText('out_data.txt')
)

p1.run()

Listing 1. Example: Patient visit count for the cardio department with core transforms

Initially, we import the beam namespace and create an alias to it. Then we create the pipeline 'p1' and a PCollection visit_count, which will store the results of all the transformations.

Pipeline p1 applies all the transforms in sequential order. It starts with reading the data using read transform into a collection and then splitting the row into columns, applying the filter for the cardio department, and creating a collection of records for that department only. Next, we make a pair of each patient, which then can be combined and summed up. Finally, a group by operation on patient_id with a sum operation will give us the required results.

Transforms

Transforms are operations that are applied to individual elements of the pipeline and provided as a function. There are quite a few transforms provided by the Apache Beam. Some predefined transforms are Map, FlatMap, and Filter, but one can create custom transforms from the “PTransform” class.

I/O Transform

Apache Beam provides built-in support for various I/O formats like file formats, messaging systems, databases, and file systems. The supported file formats are text, Avro, Parquet, S3, and Google cloud storage. All read/write transforms are under the IO namespace.

It is possible to create a PCollection using “beam.Create()”, which lets you create various types, e.g., lists, sets, or dictionaries. The following example creates a list of integers and then writes them into a file.

ilines = (
  p1
  |beam.Create([1,2,3])
  |beam.io.WriteToText('outfile.txt')
)

Listing 2. Creating a PCollection using the beam.Create()

Finally, similar to a read transform, we have a corresponding write transform, which writes the output of the Apache Beam pipeline.

Filter Transform

This is self-explanatory: Given a condition, the transform filters out all elements that don't satisfy it. A Filter transform can also be used to filter based on inequality with a given value based on the comparison ordering of the element. Here is an excellent Google Colab showing different variations of the Filter transform.

Map/FlatMap Transform

A Map transform converts a PCollection of N elements into another PCollection of N elements. A FlatMap can have zero or more outputs—it takes a PCollection of N elements and returns N collections of zero or more elements that are flattened into a final PCollection.

For example, if we have a list of three elements, applying the same lambda operation using Map will produce a collection of three lists, while FlatMap will produce a collection of six individual elements, as follows:

beam.Create([1, 2, 3]) | beam.Map(lambda x: [x, 'macroMeta']) 

# Three lists output: [[1, 'macroMeta'], [2, 'macroMeta'], [3, 'macroMeta']]


beam.Create([1, 2, 3]) | beam.FlatMap(lambda x: [x, 'macroMeta']) 

#Six single elements output: [1, 'macroMeta', 2, 'macroMeta', 3, 'macroMeta']

Listing 3. Differences in output from equivalent Map and FlatMap transforms.

ParDo Transform

ParDo is a generic parallel processing transform. It is similar to Map because it takes a collection as input, does some processing, and returns the output. However, It can return zero, one, or multiple elements of different types, which is different than Map/FlatMap, where one cannot change the element type. It is also used for slicing and dicing collections.

As ParDo is a generic transform, it is used to write more complex and customized transforms. So, we can achieve Map, FlatMap, and filtering functionality using ParDo as well.

The following code listing shows how to calculate the number of visitors to the cardio department, like in our previous example, but using the ParDo transform this time.

import apache_beam as beam

class SplitRow(beam.DoFn):
def process(self, element):
  return  [element.split(',')]

class FilterCardioPatients(beam.DoFn):
def process(self, element):
  if element[3] == 'cardio':
    return [element]
 
class PairPatients(beam.DoFn):
def process(self, element):
  return [(element[1], 1)]
class Counting(beam.DoFn):
def process(self, element):
  (key, values) = element
  return [(key, sum(values))]
 

p1 = beam.Pipeline()

visit_count = (
 
  p1
  |beam.io.ReadFromText('dept_data.txt')
 
  |beam.ParDo(SplitRow())
  |beam.ParDo(FilterCardioPatients())
  |beam.ParDo(PairPatients())
  |beam.GroupByKey()
  | beam.ParDo(Counting())
 
  |beam.io.WriteToText('parddo_output.txt')
)

p1.run()

Listing 4. Patient visit count for the cardio department using the ParDo transform

ParDo accepts a DoFn object, which contains the processing logic for each element of the input collection. It executes the user function on the distributed systems in parallel. We can define custom logic by overriding the process method of the DoFn class. Map and FlatMap also inherit from the DoFn, so we can use ParDo for any use cases to build more complex transforms.

In the above code listing, Map and Filter are replaced with ParDo, which uses the same code logic. You will notice that we are returning a list from the process method because it can produce more than one element. If we return a list of elements, it will act like Map; otherwise, it will behave like FlatMap and create multiple elements.

import apache_beam as beam


p = beam.Pipeline()


input_collection = (

    p
    | "Read input data" >> beam.io.ReadFromText('dept_data.txt')
    | "Split rows into columns" >> beam.Map(lambda record: record.split(','))
                  )

cardio_count = (
    input_collection
    | 'Filter cardio patients' >> beam.Filter(lambda record: record[3] == 'cardio')
    | 'Pair each patient with 1' >> beam.Map(lambda record: ("cardio, " +record[1], 1))
    | 'Aggegate cardio' >> beam.CombinePerKey(sum)
    )

ortho_count = (
              input_collection
              | 'Get all ortho patients' >> beam.Filter(lambda record: record[3] == 'ortho')
              | 'Pair each ortho patient with 1' >> beam.Map(lambda record: ("ortho, " +record[1], 1))
              | 'Aggregate Ortho' >> beam.CombinePerKey(sum)
          )

output =(
        (cardio_count,ortho_count)
  | beam.Flatten()
  | beam.io.WriteToText('data/both')
)

Listing 5. Patient visit count for multiple departments with pipeline branching

We can store the output of both transforms in separate files, but Apache Beam provides a transform called flatten that can combine multiple collections of the same data type. It merges various PCollectons into a single logical PCollection, similar to a union operation. In our case, both collections have the same data types and number of columns, allowing them to be flattened together. Flatten works on a tuple, so we create a tuple of collections and then call the transform and write the results to the file.

Composite Transforms

Apache Beam allows nested transformations, where a transform can perform multiple other transforms, such as ParDo, Combine, etc. These are called composite transforms. Hence, a composite transforms is a transform that has a series of transforms combined in a helpful processing pattern.

In our example, we calculated the number of visits for two different departments. If we want to add additional criteria, such as filtering the patients who have visited more than once, we will have to add this extra code in both places. Such changes are error-prone, and the resulting code is hard to maintain due to code duplication, leading to bugs. Instead, we will create a composite transform, like a reusable method for both use cases.

Our code after using composite transform will be as follows:

import apache_beam as beam

class CustomTransform(beam.PTransform):
  def expand(self, input_coll):
 
  a = (
      input_coll
                      | 'Group and sum' >> beam.CombinePerKey(sum)
                      | 'filter regular patients' >> beam.Filter(filter_on_count)
                      | 'Regular patient' >> beam.Map(format_output)
           
  )
  return a

def SplitRow(element):
  return element.split(',')

def filter_on_count(element):
 name, count = element
 if count > 1:
  return element
 def format_output(element):
 name, count = element
 #return ', '.join((name.encode('ascii'),str(count),'Regular employee'))
 return ', '.join((name,str(count),'Regular Patient'))

p = beam.Pipeline()

input_collection = (

              p
                    | "Read from text file" >> beam.io.ReadFromText('dept_data.txt')
                    | "Split rows" >> beam.Map(SplitRow)
                  )

cardio_count = (
                    input_collection
                    | 'Get all cardio patients' >> beam.Filter(lambda record: record[3] == 'cardio')
                    | 'Pair each patient with 1' >> beam.Map(lambda record: ("cardio, " +record[1], 1))
                    | 'composite cardio Patients' >> CustomTransform()
                    | 'Write results for cardio' >> beam.io.WriteToText('cardio_output.txt')
                )

ortho_count = (
              input_collection
              | 'Get all HR dept persons' >> beam.Filter(lambda record: record[3] == 'ortho')
              | 'Pair each hr employee with 1' >> beam.Map(lambda record: ("ortho, " +record[1], 1))
              | 'composite Ortho Patients' >> CustomTransform()
              | 'Write results for Ortho' >> beam.io.WriteToText('ortho_output.txt')
          )
p.run()

Listing 6. Patient visit count for multiple departments using a composite transform.

We have inherited the “CustomTransform” from the PTransform base class. All the core transforms, like Map and FlatMap, inherit from the PTransform class as well. “CustomTransform” overrides the expand method of the base class and defines the processing logic, which applies a sequence of core transforms.

This content was produced by Inbound Square

Compare