如何在Apache Beam Python流水线中传参并模块化自定义函数？

Let's break down your problem step by step — you're running into two key issues: modularizing your custom Beam functions, and correctly passing parameters to them in the pipeline. Let's fix both.

First: The Core Parameter Passing Problem #

The mistake in your current code is how you're passing arguments to check_null and del_col in beam.Filter and beam.Map. Apache Beam transforms like Filter and Map expect a function that takes the pipeline element as its first argument. When you write beam.Filter(check_null('userid', 'itemid', 'ratings')), you're immediately calling check_null with those strings as arguments (instead of passing the function itself), and Beam tries to use the return value (a boolean) as the filter function — which breaks because it expects a function that takes the data element.

Fix 1: Use Lambda Wrappers (Simplest Approach) #

Wrap your custom functions in a lambda that accepts the pipeline element and passes it along with your parameters:

# Replace your Filter and Map steps with these
| 'DeleteNullData' >> beam.Filter(lambda elem: check_null(elem, 'userid', 'itemid', 'ratings'))
| 'DeleteUnwantedData' >> beam.Map(lambda elem: del_col(elem, col_name='timestamp'))

Fix 2: Use functools.partial (Cleaner for Reusable Functions) #

For a more elegant solution, use functools.partial to bind your parameters to the function, leaving the first argument (the data element) to be filled in by Beam:

First import partial:

from functools import partial

Then update your pipeline steps:

| 'DeleteNullData' >> beam.Filter(partial(check_null, 'userid', 'itemid', 'ratings'))
| 'DeleteUnwantedData' >> beam.Map(partial(del_col, col_name='timestamp'))

Pro tip: Update your check_null function to accept variable arguments for flexibility, instead of fixed column names:

def check_null(data, *col_names):
    # Check all specified columns have non-empty values
    return all(len(data[col]) > 0 for col in col_names)

Second: Modularizing Custom Functions #

To keep your code organized, move all custom transform functions into a separate module (e.g., beam_utils.py):

Step 1: Create beam_utils.py #

# beam_utils.py
def del_col(data, col_name: str):
    del data[col_name]
    return data

def check_null(data, *col_names):
    return all(len(data[col]) > 0 for col in col_names)

def format_data(data):
    return ','.join([data['userid'], data['itemid'], data['ratings']])

def print_row(data):
    print(data)

Step 2: Import in Your Main Script #

In your main pipeline file, import the functions:

import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
from functools import partial
from beam_utils import del_col, check_null, format_data, print_row

Bonus: Fix Data Splitting Issue #

Looking at your sample data, it's space-separated, not comma-separated. Your current SplitData step uses x.split(',') which will split each row into a single element (since there are no commas). Fix that with:

| 'SplitData' >> beam.Map(lambda x: x.split())  # Splits on any whitespace

Full Corrected Pipeline Code #

Here's the complete working code with all fixes:

import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
from functools import partial
from beam_utils import del_col, check_null, format_data, print_row

if __name__ == '__main__':
    options = PipelineOptions()
    input_file = 'data.csv'
    with beam.Pipeline(options=options) as pipeline:
        (pipeline 
         | 'ReadData' >> beam.io.ReadFromText(input_file, skip_header_lines=1)  # Skip the header row
         | 'SplitData' >> beam.Map(lambda x: x.split())
         | 'FormatToDict' >> beam.Map(lambda x: {
             "userid": x[0], 
             "itemid": x[1], 
             "ratings": x[2], 
             "timestamp": x[3]
         })
         | 'DeleteNullData' >> beam.Filter(partial(check_null, 'userid', 'itemid', 'ratings'))
         | 'DeleteUnwantedData' >> beam.Map(partial(del_col, col_name='timestamp'))
         | 'FormatData' >> beam.Map(format_data)
         | 'PrintResults' >> beam.Map(print_row))  # Add this to verify output

内容的提问来源于stack exchange，提问作者wyn