如何用Polars生成变量子集全组合？实现类似tidyr::complete功能

Polars目前没有直接对应R中tidyr::complete()的内置函数，但可以通过多种方式模拟该功能，以下分场景给出解决方案，并针对性能问题做优化。

场景1：两个分组变量（country + year） #

方法1：交叉连接+左连接（通用方案） #

先提取每个变量的唯一值，生成所有可能的组合，再与原数据左连接补全缺失值：

import polars as pl

df = pl.DataFrame(
    {
        "country": ["France", "France", "UK", "UK", "Spain"],
        "year": [2020, 2021, 2019, 2020, 2022],
        "value": [1, 2, 3, 4, 5],
    }
)

balanced_df = (
    df.select("country").unique()
    .join(df.select("year").unique(), how="cross")
    .join(df, how="left", on=["country", "year"])
)

print(balanced_df)

输出与预期一致：

shape: (12, 3)
┌─────────┬──────┬───────┐
│ country ┆ year ┆ value │
│ ---     ┆ ---  ┆ ---   │
│ str     ┆ i64  ┆ i64   │
╞═════════╪══════╪═══════╡
│ France  ┆ 2020 ┆ 1     │
│ France  ┆ 2021 ┆ 2     │
│ France  ┆ 2019 ┆ null  │
│ France  ┆ 2022 ┆ null  │
│ UK      ┆ 2020 ┆ 4     │
│ UK      ┆ 2021 ┆ null  │
│ UK      ┆ 2019 ┆ 3     │
│ UK      ┆ 2022 ┆ null  │
│ Spain   ┆ 2020 ┆ null  │
│ Spain   ┆ 2021 ┆ null  │
│ Spain   ┆ 2019 ┆ null  │
│ Spain   ┆ 2022 ┆ 5     │
└─────────┴──────┴───────┘

方法2：透视+逆透视（仅适用于两个变量） #

对于两个变量的场景，可通过透视表展开所有组合，再逆透视还原结构：

balanced_df = (
    df.pivot(index="country", columns="year", values="value")
    .unpivot(index="country", variable_name="year", value_name="value")
    .sort(["country", "year"])
)

场景2：三个及以上分组变量（orig + dest + year） #

多变量场景下，交叉连接是通用方案，但可通过Lazy API和批量提取唯一值优化性能：

优化后的交叉连接方案 #

先修正你代码中的参数错误（原on参数应为["orig", "dest", "year"]），改用Lazy模式执行：

import polars as pl
import time

df = pl.DataFrame(
    {
        "orig": ["France", "France", "UK", "UK", "Spain"],
        "dest": ["Japan", "Vietnam", "Japan", "China", "China"],
        "year": [2020, 2021, 2019, 2020, 2022],
        "value": [1, 2, 3, 4, 5],
    }
)

tic = time.perf_counter()
balanced_df = (
    df.lazy()
    .select("orig", "dest", "year")
    .unique()
    .group_by("orig", maintain_order=True)
    .map_groups(lambda g: g.select("orig").unique().join(g.select("dest", "year").unique(), how="cross"))
    .join(df.lazy(), how="left", on=["orig", "dest", "year"])
    .sort(["orig", "dest", "year"])
    .collect()
)
toc = time.perf_counter()
print(f"Lazy eval: {toc - tic:0.4f} seconds")
print(balanced_df)

性能差异说明 #

你提到的R中tidyr::complete()更快的情况，主要是小数据量下的常数项开销导致：Polars的Lazy API需要额外优化和编译时间，而R的函数在小数据场景下启动成本更低。但当数据量增大（比如分组变量唯一值增多、原数据行数上万），Polars的向量化处理和Lazy优化会体现出明显性能优势。

也可以用pl.fold简化多变量交叉连接逻辑：

group_cols = ["orig", "dest", "year"]
unique_dfs = [df.select(col).unique() for col in group_cols]

full_combinations = pl.fold(
    acc=unique_dfs[0],
    function=lambda acc, df: acc.join(df, how="cross"),
    exprs=unique_dfs[1:]
)

balanced_df = full_combinations.join(df, how="left", on=group_cols)

内容的提问来源于stack exchange，提问作者bretauv