如何在Polars中正确实现返回并拼接DataFrame的UDF？

问题背景 #

现有如下Polars DataFrame，需要对每行应用自定义函数（UDF），每个UDF返回schema相同但行数可变的DataFrame，最终将所有结果合并为单个DataFrame：

import polars as pl

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

DataFrame结构：

shape: (3, 3)
┌─────┬─────┬─────┐
│ foo ┆ bar ┆ ham │
│ --- ┆ --- ┆ --- │
│ i64 ┆ f64 ┆ str │
╞═════╪═════╪═════╡
│ 1   ┆ 6.0 ┆ a   │
│ 2   ┆ 7.0 ┆ b   │
│ 3   ┆ 8.0 ┆ c   │
└─────┴─────┴─────┘

现有尝试及问题 #

最初尝试直接返回DataFrame，但结果会嵌套成object类型列，无法直接合并：

def myUDF(row_tuple):
    foo, bar, ham = row_tuple
    result = pl.DataFrame({
        "a": foo + bar,
        "b": ham
    })
    return (result,)

df.map_rows(myUDF)

输出：

shape: (3, 1)
┌────────────────┐
│ column_0       │
│ ---            │
│ object         │
╞════════════════╡
│ shape: (1, 2)  │
│ ┌─────┬─────┐  │
│ │ a …          │
│ shape: (1, 2)  │
│ ┌─────┬─────┐  │
│ │ a …          │
│ shape: (1, 2)  │
│ ┌──────┬─────┐ │
│ │ a…           │
└────────────────┘

后续改用to_dict()转换后再unnest+explode，虽然能得到正确结果，但频繁的字典转换会带来性能损耗：

def myUDF(row_tuple):
    foo, bar, ham = row_tuple
    result = pl.DataFrame({
        "a": foo + bar,
        "b": ham
    })
    return (result.to_dict(),)

df.map_rows(myUDF).unnest("column_0").explode("a", "b")

输出：

shape: (3, 2)
┌──────┬─────┐
│ a    ┆ b   │
│ ---  ┆ --- │
│ f64  ┆ str │
╞══════╪═════╡
│ 7.0  ┆ a   │
│ 9.0  ┆ b   │
│ 11.0 ┆ c   │
└──────┴─────┘

补充说明（实际场景） #

实际操作中UDF逻辑更复杂，例如从大的参考数据集中切片并添加标识列：

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

def myUDF(row_tuple):
    foo, bar, ham = row_tuple
    result = (
        reference_data
        .slice(foo, int(bar / 2))
        .with_columns(name=pl.lit(ham))
    )
    return (result.to_dict(),)

期望最终合并后的结果为：

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    │
└─────┴─────┴──────┘

正确实现方式 #

方案1：优化map_rows返回值（减少转换开销） #

让UDF直接返回各列的列表，而非DataFrame或字典，避免不必要的类型转换：

def myUDF(row_tuple):
    foo, bar, ham = row_tuple
    sliced_data = reference_data.slice(foo, int(bar / 2))
    # 直接返回各列的列表，以及重复的name列
    return (
        sliced_data["x"].to_list(),
        sliced_data["y"].to_list(),
        [ham] * len(sliced_data)
    )

# 重命名列后直接explode所有列
result = (
    df.map_rows(myUDF)
    .rename({"column_0": "x", "column_1": "y", "column_2": "name"})
    .explode(pl.all())
)
print(result)

方案2：用apply+pl.concat实现更高效的合并（推荐） #

避免逐行处理的额外开销，直接用apply返回每个行对应的DataFrame，再通过pl.concat合并所有结果：

result = (
    df
    .apply(
        lambda row: reference_data.slice(row["foo"], int(row["bar"]/2)).with_columns(name=pl.lit(row["ham"])),
        return_dtype=pl.DataFrame
    )
    .to_list()  # 将列中的DataFrame转为列表
    .pipe(pl.concat)  # 合并所有DataFrame
)
print(result)

这两种方案都能得到你期望的结果，其中方案2的性能更优，因为它利用了Polars的批量处理能力，避免了map_rows的逐行Python循环开销，尤其适合大数据量场景。

内容的提问来源于stack exchange，提问作者polars_user