Apache Arrow PyArrow Substrait API 实战指南：执行查询计划与表达式序列化-平芜编程栈

Apache Arrow PyArrow Substrait API 实战指南：执行查询计划与表达式序列化

【免费下载链接】arrowApache Arrow is a multi-language toolbox for accelerated data interchange and in-memory processing项目地址: https://gitcode.com/gh_mirrors/arrow13/arrow

Substrait 是一个开放的跨语言执行计划（execution plan）标准，用于在不同查询引擎与数据系统之间传递"计算意图"。PyArrow 在pyarrow.substrait模块中提供了 Substrait 计划的消费（Consumer）能力：你可以直接对 Arrow 数据执行序列化后的 Substrait 计划，也可以把 pyarrow 的计算表达式（Expression）序列化为 Substrait 消息在系统之间交换。读完本文，你将掌握run_query、serialize_expressions、deserialize_expressions、get_supported_functions等核心 API 的用法、参数细节、底层实现与常见错误处理，并能在自己的数据管道中落地 Substrait 互操作。

一、模块总览与能力边界

pyarrow.substrait是 PyArrow 对 Substrait 标准的一组封装，按功能可划分为三个部分（对应 docs/source/python/api/substrait.rst 的目录结构）：

类别	API	作用
查询执行（Query Execution）	`run_query`	执行一个序列化的 Substrait plan，并以`RecordBatchReader`流式读取结果
表达式序列化（Expression Serialization）	`BoundExpressions`、`serialize_expressions`、`deserialize_expressions`	在 pyarrow 计算表达式与 Substrait`ExtendedExpression`消息之间互相转换
工具（Utility）	`get_supported_functions`	返回当前引擎支持的 Substrait 函数 ID 列表

该模块的公开接口定义在 python/pyarrow/substrait.py，它从 Cython 扩展pyarrow._substrait导入四个符号：BoundExpressions、get_supported_functions、run_query、deserialize_expressions、serialize_expressions。如果导入失败，会抛出带有 "The pyarrow installation is not built with support for 'substrait'" 信息的ImportError——这提醒我们：Substrait 支持是一个可选的编译组件，并非所有预编译 PyArrow 都自带。

编译开关与安装前提

在源码构建时，Substrait 组件由 CMake 选项PYARROW_SUBSTRAIT控制。在 python/setup.py 中可以看到对应的命令行开关：

第 110 行定义了('with-substrait', None, 'build the Substrait extension')参数；
第 196 行将self.with_substrait初始化为None；
第 262 行通过append_cmake_component(self.with_substrait, 'PYARROW_SUBSTRAIT')把该选项透传给 CMake。

因此从源码构建 PyArrow 时，需要通过--with-substrait显式开启，例如：

python setup.py build_ext --with-substrait

在 docs/source/python/install.rst 的组件支持矩阵中可以看到，Substrait 功能由libarrow-substrait库提供，在基于 conda 和 wheel 的安装渠道上默认可用（表格中对应列均为 ✓），但使用前仍建议用import pyarrow.substrait验证当前环境是否具备该能力。

二、查询执行：run_query

run_query(plan, *, table_provider=None, use_threads=True)是模块中最重要的入口，它的职责是：接收一个序列化的 Substrait plan，在本地执行引擎（Acero）上运行，并以RecordBatchReader形式返回结果。其完整签名与文档字符串位于 python/pyarrow/_substrait.pyx。

参数说明

参数	类型	默认值	说明
`plan`	`Buffer`或`bytes`	必填	序列化后的 Substrait plan。其他类型（如`int`）会抛出`TypeError`
`table_provider`	可调用对象	`None`	用于将 plan 中的`NamedTable`关系解析为`pyarrow.Table`的函数；不传则只能执行基于本地文件（`local_files`）的读取计划
`use_threads`	`bool`	`True`	是否使用多线程执行；为`False`时所有 CPU 密集工作都在调用线程上串行完成

table_provider回调接收两个参数：

