Flink SQL中怎么注册python以及使用python注册的UDF中数据流是怎么流转的

背景

本文基于 Flink 1.17.0

和Spark SQL中怎么注册python以及使用python注册的UDF中数据流是怎么流转的目的一样，为了阐述 Flink SQL 对 python UDF的处理

分析

注册python udf以及调用

如create-function所示，可以用DSL进行 udf的注册，引用StreamPythonUdfSqlJob.java中的例子：

 tEnv.executeSql(
                "create temporary system function add_one as 'add_one.add_one' language python");

        tEnv.createTemporaryView("source", tEnv.fromValues(1L, 2L, 3L).as("a"));

        Iterator<Row> result = tEnv.executeSql("select add_one(a) as a from source").collect();

其中 add_one.py 为：

from pyflink.table import DataTypes
from pyflink.table.udf import udf


@udf(input_types=[DataTypes.BIGINT()], result_type=DataTypes.BIGINT())
def add_one(i):
    import pytest
    return i + 1

也是用python中注册 add_one 函数，之后在 SQL中进行调用

调用python udf的数据流

注册

如 create temporary system function add_one as 'add_one.add_one' language python 这个DSL中的定义的SQL，最终会变成
CreateTempSystemFunctionOperation 最终会走到 TableEnvironmentImpl.executeInternal 中的 createSystemFunction((CreateTempSystemFunctionOperation) operation)方法：

public void registerTemporarySystemFunction(
            String name, CatalogFunction function, boolean ignoreIfExists) {
        final String normalizedName = FunctionIdentifier.normalizeName(name);

        try {
            validateAndPrepareFunction(name, function);
        } catch (Throwable t) {
            throw new ValidationException(
                    String.format(
                            "Could not register temporary system function '%s' due to implementation errors.",
                            name),
                    t);
        }
        if (!tempSystemFunctions.containsKey(normalizedName)) {
            tempSystemFunctions.put(normalizedName, function);

最终会保存到 FunctionCatalog.tempSystemFunctions变量中, 这个变量在后续的查找函数的时候会被调用到。

调用

对于Flink来说，每一个函数，都会经过FunctionCatalog.lookupFunction方法：

 public Optional<ContextResolvedFunction> lookupFunction(UnresolvedIdentifier identifier) {
        // precise function reference
        if (identifier.getDatabaseName().isPresent()) {
            return resolvePreciseFunctionReference(catalogManager.qualifyIdentifier(identifier));
        } else {
            // ambiguous function reference
            return resolveAmbiguousFunctionReference(identifier.getObjectName());
        }
    }

对应的数据流为：

FunctionCatalog.resolveAmbiguousFunctionReference

getFunctionDefinition(normalizedName, tempSystemFunctions.get(normalizedName))

UserDefinedFunctionHelper.instantiateFunction

PythonFunctionUtils.getPythonFunction(catalogFunction.getClassName(), config, classLoader)

PythonFunctionUtils.pythonFunctionFactory(利用反射调用 getPythonFunction)

最终会调用 PythonFunctionFactory.getPythonFunction 该方法会最终调用 createPythonFunctionFactory 方法，

该方法会调用python -m pyflink.pyflink_callback_server P动，这里启动相关的都是跟Py4j有关，其中这里就会把python中的PythonFunctionFactory 放到 java中的gatewayServer 的hashMap中，而这里启动的Py4j客户端就在 startGatewayServer方法中，这个命令 python -m pyflink.pyflink_callback_server会把 python 的PythonFunctionFactory()对象放入 Py4j 的客户端中，

PythonFunctionFactory 代码如下：

class PythonFunctionFactory(object):
           """
           Used to create PythonFunction objects for Java jobs.
           """

           def getPythonFunction(self, moduleName, objectName):
               udf_wrapper = getattr(importlib.import_module(moduleName), objectName)
               return udf_wrapper._java_user_defined_function()

           class Java:
               implements = ["org.apache.flink.client.python.PythonFunctionFactory"]

所以createPythonFunctionFactory方法中：

pythonProcess =
                        launchPy4jPythonClient(
                                gatewayServer, config, commands, null, tmpDir, false);
                entryPoint = (Map<String, Object>) gatewayServer.getGateway().getEntryPoint();
...
return new PythonFunctionFactoryImpl(
                (PythonFunctionFactory) entryPoint.get("PythonFunctionFactory"), shutdownHook);

最终返回的 PythonFunctionFactoryImpl是包含了python的 PythonFunctionFactory 对象，所以前面返回的PythonFunctionUtils.getPythonFunctio都是包裹了python的java对象，所以后续的调用都是基于 Py4j 的进程间的调用了

总结

所以说 Flink SQL 调用 python UDF 还是采用了 Py4j ，这种方式也是采用了进程间通信的方式，在效率上还是比不了基于 java/scala 而写的UDF，这种方式和Spark SQL中怎么注册python以及使用python注册的UDF中数据流是怎么流转的类似。