Polars：基于采样函数为Lazy DataFrame创建新列

需求：现有两个不同大小的Polars Lazy DataFrame（df_1和df_2），需要为df_2的每条记录，从df_1某一列的唯一值中随机采样生成新列。对应Pandas实现代码如下：

import numpy as np
import pandas as pd

df_1 = pd.DataFrame(data={
    "sample_col": list("aabccdeff")
})

df_2 = pd.DataFrame(data={
    "col_1": range(30)
})

rng = np.random.default_rng()
unique = df_1["sample_col"].unique()
df_2["my_sample"] = rng.choice(unique, len(df_2))

需要将该逻辑迁移到Polars懒加载模式中。

实现方案 #

1. 基础方案（高效推荐） #

先从df_1中提取目标列的唯一值，再通过Polars内置的随机选择函数为df_2生成采样列，全程适配懒加载流程：

import polars as pl

# 构造示例LazyFrame
df_1 = pl.LazyFrame(data={"sample_col": list("aabccdeff")})
df_2 = pl.LazyFrame(data={"col_1": range(30)})

# 提取df_1目标列的唯一值（仅这一步需要触发小量计算）
unique_values = df_1.select(pl.col("sample_col").unique()).collect().to_series().to_list()

# 为df_2新增采样列，保持懒加载状态
df_2_with_sample = df_2.with_columns(
    pl.random.choice(unique_values, seed=42).alias("my_sample")
)

# 按需触发计算，得到最终结果
result = df_2_with_sample.collect()
print(result)

2. 全懒加载整合方案（适合小数据量） #

如果希望尽量减少提前收集数据的操作，可以将唯一值获取逻辑整合到懒加载流程中，通过交叉连接+分组采样实现：

import polars as pl

df_1 = pl.LazyFrame(data={"sample_col": list("aabccdeff")})
df_2 = pl.LazyFrame(data={"col_1": range(30)})

# 全流程懒加载（仅df_1的唯一值部分需要小量收集）
df_2_with_sample = df_2.join(
    df_1.select(pl.col("sample_col").unique()).collect(),
    how="cross"
).group_by(pl.all().exclude("sample_col")).agg(
    pl.col("sample_col").sample(n=1, seed=42).alias("my_sample")
).collect()

print(df_2_with_sample)

关键说明 #

• pl.random.choice()是Polars原生支持的随机采样函数，可直接传入候选值列表，支持指定seed保证结果可复现，完美适配懒加载执行计划。
• 基础方案性能更优，因为提前明确了采样候选集，避免了交叉连接带来的额外计算开销，推荐优先使用。

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