PyFlink中定义函数的3种方式

函数(Functions)

转换算子接受用户定义的函数作为输入,以定义转换的功能。本节将描述Python DataStream API中定义Python用户定义函数的不同方式。

1、实现函数接口

Python DataStream API中针对不同的转换算子提供了不同的函数接口。例如,map 转换提供了MapFunction接口,filter转换提供了FilterFunction接口等。用户可以根据转换的类型实现对应的函数接口。以MapFunction为例:

# Implementing MapFunction
class MyMapFunction(MapFunction):
    
    def map(self, value):
        return value + 1
        
data_stream = env.from_collection([1, 2, 3, 4, 5], type_info=Types.INT())
mapped_stream = data_stream.map(MyMapFunction(), output_type=Types.INT())

这里实现了MapFunction接口,并重写了map()方法来定义map转换逻辑。然后实例化函数并传入data_stream的map转换中。

类似地,可以针对其他转换,实现对应的函数接口:

• FilterFunction: 实现filter()
• FlatMapFunction: 实现flat_map()
• ReduceFunction: 实现reduce()
• CoMapFunction: 实现map1()和map2()
• …

实现函数接口可以使代码更结构化,也方便Flink调用和处理函数逻辑。

2、Lambda函数

如以下示例所示,转换算子也可以接受lambda函数来定义转换的功能:

data_stream = env.from_collection([1, 2, 3, 4, 5], type_info=Types.INT())
mapped_stream = data_stream.map(lambda x: x + 1, output_type=Types.INT())

这里直接使用lambda表达式作为map()的函数参数。

lambda函数的优点是非常方便进行简单逻辑的定义。但复杂逻辑时还是推荐实现函数接口,结构会更清晰。

注意: ConnectedStreams.map() 和 ConnectedStreams.flat_map() 不支持lambda函数,必须分别接受 CoMapFunction 和 CoFlatMapFunction。

例如:

connected_stream = stream1.connect(stream2)

# 正确
connected_stream.map(MyCoMapFunction()) 

# 错误
connected_stream.map(lambda x, y: ...)

# 正确  
connected_stream.flat_map(MyCoFlatMapFunction())

# 错误
connected_stream.flat_map(lambda x, y: ...)

原因是 ConnectedStreams 的 map 和 flat_map 需要处理两个流的数据,所以必须实现 CoMapFunction 或 CoFlatMapFunction 接口,而不能使用lambda函数。

3、Python函数

用户也可以使用普通的Python函数来定义转换的功能:

def add_one(value):
    return value + 1

data_stream = env.from_collection([1, 2, 3, 4, 5], type_info=Types.INT())
mapped_stream = data_stream.map(add_one, output_type=Types.INT())

这里定义了一个普通的 Python 函数 add_one,然后将其传入 map 转换。

Python 函数的优点是结构清晰,易于重用和测试。复杂函数逻辑时,实现 Python 函数会更加合适。

值得注意的是,普通 Python 函数不支持富函数的特性,比如获取函数的运行时上下文等。如果需要这些特性,还是需要实现函数接口(如 MapFunction)。

所以,一般建议如下:

• 简单逻辑: lambda 函数
• 复杂逻辑: Python 函数
• 需要富函数特性: 实现函数接口

根据不同的场景选择合适的函数定义方式。Python 函数提供了一种灵活的代码重用和组织方式。

PyFlink 中没有RichMapFunction这种写法，MapFunction 天生自带生命周期方法。

源码：

# Function 自带 open、close 方法
class Function(ABC):
    """
    The base class for all user-defined functions.
    """
    def open(self, runtime_context: RuntimeContext):
        pass

    def close(self):
        pass

# map、fliter等继承自Function
class MapFunction(Function):
    """
    Base class for Map functions. Map functions take elements and transform them, element wise. A
    Map function always produces a single result element for each input element. Typical
    applications are parsing elements, converting data types, or projecting out fields. Operations
    that produce multiple result elements from a single input element can be implemented using the
    FlatMapFunction.
    The basic syntax for using a MapFunction is as follows:

    ::
        >>> ds = ...
        >>> new_ds = ds.map(MyMapFunction())
    """

    @abstractmethod
    def map(self, value):
        """
        The mapping method. Takes an element from the input data and transforms it into exactly one
        element.

        :param value: The input value.
        :return: The transformed value.
        """
        pass

4、输出类型（Output Type）

在 Python DataStream API 中,用户可以显式指定转换的输出类型信息。如果不指定,默认输出类型会是 Types.PICKLED_BYTE_ARRAY,结果数据使用 pickle 序列化器序列化。

通常在以下场景需要指定输出类型:

• 将数据转换为非字节数组类型后输出,如字符串、数字等。这时需要明确指定输出类型,否则默认会使用 pickle 序列化,无法得到正确的结果。

• 进行基于类型的流优化,如基于字符串 Hash 分区等。这时需要指定类型以发挥优化效果。

• 数据写入外部系统需要进行特定编码或序列化。

指定输出类型的方式是传递一个 output_type 参数,例如:

data_stream.map(lambda x: x * 2, output_type=Types.INT)

此外,在实现函数接口时,也可以通过返回类型提示输出类型,例如:

class MyMapFunction(MapFunction):
    
    def map(self, value) -> str:
        ...
        
data_stream.map(MyMapFunction())

明确指定输出类型可以减少不必要的默认序列化,并启用基于类型的优化,是提升性能的重要手段。

5、算子链

默认情况下,多个 non-shuffle 的 Python 函数会被链在一起,以避免序列化和反序列化,提高性能。在某些情况下,可能需要禁用链,例如,有一个 flat_map 函数为每个输入元素生成大量元素,禁用链允许以不同的并行性来处理其输出。

可以通过以下方式之一禁用算子链:

• 在当前算子后添加 key_by、shuffle、rescale、rebalance 或 partition_custom 操作,禁用与后续算子链接。

• 为当前算子应用 start_new_chain 操作,禁用与前面的的算子链接。

• 为当前算子应用 disable_chaining 操作,禁用与前面和后面的算子链接。

• 为两个算子设置不同的并行度或不同的 slot 共享组,禁用两者之间的链。

• 通过配置 python.operator-chaining.enabled 禁用全部算子链。

禁用算子链可以更细粒度地控制并行度和资源利用。但也会损耗性能,需要根据实际情况进行权衡。

6、在 Python 函数中加载资源

有时希望先在 Python 函数中加载一些资源,然后重复运行计算,而无需重新加载资源。例如,您可能只想加载一个大的深度学习模型一次,然后针对该模型运行多次批预测。

此时需要重写从基类Function继承的open()方法。

例如:

class LoadModelOnceFunction(MapFunction):

    def open(self, runtime_context):
        print("Loading model")
        self.model = load_model() 

    def map(self, value):
        return self.model.predict(value)

stream.map(LoadModelOnceFunction())

这里在open()方法中加载了模型,这样模型只会被加载一次,但map()方法可以重复使用该模型进行预测。

open()方法会在函数实例初始化时调用一次,可以在其中进行只需要加载一次的操作。这是避免每次调用都加载的一种简单高效的方式。

除此之外,也可以重写close()方法进行清理工作。close()会在函数实例销毁前调用。

详细资料关注微信公众号