names：list[str]，NamedTable 的名称路径（如["t1"]）；
schema：pyarrow.Schema，计划中声明的期望 schema。

回调必须返回一个pyarrow.Table。该函数名必须能解析到实际数据，否则执行会失败（见下文错误处理）。从 Cython 实现看，该回调通过_create_named_table_provider被包装为 C++ 层的NamedTableProvider，最终在ExecuteSerializedPlan调用中通过named_table_provider配置传入（见 python/pyarrow/_substrait.pyx）。

完整示例：从 NamedTable 读取

这是模块文档字符串中的官方示例（python/pyarrow/_substrait.pyx）：

import pyarrow as pa from pyarrow.lib import tobytes import pyarrow.substrait as substrait # 准备两张表，供 table_provider 按名字返回 test_table_1 = pa.Table.from_pydict({"x": [1, 2, 3]}) test_table_2 = pa.Table.from_pydict({"x": [4, 5, 6]}) def table_provider(names, schema): if not names: raise Exception("No names provided") elif names[0] == "t1": return test_table_1 elif names[1] == "t2": return test_table_2 else: raise Exception("Unrecognized table name") substrait_query = ''' { "relations": [ {"rel": { "read": { "base_schema": { "struct": { "types": [ {"i64": {}} ] }, "names": [ "x" ] }, "namedTable": { "names": ["t1"] } } }} ] } ''' # 将 JSON 形式的 plan 解析为等价的序列化 Protobuf buf = pa._substrait._parse_json_plan(tobytes(substrait_query)) # 执行计划并读取结果 reader = pa.substrait.run_query(buf, table_provider=table_provider) print(reader.read_all())

输出结果是一个pyarrow.Table，列x: int64，数据[[1, 2, 3]]，与test_table_1完全一致。

关于上面用到的_parse_json_plan：它并不是pyarrow.substrait的公开 API，而是一个模块内辅助函数（定义于 python/pyarrow/_substrait.pyx），其作用是把 Substrait 的 JSON 表示转换成等价的 Protobuf 二进制Buffer。正式场景中，更常见的做法是直接使用其他系统（如 Java 端）序列化好的 protobuf plan，再以Buffer形式传入run_query。

完整示例：从本地 Arrow 文件读取

如果不提供table_provider，则 plan 中的read关系必须使用local_files指明数据文件位置。测试用例 test_run_serialized_query 展示了这一用法：

import os import pathlib import pyarrow as pa from pyarrow.lib import tobytes table = pa.table([[1, 2, 3, 4, 5]], names=['foo']) # 将表写入 IPC 文件（RecordBatchFileWriter） path = os.path.join(tmpdir, "read_data.arrow") with pa.ipc.RecordBatchFileWriter(path, schema=table.schema) as writer: writer.write_table(table) substrait_query = """ { "version": { "major": 9999 }, "relations": [ {"rel": { "read": { "base_schema": { "struct": { "types": [{"i64": {}}] }, "names": ["foo"] }, "local_files": { "items": [ { "uri_file": "FILE:///path/to/read_data.arrow", "arrow": {} } ] } } }} ] } """ query = tobytes(substrait_query.replace( "FILE:///path/to/read_data.arrow", pathlib.Path(path).as_uri())) buf = pa._substrait._parse_json_plan(query) reader = substrait.run_query(buf, use_threads=False) res_tb = reader.read_all() assert table.select(["foo"]) == res_tb.select(["foo"])

注意这里的"arrow": {}声明了文件为 Arrow IPC 格式。需要说明的是：C++ 层（Acero Substrait consumer，见 docs/source/cpp/acero/substrait.rst）目前只支持 parquet 与 arrow 两种文件格式，且所有 URI 必须使用filescheme；projection、partition_index、start、length等属性暂不支持。计划中若包含不支持的属性或关系，整个计划会被拒绝执行。

多线程与执行模型

use_threads参数最终以c_bool形式传入 C++ 的ExecuteSerializedPlan（python/pyarrow/_substrait.pyx）：

