已开启
codex_optimize_rowwise_udf_conversion #51
codex_optimize_rowwise_udf_conversion #51
已开启
3stone创建于 7月24日
3stone
3stone成员
7月24日

1. 批量 Python UDF ABI

1.1 优化点名称

批量 Python UDF ABI(Row-wise Python UDF 批量跨语言调用)。

1.2 背景与收益原理

逐行 Python UDF 的传统路径会为每一行执行一次 Rust → Python 调用,并重复完成参数列表创建、Python 对象转换、函数绑定、指标上下文进入退出和返回值转换。对于单行计算量较小的 UDF,跨语言边界和 Python 调度成本可能高于 UDF 本身。

该优化尝试将一个 Daft Series 中的多行参数一次性传给 Python,在一次 ABI 调用中循环执行用户函数,再一次性构造返回 Series。理论上可减少:

  • Rust/Python ABI 切换次数;
  • 重复函数绑定和模块查找;
  • 指标上下文创建次数;
  • 每行返回值单独转换的固定成本。

1.3 实现方式、分支与提交

分支:codex/ai-de-batch-python-udf-abi

实现提交:

  • 3d89835b8:codex_optimize_rowwise_udf_conversion

主要实现:

  • 在 daft/udf/execution.py 增加 call_func_many_as_series,一次接收多行已经求值的参数,在单个 metrics context 中调用绑定后的 Python 方法,并统一构造 PySeries。
  • 在 src/daft-dsl/src/python_udf/row_wise.rs 增加 call_serial_batch_as_series,预先收集每行参数后只调用一次 Python helper。
  • 仅在 on_error=Raise 且 max_retries=0 时进入快速路径,避免改变逐行错误处理和重试语义。

该分支只包含批量 ABI 优化;测试时没有叠加其他实验代码。

1.4 测试结果

Pipeline 基线中位数 优化后中位数 优化效果 加速比
pipeline_pdf_full_min 12.076 s 12.288 s -1.76% 0.983x
pipeline_text_fineweb_full_min 9.381 s 9.354 s +0.28% 1.003x
pipeline_text_vectorize_full_min 17.093 s 17.602 s -2.98% 0.971x
pipeline_image_full_min 12.031 s 12.221 s -1.57% 0.984x
pipeline_video_full_min 12.391 s 12.530 s -1.12% 0.989x
五条中位数之和 62.972 s 63.994 s -1.62% 0.984x

结论:新 base 上只有 FineWeb 提升 0.28%,其余 4 条均回退,五条汇总回退 1.62%。减少 ABI 调用的收益不足以覆盖批量参数列表和 Python 对象的集中构造成本,不建议合入当前实现。

likedislike
合并受阻
3stone3stone成员
7月24日 强制推送  1 个提交:0d3a70ac-codex_optimize_rowwise_udf_conversion
3stone3stone成员
7月24日 修改了pull request 的描述