迁移至Polars时自定义函数无法在map_elements中运行的问题

问题核心 #

你遇到的两个关键问题：

1. 原Pandas风格的fetch_price函数无法适配Polars语法——Polars不支持df[date, self.name]这种按值索引的方式，且逐元素映射并非Polars的推荐高效用法。
2. 后续改用pl.when的逻辑完全错误：pl.col(self.name)取的是当前行的价格，而非Maturity_Date对应日期的价格，不符合需求。

推荐解决方案：向量化Join实现价格匹配（高效符合Polars设计） #

Polars的核心优势是向量化操作，我们可以通过创建价格查找表，再与主表做左连接来实现需求，避免低效的逐元素映射：

def get_rate(self, direction, margin, test_cycle):
    # 1. 准备原数据的价格查找表（去重确保每个日期唯一）
    price_lookup = df[['time', self.name]].drop_nulls().unique(subset='time')
    
    # 2. 截取目标周期的主数据
    self.df = (df[['time', self.name]].drop_nulls())[test_cycle * -1:]
    
    # 3. 生成Maturity_Date列
    self.df = self.df.with_columns(
        pl.col('time').map_elements(lambda x: next_trade_day(x, self.life_cycle)).alias('Maturity_Date')
    )
    
    # 4. 检查日期是否存在于原数据中
    self.df = self.df.with_columns(
        pl.col('Maturity_Date').is_in(price_lookup['time']).cast(bool).alias('Data_Accessibility')
    )
    
    # 5. 通过左连接匹配Maturity_Date对应的价格
    self.df = self.df.join(
        price_lookup,
        left_on='Maturity_Date',
        right_on='time',
        how='left',
        suffix='_maturity'
    ).with_columns(
        # 重命名匹配到的价格列，无匹配数据时自动为None
        pl.col(f'{self.name}_maturity').alias('Maturity_Price')
    ).drop(f'{self.name}_maturity', 'time_right')  # 清理冗余列
    
    print(self.df)

若坚持使用map_elements：修正自定义函数 #

如果一定要保留map_elements的用法，需要修改fetch_price适配Polars的索引逻辑：

# 修正后的fetch_price函数
def fetch_price(self, date):
    # Polars中按日期过滤并提取价格
    price_result = df.filter(pl.col('time') == date).select(self.name)
    # 存在数据则返回值，否则返回None
    return price_result.item() if len(price_result) > 0 else None

# 在get_rate中调用修正后的函数
def get_rate(self, direction, margin, test_cycle):
    # ... 其他原有代码保持不变 ...
    
    # 生成Maturity_Price列，指定元素类型避免Polars推断错误
    self.df = self.df.with_columns(
        pl.col('Maturity_Date').map_elements(
            lambda x: self.fetch_price(x),
            element_type=pl.Float64  # 根据实际价格类型调整
        ).alias('Maturity_Price')
    )
    
    print(self.df)

关键提示 #

• 优先选择Join方案：Polars的向量化操作比map_elements效率高几个数量级，尤其适合大数据集。
• 摒弃Pandas式索引：Polars中按值查找需使用.filter()或.select()，不能直接用df[value, column]的语法。

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