用 Rust 为 Daft 构建第一个原生扩展：Hello Native Extension 实战指南

用 Rust 为 Daft 构建第一个原生扩展Hello Native Extension 实战指南【免费下载链接】DaftHigh-performance data engine for AI and multimodal workloads. Process images, audio, video, and structured data at any scale项目地址: https://gitcode.com/GitHub_Trending/da/Daft本指南以 Daft 仓库中的最小原生扩展示例 examples/hello 为核心完整讲解如何用 Rust 编写一个可被 Daft 分布式执行引擎加载的原生扩展Native ABI Extension包括标量函数、聚合函数的实现、Python 包装层、构建安装与测试全流程。读完本文你将掌握 Daft 扩展机制的底层原理Arrow C Data Interface ABI、Session 注册模型并能独立从零搭建一个可 pip 安装、可在 DataFrame 表达式中直接调用的原生扩展包。从 README 出发hello 示例解决什么问题仓库中 examples/hello/README.md 开宗明义这是一个minimal Daft native extension example最小 Daft 原生扩展示例其快速开始只有两步# Install the extension in the project .venv uv pip install -e . # Run the tests! pytest -vREADME 同时将读者引导至扩展指南。需要说明的是指南实际位于 docs/extensions/overview.md扩展模型总览与 docs/extensions/authoring.mdRust 原生扩展编写指南。根据 overview 的划分Daft 扩展有两条技术路线Python UDF 型扩展基于daft.func、daft.cls等自定义代码 API见 docs/custom-code/func.md适合编排 Python 生态、外部服务、ML 模型等场景无需编译原生库原生 ABI 型扩展以共享库shared library为载体基于 Arrow C Data Interface 的稳定 C ABI获得底层的向量化原生执行性能。hello示例正是这条路线的最小实现。项目骨架Rust 工程与 Python 包如何组织从仓库目录结构可以还原出hello扩展的完整布局它由三部分拼装而成examples/hello/ ├── Cargo.toml # Rust 工程配置cdylib 产物 ├── pyproject.toml # Python 打包配置setuptools-rust ├── setup.py # RustExtension 接线文件 ├── hello/ │ ├── __init__.py # Python 表达式包装层 │ └── py.typed # PEP 561 类型标记 ├── src/ │ └── lib.rs # Rust 原生实现模块 函数 └── tests/ └── test_hello.py # pytest 测试Cargo.toml必须以 cdylib 编译examples/hello/Cargo.toml 的关键点在于crate-type [cdylib]——只有动态库才能在运行时被dlopen加载进 Daft 进程[workspace] [package] name hello edition 2024 version 0.1.0 [lib] name hello crate-type [cdylib] [dependencies] daft-ext {path ../../src/daft-ext, features [arrow-58]} arrow-array {version 58, features [chrono-tz]} arrow-schema 58依赖方面有两点值得注意daft-ext是官方扩展 SDK位于 src/daft-ext提供类型、trait 与宏这里以仓库相对路径path方式引用正式发布时应改为版本号依赖feature flagarrow-58必须与你的 arrow-rs 版本匹配。ABI 边界使用纯 C 结构体ArrowSchema、ArrowArray因此扩展并不绑定 Daft 内部的 arrow-rs 版本在daft-ext上启用与自身 arrow-rs 一致的特性arrow-56、arrow-57或arrow-58即可获得安全的.into()转换不支持的版本则可使用from_owned/into_owned/from_raw/as_raw逃生舱接口。pyproject.tomlsetuptools-rust 构建系统examples/hello/pyproject.toml 使用setuptoolssetuptools-rust作为构建后端把 Rust cdylib 的编译集成进 Python 打包流程[build-system] requires [setuptools, setuptools-rust] build-backend setuptools.build_meta [project] name hello version 0.1.0 requires-python 3.10 dependencies [daft] [project.optional-dependencies] test [pytest]包本身依赖daft运行时在仓库内开发时通过[tool.uv.sources]将其指向仓库根目录的可编辑安装[tool.uv.sources] daft {path ../.., editable true}setup.py把编译产物装进 Python 包目录examples/hello/setup.py 是 Python 与 Rust 的接线文件from setuptools import find_packages, setup from setuptools_rust import Binding, RustExtension setup( packagesfind_packages(), rust_extensions[ RustExtension( hello.libhello, pathCargo.toml, bindingBinding.NoBinding, stripTrue, ) ], )RustExtension(hello.libhello, ...)的含义是把编译出的libhello.so放进hello/包目录内这样Session.load_extension才能定位到它bindingBinding.NoBinding是因为 Daft 扩展导出的是裸 C 符号daft_module_magic而非 PyO3 绑定stripTrue用于裁剪符号表减小体积。扩展入口#[daft_extension]模块与 install 钩子原生扩展的 Rust 侧入口在 examples/hello/src/lib.rs 中。一个扩展由两部分组成模块module入口点与一个或多个函数标量或聚合。use std::{ffi::CStr, sync::Arc}; use arrow_array::{Array, ArrayRef}; use arrow_schema::{DataType, Field}; use daft_ext::{daft_extension, prelude::*}; // ── Module ────────────────────────────────────────────────────────── #[daft_extension] struct HelloExtension; impl DaftExtension for HelloExtension { fn install(session: mut dyn DaftSession) { session.define_function(Arc::new(Greet)); session.define_aggregate_function(Arc::new(StringCount)); } }机制上详见 docs/extensions/authoring.md#[daft_extension]宏会生成daft_module_magic这个 C 符号Daft 运行时dlopen加载共享库时就是靠它发现扩展入口的impl DaftExtension中的install(session: mut dyn DaftSession)是扩展安装钩子在扩展被加载进某个会话时调用一次所有函数标量用define_function、聚合用define_aggregate_function都在这里注册函数名在会话内是全局的authoring 指南建议使用前缀如myext_greet避免多扩展共存时冲突。标量函数greet 的 Rust 实现当前仓库的hello示例使用#[daft_func]宏实现标量函数写法非常简洁// ── Scalar Function ──────────────────────────────────────────────── #[daft_func] fn greet(name: str) - String { format!(Hello, {}!, name) }authoring 指南则展示了更底层的等价写法每个标量函数是一个实现DaftScalarFunctiontrait 的结构体包含三个方法name(self) - CStr函数名Python 侧通过它查找注意使用c...字符串字面量return_field(self, args: [ArrowSchema]) - DaftResultArrowSchema类型检查与输出类型声明。输入字段以 C Data Interface 的ArrowSchema形式传入可用.as_raw()零拷贝借用为 arrow-rs 的FFI_ArrowSchema做校验如检查参数个数、是否字符串类型再.into()转换回 ABI 类型返回call(self, args: VecArrowData) - DaftResultArrowData真正的执行逻辑接收整列ArrowData输出整列结果——所有数据都走 Arrow 数组没有任何逐行 Python 开销。fn call(self, args: VecArrowData) - DaftResultArrowData { // 1. 取出输入列并转换为 arrow-rs FFI 类型 // 2. 用 StringBuilder 逐行构建输出null 也要保留 // 3. 通过 arrow::ffi::to_ffi 转回 ABI 类型返回 }实现中有两个容易踩坑的点authoring 指南明确提醒字符串类型必须是LargeUtf8i64 偏移Daft 内部字符串使用 i64 offsets向下转型时必须用as_string::i64()用i32会在运行时 panicreturn_field做类型校验时也要接受DataType::LargeUtf8错误分类return_field中的 schema 违规返回Err(DaftError::TypeError(...))call中的执行失败返回Err(DaftError::RuntimeError(...))。聚合函数string_count 的三阶段管线除了标量函数扩展还可以注册聚合函数UDAF。仓库中的StringCount统计非空字符串个数是一个完整的实现范例examples/hello/src/lib.rs。聚合函数遵循三阶段管线聚合aggregate把输入数组加工成部分状态partial state合并combine把多个部分状态合并成一个终结finalize从合并后的状态产出最终标量结果。状态以VecArrowData交换每个元素对应一个状态字段FFI 层会把这些字段透明地打包成一个 Struct 数组。对应 trait 的实现要点struct StringCount; impl DaftAggregateFunction for StringCount { fn name(self) - CStr { cstring_count } // 输出类型声明1 个参数否则 TypeError fn return_field(self, args: [ArrowSchema]) - DaftResultArrowSchema { // ... 校验 args.len() 1返回 Int64 字段 string_count } // 中间状态字段声明单个 Int64 字段 count fn state_fields(self, _args: [ArrowSchema]) - DaftResultVecArrowSchema { // ... } // 阶段一统计输入列的非空个数产出单行状态 fn aggregate(self, inputs: VecArrowData) - DaftResultVecArrowData { let non_null_count input.len() - input.null_count(); // 返回 Int64Array 状态 } // 阶段二把多行状态求和 fn combine(self, states: VecArrowData) - DaftResultVecArrowData { // 遍历 counts 数组求和跳过 null } // 阶段三取出最终 count fn finalize(self, states: VecArrowData) - DaftResultArrowData { // ... } }register 时使用session.define_aggregate_function(Arc::new(StringCount))与标量函数并存于同一个install钩子中。这种三阶段设计天然适配分布式执行每个分区先在本地aggregate出部分状态再在 shuffle 后combine最后finalize产出结果。Python 包装层把原生函数接入 Expression DSL原生函数需要在 Python 侧提供符号才能在 DataFrame 表达式 DSL 中使用。包装层在 examples/hello/hello/init.pyfrom __future__ import annotations from typing import TYPE_CHECKING import daft if TYPE_CHECKING: from daft.expressions import Expression def greet(name: Expression) - Expression: Greet someone by name. return daft.get_function(greet, name) def string_count(name: Expression) - Expression: Count non-null strings. return daft.get_aggregate_function(string_count, name)标量函数通过daft.get_function(greet, name)解析——它调用当前会话的get_function按DaftScalarFunction::name()注册的名字查找函数并绑定参数模块级 API 见 daft/session.py聚合函数通过daft.get_aggregate_function(...)走同样的解析逻辑对应DaftAggregateFunction::name()加一层 Python 包装的意义在于可以用类型注解Expression - Expression、docstring 提供自动补全与文档还可以在链接到原生符号之前做参数预处理。authoring 指南特别指出SQL 解析函数并不需要这层 Python所以扩展不依赖 PyO3但 Python 函数让 Expression DSL 用起来更符合 Python 习惯。在 Python 包内放置空的py.typed标记文件本示例即如此即可获得类型检查器的完整支持。构建、安装与测试从源码到可运行安装按照 README 的快速开始在项目虚拟环境中执行uv pip install -e .setuptools-rust会先编译 Rust cdylib再以可编辑模式安装 Python 包开发时通过[tool.uv.sources]将daft指向仓库根目录保证扩展依赖的 Daft 与当前源码一致。加载与使用import daft # Step 1. Import your extension module import hello # Step 2. Load the extension into the current daft session daft.load_extension(hello) # Step 3. Use in your dataframe! df daft.from_pydict({name: [John, Paul]}) df df.select(hello.greet(df[name])) df.show()输出╭──────────────╮ │ greet │ │ --- │ │ String │ ╞══════════════╡ │ Hello, John! │ ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ │ Hello, Paul! │ ╰──────────────╯daft.load_extension的模块级入口定义在 daft/session.py它接受字符串、模块对象或路径最终委托给会话的Session.load_extensiondaft/session.py。测试验证examples/hello/tests/test_hello.py 提供了完整的 pytest 覆盖直接pytest -v即可运行test_greet/test_greet_null验证标量函数对普通值与None的处理null 必须原样透传test_greet_show验证.show()的可视化输出test_string_count/test_string_count_with_nulls验证聚合函数对含 null 数据的非空计数test_aggregate_not_available_without_extension验证未加载扩展的会话中调用string_count会抛异常——这从反面印证了函数是按会话隔离注册的。测试里使用了独立作用域会话而非全局活动会话sess Session() sess.load_extension(hello) with sess: result df.select(greet(col(name))).collect().to_pydict()会话作用域与最佳实践作者指南对会话机制有几个关键说明直接关系到扩展的正确使用扩展在进程内只加载一次重复调用load_extension对同一共享库只会dlopen一次会话只是名字解析的作用域机制函数仅在加载了该扩展的会话中可用with sess:上下文管理器可以把查询限定到特定会话避免污染全局会话命名前缀函数名在会话内全局唯一定义多个函数或可能与其他扩展共存时建议使用extension_name_fn_name前缀防止冲突类型检查器Python 包装函数统一使用TYPE_CHECKING守卫导入Expression并给每个函数加上类型注解与 docstring。延伸更多原生扩展范式hello是理解 Daft 原生扩展机制的最佳起点仓库中还有两个同构但更具实战价值的示例可以继续研读examples/dvectorpgvector 风格的原生扩展提供l2_distance、cosine_distance、inner_product、jaccard_distance等向量距离函数展示了多个函数并存注册的组织方式examples/hello_cpp纯 C 原生扩展基于 Apache Arrow C 直接使用 Daft 的裸 C ABI说明只要语言能产出共享库、导出约定的 C ABI 并能读写 Arrow C Data Interface 数组就可以接入 Daft。此外扩展 SDK 的完整定义位于 src/daft-ext/src含abi/、ffi/、function.rs、aggregate.rs、session.rs等模块阅读它可以进一步理解 ABI 结构体与 trait 的底层契约。对于想直接用 Python 快速构建领域扩展AI 推理、文件处理等的场景则可以参考 docs/extensions/overview.md 中介绍的 UDF 型扩展路线。【免费下载链接】DaftHigh-performance data engine for AI and multimodal workloads. Process images, audio, video, and structured data at any scale项目地址: https://gitcode.com/GitHub_Trending/da/Daft创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考