Polars中日内特定时段扩展最大值计算后无法合并回原表问题

需求与问题场景 #

现有存储在Polars DataFrame中的OHLC股票数据，需针对每个交易日，计算09:15到10:15之间Close价格的扩展最大值（cum_max）。现有代码能算出每日该时段的扩展最大值列表，但因提前过滤时间范围后聚合丢失了对应时间戳，无法将结果正确合并回原始DataFrame。

现有代码：

start = '09:15'
end = '10:15'
start_time = datetime.time.fromisoformat(start)
end_time = datetime.time.fromisoformat(end)

# Filter and calculate the expanding max within the time range for each day
max_values = und_df2.filter((und_df2['timestamp'].dt.time() >= start_time) & (und_df2['timestamp'].dt.time() <= end_time)) \
        .group_by_dynamic('timestamp', every='1d', closed='left') \
        .agg(pl.col('Close').cum_max().alias('fbarh3')) 

# Join the maximum values back to the original dataframe
und_df2 = und_df2.with_columns(pl.col('timestamp').dt.truncate('1d').alias('date')).join(max_values, left_on='date', right_on='timestamp' ,  how='left' ).drop('date')

生成的max_values结构（丢失时段内具体时间戳，仅保留日期和扩展最大值列表）：

shape: (824, 2)
timestamp   fbarh3
datetime[μs]    list[f64]
2020-03-02 00:00:00 [11356.8, 11358.45, … 11388.65]
2020-03-03 00:00:00 [11293.65, 11293.95, … 11310.7]
2020-03-04 00:00:00 [11296.95, 11296.95, … 11325.25]
2020-03-05 00:00:00 [11312.65, 11312.65, … 11312.65]
2020-03-06 00:00:00 [10894.8, 10914.85, … 10965.3]

问题原因 #

原代码提前过滤时间范围后按日聚合，将时段内所有cum_max值打包为列表，丢失了每个值对应的具体时间戳。后续按日期合并时，无法将列表中的单个值匹配到原表对应时间点，导致结果无法正确映射。

解决方案 #

调整思路：保留时段内每个行的原始时间戳，计算对应cum_max后，直接按时间戳合并回原表。

实现代码 #

import polars as pl
from datetime import datetime

start = '09:15'
end = '10:15'
start_time = datetime.time.fromisoformat(start)
end_time = datetime.time.fromisoformat(end)

# 1. 筛选目标时间窗口内的数据，计算每个时间点对应的扩展最大值（保留原始timestamp）
window_cum_max = und_df2.filter(
    pl.col('timestamp').dt.time().is_between(start_time, end_time)
).with_columns(
    # 按日期分组，计算组内Close的扩展最大值
    pl.col('Close').cum_max().over(pl.col('timestamp').dt.truncate('1d')).alias('fbarh3')
)

# 2. 将结果左连接回原DataFrame，匹配timestamp
und_df2 = und_df2.join(window_cum_max.select('timestamp', 'fbarh3'), on='timestamp', how='left')

代码说明 #

• 第一步：仅筛选09:15-10:15的行，通过over(pl.col('timestamp').dt.truncate('1d'))按日期分组，对每组内的Close价格计算cum_max，同时保留原始timestamp，确保每个cum_max值对应具体时间点。
• 第二步：直接以timestamp为连接键，将计算结果左连接回原表。原表中不在目标时间窗口内的行，fbarh3字段会自动填充为null，符合需求。

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