如何优化Polars懒加载DataFrame的频率编码实现？

希望将懒加载（lazy）状态下的DataFrame中的类别列替换为其出现频率，现有实现需要3次collect()操作，对全量数据多次扫描，寻求更高效的实现方式。

输入代码 #

df = pl.DataFrame({"a": [1, 8, 3], "b": [4, 5, None], "c": ["foo", "bar", "bar"]}).lazy()
print(df.collect())

输入输出 #

shape: (3, 3)
┌─────┬──────┬─────┐
│ a   ┆ b    ┆ c   │
│ --- ┆ ---  ┆ --- │
│ i64 ┆ i64  ┆ str │
╞═════╪══════╪═════╡
│ 1   ┆ 4    ┆ foo │
│ 8   ┆ 5    ┆ bar │
│ 3   ┆ null ┆ bar │
└─────┴──────┴─────┘

期望输出 #

shape: (3, 3)
┌─────┬──────┬────────────────────┐
│ a   ┆ b    ┆ c                  │
│ --- ┆ ---  ┆ ---                │
│ i64 ┆ i64  ┆ str                │
╞═════╪══════╪════════════════════╡
│ 1   ┆ 4    ┆ 0.3333333333333333 │
│ 8   ┆ 5    ┆ 0.6666666666666666 │
│ 3   ┆ null ┆ 0.6666666666666666 │
└─────┴──────┴────────────────────┘

当前实现代码 #

l = df.select("c").collect().shape[0]
rep = df.group_by("c").len().collect().with_columns(pl.col("len")/l).lazy()
df_out = df.with_context(rep.select(pl.all().name.prefix("context_"))).with_columns(pl.col("c").replace(pl.col("context_c"), pl.col("context_len"))).collect()
print(df_out)

以下两种方案均只需一次collect()操作，全程保持懒加载执行计划，Polars会自动优化底层数据扫描逻辑，避免多次全量扫描：

方法1：窗口函数直接计算频率 #

利用窗口函数over分别计算每个类别的出现次数和总行数，直接得到频率并替换原列：

df_out = df.with_columns(
    c=(pl.count().over("c") / pl.count().over()).cast(pl.Float64)
).collect()
print(df_out)

方法2：分组计算频率后关联 #

先分组统计每个类别的频率，再通过join关联回原DataFrame，替换目标列：

# 先计算每个类别的频率
freq_df = df.group_by("c").agg(
    freq=pl.count() / pl.select(pl.count()).first()
)
# 关联后替换原列
df_out = df.join(freq_df, on="c").drop("c").rename({"freq": "c"}).collect()
print(df_out)

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