DoFn Deep Dive

1. Introduction

2. Why This Concept Exists

3. Key Terminology

• DoFn Class: A Python class that inherits from beam.DoFn.
• process method: The core method called for every element in the input PCollection.
• Yielding vs Returning: A DoFn yields elements using yield (or returns an iterable) instead of using return to support one-to-many outputs.

4. How It Works: The Lifecycle Hooks

1. setup(): Called once per worker initialization (ideal for establishing database or API connections).
2. start_bundle(): Called at the beginning of processing a bundle of elements.
3. process(element): Called for every element in the PCollection.
4. finish_bundle(): Called after a bundle of elements is processed.
5. teardown(): Called once when the worker instance is destroyed (ideal for closing connections).

5. Visual Diagram

6. Code Example

Creating a basic DoFn that yields elements:

pythonCopy
import apache_beam as beam

class FilterNumbersDoFn(beam.DoFn):
    def process(self, element):
        if element >= 10:
            yield element

# Applying the DoFn inside your pipeline
# output = input_pcoll | beam.ParDo(FilterNumbersDoFn())

7. Code Explanation

• FilterNumbersDoFn inherits from beam.DoFn.
• The process method takes element as an argument.
• It checks the condition and uses yield element to output the record.
• We apply it inside a pipeline using beam.ParDo(FilterNumbersDoFn()).

8. Real Production Example

pythonCopy
class DatabaseLookupDoFn(beam.DoFn):
    def setup(self):
        # Establish connection once per worker initialization
        self.db_client = my_database_library.connect("host-name")

    def process(self, element):
        # Query database and emit record
        user_info = self.db_client.query(element["user_id"])
        yield {**element, "user_name": user_info.name}

    def teardown(self):
        # Close connection when worker shuts down
        self.db_client.close()

9. Common Mistakes

• Instantiating non-serializable objects in __init__: Do not create database connections inside the class constructor. The constructor runs on the developer's computer, not the worker machines. Create connections inside the setup() method.
• Returning values instead of iterables: The process method must return an iterable (such as a list) or yield values. Returning a primitive type (like an integer) will throw a runtime error.

10. Best Practices

• Never write local file outputs inside process(). Worker local disks are ephemeral and isolated.
• Catch connection errors in setup() and implement retry strategies to prevent worker crashes.

11. Summary

• DoFn is the processing class executed inside ParDo.
• Features lifecycle hooks: setup, start_bundle, process, finish_bundle, and teardown.
• Must yield elements or return an iterable.

12. Interactive Challenges

13. Related Content

• ParDo Processing
• Map Transform