with nogil: c_res_reader = ExecuteSerializedPlan( deref(c_buf_plan), default_extension_id_registry(), GetFunctionRegistry(), c_conversion_options, c_use_threads)

这里可以看到计划执行同时依赖两个关键组件：

default_extension_id_registry()：默认的扩展函数 ID 注册表，用于把 Substrait 计划中的函数 ID（如functions_arithmetic.yaml#add）映射到具体实现；
GetFunctionRegistry()：pyarrow 自身的函数注册表。

在多线程模式下，引擎会并行处理可并行的工作；use_threads=False则退化为单线程串行，便于调试和保证执行顺序的可复现性（测试用例 test_hash_aggregate_udf_basic 的注释正是"Ordering of k is deterministic because this is running with serial execution"）。

常见错误与边界

非法 plan 输入类型：run_query只接受pyarrow.Buffer或bytes。测试 test_run_query_input_types 验证：传入int会抛出TypeError: Expected 'pyarrow.Buffer' or bytes, got '<class 'int'>'；而传入合法类型但内容是无效 protobuf 时，会抛出OSError，错误信息为ParseFromZeroCopyStream failed for substrait.Plan。
空计划：如果 plan 没有任何 relation，执行时会报ArrowInvalid: Plan has no relations（见 test_invalid_plan）。
NamedTable 名称无法解析：当table_provider无法识别表名时抛出ArrowInvalid: Invalid NamedTable Source；当名称列表为空时抛出ArrowInvalid: names for NamedTable not provided（分别对应 test_named_table_invalid_table_name 与 test_named_table_empty_names）。
未知函数：如果计划引用了未注册的函数，会抛出ArrowKeyError，信息形如No function registered（见 test_udf_via_substrait_wrong_udf_name）。

三、表达式序列化：serialize_expressions 与 deserialize_expressions

Substrait 表达式是"绑定"（bound）到 schema 的：同样的a + b，在a:i32与a:i64两种 schema 下是两条不同的 Substrait 表达式；而 pyarrow 的Expression通常是"未绑定"的。因此序列化表达式时必须同时提供 schema。

serialize_expressions

签名（python/pyarrow/_substrait.pyx）：

serialize_expressions(exprs, names, schema, *, allow_arrow_extensions=False)

参数	类型	说明
`exprs`	`list[Expression]`	待序列化的 pyarrow 计算表达式
`names`	`list[str]`	与`exprs`一一对应的表达式名称（长度必须一致）
`schema`	`Schema`	表达式将要绑定到的 schema
`allow_arrow_extensions`	`bool`，默认`False`	为`False`时只允许核心 Substrait 函数定义；为`True`时允许 pyarrow 特有函数与用户自定义函数，但结果可能不被其他计算引擎接受

返回一个Buffer，内容为 Substrait 的ExtendedExpression消息。校验逻辑如下（Cython 层）：

len(exprs) != len(names)时抛出ValueError: exprs and names need to have the same length；
exprs中存在非Expression对象时抛出TypeError；
表达式引用了 schema 中不存在的字段时，序列化失败并抛出ValueError: No match for FieldRef（见 test_invalid_expression_ser_des）。

基本用法：

import pyarrow as pa import pyarrow.compute as pc import pyarrow.substrait as substrait schema = pa.schema([ pa.field("x", pa.int32()), pa.field("y", pa.int32()) ]) # 构造一个未绑定表达式：x == 7 expr = pc.equal(pc.field("x"), 7) buf = substrait.serialize_expressions([expr], ["test_expr"], schema)

deserialize_expressions 与 BoundExpressions

deserialize_expressions(buf) -> BoundExpressions

接受Buffer或bytes，将ExtendedExpression消息还原为一个BoundExpressions对象。BoundExpressions是一个只读视图类（不可直接实例化，见 python/pyarrow/_substrait.pyx），它对应 Substrait 的ExtendedExpression消息，暴露两个属性：

