用 Python 函数处理#
Polars 的内置函数和运算符覆盖了大部分常见操作，但总有些处理它表达不了：查询外部表、调用 API、复用现成的 Python 算法。这时可以把自己的函数放进表达式链或聚合里，让 Polars 负责分组、并行和内存管理，自己只写那一段业务逻辑。
这一章按“侵入程度”从低到高介绍几种手段：
pipe：只是把函数接入链里，不碰数据。map_batches/map_groups：把整组 Series 交给函数。map_elements：逐元素调用。rolling_map：把窗口交给函数。
越往后越灵活，也越慢，所以选哪种要看实际需要。能用内置表达式完成时，优先用内置表达式；只有确实需要 Python 逻辑时，才把函数接进来。
import polars as pl
import numpy as np
from helper.jupyter import row
pipe#
pipe 不改变数据，只是把函数接入调用链，让一长串操作能拆成有名字的步骤。它有两个作用点。
DataFrame.pipe()Expr.pipe()
两者都不会带来性能开销，纯粹是为了可读性。下面按从低到高的顺序分别举例。
DataFrame.pipe#
DataFrame.pipe 把函数接入链式处理里：函数接收当前的表、返回一个新表，于是复杂的处理可以按步骤写下来，每一步都有名字，出错时也容易定位。
pipe() 的第一个参数是要执行的函数，DataFrame 本身会自动作为它的第一个参数传入。其余的位置参数和关键字参数原样转给这个函数，所以带配置的函数也能直接用。
df = pl.DataFrame({
"A": [1, 2, 3],
"B": [4, 5, 6],
"C": [7, 8, 9]
})
def add_column(df, col_name, value):
return df.with_columns(pl.lit(value).alias(col_name))
def multiply_column(df, col_name, factor):
return df.with_columns((pl.col(col_name) * factor).alias(col_name))
df_res = (
df.pipe(add_column, col_name="D", value=10)
.pipe(multiply_column, col_name="A", factor=2)
)
row(df, df_res)
shape: (3, 3)
|
shape: (3, 4)
|
Expr.pipe#
表达式的 pipe 同理，只是对象换成了 Expr：把表达式交给函数，函数返回一个新的表达式。
下面定义一个计算二次式的函数 quadratic()，再用 pipe 把它套到多个表达式上。同一个公式要用在几列上时，这样写比复制代码更可靠。
def quadratic(x, a=0, b=0, c=0, suffix='_2'):
return (a * x ** 2 + b * x + c).name.suffix(suffix)
df_res = (
df.with_columns(
pl.col("A").pipe(quadratic, 1, 2, 3),
(pl.col("B") + pl.col("C")).pipe(quadratic, a=1, b=2, suffix='_q')
)
)
row(df, df_res)
shape: (3, 3)
|
shape: (3, 5)
|
map_batches#
pl.map_batches 把多列合成一组交给用户函数，函数的返回值作为新列加进表里。
它适合“一次需要读取多列才能算出一个值”的场景，比如求两列的欧氏距离、按行做业务规则判断。
pl.map_batches(
column_names,
f: Callable[[Sequence[Series]], Series | Any],
return_dtype: PolarsDataType | pl.DataTypeExpr | None = None
)
column_names：交给用户函数的列名或表达式列表。f：处理数据的用户定义函数。它接收指定列的pl.Series列表，返回新的pl.Series或其他值。return_dtype：用户函数返回值中元素的数据类型。默认None，此时自动推断。
下面的例子用 a、b 两列调用用户定义函数 hypot，把结果作为新列 c 加进数据框。
df = pl.DataFrame(
{
"a": [3, 3, 3, 4],
"b": [4, 12, 6, 7],
"g": ['A', 'B', 'A', 'B']
}
)
def hypot(args:list[pl.Series]):
a, b = args
return (a**2 + b**2)**0.5
df.with_columns(
pl.map_batches(['a', 'b'], hypot, return_dtype=pl.Float64).alias('c')
)
| a | b | g | c |
|---|---|---|---|
| i64 | i64 | str | f64 |
| 3 | 4 | "A" | 5.0 |
| 3 | 12 | "B" | 12.369317 |
| 3 | 6 | "A" | 6.708204 |
| 4 | 7 | "B" | 8.062258 |
column_names 参数也可以传表达式。下面的代码用“a 列加 1”和 b 列做计算。
df.with_columns(
pl.map_batches([pl.col('a') + 1, 'b'], hypot).alias('c')
)
| a | b | g | c |
|---|---|---|---|
| i64 | i64 | str | f64 |
| 3 | 4 | "A" | 5.656854 |
| 3 | 12 | "B" | 12.649111 |
| 3 | 6 | "A" | 7.211103 |
| 4 | 7 | "B" | 8.602325 |
除表达式和用户函数外，pl.map_batches() 还有几个参数决定 Polars 如何调用传入的函数，写错时结果可能在不报错的情况下出错。
return_dtype(默认：None)：指定返回 Series 的元素类型，不给就由 Polars 推断。is_elementwise(默认：False)：告诉 Polars 这个函数是不是逐元素处理。设为True后，不论有没有分组，输入都会整块交给函数。returns_scalar(默认：False)：函数返回标量而不是 Series 时设True。
下面的 f1()、f2()、f3() 分别演示这些参数的作用：f1 只做元素间计算，f2 需要看整列，f3 返回一个标量。
import threading
lock = threading.Lock()
def f1(args):
with lock:
a, b = args
print(f"{a.to_list()} {b.to_list()}")
return a + b
def f2(args):
with lock:
a, b = args
print(f"{a.to_list()} {b.to_list()}")
return a - b.min()
def f3(args):
with lock:
a, b = args
print(f"{a.to_list()} {b.to_list()}")
return a.max() - b.min()
不指定 return_dtype 时，Polars 会用假输入([1, 1])先调用一次函数，以此推断输出类型。
这就是为什么函数里不能假设输入有几十万行，试调用时只有两行。
df.select(
pl.map_batches(['a', 'b'], f1)
);
[1, 1] [1, 1]
[3, 3, 3, 4] [4, 12, 6, 7]
指定 return_dtype 后，这次试调用就不需要了：既省一次调用，也避免了假输入触发边界错误。
df.select(
pl.map_batches(['a', 'b'], f1, return_dtype=pl.Int64)
);
[3, 3, 3, 4] [4, 12, 6, 7]
在 agg() 上下文里，is_elementwise 的作用看得更清楚：为 False 时，各组数据按组交给用户函数，函数会被调用多次。
结果总是对的，只是调用开销大一些，分组一多，这个差别就很明显。
df.group_by('g', maintain_order=True).agg(
pl.map_batches(['a', 'b'], f1, return_dtype=pl.Int64, is_elementwise=False).alias('f1'),
pl.map_batches(['a', 'b'], f2, return_dtype=pl.Int64, is_elementwise=False).alias('f2'),
)
[3, 3] [4, 6]
[3, 3] [4, 6]
[3, 4] [12, 7]
[3, 4] [12, 7]
| g | f1 | f2 |
|---|---|---|
| str | list[i64] | list[i64] |
| "A" | [7, 9] | [-1, -1] |
| "B" | [15, 11] | [-4, -3] |
为 True 时，分组信息会被忽略，所有数据一次性交给函数。
这对“元素之间”的计算没问题，但像 f2() 这种需要看到整组数据的函数，结果就会算错，所以 is_elementwise=True 只在函数确实逐元素时才安全。
df.group_by('g', maintain_order=True).agg(
pl.map_batches(['a', 'b'], f1, return_dtype=pl.Int64, is_elementwise=True).alias('f1'),
pl.map_batches(['a', 'b'], f2, return_dtype=pl.Int64, is_elementwise=True).alias('f2'),
)
[3, 3, 3, 4] [4, 12, 6, 7]
[3, 3, 3, 4] [4, 12, 6, 7]
| g | f1 | f2 |
|---|---|---|
| str | list[i64] | list[i64] |
| "A" | [7, 9] | [-1, -1] |
| "B" | [15, 11] | [-1, 0] |
f3() 返回标量(一个数值)，不指定 returns_scalar=True 就会报错：Polars 会按“返回 Series”去接，类型不匹配。
df.group_by('g', maintain_order=True).agg(
pl.map_batches(['a', 'b'], f3, return_dtype=pl.Int64, returns_scalar=True).alias('f3'),
)
[3, 3] [4, 6]
[3, 4] [12, 7]
| g | f3 |
|---|---|
| str | i64 |
| "A" | -1 |
| "B" | -3 |
表达式也有 map_batches() 方法：把表达式的计算结果交给函数处理，函数接收的是 Series。
下面的例子里两列的数据分两次传给 square()，每次处理一列。
def square(s):
return s**2
df.select(
pl.col('a', 'b').map_batches(square)
)
| a | b |
|---|---|
| i64 | i64 |
| 9 | 16 |
| 9 | 144 |
| 9 | 36 |
| 16 | 49 |
也可以先用 pl.struct() 把多列合成一个结构体列，再用 map_batches() 套自定义函数。
这样函数只需要接一个参数，签名更简单，也不必记住列的顺序。
def square2(s):
return s.struct.field('a')**2 + s.struct.field('b')**2
df.select(
pl.struct('a', 'b').map_batches(square2)
)
| a |
|---|
| i64 |
| 25 |
| 153 |
| 45 |
| 65 |
pl.map_groups() 把各组的整块数据交给用户函数，再把每组的返回值汇总成分组结果。
实现上它和 map_batches() 是同一套机制，但在 agg() 上下文里写 map_groups() 更能表达“需要查看整组数据”的意图，代码也更好读。下面的例子按组计算 a 列最大值与 b 列最大值的比。
df = pl.DataFrame(
{
"a": [3, 3, 3, 4],
"b": [4, 12, 6, 7],
"g": ['A', 'B', 'A', 'B']
}
)
def ratio_max(args):
a, b = args
return a.max() / b.max()
df.group_by('g').agg(
pl.map_groups(['a', 'b'], ratio_max, returns_scalar=True)
)
| g | a |
|---|---|
| str | f64 |
| "B" | 0.333333 |
| "A" | 0.5 |
Expr.map_elements() 把表达式的每个值单独交给自定义函数，结果作为新列。
它是所有方案里最灵活、也是最慢的一种，因为每个元素都要进入一次 Python 层。下面的例子用 f 逐个处理取值，可以看到调用次数等于元素个数。
如果一次调用能处理一整列，优先用 map_batches()；只有逻辑确实依赖单个值、又无法用表达式表达时，才用 map_elements()。后者的调用次数等于行数，几百万行就会产生几百万次 Python 调用，开销很容易超过计算本身。
def f(x):
if x > 10:
return 'large'
elif x > 5:
return 'middle'
else:
return 'small'
df.with_columns(
pl.col.b.map_elements(f, return_dtype=pl.String).alias('b_category')
)
| a | b | g | b_category |
|---|---|---|---|
| i64 | i64 | str | str |
| 3 | 4 | "A" | "small" |
| 3 | 12 | "B" | "large" |
| 3 | 6 | "A" | "middle" |
| 4 | 7 | "B" | "middle" |
map_elements() 逐个处理取值，数据量大时会慢一个数量级，所以只在向量化运算和内置函数都做不到时才用。
上面的例子可以改写成向量化写法：把“判断 + 取值”写成 pl.when().then().otherwise()，下面就是等价的实现，速度会快很多。
df.with_columns(
pl.when(pl.col('b') > 10).then(pl.lit('large'))
.when(pl.col('b') > 5).then(pl.lit('middle'))
.otherwise(pl.lit('small'))
.alias('b_category')
)
| a | b | g | b_category |
|---|---|---|---|
| i64 | i64 | str | str |
| 3 | 4 | "A" | "small" |
| 3 | 12 | "B" | "large" |
| 3 | 6 | "A" | "middle" |
| 4 | 7 | "B" | "middle" |
在 agg() 或 over() 上下文里用 map_elements()，传给用户函数的仍然是标量值，不是整组数据，这里很容易产生误解。
下面的例子里 f 对每个元素各调用一次，值依次传进去。
def f(x):
print("---")
print(x)
return x
g = df.group_by('g', maintain_order=True)
g.agg(pl.col('a').map_elements(f, return_dtype=pl.Int64))
---
3
---
3
---
3
---
4
| g | a |
|---|---|
| str | list[i64] |
| "A" | [3, 3] |
| "B" | [3, 4] |
想把整组数据交给用户函数，需要用 map_batches()，或者 .implode().map_elements() 先把一列合并成一个列表再处理。
两者的 return_dtype 含义不同：map_batches() 指定返回值中元素的类型，map_elements() 指定返回值本身的类型。比如 .implode().map_elements() 返回的是 Series，return_dtype 就要给 pl.List(...)。
g.agg(pl.col('a').implode().map_elements(f, return_dtype=pl.List(pl.Int64)))
---
shape: (2,)
Series: '' [i64]
[
3
3
]
---
shape: (2,)
Series: '' [i64]
[
3
4
]
| g | a |
|---|---|
| str | list[i64] |
| "A" | [3, 3] |
| "B" | [3, 4] |
g.agg(pl.col('a').map_batches(f, return_dtype=pl.Int64))
------
shape: (2,)
Series: '' [i64]
[
3
4
]
shape: (2,)
Series: '' [i64]
[
3
3
]
| g | a |
|---|---|
| str | list[i64] |
| "A" | [3, 3] |
| "B" | [3, 4] |
rolling_map#
pl.rolling_mean() 这类移动窗口函数覆盖不到的情况，可以用 rolling_map() 把窗口内的数据交给自定义函数，它是“滚动窗口 + 任意 Python 逻辑”的组合。
下面的代码把每个窗口内的数据传给用户函数，结果作为新列。速度自然比不上原生窗口函数，所以只在确实需要自定义逻辑时用。
df = pl.DataFrame(
{
"a": [2, 4, 0, 3, 1, 6]
}
)
arguments = []
def f(s):
arguments.append((id(s._s), s.to_list()))
return s.mean()
df.select(
pl.col('a').rolling_map(f, window_size=3, center=True, min_samples=1)
)
for id_, arg in arguments:
print(id_, arg)
1492088837328 [2, 4]
1492088837328 [2, 4, 0]
1492088837328 [4, 0, 3]
1492088837328 [0, 3, 1]
1492088837328 [3, 1, 6]
1492088837328 [1, 6]
ufunc#
可以直接用 NumPy 的 ufunc 对列或表达式做运算。比如 np.hypot 计算两数平方和的平方根 \(\sqrt{a^2+b^2}\)，在 Polars 里直接把列传给它即可。
ufunc 会被包装成一个 python_udf() 表达式，因此仍然走 Python 调用，速度不如原生表达式；适合一次算完、调用次数不多的场景。
df = pl.DataFrame(
{
"a": [3, 3, 3, 4],
"b": [4, 12, 6, 7],
"g": ['A', 'B', 'A', 'B']
}
)
df.with_columns(
np.hypot(pl.col("a"), pl.col("b")).alias("c")
)
| a | b | g | c |
|---|---|---|---|
| i64 | i64 | str | f64 |
| 3 | 4 | "A" | 5.0 |
| 3 | 12 | "B" | 12.369317 |
| 3 | 6 | "A" | 6.708204 |
| 4 | 7 | "B" | 8.062258 |
用 np.hypot() 对列做运算时，Polars 的处理方式可以从表达式本身看清：它把 a 列转成结构体列，再加入 b 列作为 "argument_1"，然后对这个结构体列调用 python_udf()，由它调用 NumPy 的 np.hypot()。
np.hypot(pl.col("a"), pl.col("b"))
列名与字段名改写#
Expr.name.map()
用用户函数改写列名。
Expr.name.map_fields()
用用户函数改写 Struct 的 field 名。
def rename(n):
return n.upper() + n
df.select(pl.all().name.map(rename))
| Aa | Bb | Gg |
|---|---|---|
| i64 | i64 | str |
| 3 | 4 | "A" |
| 3 | 12 | "B" |
| 3 | 6 | "A" |
| 4 | 7 | "B" |
DataFrame.map_rows()
把每行作为 Tuple 交给用户函数。
def f(row):
a, b, g = row
return f"{g}:{a + b}", f"{g}:{a - b}"
df.map_rows(f).rename({"column_0": "add", "column_1": "sub"})
| add | sub |
|---|---|
| str | str |
| "A:7" | "A:-1" |
| "B:15" | "B:-9" |
| "A:9" | "A:-3" |
| "B:11" | "B:-3" |
需要按字典逐行处理时，用 DataFrame.iter_rows()。
for row in df.iter_rows(named=True):
print(row)
{'a': 3, 'b': 4, 'g': 'A'}
{'a': 3, 'b': 12, 'g': 'B'}
{'a': 3, 'b': 6, 'g': 'A'}
{'a': 4, 'b': 7, 'g': 'B'}