Polars UDF 行级映射与多行结果 DataFrame 的高效拼接

本文详解如何在 Polars 中正确实现「每行输入生成一个同结构但行数可变的 DataFrame」的 UDF 模式，对比 map_rows、iter_rows + concat 与表达式向量化三种方案，强调性能陷阱与最佳实践。

本文详解如何在 polars 中正确实现「每行输入生成一个同结构但行数可变的 dataframe」的 udf 模式，对比 `map_rows`、`iter_rows + concat` 与表达式向量化三种方案，强调性能陷阱与最佳实践。

在 Polars 中，当需要对每一行输入执行复杂逻辑（如基于参数切片参考数据、生成变长结果集）并最终合并为单个 DataFrame 时，开发者常误用 map_rows 返回嵌套 DataFrame，导致性能劣化与类型混乱。核心原则是：map_rows 本质是逐行标量映射，不适用于返回多行结构；真正支持“1→N”行扩展的操作应使用 iter_rows + pl.concat() 或（更优）完全向量化表达式重构。

✅ 推荐方案一：避免 map_rows，改用 iter_rows + pl.concat()

这是最直接、语义清晰且可控性高的方式，尤其适合逻辑无法轻易向量化（如需调用外部库、动态切片大表等）的场景：

import polars as pl

df = pl.DataFrame({
    "foo": [1, 2, 3],
    "bar": [6.0, 7.0, 8.0],
    "ham": ["a", "b", "c"],
})

reference_data = pl.DataFrame({
    "x": list(range(10)),
    "y": [chr(ord("a") + i % 26) for i in range(10)],
})

def myUDF(row_tuple):
    foo, bar, ham = row_tuple
    # 动态切片：从 foo 开始取 floor(bar/2) 行
    n_rows = int(bar / 2)
    return (
        reference_data
        .slice(foo, n_rows)
        .with_columns(name=pl.lit(ham))
    )

# ✅ 正确做法：逐行调用 UDF 并拼接
result = pl.concat([myUDF(row) for row in df.iter_rows()], how="vertical")
print(result)

输出：

shape: (10, 3)
┌─────┬─────┬──────┐
│ x   ┆ y   ┆ name │
│ --- ┆ --- ┆ ---  │
│ i64 ┆ str ┆ str  │
╞═════╪═════╪══════╡
│ 1   ┆ b   ┆ a    │
│ 2   ┆ c   ┆ a    │
│ 3   ┆ d   ┆ a    │
│ 2   ┆ c   ┆ b    │
│ 3   ┆ d   ┆ b    │
│ 4   ┆ e   ┆ b    │
│ 3   ┆ d   ┆ c    │
│ 4   ┆ e   ┆ c    │
│ 5   ┆ f   ┆ c    │
│ 6   ┆ g   ┆ c    │
└─────┴─────┴──────┘

⚠️ 注意事项：

Calliper 文档对比神器

文档内容对比神器

• pl.concat(..., how="vertical") 要求所有子 DataFrame 具有完全一致的 schema（列名、类型、顺序），否则会报错；
• iter_rows() 返回 Python tuple，默认按列顺序解包，确保 UDF 内部解构顺序与 DataFrame 列序严格一致；
• 若 reference_data 极大（如千万级），建议预先 .lazy() 化并在 UDF 中构建 LazyFrame，最后统一 .collect()，避免重复 eager 计算。

✅ 推荐方案二：彻底向量化 —— 用 Polars 表达式替代 UDF（最高性能）

若逻辑可表达为列运算（如本例中“切片范围由 foo 和 bar 决定”），应优先放弃行级循环，转为纯表达式操作：

# ✅ 向量化实现（无 Python 循环，充分利用 Polars 引擎）
result_vec = (
    df
    .with_row_index("row_idx")  # 添加索引用于后续 join 或 expand
    .with_columns(
        start = pl.col("foo"),
        stop = (pl.col("bar") / 2).cast(pl.Int64),
        name = pl.col("ham")
    )
    # 生成每个 row 对应的 range(start, stop)，再 explode 展开
    .with_columns(
        range_list = pl.arange(pl.col("start"), pl.col("stop"), eager=False)
    )
    .explode("range_list")
    .join(
        reference_data.rename({"x": "range_list"}),
        on="range_list",
        how="left"
    )
    .select("x", "y", "name")
)

该方式零 Python 解释器开销，支持多线程并行，是 Polars 的设计哲学首选。

❌ 不推荐方案：滥用 map_rows 返回 DataFrame

map_rows 的设计目标是返回标量元组（如 (a, b)），其输出会被自动构造成新 DataFrame 的一行。若在函数内创建 DataFrame 并返回（如 return (result,)），Polars 仅将其作为 object 类型存储，后续需 .unnest().explode() 等昂贵操作，丧失 Polars 的内存与计算优势：

# ❌ 反模式：低效、类型丢失、难以调试
df.map_rows(lambda r: (pl.DataFrame({"a": [r[0]+r[1]], "b": [r[2]]}),))

总结：选择路径的决策树

记住：Polars 的强大源于列式计算与查询优化器，而非模拟 Pandas 的 .apply(axis=1)。拥抱向量化，谨慎使用行级迭代——这才是“正确的 Polars 方式”。