多表合并#
把多个数据框合并成一个，是数据预处理的核心环节：日志按天分成不同的文件，指标分散在若干维表里，外部数据要并进主表，最后总要在某一步把它们连接起来。Polars 的合并方式有几种，区别在于以什么为基准对齐，以及对不上时怎么处理。这一章讲按位置堆叠的 concat、按键对齐的 join、按最近的键对齐的 join_asof、按条件对齐的 join_where，以及就地追加行和列的几个方法。
import numpy as np
import polars as pl
from helper.jupyter import row
concat#
pl.concat() 把多个数据框纵向或横向合并，合并方式由 how 指定。
vertical/vertical_relaxed：纵向合并，relaxed会把类型统一到上位类型horizontal：横向合并，要求行数相同horizontal_extend：横向合并，不足的行用 NULL 填补diagonal/diagonal_relaxed：列只有一部分重叠时，纵向与横向两个方向一起合并align：先用共同的键列把行对齐，再合并
下面用两个数据框逐一确认各种方法的结果，先看纵向合并。
df1 = pl.DataFrame({"x":[1, 2, 3], "y":[2, 3, 1]})
df2 = pl.DataFrame({"x":[6, 2, 1, 5], "y":[12, 3, 2, 4]})
row(df1, df2)
shape: (3, 2)
|
shape: (4, 2)
|
纵向合并#
列名和数据类型都一致时，用 vertical 直接纵向堆叠即可，这是最快也最不容易出错的情形。
pl.concat([df1, df2], how='vertical')
| x | y |
|---|---|
| i64 | i64 |
| 1 | 2 |
| 2 | 3 |
| 3 | 1 |
| 6 | 12 |
| 2 | 3 |
| 1 | 2 |
| 5 | 4 |
列名相同而数据类型不同时需要用 vertical_relaxed，它会把结果统一到上位类型。下面的代码先把 df2 的 x 转成 Float64 再合并，因此合并后 x 的类型是 Float64，而 df1 里原本的整数值并不改变，只是整列被提升为浮点类型。
pl.concat([
df1,
df2.with_columns(pl.col('x').cast(pl.Float64))
],
how='vertical_relaxed')
| x | y |
|---|---|
| f64 | i64 |
| 1.0 | 2 |
| 2.0 | 3 |
| 3.0 | 1 |
| 6.0 | 12 |
| 2.0 | 3 |
| 1.0 | 2 |
| 5.0 | 4 |
要区分某一行来自哪个数据框，需要在合并之前加一列键，因为合并之后原始出处就已经丢失，无法再从行号或数值反推。下面的代码给 df1 和 df2 各加一列 key，再纵向堆叠。
data = {"A":df1, "B":df2}
pl.concat([
df.select(pl.lit(key).alias("key"), pl.all())
for key, df in data.items()
])
| key | x | y |
|---|---|---|
| str | i64 | i64 |
| "A" | 1 | 2 |
| "A" | 2 | 3 |
| "A" | 3 | 1 |
| "B" | 6 | 12 |
| "B" | 2 | 3 |
| "B" | 1 | 2 |
| "B" | 5 | 4 |
横向合并#
列名不同的数据框要横向并排时用 horizontal，它要求两边的行数相同，因此这里给列名加后缀区分。Polars 2.0 中 horizontal 遇到行数不同会直接报错，想把较短的一边用 NULL 填满，需要显式指定 horizontal_extend。
# 列名分别加 1 / 2 后并排（df2 多一行，所以用 horizontal_extend）
pl.concat(
[
df1.rename(lambda name: f"{name}1"),
df2.select(pl.all().name.suffix("2")),
],
how="horizontal_extend",
)
| x1 | y1 | x2 | y2 |
|---|---|---|---|
| i64 | i64 | i64 | i64 |
| 1 | 2 | 6 | 12 |
| 2 | 3 | 2 | 3 |
| 3 | 1 | 1 | 2 |
| null | null | 5 | 4 |
只有一部分列名重叠时，需要用 diagonal 或 diagonal_relaxed，其中 diagonal_relaxed 会把类型统一到上位类型。下面的代码给 df1 加上 u、给 df2 加上 v，重叠的列纵向合并，不重叠的列用 NULL 填补。
纵向与横向合并#
dfs = [
df1.with_columns(u=pl.col('x') + pl.col('y')),
df2.with_columns(v=pl.col('x') * pl.col('y'))
]
pl.concat(dfs, how='diagonal')
| x | y | u | v |
|---|---|---|---|
| i64 | i64 | i64 | i64 |
| 1 | 2 | 3 | null |
| 2 | 3 | 5 | null |
| 3 | 1 | 4 | null |
| 6 | 12 | null | 72 |
| 2 | 3 | null | 6 |
| 1 | 2 | null | 2 |
| 5 | 4 | null | 20 |
按列对齐合并#
align 与 diagonal 相似，但它先用共同的列作为键把行对齐，再合并。下面的例子用 x 和 y 作键。因此 x=1, y=2 那一行里，u=3 与 v=2 会合并成一行。
pl.concat(dfs, how="align")
| x | y | u | v |
|---|---|---|---|
| i64 | i64 | i64 | i64 |
| 1 | 2 | 3 | 2 |
| 2 | 3 | 5 | 6 |
| 3 | 1 | 4 | null |
| 5 | 4 | null | 20 |
| 6 | 12 | null | 72 |
用 pl.align_frames() 可以把多个数据框按指定的列对齐，对齐之后各数据框的行数与键的顺序都一致，可以直接逐行比较，也可以按行横向拼接。
row(*pl.align_frames(*dfs, on=['x', 'y']))
shape: (5, 3)
|
shape: (5, 3)
|
join#
join() 按键把两个数据框对上，主要参数的用途写在下面的注释里。
df.join(
other, # 要合并的对方
on=None, # 两边共同的键列，或表达式
how="inner", # 合并方法
left_on=None, # 只有左表使用的键列
right_on=None, # 只有右表使用的键列
suffix="_right", # 加在重名列上的后缀
validate="m:m", # 校验键的对应关系，取值 m:m / m:1 / 1:m / 1:1
nulls_equal=False, # NULL 之间是否算一致，2.0 中由 join_nulls 改名
coalesce=None, # 是否把键列合成一列
maintain_order=None, # 是否保持行序，取值 left / right / none 等
)
how 决定对不上的行怎么处理，这也是一组最容易选错的取值。
inner：只返回两边都一致的行
left：返回左表全部行与右表一致的行
right：返回右表全部行与左表一致的行
full：返回任意一边一致的所有行
semi：返回左表中与右表一致的行，只含左表的列
anti：返回左表中与右表不一致的行
Polars 2.0 中 collect() 默认使用流式引擎，因此 join 的行序不被保证。顺序有意义时，需要指定 maintain_order，或者在最后调用 sort()。
df_left = pl.DataFrame({
"id": [1, 2, 3, 4],
"name": ["Alice", "Bob", "Charlie", "David"]
})
df_right = pl.DataFrame({
"id": [3, 4, 5],
"age": [23, 30, 40]
})
按两边都存在的 id 做内连接：下面的例子中只有 id 为 3 和 4 的行与对方一致，其余行不返回。
row(
df_left, df_right,
df_left.join(df_right, on="id", how="inner")
)
shape: (4, 2)
|
shape: (3, 2)
|
shape: (2, 3)
|
left 保留左表的全部行。下面的例子中 id 为 1 和 2 的行在右表没有对应数据，因此 age 是 NULL，right 则相反。
row(
df_left, df_right,
df_left.join(df_right, on="id", how="left")
)
shape: (4, 2)
|
shape: (3, 2)
|
shape: (4, 3)
|
row(
df_left, df_right,
df_left.join(df_right, on="id", how="right")
)
shape: (4, 2)
|
shape: (3, 2)
|
shape: (3, 3)
|
full 返回两边的全部行。为了区分行来自左边还是右边，合并键会变成 id 与 id_right 两列。由此可以看出 id 为 1 和 2 的行只存在于左表，5 只存在于右表。
row(
df_left, df_right,
df_left.join(df_right, on="id", how="full")
)
shape: (4, 2)
|
shape: (3, 2)
|
shape: (5, 4)
|
指定 coalesce=True 后，这两列键会合并成一列，代价是看不出某一行究竟来自哪一边。
row(
df_left, df_right,
df_left.join(df_right, on="id", how="full", coalesce=True)
)
shape: (4, 2)
|
shape: (3, 2)
|
shape: (5, 3)
|
semi 与 anti#
semi 只返回左表中能在右表找到匹配的行，anti 只返回左表中在右表找不到匹配的行。两者都不包含右表的列，因此适合用来筛选，而不是补充信息。
row(
df_left, df_right,
df_left.join(df_right, on="id", how="semi")
)
shape: (4, 2)
|
shape: (3, 2)
|
shape: (2, 2)
|
row(
df_left, df_right,
df_left.join(df_right, on="id", how="anti")
)
shape: (4, 2)
|
shape: (3, 2)
|
shape: (2, 2)
|
cross#
cross 返回两个数据框的笛卡尔积，结果是左表全部行与右表全部行的所有组合，因此不需要指定 on。行数是两边行数的乘积，使用前要先确认数据量。
row(
df_left, df_right,
df_left.join(df_right, how="cross")
)
shape: (4, 2)
|
shape: (3, 2)
|
shape: (12, 4)
|
合并键的对应关系可以分成四类，事先在代码里写清楚，比事后发现行数翻倍再排查要好得多。
1:1：左右两边的键都唯一，例如员工 ID 与员工证号
1:m：左边唯一，右边可以重复，例如部门与所属员工
m:1：左边可以重复，右边唯一，例如订单明细与商品主数据
m:m：两边都允许重复，这是默认值
一致的行有多个时，默认会把所有组合，即笛卡尔积放进结果。指定 validate 可以强制预期的关系，违反时直接报错。下面的两个数据框中，id == 3 的行各有两条，恰好可以用来观察这种膨胀。
df_left = pl.DataFrame({
"id": [1, 2, 3, 4, 3],
"name": ["Alice", "Bob", "Charlie", "David", "charlie"]
})
df_right = pl.DataFrame({
"id": [3, 4, 5, 3],
"age": [23, 30, 40, 46]
})
直接用 id 合并时，id == 3 会在两边各自匹配到对方的 2 行，因此一共得到 4 行。
df_left.join(df_right, on="id")
| id | name | age |
|---|---|---|
| i64 | str | i64 |
| 3 | "Charlie" | 23 |
| 3 | "Charlie" | 46 |
| 4 | "David" | 30 |
| 3 | "charlie" | 23 |
| 3 | "charlie" | 46 |
指定 validate="1:1" 后，键不是一对一就会报错，数据的重复因此能在合并这一步被发现。
%%capture_except
df_left.join(df_right, on="id", validate='1:1')
ComputeError: join keys did not fulfill 1:1 validation
DataFrame.update() 用另一个数据框的值覆盖键一致的行。下面的代码按 id 匹配，用 df_update 里的值覆盖 x、y，其余行保持原样。
df = pl.DataFrame({
"id": [1, 2, 3, 4],
"x": [10, 20, 30, 40],
"y": [40, 30, 20, 10]
})
df_update = pl.DataFrame({
"id": [1, 3],
"x": [100, 300],
"y": [-10, -20],
})
df_res = df.update(df_update, on='id')
row(df, df_update, df_res)
shape: (4, 3)
|
shape: (2, 3)
|
shape: (4, 3)
|
join_asof#
join_asof 用时间或数值这类连续值对齐最近的一行，它不要求完全一致，而是在一方的值上寻找最近的行，时序数据的预处理中经常出现。
一方是某个时点的快照，另一方持有最近时点的值
可以指定查找方向，向前找或向后找
import polars as pl
df1 = pl.DataFrame(
{
"time": [1, 5, 10, 15, 20],
"event": ["A", "B", "C", "D", "E"],
}
)
df2 = pl.DataFrame(
{
"time": [2, 6, 12, 18],
"price": [100, 105, 110, 115],
}
)
result = df1.join_asof(df2, on="time", strategy="backward")
row(df1, df2, result)
shape: (5, 2)
|
shape: (4, 2)
|
shape: (5, 3)
|
对 df1 的每一行，合并 df2 中 time 在过去或同一时刻最近的行，strategy="backward" 就是由这个参数指定的。把它改成 "forward"，则改为选择未来的值。
join_where#
join_where() 可以把列与列之间的比较条件直接作为合并条件。下面的例子中，df1 的 time 落在 df2 的 time_span 范围内时两行才会合并，传入多个条件时只有全部满足的行才会留下。
df1 = pl.DataFrame(
{
"id": [100, 101, 102],
"time": [120, 140, 160],
}
)
df2 = pl.DataFrame(
{
"t_id": [404, 498, 676, 742],
"time_span": [(100, 110), (110, 130), (90, 100), (150, 170)],
}
)
df_res = df1.join_where(
df2,
pl.col('time') >= pl.col('time_span').list.get(0),
pl.col('time') <= pl.col('time_span').list.get(1)
)
row(df1, df2, df_res)
shape: (3, 2)
|
shape: (4, 2)
|
shape: (2, 4)
|
布尔运算和 is_between() 也可以用在条件里，写法与普通的列表达式相同。
df_res1 = df1.join_where(
df2,
(pl.col('time') >= pl.col('time_span').list.get(0)) & (pl.col('time') <= pl.col('time_span').list.get(1))
)
df_res2 = df1.join_where(
df2,
pl.col('time').is_between(pl.col('time_span').list.get(0), pl.col('time_span').list.get(1))
)
row(df_res1, df_res2)
shape: (2, 4)
|
shape: (2, 4)
|
join_where() 内部先做出两个数据框的笛卡尔积，再用条件筛选，因此行数一多就会明显变慢。下面的例子从两个点群中寻找距离小于 0.01 的点对，计算量与 \(N^2\) 成正比。
n = 4000
np.random.seed(42)
df1 = pl.DataFrame(dict(
x1=np.random.normal(size=n),
y1=np.random.normal(size=n),
))
df2 = pl.DataFrame(dict(
x2=np.random.normal(size=n),
y2=np.random.normal(size=n),
))
%%time
df_res1 = (
df1
.join_where(
df2,
(pl.col('x2') - pl.col('x1'))**2 + (pl.col('y2') - pl.col('y1'))**2 < 0.0001
)
.sort('x1', 'y1')
)
CPU times: total: 734 ms
Wall time: 46 ms
这类处理改用 SciPy 的 KDTree 会快得多：用 query_ball_tree() 求出范围内的点对，再以行号为键把 df1 与 df2 合并。
from scipy.spatial import KDTree
%%time
tree1 = KDTree(df1.to_numpy())
tree2 = KDTree(df2.to_numpy())
res = tree1.query_ball_tree(tree2, 0.01)
df_pair = (
pl.DataFrame(pl.Series('index2', res, dtype=pl.List(pl.UInt32))) # 把 res 转成数据框
.with_row_index('index1')
.filter(pl.col('index2').list.len() > 0) # 只保留有近距离点的行
.explode('index2') # 展开列表，一行一对
)
df_res2 = (
df1
.with_row_index('index1')
.join(df_pair, on='index1')
.join(
df2.with_row_index('index2'), on='index2')
.drop('index1', 'index2')
.sort('x1', 'y1')
)
CPU times: total: 0 ns
Wall time: 13.5 ms
row(df_res1, df_pair, df_res2)
shape: (420, 4)
|
shape: (420, 2)
|
shape: (420, 4)
|
res 是双层列表，对 tree1 的每个点，其中存放着 tree2 中距离在 0.01 以内的点的编号。下面的代码遍历前 50 个点，取出其中非空的部分。
for i, index2 in enumerate(res[:50]):
if index2:
print(i, index2)
10 [1832, 2789]
41 [1502]
44 [739]
45 [2018]
extend 与 vstack#
DataFrame.extend() 与 DataFrame.vstack() 都是往已有数据框追加另一个数据框各行的旧方法，两者都已弃用，现在推荐使用 pl.concat()，这里只确认行为上的差异。
extend()：就地修改原数据框，同时返回修改后的数据框vstack()：用in_place选择就地修改，还是返回新的数据框
df1 = pl.DataFrame({"a": [1, 2], "b": [3, 4]})
df2 = pl.DataFrame({"a": [5, 6], "b": [7, 8]})
# 用 clone() 复制原表
df_extend = df1.clone()
df_extend.extend(df2) # extend 就地修改
df_vstack = df1.clone()
df_vstack.vstack(df2, in_place=True) # vstack 就地修改
row(df_extend, df_vstack)
shape: (4, 2)
|
shape: (4, 2)
|
两者的区别在内存的持有方式。extend() 把要追加的数据复制到原内存的末尾，合并后的内存是连续的，但这一步会发生重新分配，因此在循环里多次调用并不可取。vstack() 保留两个数据框各自的内存，用链接把它们接起来，结果的内存不连续，之后再计算时可能变慢。在循环里 vstack()、最后调用一次 rechunk()，曾是标准做法。现在请使用 pl.concat()，必要时在最后调用 rechunk()。
print(f'{df_extend.n_chunks() = }') # extend 后的 chunk 数
print(f'{df_vstack.n_chunks() = }') # vstack 后的 chunk 数
print(f'{df_vstack.rechunk().n_chunks() = }') # rechunk 后的 chunk 数
df_extend.n_chunks() = 1
df_vstack.n_chunks() = 2
df_vstack.rechunk().n_chunks() = 1
DataFrame.hstack() 把两个数据框横向并排，指定 in_place=True 时往原数据框里追加列。这个方法也已弃用，新代码请使用 pl.concat(..., how="horizontal") 或 with_columns()。
df1 = pl.DataFrame({"a": [1, 2], "b": [3, 4]})
df2 = pl.DataFrame({"c": [5, 6], "d": [7, 8]})
# 不就地修改，返回新表
df3 = df1.hstack(df2)
# 就地修改
df1.hstack(df2, in_place=True)
row(df3, df1)
shape: (2, 4)
|
shape: (2, 4)
|
merge_sorted() 把已经按键列排序的两个数据框纵向合并，返回时保持键的顺序。它要求输入确实已经排序，只做归并而不检查顺序。
df1 = pl.DataFrame({"a": [3, 1, 5], "b": [6, 4, 2]})
df2 = pl.DataFrame({"a": [6, 2, 4], "b": [7, 5, 3]})
# 按 a 列排序后再合并
df3 = df1.sort("a").merge_sorted(df2.sort("a"), key="a")
# 按 b 列排序后再合并
df4 = df1.sort("b").merge_sorted(df2.sort("b"), key="b")
row(df3, df4)
shape: (6, 2)
|
shape: (6, 2)
|