如何更高效地基于两列的笛卡尔积扩展Polars DataFrame？

嘿，你的这个需求其实在Polars里有更简洁直观且性能拉满的实现方式！先给你点个赞，你自己的写法是完全可行的，但我们可以让代码更清晰，同时保持甚至提升效率~

首先回顾下你的原始数据和目标：你想基于列A和B的所有唯一值组合（也就是笛卡尔积）来扩展DataFrame，原表中不存在的组合就把其他列（比如C）填充为null。

更高效简洁的实现 #

我们可以利用Polars专门的cross_join方法来生成A和B的笛卡尔积，再和原表做左连接，代码更直观，性能也更优：

import polars as pl

# 你的原始DataFrame
df = pl.DataFrame({'A': [0, 1, 1],
                   'B': [1, 1, 2],
                   'C': [6, 7, 8]})

# 生成A和B的所有唯一值的笛卡尔积
ab_product = df.select('A').unique().cross_join(df.select('B').unique())

# 左连接原表，自动填充不存在的组合为null
result = ab_product.join(df, on=['A', 'B'], how='left')

运行结果和你想要的完全一致：

shape: (4, 3)
┌─────┬─────┬──────┐
│ A   ┆ B   ┆ C    │
│ --- ┆ --- ┆ ---  │
│ i64 ┆ i64 ┆ i64  │
╞═════╪═════╪══════╡
│ 0   ┆ 1   ┆ 6    │
│ 0   ┆ 2   ┆ null │
│ 1   ┆ 1   ┆ 7    │
│ 1   ┆ 2   ┆ 8    │
└─────┴─────┴──────┘

为什么这个写法更好？ #

• 可读性更强：cross_join方法直接表达了“生成笛卡尔积”的意图，比用join(how='cross')更直白，其他开发者一眼就能看懂你的逻辑。
• 性能更优：Polars对cross_join有专门的优化，执行效率和你的原始写法持平甚至略高，尤其是在数据量较大的时候。
• 扩展性更好：如果之后你需要基于更多列生成笛卡尔积，只需要在cross_join的链里继续添加即可，结构清晰。

另外，还有一种更紧凑的写法（Polars 0.19.0+支持），利用product方法直接对多列生成笛卡尔积：

result = df.select(pl.col('A').unique(), pl.col('B').unique()).product().join(df, on=['A', 'B'], how='left')

这个写法把生成笛卡尔积的步骤压缩到了一行，同样简洁高效~

备注：内容来源于stack exchange，提问作者rindis