高效数据交互 Arrow 插件使用教程

本例需要使用 2.00.12 及以上版本的 DolphinDB server 和对应的 DolphinDB Arrow 插件。使用前可以分别参考 DolphinDB 部署教程和章节 2 中的内容进行安装和加载。

本例需要使用 1.30.21.1 及以上版本的 DolphinDB Python API。详细安装说明请参见文档。为保证本例代码的正常运行，请安装如下库：

• NumPy：测试代码需要，建议安装最新版本。
• pandas：测试代码和全流程使用 Arrow 格式需要，请安装 2.0 及以上版本。
• pyarrow：为测试场景中全流程可以使用 Arrow 作为中间格式，请安装 9.0.0 及以上版本。

本例中使用 PySpark 作为数据计算组件，需安装 PySpark，详情请参考 PySpark 下载指南。启动 PySpark 前，需要指定环境变量 PYSPARK_PYTHON 的值为 Python 可执行文件的路径，或在 Python 脚本中进行如下配置：

import os
os.environ["PYSPARK_PYTHON"] = "/path_to_python_env/python"

本例中，定义 _SparkBase 类用于管理 PySpark 环境。使用时，需手动调用 start 方法启动 PySpark 环境。

• 如果指定 use_arrow=True，表示启用 Arrow 作为 PySpark 内部的中间格式。
• 如果指定 use_arrow=False，表示使用 pandas DataFrame 作为 PySpark 内部的中间格式。

结束时需调用 shutdown 方法关闭 PySpark 环境。

class _SparkBase(object):
    @classmethod
    def start(cls, use_arrow=True):
        conf = SparkConf()
        conf.set("spark.sql.execution.arrow.pyspark.enabled", "true" if use_arrow else "false")
        conf.set("spark.default.parallelism", "100")
        cls.context: SparkContext = SparkContext(master='local[*]', appName=cls.__name__, conf=conf)
        cls.session: SparkSession = SparkSession.builder.getOrCreate()
        cls.use_arrow = use_arrow

    @classmethod
    def shutdown(cls):
        cls.session.stop()
        cls.session = None
        cls.context.stop()
        cls.context = None

本例中使用包含一列 STRING 和 三列 DOUBLE ，有 10000000 行随机数据的共享内存表（约 300 MB）进行计算演示。数据准备脚本如下，其中 HOST/PORT/USER/PASSWD 可参考 Python API 使用教程配置为数据库连接信息。

import dolphindb as ddb
script = """
    n = 10000000;
    id = take(`AAA`BBB`CCC, n);
    v1 = rand(10.0, n);
    v2 = rand(20.0, n);
    v3 = rand(5.0, n);
    share table(id as id, v1 as v1, v2 as v2, v3 as v3) as t;
"""

conn = ddb.Session()
conn.connect(HOST, PORT, USER, PASSWD)
conn.run(script)

定义 cal 函数，根据 id 列值划分的各组，分别计算 (v1 + v2) * v3 的均值，并将结果返回。

def cal(df: sql.DataFrame):
    df = df.withColumn("v", (F.col("v1") + F.col("v2")) * F.col("v3"))
    df = df.groupby("id").avg("v")
    return df

参考上图，使用 pandas 作为 PySpark 的中间格式，整体计算流程如下：

1. 从 API 获取 pandas DataFrame 数据，该过程中，会将 DolphinDB 数据转换为 pandas DataFrame 。
2. Python 通过 Socket 调用 Spark API，在计算时，将 Python 端 pandas DataFrame 数据经过多次序列化后发送至 Spark 处理。
3. 计算完成后，再经过多次反序列化（pickle）将 Spark 数据反序列化至 Python 端。

在该流程中，第二步和第三步将会在 PySpark 内部经过多次序列化和反序列化，且默认使用行式处理，额外开销较大。

具体实现如下，详细代码可参考附录：

import dolphindb as ddb
import dolphindb.settings as keys
class PySparkTest(_SparkBase):
    def testPROTOCOL_PICKLE(self):
        s = ddb.Session(protocol=keys.PROTOCOL_PICKLE)
        s.connect(HOST, PORT, USER, PASSWD)
        df = s.run("t")
        df = self.session.createDataFrame(df)
        cal(df)

PySparkTest.start(use_arrow=False)
PySparkTest().testPROTOCOL_PICKLE()
PySparkTest.shutdown()

上述代码中，启动 PySpark 时指定 use_arrow=False，表示不使用 Arrow 作为 Spark 内部中间格式。指定 Python API 使用 PROTOCOL_PICKLE 协议（默认协议）从 DolphinDB 获取 pandas DataFrame 数据，并调用 createDataFrame 方法将其转换为 spark 格式的 DataFrame，传入 cal 函数进行计算。

参考上图，使用 Arrow 作为 PySpark 的中间格式，将会经过以下步骤：

1. 从 API 获取 Arrow 格式数据，该过程中会将 DolphinDB 表数据转换为 Arrow 格式数据。
2. Python 端将 Arrow 格式数据转换为同样使用 Arrow 作为数据后端的 pandas DataFrame，因为都使用 Arrow 格式，转换开销极小。
3. Python 通过 Socket 调用 Spark API 将使用 Arrow 作为数据后端的 pandas DataFrame 传入 Spark 系统，和原流程不同，Arrow 数据格式在序列化和反序列化时能够以近乎零成本的方式进行，启用 Arrow 作为中间格式的 PySpark，可以以最高效的方式接收数据。
4. 计算时，Arrow 作为列式数据格式，对于计算也有一定优化，计算完成后通过零成本序列化和反序列化将结果发送回 Python 端。

该流程中，Arrow 能够在 Python 系统和 Spark 系统之间以较低开销进行传输，Arrow 插件使得 DolphinDB 可以直接提供 Arrow 格式的数据，并通过 Arrow 后端的 pandas DataFrame 作为桥梁，省去流程中繁琐和开销极大的类型转换，从而提供整体流程效率。

具体实现如下，详细代码可参考附录：

class PySparkTest(_SparkBase):
    def testPROTOCOL_ARROW(self):
        s = ddb.Session(protocol=keys.PROTOCOL_ARROW)
        s.connect(HOST, PORT, USER, PASSWD)
        pt = s.run("t")
        df = pd.DataFrame({
            'id': pd.Series(pt["id"].combine_chunks(), dtype=pd.ArrowDtype(pa.utf8())),
            'v1': pd.Series(pt["v1"].combine_chunks(), dtype=pd.ArrowDtype(pa.float64())),
            'v2': pd.Series(pt["v2"].combine_chunks(), dtype=pd.ArrowDtype(pa.float64())),
            'v3': pd.Series(pt["v3"].combine_chunks(), dtype=pd.ArrowDtype(pa.float64())),
        })
        df = self.session.createDataFrame(df)
        cal(df)
PySparkTest.start(use_arrow=True)
PySparkTest().testPROTOCOL_ARROW()
PySparkTest.shutdown()

上述代码中，启动 PySpark 时指定 use_arrow=True，表示使用 Arrow 作为 Spark 内部中间格式。指定 Python API 使用 PROTOCOL_ARROW 协议从 DolphinDB 获取 Arrow 格式数据，将其转为 PyArrow 后端的 pandas DataFrame（pandas 2.0 支持使用 PyArrow 替代 NumPy 作为数据后端）。再调用 createDataFrame 方法，将其转为 PySpark DataFrame，最后调用 cal 进行计算。