Polars DataFrame逐行应用Python UDF时抛出‘expected tuple, got list’异常

问题背景 #

有如下Polars DataFrame：

import polars as pl

df = pl.DataFrame({
    "user_movies": [[7064, 7153, 78009], [6, 7, 1042], [99, 110, 3927], [2, 11, 152081], [260, 318, 195627]],
    "user_ratings": [[5.0, 5.0, 5.0], [4.0, 2.0, 4.0], [4.0, 4.0, 3.0], [3.5, 3.0, 4.0], [1.0, 4.5, 0.5]],
    "common_movies": [[7064, 7153], [7], [110, 3927], [2], [260, 195627]]
})

需要新增common_movie_ratings列，从每行的user_ratings列表中提取common_movies对应影片的评分。例如第一行返回[5.0, 5.0]，第二行返回[2.0]等。

用户通过自定义函数（UDF）结合df.apply实现，但处理第三行时抛出错误：

expected tuple, got list

错误原因 #

Polars的apply函数传递给UDF的row参数类型取决于列类型：当列是List类型时，row实际是列表而非元组。用户代码中错误假设row为元组，导致索引访问时类型不匹配。

解决方案 #

方案1：修复UDF的参数兼容 #

无需假设row是元组，直接按列表处理，或通过结构体明确指定列，避免依赖索引顺序：

def get_common_movie_ratings(row):
    # 建立影片与评分的映射，避免重复查找索引
    movie_rating_map = dict(zip(row["user_movies"], row["user_ratings"]))
    return [movie_rating_map[movie] for movie in row["common_movies"]]

df = df.with_columns(
    pl.struct(["user_movies", "user_ratings", "common_movies"])
    .apply(get_common_movie_ratings, return_dtype=pl.List(pl.Float64))
    .alias("common_movie_ratings")
)

方案2：Polars原生向量化操作（推荐） #

UDF在大数据量下效率较低，推荐使用Polars内置列表操作实现，性能更优：

df = df.with_columns(
    # 将影片与评分配对成结构体列表
    pl.zip("user_movies", "user_ratings").alias("movie_rating_pairs"),
    # 匹配common_movies对应的评分并聚合
    pl.col("common_movies").list.eval(
        pl.element().map_batches(lambda x: pl.col("movie_rating_pairs").list.eval(
            pl.when(pl.element().struct.field("user_movies") == x)
            .then(pl.element().struct.field("user_ratings"))
            .drop_nulls()
        ).flatten()
    ).alias("common_movie_ratings")
).drop("movie_rating_pairs")

验证结果 #

执行后common_movie_ratings列的结果为：

shape: (5,)
Series: 'common_movie_ratings' [list[f64]]
[
        [5.0, 5.0]
        [2.0]
        [4.0, 3.0]
        [3.5]
        [1.0, 0.5]
]

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