schema：所有表达式共同绑定到的pyarrow.Schema；
expressions：一个字典，键为表达式名称，值为pyarrow.compute.Expression。

完整往返（round-trip）示例：

returned = substrait.deserialize_expressions(buf) assert schema == returned.schema assert len(returned.expressions) == 1 assert "test_expr" in returned.expressions

多表达式序列化同样受支持：传入多个表达式时，返回的expressions字典会包含全部名称到表达式的映射（见 test_serializing_multiple_expressions）。

类型覆盖：Arrow 特有类型与"单向"类型

测试 test_arrow_specific_types 验证了以下 Arrow 类型可以被序列化并完整往返（schema 保持完全一致）：time32("s")、time32("ms")、time64("ns")、date64()、large_string、large_binary。

而测试 test_arrow_one_way_types 揭示了一类重要的行为：binary_view、string_view、dictionary(int32, string)、run_end_encoded(int32, string)这四种类型可以单向序列化，但反序列化回来后 schema 中的类型会被"降级"为对应的普通类型（binary、string、string、string）。也就是说：

对于time32/time64/date64/large_*等类型，往返是无损的；
对于*_view、dictionary、run_end_encoded等新类型，往返是有损的（类型被映射为 Substrait 能表达的基础类型）。

这在做跨引擎互操作时需要特别留意：不要假设 schema 一定能原样往返。

在 pyarrow.compute 中的便捷入口

pyarrow.compute.Expression也提供了两个便捷方法（定义于 python/pyarrow/_compute.pyx）：

Expression.to_substrait(schema, allow_arrow_extensions=False)：把单个表达式序列化为ExtendedExpression消息，表达式固定命名为"expression"；
Expression.from_substrait(buffer)：从消息中反序列化出单个Expression。它内部调用deserialize_expressions后只取唯一表达式——如果消息包含 0 个表达式会抛ValueError: Substrait message did not contain any expressions，包含多个表达式会抛ValueError: Substrait message contained multiple expressions. Use pyarrow.substrait.deserialize_expressions instead；表达式名称与绑定 schema 会被忽略。

# 单表达式便捷入口 buf = expr.to_substrait(schema) expr2 = pc.Expression.from_substrait(buf) assert str(expr2) == str(pc.equal(pc.field(0), 7)) # 字段引用被归一化为位置索引

测试 test_serializing_with_compute 还验证了反序列化时表达式名称不影响结果（"weirdname"同样可被from_substrait接受）。

自定义函数与 allow_arrow_extensions

Substrait 有标准函数定义集（functions_arithmetic.yaml、functions_comparison.yaml等）。当一个表达式使用不在 Substrait 标准函数集中的函数时（例如shift_left，它是 pyarrow 特有函数），默认会抛出ArrowNotImplementedError；只有设置allow_arrow_extensions=True才能成功序列化（见 test_serializing_udfs）：

schema = pa.schema([pa.field("x", pa.uint32())]) exprs = [pc.shift_left(pc.scalar(10), pc.scalar(4))] # 默认失败 with pytest.raises(ArrowNotImplementedError): pa.substrait.serialize_expressions(exprs, ["expr"], schema) # 允许 Arrow 扩展后成功，且能完整往返 buf = pa.substrait.serialize_expressions( exprs, ["expr"], schema, allow_arrow_extensions=True) returned = pa.substrait.deserialize_expressions(buf) assert str(returned.expressions["expr"]) == str(exprs[0])

这里"UDF"一词在 Substrait 语境中指的是"Substrait 无法识别的函数"，它可能是 pyarrow 的内置函数而非用户自定义函数。需要权衡的是：开启扩展后，生成的ExtendedExpression消息可能不被其他 Substrait 引擎接受，因此仅在确知目标端支持时才应开启。

四、工具函数：get_supported_functions

get_supported_functions()返回当前底层引擎支持的 Substrait 函数 ID 列表，每个 ID 的格式为'{uri}#{name}'（python/pyarrow/_substrait.pyx），例如.../functions_arithmetic.yaml#add。实现上它直接调用 C++ 层default_extension_id_registry().GetSupportedSubstraitFunctions()。

