PyArrow Substrait API 实战：在 Arrow 数据上执行 Substrait 查询计划并互算表达式

PyArrow Substrait API 实战在 Arrow 数据上执行 Substrait 查询计划并互算表达式【免费下载链接】arrowApache Arrow is the universal columnar format and multi-language toolbox for fast data interchange and in-memory analytics项目地址: https://gitcode.com/GitHub_Trending/arrow3/arrow本文基于 Apache Arrow 仓库的 Python API 参考文档 docs/source/python/api/substrait.rst系统讲解pyarrow.substrait模块提供的 8 个公开接口如何把序列化的 Substrait 计划交给run_query在 Arrow 表上执行、如何用serialize_expressions/deserialize_expressions在 SubstraitExtendedExpression与 PyArrow 计算表达式之间互转、如何用serialize_schema/deserialize_schema处理 Schema以及如何通过get_supported_functions查询引擎支持的函数清单。读完本文你可以独立完成 Substrait 计划的执行、表达式/Schema 的序列化往返并理解其背后的 C 引擎实现与构建前提。模块定位与构建前提pyarrow.substrait是 PyArrow 面向 Substrait 查询交换标准的接口层。Substrait 采用 Protobuf 编码描述关系代数读取、投影、过滤、聚合等PyArrow 的 C 引擎acero/substrait 组件负责解析并执行这些计划。模块公开了如下三类 API与参考文档中的 autosummary 分组一一对应分组API作用Query Executionrun_query执行序列化的 Substrait 计划返回RecordBatchReaderExpression SerializationBoundExpressions、serialize_expressions、deserialize_expressions、serialize_schema、deserialize_schema、SubstraitSchemaPyArrow 计算表达式与 Schema 和 Substrait 消息之间的序列化/反序列化Utilityget_supported_functions列出底层引擎支持的 Substrait 函数 ID构建与安装层面有两点需要注意可选组件Substrait 支持由 CMake 构建开关ARROW_SUBSTRAIT控制定义见 cpp/cmake_modules/DefineOptions.cmake第三方依赖拉取逻辑见 cpp/cmake_modules/ThirdpartyToolchain.cmake。若当前安装未启用该组件导入pyarrow.substrait会抛出带明确提示的ImportError见 python/pyarrow/substrait.pyraise ImportError( The pyarrow installation is not built with support ffor substrait ({str(exc)}) )可选的 Python 包substrait_substrait.pyx会尝试import substrait即 substrait-python 包但它只在SubstraitSchema.to_pysubstrait()中用到属于可选项未安装时不影响其余功能见 python/pyarrow/_substrait.pyx。模块内部实现全部由 Cython 层 python/pyarrow/_substrait.pyx 完成并通过 python/pyarrow/substrait.py 重新导出 8 个公开符号run_query、BoundExpressions、get_supported_functions、serialize_expressions、deserialize_expressions、serialize_schema、deserialize_schema、SubstraitSchema。run_query执行序列化的 Substrait 计划run_query的函数签名与参数语义如下源码 docstringpython/pyarrow/_substrait.pyxrun_query(plan, *, table_providerNone, use_threadsTrue)参数类型说明planBuffer或bytes序列化的 Substrait 计划binary protobuf 格式table_provider可调用对象可选解析计划中NamedTable关系。回调接收两个参数表名列表list[str]与期望的pyarrow.Schema须返回一个pyarrow.Tableuse_threadsbool默认TrueTrue时使用多线程执行False时所有 CPU 密集工作都在调用线程上完成返回值为RecordBatchReader用reader.read_all()即可拿到结果Table。最小可运行示例NamedTable 场景以下示例完整继承自源码 docstring展示了用 JSON 形式的 Substrait 计划读取名为t1的表执行查询的全过程 import pyarrow as pa from pyarrow.lib import tobytes import pyarrow.substrait as substrait 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] ... } ... } ... }} ... ] ... } ... buf pa._substrait._parse_json_plan(tobytes(substrait_query)) reader pa.substrait.run_query(buf, table_providertable_provider) reader.read_all() pyarrow.Table x: int64 ---- x: [[1,2,3]]几个值得注意的细节_parse_json_plan是测试辅助函数。run_query只接受 binary protobufdocstring 中使用的pa._substrait._parse_json_plan是 _substrait.pyx 中的私有辅助函数其 docstring 明确写着 Parse a JSON plan into equivalent serialized Protobuf输入 JSON 字节串、输出Buffer。在真实场景中通常由计划生成端编译器或 substrait-python直接产出 Protobuf 字节串。table_provider的回调协议C 侧通过 Cython 辅助函数_create_named_table_providerpython/pyarrow/_substrait.pyx把 Python 回调绑定为CNamedTableProvider它把计划中的表名序列与base_schema转成 Python 对象后调用你的回调再把返回的pyarrow.Table包装成table_source声明交给执行引擎。use_threadsFalse保证串行确定性测试 test_substrait.py 中的哈希聚合用例特意以use_threadsFalse运行注释说明 Ordering of k is deterministic because this is running with serial execution——如果你的计划输出顺序对调试很重要可关闭线程。本地文件读取local_files除了NamedTable计划还可以用local_files指定磁盘上的 Arrow IPC 文件。测试 test_substrait.py 展示了完整流程先把数据写成 Arrow 文件再把文件 URI 填进 JSON 计划最后run_query时不需要table_providersubstrait_query { version: { major: 9999 }, relations: [ {rel: { read: { base_schema: { struct: { types: [ {i64: {}} ] }, names: [ foo ] }, local_files: { items: [ { uri_file: FILENAME_PLACEHOLDER, arrow: {} } ] } } }} ] } file_name read_data.arrow table pa.table([[1, 2, 3, 4, 5]], names[foo]) path _write_dummy_data_to_disk(tmpdir, file_name, table) # 写 Arrow IPC 文件 query tobytes(substrait_query.replace( FILENAME_PLACEHOLDER, pathlib.Path(path).as_uri())) buf pa._substrait._parse_json_plan(query) reader substrait.run_query(buf, use_threadsuse_threads) res_tb reader.read_all() assert table.select([foo]) res_tb.select([foo])测试 test_binary_conversion_with_json_options 还验证了在local_files条目上附加metadata如created_by等 JSON 扩展字段不影响二进制转换。错误处理边界测试文件 python/pyarrow/tests/test_substrait.py 覆盖了run_query的主要失败路径可作排错参考场景异常与消息测试位置plan传入非法类型如intTypeError: Expected pyarrow.Buffer or bytes, got ...L96-L104字节串不是合法 ProtobufArrowInvalid: ParseFromZeroCopyStream failed for substrait.PlanL106-L109计划缺少relationsArrowInvalid: Plan has no relationsL112-L122NamedTable名字无法被 provider 识别ArrowInvalid: Invalid NamedTable SourceL283-L281NamedTable的names为空ArrowInvalid: names for NamedTable not providedL284-L323扩展函数名未注册UDF 写错ArrowKeyError: No function registered ...L563-L565执行链路从 Python 到 C 引擎run_query的实现python/pyarrow/_substrait.pyx流程为校验plan类型并解包为 CBuffer→ 若提供table_provider通过BindFunction绑定命名表回调写入CConversionOptions→ 在nogil块中调用 C 层ExecuteSerializedPlan。该 C API 声明于 cpp/src/arrow/engine/substrait/util.h其完整参数集比 Python 暴露的更丰富ARROW_ENGINE_EXPORT Resultstd::shared_ptrRecordBatchReader ExecuteSerializedPlan( const Buffer substrait_buffer, const ExtensionIdRegistry* registry NULLPTR, compute::FunctionRegistry* func_registry NULLPTR, const ConversionOptions conversion_options {}, bool use_threads true, MemoryPool* memory_pool default_memory_pool());Python 层固定传入默认扩展 ID 注册表与GetFunctionRegistry()而registry、memory_pool等参数在 C 侧仍可供原生调用方定制。用户定义函数UDF经 Substrait 调用测试 test_udf_via_substrait 展示了一个进阶场景通过计划的extensionUris/extensions声明扩展函数如uri: urn:arrow:substrait_simple_extension_function、name: yx1在project节点的scalarFunction中以functionReference引用即可在执行时对每行调用注册的 Python UDF。Arrow 还内置了扩展关系类型测试 test_scalar_aggregate_udf_basic 中的type: /arrow.substrait_ext.SegmentedAggregateRel分区分段聚合以及 test_hash_aggregate_udf_basic 中带groupingKeys的哈希聚合属于 Arrow 对 Substrait 的扩展而非核心标准跨引擎使用需谨慎。表达式序列化BoundExpressions、serialize_expressions 与 deserialize_expressions为什么必须绑定 Schemaserialize_expressions 的 docstring 解释了这一组 API 的核心概念Substrait 表达式必须绑定到 Schema。例如 Substrait 表达式a:i32 b:i32与a:i64 b:i64是不同的表达式而 PyArrow 表达式通常是未绑定的两者都写作a b。因此序列化表达式时必须提供 Schema且当找不到匹配的函数调用时序列化可能失败。函数签名与参数serialize_expressions(exprs, names, schema, *, allow_arrow_extensionsFalse)参数说明exprslist[Expression]待序列化的 PyArrow 计算表达式长度必须与names一致元素必须是Expression实例见 L336-L342 的校验逻辑nameslist[str]各表达式的名称schema表达式所绑定的pyarrow.Schemaallow_arrow_extensionsbool默认False。False时仅允许核心 Substrait 函数定义中的函数True时允许 PyArrow 专有函数与用户定义函数但结果可能不被其他计算引擎接受返回值为包含ExtendedExpression消息的Buffer。反序列化方向由 deserialize_expressions 完成deserialize_expressions(buf) # buf: Buffer 或 bytes - BoundExpressionsBoundExpressionspython/pyarrow/_substrait.pyx是抽象容器类等价于 Substrait 的ExtendedExpression消息提供schema属性所有表达式共同绑定的 Schemaexpressions属性{名称: Expression}字典类方法from_substrait(message)接受Buffer、bytes或带SerializeToString()的 Protobuf 消息对象见测试 test_bound_expression_from_Message。往返示例继承自测试最典型的用法是序列化—反序列化往返验证test_serializing_expressionsimport pyarrow as pa import pyarrow.compute as pc schema pa.schema([ pa.field(x, pa.int32()), pa.field(y, pa.int32()) ]) expr pc.equal(pc.field(x), 7) buf pa.substrait.serialize_expressions([expr], [test_expr], schema) returned pa.substrait.deserialize_expressions(buf) assert schema returned.schema assert test_expr in returned.expressions多个表达式可一次性序列化反序列化后名称保持对应test_serializing_multiple_expressions。一个值得注意的现象反序列化回来的表达式中字段引用会归一化为按序号引用——pc.field(x)往返后打印为pc.field(0)测试中以norm_exprs断言了这一点。pyarrow.compute.Expression还内建了与 Substrait 互通的便捷方法expr.to_substrait(schema)序列化、pc.Expression.from_substrait(buf)反序列化见 test_serializing_with_compute。后者有两个限制只能处理单表达式消息多表达式会抛ValueError: ... contained multiple expressions且不依赖表达式序列化名。类型覆盖与 UDF 序列化边界测试文件对类型互转边界做了明确验证可双向往返的 Arrow 特有类型test_arrow_specific_typestime32(s/ms)、time64(ns)、date64、large_string、large_binary等在默认allow_arrow_extensionsFalse下即可往返仅单向的类型test_arrow_one_way_typesbinary_view/string_view/dictionary/run_end_encoded序列化后往返时 Schema 会降级为binary/string/string/string非标准函数的 UDF 序列化test_serializing_udfs像pc.shift_left(a, b)这类未被 Substrait 核心函数集识别的函数默认序列化会抛ArrowNotImplementedError加allow_arrow_extensionsTrue后成功往返。失败路径方面test_invalid_expression_ser_desexprs与names长度不一致抛ValueErrorneed to have the same length表达式引用了 Schema 中不存在的字段抛ValueErrorNo match for FieldRef。Schema 序列化serialize_schema、deserialize_schema 与 SubstraitSchemaSchema 层面提供三个 APIserialize_schema(schema)python/pyarrow/_substrait.pyx把pyarrow.Schema包装为SubstraitSchema对象。其实现同时产出两份表示schema属性CSerializeSchema产出的 SubstraitNamedStruct二进制见_serialize_namedstruct_schemaL241-L252expression属性以空表达式列表、allow_arrow_extensionsTrue调用的serialize_expressions结果即ExtendedExpression消息。之所以两者并存正如 SubstraitSchema 的 docstring 所述当 Schema 中使用了需要扩展信息才能解码的类型时ExtendedExpression形式base_schema 全部 extensions才能完整还原。SubstraitSchema还提供to_pysubstrait()方法把expression转成 substrait-python 的ExtendedExpression对象需要安装substrait包否则抛ImportError。deserialize_schema(buf)python/pyarrow/_substrait.pyx接受三种输入SubstraitSchema对象直接返回deserialize_expressions(buf.expression).schemabytes/memoryview/Buffer按 SubstraitNamedStruct消息调用 CDeserializeSchema其他类型抛TypeError: Expected pyarrow.Buffer or bytes, got ...。往返测试test_serializing_schema展示了全部三条路径substrait_schema b\n\x01x\n\x01y\x12\x0c\n\x04*\x02\x10\x01\n\x04b\x02\x10\x01 returned pa.substrait.deserialize_schema(substrait_schema) # - pa.schema([pa.field(x, pa.int32()), pa.field(y, pa.string())]) arrow_substrait_schema pa.substrait.serialize_schema(returned) assert arrow_substrait_schema.schema substrait_schema # NamedStruct 字节一致 assert pa.substrait.deserialize_schema(arrow_substrait_schema) returned assert pa.substrait.deserialize_schema(arrow_substrait_schema.schema) returned returned pa.substrait.deserialize_expressions(arrow_substrait_schema.expression) assert returned.schema expected_schemaget_supported_functions查询引擎能力边界get_supported_functions 返回底层引擎当前支持的 Substrait 函数 ID 列表每项编码为{uri}#{name}字符串。实现上它从默认扩展 ID 注册表取GetSupportedSubstraitFunctions()无需构造任何计划即可自省引擎能力——在做计划移植或兼容性检查时非常有用。测试 test_get_supported_functions 说明了一个现实约束Substrait 尚未为各标准函数文件敲定最终 URI因此测试只校验后缀如functions_arithmetic.yaml#add、functions_arithmetic.yaml#sum而不是完整 IDdef has_function(fns, ext_file, fn_name): suffix f{ext_file}#{fn_name} for fn in fns: if fn.endswith(suffix): return True return False版本兼容性与工程化使用Substrait 版本约束C 引擎头文件 cpp/src/arrow/engine/substrait/util.h 中硬编码了版本常量——当前对应 Substrait 版本为0.44.0最低支持版本为0.20CheckVersion会在计划校验阶段执行该检查。计划 JSON 里的version.major字段测试中统一填9999表示未来版本即参与此协商。测试与 CI 集成全部 Substrait 测试集中在 python/pyarrow/tests/test_substrait.py模块级标记pytestmark pytest.mark.substrait文件头注释说明未构建 Substrait 支持的环境可用pytest -m not substrait整体跳过涉及 numpy 的 UDF 聚合用例额外标记pytest.mark.numpy。典型集成路径小结计划来源为 Protobuf 字节串或Buffer时直接pa.substrait.run_query(buf, table_provider..., use_threads...)需要把 PyArrow 计算表达式过滤条件、投影表达式发给远端 Substrait 计划生成器时用serialize_expressions(exprs, names, schema)跨引擎场景保持allow_arrow_extensionsFalse接收远端传来的ExtendedExpression/NamedStruct消息时用BoundExpressions.from_substrait与deserialize_schema还原为Expression字典和pyarrow.Schema。小结pyarrow.substrait把 PyArrow 同时置于 Substrait 生态的两端作为执行端run_query借助 C 引擎ExecuteSerializedPlancpp/src/arrow/engine/substrait/util.h直接执行序列化的关系计划支持NamedTable回调供数与local_files直读并可用get_supported_functions自省函数能力作为互操作端serialize_expressions/deserialize_expressions/serialize_schema/deserialize_schema提供了带 Schema 绑定的表达式与 Schema 双向序列化并以allow_arrow_extensions开关和SubstraitSchema双表示设计处理 Arrow 专有类型的跨引擎兼容问题。所有行为边界类型降级、未注册函数、非法计划等均可在 python/pyarrow/tests/test_substrait.py 中找到对应的断言级证据可作为生产环境排错的可靠参照。【免费下载链接】arrowApache Arrow is the universal columnar format and multi-language toolbox for fast data interchange and in-memory analytics项目地址: https://gitcode.com/GitHub_Trending/arrow3/arrow创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考