玩转PyFlink：自定义SourceFunction在实时数据处理中的运用

简介

在PyFlink中，SourceFunction与其他算子（如Map、FlatMap等）的实现方式有所不同。本文将介绍PyFlink中如何自定义SourceFunction，并通过示例代码展示如何在实时数据处理中灵活应用自定义SourceFunction。

自定义源算子

1. 理解自定义源算子的必要性和作用。
2. 创建实现了 SourceFunction 接口的自定义源算子类。
3. 实现核心方法：run() 和 cancel()。
4. 在 run() 方法中定义数据生成和发送的逻辑。
5. 使用自定义源算子，执行任务并观察输出。

PyFlink 中 SourceFunction 的定义

PyFlink中的SourceFunction用于实现数据源，但与其他算子不同，无法通过Python UDF直接实现。而是需要先在Java中实现SourceFunction，然后在Python作业中引入。

class SourceFunction(JavaFunctionWrapper):
    """
    Base class for all stream data source in Flink.
    """

    def __init__(self, source_func: Union[str, JavaObject]):
        """
        Constructor of SinkFunction.

        :param source_func: The java SourceFunction object.
        """
        super(SourceFunction, self).__init__(source_func)

PyFlink 中 MapFunction 的定义

class Function(ABC):
    """
    The base class for all user-defined functions.
    """
    def open(self, runtime_context: RuntimeContext):
        pass

    def close(self):
        pass


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

由以上源码可见，SourceFunction 和 MapFunction 等这些算子的定义是不同的。

在PyFlink中MapFunction、FlatMapFunction 等天生就存在生命周期方法。

编写 Java 代码

/*
* 自定义一个简单的 SourceFunction
* */
public class MyJavaCustomSourceFunction implements SourceFunction<Row> {

    // 自定义一个数组当作数据源
    private static final String[] NAMES = {"Flink", "Spark", "Hadoop", "Hive", "Doris", "PyFlink"};

    public void run(SourceContext<Row> sourceContext) throws Exception {
        Random random = new Random();
        for (int i = 0; i < NAMES.length; i++) {
            Row row = Row.of(i, NAMES[i]);
            sourceContext.collect(row);
        }
    }

    public void cancel() {

    }
}

简单的SourceFunction

以下示例展示了如何在Java中编写自定义的SourceFunction。

def get_customer_source(env):
    """
    自定义一个简单的数据源
    :param env:
    :return:
    """
    # Java 代码打包的jar包地址
    environment.add_jars('file:///Users/oscar/software/jars/customer_source-1.0-SNAPSHOT.jar')
    # 直接填写类路径，实现对java source的引用
    custom_source = SourceFunction("org.tutorial.MyJavaCustomSourceFunction")

    # add_source
    ds = env.add_source(custom_source, type_info=Types.ROW([Types.INT(), Types.STRING()]))

    ds.print()

以上这种方法缺乏灵活性，无法传递参数，这在生产中极其麻烦。、

自定义可传递参数的SourceFunction

以下示例展示如何在Java中编写可传递参数的SourceFunction，并如何在Python中引用和使用它。

/*
* 自定义Redis 的 SourceFunction，实现 Python 端可传递参数
* */
public class MyJavaRedisSourceFunction implements SourceFunction<String> {

    private final String host;
    private final int port;
    private final String password;

    public MyJavaRedisSourceFunction(String host, int port, String password){
        this.host = host;
        this.port = port;
        this.password = password;
    }


    public void run(SourceContext<String> sourceContext) throws Exception {
        // 初始化 Redis 连接
        Jedis jedis = new Jedis(this.host, this.port);
        jedis.auth(this.password);

        jedis.set("key", "测试PyFlink自定义Source");
        int i = 0;
        while (i < 100){
            String s = Integer.toString(i);
            String value = jedis.get("key");
            String result = value + "-->" + s;
            i++;
            sourceContext.collect(result);
            System.out.println(result);
        }
    }

    public void cancel() {

    }
}

python 端代码

class MyRedisSourceFunction(SourceFunction):
    """
    实现Python SourceFunction，内部初始化java SourceFunction
    """
    def __init__(self, host, port, password):
        # 获取Java RedisSource类，也就是自定义的Java SourceFunction 类
        JVMRedisSourceFunction = get_gateway().jvm.org.tutorial.MyJavaRedisSourceFunction

        # 初始化Java RedisSource，传递参数 host、port
        j_redis_source = JVMRedisSourceFunction(host, port, password)

        # 调用父类__init__
        super(MyRedisSourceFunction, self).__init__(source_func=j_redis_source)

def get_customer_source_redis(env):
    """
    自定义Redis为数据源
    :param env:
    :return:
    """
    host = '127.0.0.1'
    port = 6379
    password = 'Oscar&0503'
    # jar包地址
    environment.add_jars('file:///Users/oscar/software/jars/customer_source-1.0-SNAPSHOT.jar')
    redis_source = MyRedisSourceFunction(host, port, password)

    ds = env.add_source(redis_source, type_info=Types.STRING())
    ds.print()

本文深入介绍了PyFlink中自定义SourceFunction的实现与应用。通过Java实现SourceFunction，然后在Python作业中引用，我们可以灵活地构建不同类型的数据源，实现更丰富的实时数据处理应用。这种跨语言的灵活性为PyFlink带来了更多的可能性，让我们在实际应用中能够更好地发挥其强大的功能。

完整示例代码：

#!/usr/bin/python
# -*- coding:UTF-8 -*-
import json

from pyflink.common import Types
from pyflink.datastream import StreamExecutionEnvironment, SourceFunction
from pyflink.java_gateway import get_gateway


class MyRedisSourceFunction(SourceFunction):
    """
    实现Python SourceFunction，内部初始化java SourceFunction
    """

    def __init__(self, host, port, password):
        # 获取Java RedisSource类，也就是自定义的Java SourceFunction 类
        JVMRedisSourceFunction = get_gateway().jvm.org.tutorial.MyJavaRedisSourceFunction

        # 初始化Java redis source，传递参数 host、port
        j_redis_source = JVMRedisSourceFunction(host, port, password)

        # 调用父类__init__
        super(MyRedisSourceFunction, self).__init__(source_func=j_redis_source)


def get_customer_source(env):
    """
    自定义一个简单的数据源
    :param env:
    :return:
    """
    # 直接填写类路径，实现对java source的引用
    custom_source = SourceFunction("org.tutorial.MyJavaCustomSourceFunction")

    # add_source
    ds = env.add_source(custom_source, type_info=Types.ROW([Types.INT(), Types.STRING()]))

    ds.print()


def get_customer_source_redis(env):
    """
    自定义Redis为数据源
    :param env:
    :return:
    """
    host = '127.0.0.1'
    port = 6379
    password = 'Oscar&0503'
    # 初始化Python class
    redis_source = MyRedisSourceFunction(host, port, password)

    # add_source
    ds = env.add_source(redis_source, type_info=Types.STRING())
    ds.print()


if __name__ == '__main__':
    """
    PyFlink DataStream API 自定义 SourceFunction
    """
    # 执行环境
    environment = StreamExecutionEnvironment.get_execution_environment()
    environment.set_parallelism(1)

    # 添加依赖的jar包
    environment.add_jars('file:///Users/oscar/software/jars/flink-sql-connector-kafka-1.17.1.jar')
    environment.add_jars('file:///Users/oscar/software/jars/customer_source-1.0-SNAPSHOT.jar')

    get_customer_source(env=environment)

    get_customer_source_redis(env=environment)

    environment.execute()

详细资料关注微信公众号