Polars 使用技巧#
Polars 的 API 本身已经比较简洁，但构造复杂处理时，还是有三类问题会出现：条件分支一多就写成一串 when().then()，方法链一长就无法看清中间数据的状态，表达式写起来离纸上的公式越来越远。
这一章讲的是本书提供的小工具和几个使用习惯，目的只有一个：让调试和复用变容易。里面的细节不必记住住，遇到对应场景时回来查即可。
这些工具都不是 Polars 核心的一部分，而是为了补上“表达能力强但调试困难”这块短板：match() 压缩条件分支，DataCapturer 和 PipeLogger 记录方法链，polars_exprs 让公式更接近数学写法。理解它们解决的问题，比记住用法更重要。
import panel as pn
pn.extension()
import polars as pl
from helper.jupyter import row
条件分支#
Polars 写条件分支一般用 pl.when().then().when().then().otherwise() 这样的方法链。逻辑一多，括号层层嵌套，改一个条件就要重排整段。
本书用 match() 把这类方法链写得简洁些：奇数位是条件、偶数位是结果，最后多出来的那个参数就是默认值。下面的代码里 rank_expr1 用 pl.when() 链写分级逻辑，rank_expr2 用 match() 写同样的处理，对比一下长度就能看出差别。match() 的参数规则如下：
奇数位参数是条件式。
偶数位参数是条件成立时的结果式。
参数个数为奇数时，最后一个参数作为
otherwise()的默认值。
from helper.polars import match
df = pl.DataFrame({
"name": ["Alice", "Bob", "Charlie", "Diana"],
"score": [95, 85, 70, 50],
})
rank_expr1 = (
pl.when(pl.col("score").is_between(90, 100)).then(pl.lit("A"))
.when(pl.col("score").is_between(80, 89)).then(pl.lit("B"))
.when(pl.col("score").is_between(70, 79)).then(pl.lit("C"))
.otherwise(pl.lit("F"))
)
rank_expr2 = match(
pl.col("score").is_between(90, 100), pl.lit("A"),
pl.col("score").is_between(80, 89), pl.lit("B"),
pl.col("score").is_between(70, 79), pl.lit("C"),
pl.lit("F")
)
df2 = df.with_columns(
rank_expr1.alias('rank1'),
rank_expr2.alias("rank2")
)
df2
| name | score | rank1 | rank2 |
|---|---|---|---|
| str | i64 | str | str |
| "Alice" | 95 | "A" | "A" |
| "Bob" | 85 | "B" | "B" |
| "Charlie" | 70 | "C" | "C" |
| "Diana" | 50 | "F" | "F" |
这样写就把重复的结构去掉了，条件增删也只是加一对参数，不用动缩进，可读性和可维护性都优于手写的方法链。
方法链#
方法链能把处理写得像流水线，但链一长就有一个问题：无法只凭代码判断中间数据的状态，出错时也不知道是哪一步引入的。
调试时可以用下面两种工具，它们都不改变计算结果，只在旁边记录。
helper.polars.DataCapturer：在方法链的任意位置捕获数据，查看当时的状态。helper.polars.PipeLogger：捕获整条方法链的输入输出和参数，并可视化。
DataCapturer#
DataCapturer 用来在方法链的指定位置捕获数据状态，相当于在流水线上设置了几个“取样口”。下面是用法示例。
from helper.polars import DataCapturer
df = pl.DataFrame({
"name": ["Alice", "Bob", "Charlie", "David", "Eve"],
"age": [25, 30, 35, 40, 45],
"score": [85, 90, 95, 80, 70],
"department": ["HR", "IT", "HR", "IT", "HR"]
})
cap = DataCapturer() #❶
result = (
df
.filter(pl.col("age") > 30)
.with_columns(
(pl.col("score") * 1.1).alias("adjusted_score")
)
.pipe(cap.before_group) #❷
.group_by("department")
.agg([
pl.col("adjusted_score").mean().alias("average_score"),
pl.col("age").max().alias("max_age")
])
.pipe(cap.before_sort) #❷
.sort("average_score", descending=True)
)
row(cap.before_group, cap.before_sort) #❸
shape: (3, 5)
|
shape: (2, 3)
|
这段代码在分组前用 cap.before_group 捕获数据，在排序前用 cap.before_sort 捕获数据。
❶ 在开始方法链之前先建 DataCapturer() 的实例 cap。
❷ 在链的指定位置插入 .pipe(cap.name)，把当时的数据以 name 为名存下来。
❸ 捕获的数据用 cap.name 取。
PipeLogger#
PipeLogger 把方法链每一步的输入输出可视化：它不是打印日志，而是在 JupyterLab 里画出每步的表，逐步确认数据怎么变。
链特别长、或者接手他人编写的方法链时，这比逐行读代码快得多。下面的例子用 PipeLogger 操作一个数据框，逐步记录每一步的形状。
from helper.polars import PipeLogger
result = (
PipeLogger(df) #❶
.filter(pl.col("age") > 30)
.with_columns(
(pl.col("score") * 1.1).alias("adjusted_score")
)
.group_by("department")
.agg([
pl.col("adjusted_score").mean().alias("average_score"),
pl.col("age").max().alias("max_age")
])
.sort("average_score", descending=True)
)
result #❷
用 PipeLogger 可以把方法链里数据的变化可视化，调试时比逐行阅读代码更快。
❶ 把最初的 df 用 PipeLogger() 包起来。这样结果不再是 DataFrame，而是捕获全部输入输出的 PipeLogger 对象，它会保存链上每一步的数据状态。
❷ 最后把 PipeLogger 对象放在单元格末尾，以交互式 widget 显示结果，各个方法的输入输出就能看到了。
ExprCapturer#
想看表达式中间流过的数据，最简单的是插一个 Expr.inspect()：值会立即打印，表达式的结果照原样往后传，不影响计算。
缺点是它只打印、不保存，事后想再拿出来分析就不行了。
# 想看中间数据时插一个 inspect()（一次处理一列）
df.select(
pl.col("name").str.slice(0, 1).str.to_lowercase().value_counts().inspect().implode()
)
shape: (5,)
Series: 'name' [struct[2]]
[
{"c",1}
{"b",1}
{"a",1}
{"d",1}
{"e",1}
]
| name |
|---|
| list[struct[2]] |
| [{"c",1}, {"b",1}, … {"e",1}] |
不过，想把表达式链里流过的数据存下来以后再分析，建议用本书提供的 ExprCapturer。下面的代码建实例 cap，用 .pipe(cap.counts) 把数据以 counts 为名捕获。
from helper.polars import ExprCapturer
cap = ExprCapturer()
df_res = df.select(
pl.col('name', 'department').str.slice(0, 1).str.to_lowercase().value_counts().pipe(cap.counts).implode()
)
df_res
| name | department |
|---|---|
| list[struct[2]] | list[struct[2]] |
| [{"a",1}, {"b",1}, … {"e",1}] | [{"h",3}, {"i",2}] |
因为是对多列计算，捕获的数据会存成多个 Series 对象的列表。比如可以这样查看结果的前几行。
row(cap.counts[0], cap.counts[1])
shape: (2,)
|
shape: (2,)
|
复杂表达式#
同一个中间结果用多次时，方法链的好处就不明显了：要么重复写表达式，要么拆成好几步、为每个中间变量命名。
下面的例子把三维向量归一化，向量长度要用好几次。看上去 length 的表达式会算多遍，但 Polars 会缓存运算结果，同样的计算不会执行两次。这也是为什么可以直接写“重复”的表达式，而不必为了性能事先拆变量。
df = pl.DataFrame(
dict(x=[1.0, 2.0, 3.0], y=[3.0, 2.0, 1.0], z=[2.0, 1.0, 3.0])
)
x, y, z = pl.col('x'), pl.col('y'), pl.col('z')
length = (x**2 + y**2 + z**2).sqrt()
df1 = df.select(
x = x / length,
y = y / length,
z = z / length
)
row(df, df1)
shape: (3, 3)
|
shape: (3, 3)
|
下面的代码用 LazyFrame.explain() 显示执行计划，可以直接看到缓存：计算向量长度的表达式被提取为一列，列名是 __POLARS_CSER_0x30925ccb05afdfe7 这样的临时名。
print(
df
.lazy()
.with_columns(
x = x / length,
y = y / length,
z = z / length
)
.explain()
)
simple π 3/3 ["x", "y", "z"]
WITH_COLUMNS:
[(col("x") / col("__POLARS_CSER_0x10000000a")), (col("y") / col("__POLARS_CSER_0x10000000a")), (col("z") / col("__POLARS_CSER_0x10000000a"))]
WITH_COLUMNS:
[(((col("x").pow([dyn int: 2]) + col("y").pow([dyn int: 2])) + col("z").pow([dyn int: 2]))).sqrt().alias("__POLARS_CSER_0x10000000a")]
DF ["x", "y", "z"]; PROJECT */3 COLUMNS
直接用表达式写复杂计算很费力：求平方根不能写 sqrt(x)，需要写成 x.sqrt()，一个公式拆成方法链之后，就与纸上的写法差别很大。
本书提供 polars_exprs 装饰器，把普通函数转成 Polars 表达式，函数体里可以照数学式子写，同时照样享受表达式的好处(延迟求值、CSE)。另外还有两个便利。.list.first() 这类子命名空间的操作也能以函数形式(如 list_first())调用。以 c_ 开头的变量则会解释成列表达式，比如 c_x 等同于 pl.col('x')。
from helper.polars import polars_exprs
@polars_exprs
def norm():
length = sqrt(c_x**2 + c_y**2 + c_z**2)
return dict(x=x / length, y=y / length, z=z / length)
df2 = df.select(**norm())
row(df, df2)
shape: (3, 3)
|
shape: (3, 3)
|
不用 c_ 开头的变量名时，可以用关键字参数指定列名。下面的代码是一个完整的例子。
@polars_exprs
def norm2():
length = sqrt(x**2 + y**2 + z**2)
return dict(x=x / length, y=y / length, z=z / length)
df = pl.DataFrame(
dict(a=[1.0, 2.0, 3.0], b=[3.0, 2.0, 1.0], c=[2.0, 1.0, 3.0])
)
df1 = df.select(**norm2(x='a', y='b', z='c'))
row(df, df1)
shape: (3, 3)
|
shape: (3, 3)
|
被 polars_exprs 装饰的函数执行时，符号可以按下面的顺序解释。
关键字参数：按关键字参数原样解释，参数值是字符串时，自动换成用
pl.col()表示列的表达式。以
c_开头的符号：转成pl.col()，比如c_name变成pl.col('name')。Polars 库内与表达式相关的函数和方法。
函数定义时的全局变量。
Python 内置函数。
下面的例子里，各个符号这样解释：
A和p：关键字参数指定的列，对应pl.col('Amp')和pl.col('p')。c_f：转成列pl.col('f')。math：当作全局变量。sin()：用 Polars 表达式的sin()方法。print()：用 Python 内置的print()。
import math
@polars_exprs
def my_expr():
phase = 2 * math.pi * c_f + p
print(phase)
value = A * sin(phase)
return value
df = pl.DataFrame(dict(
Amp=[1, 2, 3],
f=[10, 20, 30],
p=[0.1, 0.2, 0.3]
))
df.select(my_expr(A="Amp", p="p"))
[([(dyn float: 6.283185307179586) * (col("f"))]) + (col("p"))]
| Amp |
|---|
| f64 |
| 0.099833 |
| 0.397339 |
| 0.886561 |
用 polars_exprs 装饰器，公式可以按数学写法书写，列名由关键字参数或 c_ 前缀给出，同一段代码也能反复用在结构不同的表上。
处理数组的用户函数#
有些计算用表达式很难表达，而 NumPy 一行就能写完，比如“a 小于 b 的位置换成指定值”这种带掩码的赋值。
下面的代码用 map_batches() 把两列交给 NumPy 处理。代价是数据要进出 Python 一次，所以适合“表达式难写、但数据量可接受”的场景。
import numpy as np
df = pl.DataFrame({
"a": [1, 2, 3, 10],
"b": [1, 3, 4, 5],
})
def replace_where_smaller(args, values):
a, b = [s.to_numpy() for s in args]
c = np.copy(a)
c[a < b] = values
return pl.Series(c)
df.with_columns(
pl_result=pl.map_batches(
["a", "b"],
lambda args: replace_where_smaller(args, [100, 200]),
return_dtype=pl.Int64,
)
)
| a | b | pl_result |
|---|---|---|
| i64 | i64 | i64 |
| 1 | 1 | 1 |
| 2 | 3 | 100 |
| 3 | 4 | 200 |
| 10 | 5 | 10 |
map_batches() 有几处不方便的地方，用起来像在迁就接口。
多列需要合并成一个列表交给用户函数。
用户函数无法接收列以外的参数。
输入输出不是
Series就会出错。
Polars 2.0 里还要求显式给出 return_dtype，因为推断不出输出类型。
为了把这些一并解决，本书提供 series_batch、pyarrow_batch、numpy_batch 三个装饰器。装饰后的函数用 func(values=values) 绑好参数，再以 lambda args: f(args) 的形式传给 map_batches()，额外参数和返回类型的问题一起解决了。
from helper.polars import pyarrow_batch, numpy_batch, series_batch
@series_batch
def se_func(a, b, values):
return a.scatter((a < b).arg_true(), values)
@pyarrow_batch
def pa_func(a, b, values):
from pyarrow.compute import less, replace_with_mask
return replace_with_mask(a, less(a, b), values)
@numpy_batch
def np_func(a, b, values):
c = a.copy()
c[a < b] = values
return c
values = [100, 200]
dtype = pl.Int64 # 2.0 里要显式给出输出类型
# 先绑好参数，再以函数形式传给 map_batches
se = se_func(values=values)
pa = pa_func(values=values)
npf = np_func(values=values)
df.with_columns(
se_result=pl.map_batches(["a", "b"], lambda args: se(args), return_dtype=dtype),
pa_result=pl.map_batches(["a", "b"], lambda args: pa(args), return_dtype=dtype),
np_result=pl.map_batches(["a", "b"], lambda args: npf(args), return_dtype=dtype),
)
| a | b | se_result | pa_result | np_result |
|---|---|---|---|---|
| i64 | i64 | i64 | i64 | i64 |
| 1 | 1 | 1 | 1 | 1 |
| 2 | 3 | 100 | 100 | 100 |
| 3 | 4 | 200 | 200 | 200 |
| 10 | 5 | 10 | 10 | 10 |
表达式替换#
下面的 calc_rotate_expr() 生成做三维坐标旋转的表达式：x、y、z 是三维坐标，nx、ny、nz 是旋转轴向量，angle 是旋转角(度)。
它把列名写死在表达式里，好处是可以直接使用，代价是只能处理结构固定的表。
def calc_rotate_expr():
x, y, z, nx, ny, nz, angle = [pl.col(name) for name in 'x y z nx ny nz angle'.split()]
magnitude = (nx**2 + ny**2 + nz**2).sqrt()
nx /= magnitude
ny /= magnitude
nz /= magnitude
# 把角度换成弧度
angle_rad = angle.radians()
# 计算角度的余弦和正弦
cos_theta = angle_rad.cos()
sin_theta = angle_rad.sin()
# 用罗德里格斯旋转公式计算旋转后的坐标
dot_product = nx * x + ny * y + nz * z
cross_product_x = ny * z - nz * y
cross_product_y = nz * x - nx * z
cross_product_z = nx * y - ny * x
rx = (cos_theta * x +
sin_theta * cross_product_x +
(1 - cos_theta) * dot_product * nx)
ry = (cos_theta * y +
sin_theta * cross_product_y +
(1 - cos_theta) * dot_product * ny)
rz = (cos_theta * z +
sin_theta * cross_product_z +
(1 - cos_theta) * dot_product * nz)
return pl.struct(rx.alias('rx'), ry.alias('ry'), rz.alias('rz'))
这个函数生成的表达式这样使用。
df = pl.DataFrame(dict(
x=[1, 2, 3, 4],
y=[2, 1, 3, 2],
z=[1, 0, 0, 5],
nx=[1, 1, 1, 1],
ny=[0, 0, 0, 0],
nz=[0, 0, 0, 0],
angle=[10, 20, 30, 40],
))
rot_expr = calc_rotate_expr()
df_rot = df.select(rot_expr.struct.unnest())
row(df, df_rot)
shape: (4, 7)
|
shape: (4, 3)
|
这个表达式有两个问题：
列名写死在表达式里，列名不同的数据框没法直接处理。
想把某些列当常数时，处理会比较繁琐：需要先添加一列常数，用完再删。
为解决这些问题，本书提供 expression_replace()，可以替换表达式里的列名，也能把列换成常数或别的表达式。
expression_replace(expr, mapper=None, **kw)
这个函数用 mapper 或关键字参数替换表达式 expr 里的列名。下面的代码把 rot_expr 里的 nx、ny、nz、angle 四列换成常数，再对 df2 计算，同一个表达式因此能用在结构不同的表上。
from helper.polars import expression_replace
rot_expr2 = expression_replace(rot_expr, nx=pl.lit(1), ny=pl.lit(0), nz=pl.lit(0), angle=pl.lit(45))
df2 = df.select('x', 'y', 'z')
df2_rot = df2.select(rot_expr2.struct.unnest())
row(df2, df2_rot)
shape: (4, 3)
|
shape: (4, 3)
|
下面的代码把列名换成对应的大写。
rot_expr3 = expression_replace(
rot_expr,
x='X', y='Y', z='Z',
angle=pl.lit(30)
)
df3 = df.select(pl.col('x', 'y', 'z').name.to_uppercase(), pl.col('nx', 'ny', 'nz'))
df3_rot = df3.select(rot_expr3.struct.unnest())
row(df3, df3_rot)
shape: (4, 6)
|
shape: (4, 3)
|
从 SymPy 式子转换成 Polars 表达式#
本书提供的 helper.polars.sympy_to_polars 能把 SymPy 的公式自动转成 Polars 表达式：符号推导交给 SymPy，高速计算交给 Polars。
比如下面的代码用 SymPy 的函数，把三维坐标 \((x, y, z)\) 依次绕 Z 轴、Y 轴、X 轴旋转 \(\gamma\)、\(\beta\)、\(\alpha\) 度。得到的式子相当长，手写几乎一定会出错。
import sympy as sp
x, y, z, a, b, r = sp.symbols("x y z alpha beta gamma")
xr, yr, zr = sp.rot_axis1(sp.rad(a)) * sp.rot_axis1(sp.rad(b)) * sp.rot_axis2(sp.rad(r)) * sp.Matrix([x, y, z])
yr
接着用 sympy_to_polars() 把式子 \(x_r\)、\(y_r\)、\(z_r\) 转成 Polars 表达式。想自定义列名时，给 symbol_map 参数传一个“SymPy 符号 → 列名”的字典。
from helper.polars import sympy_to_polars
xr_expr, yr_expr, zr_expr = [sympy_to_polars(expr) for expr in (xr, yr, zr)]
print(yr_expr)
[([([(col("y")) * ([([([([(dyn float: 0.005555555555555556) * (dyn float: 3.141592653589793)]) * (col("alpha"))].cos()) * ([([(dyn float: 0.005555555555555556) * (dyn float: 3.141592653589793)]) * (col("beta"))].cos())]) + ([([(dyn float: -1) * ([([(dyn float: 0.005555555555555556) * (dyn float: 3.141592653589793)]) * (col("alpha"))].sin())]) * ([([(dyn float: 0.005555555555555556) * (dyn float: 3.141592653589793)]) * (col("beta"))].sin())])])]) + ([([(col("x")) * ([([([([(dyn float: 0.005555555555555556) * (dyn float: 3.141592653589793)]) * (col("alpha"))].cos()) * ([([(dyn float: 0.005555555555555556) * (dyn float: 3.141592653589793)]) * (col("beta"))].sin())]) + ([([([(dyn float: 0.005555555555555556) * (dyn float: 3.141592653589793)]) * (col("beta"))].cos()) * ([([(dyn float: 0.005555555555555556) * (dyn float: 3.141592653589793)]) * (col("alpha"))].sin())])])]) * ([([(dyn float: 0.005555555555555556) * (dyn float: 3.141592653589793)]) * (col("gamma"))].sin())])]) + ([([(col("z")) * ([([([([(dyn float: 0.005555555555555556) * (dyn float: 3.141592653589793)]) * (col("alpha"))].cos()) * ([([(dyn float: 0.005555555555555556) * (dyn float: 3.141592653589793)]) * (col("beta"))].sin())]) + ([([([(dyn float: 0.005555555555555556) * (dyn float: 3.141592653589793)]) * (col("beta"))].cos()) * ([([(dyn float: 0.005555555555555556) * (dyn float: 3.141592653589793)]) * (col("alpha"))].sin())])])]) * ([([(dyn float: 0.005555555555555556) * (dyn float: 3.141592653589793)]) * (col("gamma"))].cos())])]
下面的代码把自动生成的表达式和 SymPy 式子的结果做对比，逐个元素验证转换没有出错。
df = pl.DataFrame(dict(
x=[1, 2, 3, 4],
y=[2, 1, 3, 2],
z=[1, 0, 0, 5],
alpha=[30, 40, 50, 60],
beta=[10, 20, -30, -60],
gamma=[5, -5, 5, -5],
)).select(pl.all().cast(pl.Float64))
df_pl = df.select(xr=xr_expr, yr=yr_expr, zr=zr_expr)
df_sp = pl.DataFrame([
{
'xr':xr.evalf(subs=row),
'yr':yr.evalf(subs=row),
'zr':zr.evalf(subs=row)
}
for row in df.rows(named=True)
])
row(df_sp, df_pl)
shape: (4, 3)
|
shape: (4, 3)
|
Polars 算多个表达式时会自动优化，重复的计算不会做两遍。
比如下面打印优化后的表达式，可以看到部分中间结果被提成临时列，后面的计算直接引用它们。这种优化叫 CSE(Common Subexpression Elimination，公共子表达式消除)，也是长公式可以放心写成表达式的原因：写两遍不等于算两遍。
print(df.lazy().select(xr=xr_expr, yr=yr_expr, zr=zr_expr).explain())
SELECT [col("x").fma([col("__POLARS_CSER_0x100000009").alias("literal"), (-1.0 * col("z")) * col("__POLARS_CSER_0x100000004")]).alias("xr"), (col("z") * col("__POLARS_CSER_0x100000015")).fma([col("__POLARS_CSER_0x100000009").alias("literal"), col("y").fma([col("__POLARS_CSER_0x10000001a").alias("literal"), (col("x") * col("__POLARS_CSER_0x100000015")) * col("__POLARS_CSER_0x100000004")])]).alias("yr"), (col("z") * col("__POLARS_CSER_0x10000001a")).fma([col("__POLARS_CSER_0x100000009").alias("literal"), col("y").fma([(-1.0 * col("__POLARS_CSER_0x100000014")).fma([col("__POLARS_CSER_0x100000013").alias("literal"), (-1.0 * col("__POLARS_CSER_0x100000011")) * col("__POLARS_CSER_0x10000000e")]), (col("x") * col("__POLARS_CSER_0x10000001a")) * col("__POLARS_CSER_0x100000004")])]).alias("zr")]
WITH_COLUMNS:
[((0.017453 * col("gamma"))).sin().alias("__POLARS_CSER_0x100000004"), ((0.017453 * col("gamma"))).cos().alias("__POLARS_CSER_0x100000009"), ((0.017453 * col("alpha"))).cos().fma([((0.017453 * col("beta"))).sin(), (((0.017453 * col("beta"))).cos() * ((0.017453 * col("alpha"))).sin())]).alias("__POLARS_CSER_0x100000015"), ((0.017453 * col("alpha"))).cos().fma([((0.017453 * col("beta"))).cos(), (-1.0 * ((0.017453 * col("alpha"))).sin()) * ((0.017453 * col("beta"))).sin()]).alias("__POLARS_CSER_0x10000001a"), ((0.017453 * col("alpha"))).sin().alias("__POLARS_CSER_0x10000000e"), ((0.017453 * col("beta"))).cos().alias("__POLARS_CSER_0x100000011"), ((0.017453 * col("beta"))).sin().alias("__POLARS_CSER_0x100000013"), ((0.017453 * col("alpha"))).cos().alias("__POLARS_CSER_0x100000014")]
DF ["x", "y", "z", "alpha", ...]; PROJECT */6 COLUMNS