Custom IO

1. Introduction

2. Why This Concept Exists

Modern business architectures use custom internally developed REST APIs or legacy mainframes. In order to process this data at scale, pipelines need custom connectors that can parallelize work, divide datasets, track reader offsets, and manage bundle writes.

3. Key Terminology

• BoundedSource: A Beam base class representing a finite dataset that can be read and split into sub-sources.
• RangeTracker: A thread-safe helper class used to track the reader's current position and prevent overlapping record reads during dynamic splits.
• Dynamic Work Rebalancing: The ability of a runner to dynamically split a source's remaining work across multiple workers based on current progress.

4. How It Works

• Source: You subclass BoundedSource and implement the read(), get_range_tracker(), and split() methods.
• Sink: You subclass FileBasedSink (for files) or implement a custom DoFn that manages batch writes in its lifecycle hook methods (setup, process, finish_bundle).

5. Visual Diagram

6. Code Example

A custom Source that generates a sequence of integers:

pythonCopy
import apache_beam as beam
from apache_beam.io import iobase
from apache_beam.io.range_trackers import OffsetRangeTracker

class RangeSource(iobase.BoundedSource):
    def __init__(self, start, end):
        self._start = start
        self._end = end

    def estimate_size(self):
        return self._end - self._start

    def get_range_tracker(self, start_position, end_position):
        if start_position is None:
            start_position = self._start
        if end_position is None:
            end_position = self._end
        return OffsetRangeTracker(start_position, end_position)

    def read(self, range_tracker):
        for i in range(range_tracker.start_position(), range_tracker.stop_position()):
            if not range_tracker.try_claim(i):
                return
            yield i

    def split(self, desired_bundle_size, start_position=None, end_position=None):
        # Allow splitting the source range into smaller segments
        pass

7. Code Explanation

• estimate_size() helps the runner decide how many workers to allocate.
• get_range_tracker() creates a position coordinator.
• read() loops through the range, claiming indices via range_tracker.try_claim(i). If a claim fails (indicating the runner split the work), execution halts.

8. Real Production Example

pythonCopy
import apache_beam as beam

class ApiWriteDoFn(beam.DoFn):
    def __init__(self, endpoint):
        self.endpoint = endpoint
        
    def setup(self):
        # Initialize client
        self.client = self.connect_to_api(self.endpoint)
        self.batch = []
        
    def process(self, element):
        self.batch.append(element)
        # Flush if batch size exceeded
        if len(self.batch) >= 100:
            self.flush()
            
    def finish_bundle(self):
        # Guarantee remaining elements are flushed when bundle completes
        if self.batch:
            self.flush()
            
    def flush(self):
        self.client.post_batch(self.batch)
        self.batch = []
        
    def connect_to_api(self, endpoint):
        class DummyClient:
            def post_batch(self, b): pass
        return DummyClient()

9. Common Mistakes

• Failing to claim positions: Modifying range-tracker values inside read() loops without calling try_claim(), which breaks pipeline work rebalancing.
• Not flushing in finish_bundle(): Buffering elements in process() but forgetting to write the remaining buffered elements in finish_bundle(), resulting in data loss.

10. Best Practices

• Implement estimate_size() carefully; poor estimates lead to bad worker provisioning.
• Always implement finish_bundle() to flush batch writers and protect against data loss.

11. Summary

• Custom IO allows connecting to proprietary storage APIs.
• BoundedSource subclasses define read bounds and sizing.
• RangeTracker claims records, supporting dynamic work rebalancing.

12. Interactive Challenges

13. Related Content

• Databases
• Others