import pyarrow as pa supported = pa._substrait.get_supported_functions() # 或通过公开模块 import pyarrow.substrait as substrait supported = substrait.get_supported_functions() # 检查某个函数是否受支持（匹配 URI#name 后缀） def has_function(fns, ext_file, fn_name): suffix = f'{ext_file}#{fn_name}' return any(fn.endswith(suffix) for fn in fns) assert has_function(supported, 'functions_arithmetic.yaml', 'add') assert has_function(supported, 'functions_arithmetic.yaml', 'sum')

测试 test_get_supported_functions 展示了它的典型用法：抽样验证"一个非聚合函数add和一个聚合函数sum"均被支持。在实际工程中，这个函数适合用于：

可行性预检：在序列化表达式或构造计划前，先确认用到的函数在目标引擎中可用；
函数映射诊断：排查No function registered类错误时，对照支持列表定位差异。

五、执行计划（Plan）的更多细节与已知限制

虽然 Python 侧 API 只有几个函数，但计划本身的能力边界由 C++ 引擎（Acero 的 Substrait consumer）决定，详见 docs/source/cpp/acero/substrait.rst。以下是需要了解的关键约束：

版本基线：当前 consumer 基于 Substrait 0.20.0，不支持的更新特性会导致计划被拒绝；由于 0.20.0 存在破坏性变更，早于 0.20.0 的旧计划也会被拒绝。
计划结构：一个 plan 应只有单个顶层 relation。
扩展：包含扩展类型变体（extension type variations）的计划会被拒绝；更高级的扩展可通过自定义arrow::engine::ExtensionProvider提供（这是 C++ 层能力）。
Read 关系：projection属性不支持；VirtualTable与ExtensionTable读取类型不支持；只支持 parquet 与 arrow 文件格式；所有 URI 必须使用filescheme；partition_index、start、length不支持非默认值。此外，Acero 对 read 关系中的 filter 只做投影下推，可能无法完全满足 filter，官方建议在 read 之后额外附加一个相同 filter 的 filter 关系。
Join 关系：不支持JOIN_TYPE_SINGLE；连接表达式必须是equal或is_not_distinct_from函数调用且参数为直接引用；只支持单个连接键；post_join_filter会被忽略。
Aggregate 关系：最多一个 grouping set；每个分组表达式与每个 measure 的参数必须是直接引用；measure 不能带 filter 或 sort；invocation 必须是AGGREGATION_INVOCATION_ALL或未指定；phase 必须是AGGREGATION_PHASE_INITIAL_TO_RESULT。

这些限制意味着：在生成 Substrait 计划时（无论使用 JSON 手写、还是借助其他工具生成），都应先对照上述约束，否则整个计划会在执行阶段被整体拒绝。另外，Python 测试中大量使用了"version": { "major": 9999 }这类"宽松版本号"，因为当前实现对版本校验相对宽容，但正式环境仍建议声明明确的 Substrait 版本。

六、端到端综合示例：UDF 投影计划

下面综合前面所有知识点，展示一个通过 Substrait 计划调用 pyarrow 自定义一元函数（UDF）并投影出新列的完整流程，取自测试 test_udf_via_substrait 的思路：

