如何优雅高效地在Polars中生成包含所有索引列唯一值组合的数据集

我在使用Polars处理数据时，有个需求：要确保我的数据集包含指定索引列所有可能的唯一值组合，要是某个组合在原始数据里缺失，就用null来填充。

目前我是用循环交叉连接的方式来实现的，代码如下：

from typing import Sequence
import polars as pl

def ensure_uniform(df: pl.DataFrame, index_cols: Sequence[str]) -> pl.DataFrame:
    # 快速退出：只有一个索引列时直接返回原数据
    if len(index_cols) == 1:
        return df

    # 先获取第一个索引列的唯一值
    uniform_df = df.select(index_cols[0]).unique(maintain_order=True)

    # 依次和其他索引列的唯一值做交叉连接
    for i in range(1, len(index_cols)):
        unique_index_values = df.select(index_cols[i]).unique(maintain_order=True)
        uniform_df = uniform_df.join(unique_index_values, how="cross")

    # 左连接原始数据，保留已有值，缺失的自动补null
    return uniform_df.join(df, on=index_cols, how="left")

示例演示 #

我用一个简单的例子来展示这个函数的效果：

df = pl.from_repr('''
┌─────┬─────┬─────┬───────┐
│ g1  ┆ g2  ┆ g3  ┆ value │
│ --- ┆ --- ┆ --- ┆ ---   │
│ str ┆ i64 ┆ i64 ┆ i64   │
╞═════╪═════╪═════╪═══════╡
│ A   ┆ 1   ┆ 1   ┆ 10    │
│ A   ┆ 1   ┆ 2   ┆ 20    │
│ B   ┆ 2   ┆ 1   ┆ 30    │
│ B   ┆ 2   ┆ 2   ┆ 40    │
└─────┴─────┴─────┴───────┘
''')

uniform_df = ensure_uniform(df, index_cols=["g1", "g2", "g3"])
print(uniform_df)

运行后得到的结果会补全所有g1/g2/g3的组合，缺失的value列用null填充：

┌─────┬─────┬─────┬───────┐
│ g1  ┆ g2  ┆ g3  ┆ value │
│ --- ┆ --- ┆ --- ┆ ---   │
│ str ┆ i64 ┆ i64 ┆ i64   │
╞═════╪═════╪═════╪═══════╡
│ A   ┆ 1   ┆ 1   ┆ 10    │
│ A   ┆ 1   ┆ 2   ┆ 20    │
│ A   ┆ 2   ┆ 1   ┆ null  │
│ A   ┆ 2   ┆ 2   ┆ null  │
│ B   ┆ 1   ┆ 1   ┆ null  │
│ B   ┆ 1   ┆ 2   ┆ null  │
│ B   ┆ 2   ┆ 1   ┆ 30    │
│ B   ┆ 2   ┆ 2   ┆ 40    │
└─────┴─────┴─────┴───────┘

不过我总觉得这个实现不够优雅，想问问有没有更高效、更简洁的实现方式？

后续测试反馈（编辑补充） #

感谢@Dean MacGregor 和 @orlp 给出的方案！我测试后发现几种方案的性能相差不大（误差在±10%以内），其中@Dean MacGregor的方案表现略好一点。

经过多环境测试后我发现，真正的性能瓶颈其实在最后一步：用全组合数据集左连接原始数据的过程，而不是之前生成笛卡尔积的环节。这意味着不管用什么方式生成全组合，随着数据集规模增长，速度和内存峰值的表现都差不多。

另外，这些优化方案都支持Polars的延迟执行（lazy operations），这在很多场景下会非常实用。

目前Polars还没有专门的笛卡尔积生成方法，所以现有的几种实现方式都是可行的选择。

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