如何基于另一DataFrame的行筛选为LazyFrame创建列？

问题描述 #

我有一个名为hourly_data的LazyFrame，包含小时级datetime类型的time列；还有一个future_periods DataFrame，包含start（时间段开始）和end（时间段结束）两个datetime列，且所有时间段互不重叠。我想给hourly_data新增一个period整数列，值对应time落在future_periods哪一行的时间段内（索引从0开始），不在任何时间段内的则设为None。

我尝试了以下代码：

periods = pl.Series(range(future_periods.height))
hourly_data = hourly_data.with_columns(
    (
        pl.when(((future_periods.get_column('start') <= pl.col('time')) & (pl.col('time') <= future_periods.get_column('end'))).any())
        .then(periods.filter(pl.Series((future_periods.get_column('start') <= pl.col('real_time')) & (pl.col('real_time') <= future_periods.get_column('end')))).to_list()[0])
        .otherwise(None)
    ).alias('period')
)

但报错：TypeError: Series constructor called with unsupported type 'Expr' for the values parameter。

期望输入输出 #

输入：

hourly_data:
┌────────────────────┐
│ time               │
│ ---                │
│ datetime           │
╞════════════════════╡
│ 2024-01-01 00:00:00│
│ 2024-01-01 01:00:00│
│ 2024-01-01 02:00:00│
│         ...        │
│ 2024-03-31 23:00:00│
│ 2024-04-01 00:00:00│
│ 2024-04-01 01:00:00│
│         ...        │
│ 2024-06-01 00:00:00│
└────────────────────┘

future_periods:
┌─────────────────────────┬───────────────────────┐
│ start                   ┆ end                   │
│ ---                     ┆ ---                   │
│ datetime                ┆ datetime              │
╞═════════════════════════╪═══════════════════════╡
│ 2024-01-01 00:00:00     ┆ 2024-01-31 23:00:00   │
│ 2024-02-01 00:00:00     ┆ 2024-02-28 23:00:00   │
│ 2024-03-01 00:00:00     ┆ 2024-03-31 23:00:00   │
│ 2024-04-01 00:00:00     ┆ 2024-05-31 23:00:00   │
└─────────────────────────┴───────────────────────┘

输出：

hourly_data:
┌─────────────────────────┬────────┐
│ time                    ┆ period │
│ ---                     ┆ ---    │
│ datetime                ┆ int    │
╞═════════════════════════╪════════╡
│ 2024-01-01 00:00:00     ┆ 0      │
│ 2024-01-01 01:00:00     ┆ 0      │
│ 2024-01-01 02:00:00     ┆ 0      │
│          ...            ┆ ...    │
│ 2024-03-31 23:00:00     ┆ 2      │
│ 2024-04-01 00:00:00     ┆ 3      │
│ 2024-04-01 01:00:00     ┆ 3      │
│          ...            ┆ ...    │
│ 2024-06-01 00:00:00     ┆ None   │
└─────────────────────────┴────────┘

错误原因 #

你的代码报错是因为试图用Polars表达式（Expr）作为pl.Series的构造参数，而pl.Series只能接收具体的数值/数据集合，不能直接传入表达式。另外，这种直接把两个DataFrame的列做逐行匹配的逻辑不符合Polars的矢量化操作范式，效率也很低。

可行解法 #

解法1：链式pl.when条件判断 #

利用future_periods的时间段不重叠的特性，遍历每个时间段，用pl.when链式判断time落在哪个区间，匹配成功则返回对应的索引：

# 先给future_periods添加索引列
future_periods = future_periods.with_columns(period=pl.int_range(0, pl.count()))

# 构建链式条件表达式
period_expr = pl.lit(None)
for idx, row in future_periods.iter_rows(named=True):
    period_expr = pl.when(
        (pl.col('time') >= row['start']) & (pl.col('time') <= row['end'])
    ).then(pl.lit(idx)).otherwise(period_expr)

# 给hourly_data添加period列
hourly_data = hourly_data.with_columns(period_expr.alias('period'))

# 若需要立即获取结果，执行懒加载
# hourly_data = hourly_data.collect()

解法2：区间连接（Interval Join） #

Polars支持区间连接，这是更高效的矢量化操作方式，尤其适合大数据量场景：

# 先给future_periods添加索引列
future_periods = future_periods.with_columns(period=pl.int_range(0, pl.count()))

# 执行区间左连接，匹配time落在[start, end]的时间段
hourly_data = hourly_data.join(
    future_periods,
    how='left',
    left_on='time',
    right_on=['start', 'end'],
    strategy='interval'
).drop(['start', 'end'])  # 移除不需要的列

# 若需要立即获取结果，执行懒加载
# hourly_data = hourly_data.collect()

说明 #

• 解法2的区间连接是Polars原生优化操作，性能比循环判断好很多，推荐优先使用。
• 两种方法都会给不在任何时间段的time对应的period列设为None，符合需求。
• 因为hourly_data是LazyFrame，所有操作会保持懒执行状态，直到调用.collect()才会实际计算。

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