import pyarrow as pa import pyarrow.substrait as substrait from pyarrow.lib import tobytes test_table = pa.Table.from_pydict({"x": [1, 2, 3]}) def table_provider(names, _): if not names: raise Exception("No names provided") elif names[0] == "t1": return test_table else: raise Exception("Unrecognized table name") # 计划结构：读命名表 t1 -> project 计算 y=x+1 -> root 输出列 x, y substrait_query = b""" { "extensionUris": [ {"extensionUriAnchor": 1}, {"extensionUriAnchor": 2, "uri": "urn:arrow:substrait_simple_extension_function"} ], "extensions": [ { "extensionFunction": { "extensionUriReference": 2, "functionAnchor": 1, "name": "y=x+1" } } ], "relations": [ { "root": { "input": { "project": { "common": {"emit": {"outputMapping": [1, 2]}}, "input": { "read": { "baseSchema": { "names": ["t"], "struct": { "types": [{"i64": {"nullability": "NULLABILITY_REQUIRED"}}], "nullability": "NULLABILITY_REQUIRED" } }, "namedTable": {"names": ["t1"]} } }, "expressions": [ {"selection": {"directReference": {"structField": {}}, "rootReference": {}}}, { "scalarFunction": { "functionReference": 1, "outputType": {"i64": {"nullability": "NULLABILITY_NULLABLE"}}, "arguments": [ {"value": {"selection": {"directReference": {"structField": {}}, "rootReference": {}}}} ] } } ] } }, "names": ["x", "y"] } } ] } """ buf = pa._substrait._parse_json_plan(substrait_query) reader = pa.substrait.run_query(buf, table_provider=table_provider, use_threads=False) res_tb = reader.read_all() print(res_tb)

计划要点拆解：

extensionUris/extensions：声明扩展函数 URI（urn:arrow:substrait_simple_extension_function）与函数锚点（functionAnchor: 1，名y=x+1），这是 Substrait 扩展自定义函数的通用机制；
emit.outputMapping: [1, 2]：project 输出表达式索引 1（原始列t）与 2（计算列）；
scalarFunction.functionReference: 1：引用前面声明的函数锚点；
root.names: ["x", "y"]：最终输出列名。

若该函数名与已注册的 UDF 不匹配，run_query会抛出ArrowKeyError（No function registered）；若函数未在计划中正确声明扩展，则可能触发ArrowNotImplementedError。

七、最佳实践与注意事项小结

先验证构建能力：使用import pyarrow.substrait确认当前安装已包含该组件；源码构建记得加--with-substrait。
区分 plan 与表达式消息：run_query消费的是 SubstraitPlan消息（关系树）；serialize_expressions/deserialize_expressions处理的是ExtendedExpression消息（表达式集合 + 绑定 schema）。两者不可混用。
表达式必须绑定 schema：序列化前务必提供准确的输入 schema，字段名拼错或类型不符会直接失败。
留意往返有损类型：binary_view、string_view、dictionary、run_end_encoded等类型反序列化后会被降级为基础类型；对 schema 保真度有要求时需额外校验。
谨慎使用allow_arrow_extensions：仅在与本端（pyarrow）互操作或目标端已知支持时开启，否则生成的消息可能无法被其他 Substrait 引擎消费。
对照 Acero 的 conformance 限制生成计划：版本、文件格式、URI scheme、join/aggregate 属性等约束若不满足，整个计划会被拒绝。
利用get_supported_functions做预检：在跨引擎场景中提前确认函数可用性，避免运行期才发现No function registered。
善用use_threads=False做调试：串行执行下结果顺序确定，便于断言与排查。

八、延伸阅读

Substrait API 参考文档：本文对应的官方 API 文档页；
pyarrow.substrait 模块源码：公开接口导入层；
Substrait Cython 实现：run_query、表达式序列化与get_supported_functions的底层绑定；
Substrait 测试套件：覆盖查询执行、表达式往返、UDF、错误处理等全部行为的权威用例；
Using Acero with Substrait：C++ 引擎层的 Substrait 兼容性（conformance）细节；
Python 安装说明中的组件矩阵：确认pyarrow.substrait在当前渠道的可用性。

通过本文介绍的 API 组合，你可以在 PyArrow 与任何符合 Substrait 标准的数据系统之间，以标准化的方式交换并执行查询计划与计算表达式，从而在保持计算能力的同时避免被单一执行引擎锁定。

【免费下载链接】arrowApache Arrow is a multi-language toolbox for accelerated data interchange and in-memory processing项目地址: https://gitcode.com/gh_mirrors/arrow13/arrow

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考