Polars中UDF的哪些特性会导致需用map_batches()而非map_elements()？

问题背景 #

我正在搞懂Polars里什么时候该用map_batches()，什么时候用map_elements()。下面的例子里，三种计算z的方法（z_1到z_3）结果完全一样。**要让这三种方法算出不同结果，我的UDF my_sum()得改成什么样？**我知道map_batches()支持并行，但不清楚UDF的哪些特性决定了能不能用它、必须用它。

环境准备 #

pip install numpy==2.1.0
pip install polars==1.6.0

示例代码 #

import numpy as np
import polars as pl 
from polars import col, lit

TRIALS = 1
SIMS = 10

np.random.seed(42)

df = pl.DataFrame({
  'a': np.random.binomial(TRIALS, .45, SIMS),
  'b': np.random.binomial(TRIALS, .5, SIMS),
  'c': np.random.binomial(TRIALS, .55, SIMS)
})

df.head()

（生成的DataFrame前5行是a、b、c三列的0/1随机值）

def my_sum(x, y):
  z = x + y
  return(z)

df2 = df.with_columns(
  z_1 = my_sum(col("a"), col("b")),
  z_2 = pl.struct("a", "b")
    .map_elements(lambda x: my_sum(x["a"], x["b"]), return_dtype=pl.Int64),
  z_3 = pl.struct("a", "b")
    .map_batches(lambda x: my_sum(x.struct["a"], x.struct["b"])) # 注意map_batches里用x.struct[]访问字段
)

df2.head()

（新增的z_1、z_2、z_3三列值完全相同，都是对应行a和b的和）

关键答案 #

要让三种方法产生差异，my_sum()需要具备以下任意一种特性：

1. 依赖整批数据的状态或统计量 #

如果my_sum()的计算需要用到整个批次的全局信息（比如整列的均值、排序后的位置、累积和等），而不是只看单个元素：

• z_1：直接传Polars列表达式的话，Polars会优化成向量化运算，但如果UDF需要自定义的批量逻辑（比如基于整批均值调整每个元素），这种写法要么无法实现，要么逻辑和预期不符。
• z_2：map_elements()是逐行处理单个struct元素，根本拿不到整批数据的信息，直接报错或者结果完全不对。
• z_3：map_batches()接收的是整批的struct列，可以先计算整批的统计量，再基于这个统计量处理每个元素，能正确实现逻辑。

举个例子，修改my_sum()为：

def my_sum(x, y):
    # 给每个元素加上整批x的均值
    batch_mean = x.mean()
    return x + y + batch_mean

这时候map_elements()会报错（单个int值没有mean()方法），z_1如果用my_sum(col("a"), col("b"))，Polars会把x.mean()解析为整个列的均值，结果和map_batches()一致，但如果是更复杂的自定义批量逻辑（比如给每个元素加上它在整批中的排名），map_elements()就彻底无能为力了。

2. 只支持标量输入，不兼容向量化 #

如果my_sum()只能处理单个数值，不能接收Polars Series或numpy数组这类批量数据：

• z_1：直接传列表达式会报错，因为传入的是Polars Expr对象，不是单个标量。
• z_2：map_elements()会把每个struct里的a、b拆成单个标量传入my_sum，能正常运行，但速度极慢（逐行处理）。
• z_3：map_batches()传入的是整批的Series，不符合UDF的输入要求，直接报错。

比如修改my_sum()为：

def my_sum(x, y):
    if not isinstance(x, int) or not isinstance(y, int):
        raise ValueError("只能传整数标量")
    return x + y

此时z_1和z_3都会报错，只有z_2能正常出结果。

3. 需要修改输出的批量结构 #

如果my_sum()需要返回和输入批次长度不同的结果（比如过滤掉部分元素、聚合后返回更少行）：

• z_1：Polars表达式要求输出长度必须和输入一致，根本实现不了。
• z_2：map_elements()要求每个输入元素对应一个输出，长度也必须匹配，同样实现不了。
• z_3：map_batches()可以返回任意长度的Series甚至新的DataFrame，能实现这类逻辑。

比如修改my_sum()为：

def my_sum(x, y):
    # 只保留x+y大于等于1的结果
    combined = x + y
    return combined[combined >= 1]

这时候z_1和z_2都会因为输出长度不匹配报错，只有map_batches()能正常返回过滤后的结果（注意此时用with_columns会失败，需要用map_batches直接处理整个DataFrame，但逻辑上只有它支持这种批量结构修改）。

怎么判断该用map_batches()还是map_elements()？ #

看你的UDF这几个特性：

• 能处理批量数据（Series/数组）：优先用map_batches()，速度快还支持并行。
• 需要整批数据的上下文：必须用map_batches()，map_elements()拿不到全局信息。
• 只能处理单个标量：只能用map_elements()，map_batches()会报错。
• 大数据量要性能：选map_batches()，map_elements()逐行跑慢到离谱。

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