如何在Rust Polars中通过自定义函数为DataFrame创建新列？

问题核心 #

直接在with_columns中调用普通Rust函数unlevered_beta_f会失败，原因是Polars的with_columns仅接受Polars表达式（Expr类型），而普通函数处理的是具体数值，无法适配Polars的延迟计算机制。

解决方案 #

需要将自定义的无杠杆贝塔计算逻辑包装为返回Expr的函数，以下分两种场景实现：

场景1：计算逻辑可通过Polars内置表达式组合实现 #

如果unlevered_beta_f是常规公式（比如无杠杆贝塔=杠杆贝塔/[1+(1-税率)*(债务/市值)-(现金/企业价值)]），直接用Polars的Expr操作组合即可：

use polars::prelude::*;

// 定义返回Expr的无杠杆贝塔计算函数
fn unlevered_beta_expr(
    levered_beta: Expr,
    debt_to_equity: Expr,
    marginal_tax_rate: f64,
    effective_tax_rate: Expr,
    cash_to_firm_value: Expr,
) -> Expr {
    // 替换为你的实际计算逻辑，全部用Polars Expr操作实现
    levered_beta / (
        lit(1.0) + (lit(1.0) - effective_tax_rate) * debt_to_equity - cash_to_firm_value
    )
}

// 在处理DataFrame的函数中调用
fn industry_total_beta_f(marginal_tax_rate: f64) -> Result<DataFrame> {
    // 读取CSV数据
    let df = CsvReader::from_path("your_data.csv")?.finish()?;

    // 添加新列
    let df_with_unlevered_beta = df.with_columns([
        unlevered_beta_expr(
            col("Average of Beta"),
            col("Sum of Total Debt incl leases (in US $)") / col("Sum of Market Cap (in US $)"),
            marginal_tax_rate,
            col("Average of Effective Tax Rate"),
            col("Sum of Cash") / col("Sum of Firm Value (in US $)"),
        ).alias("Average Unlevered Beta")
    ]);

    Ok(df_with_unlevered_beta?)
}

场景2：计算逻辑复杂，无法用内置表达式组合实现 #

如果unlevered_beta_f包含复杂逻辑（比如分支判断、外部依赖调用），需要用map方法将多列数据打包后，逐行应用自定义函数：

use polars::prelude::*;

// 原始的数值级无杠杆贝塔计算函数
fn unlevered_beta_f(
    levered_beta: f64,
    debt_to_equity: f64,
    marginal_tax_rate: f64,
    effective_tax_rate: f64,
    cash_to_firm_value: f64,
) -> f64 {
    // 这里保留你的原始计算逻辑
    levered_beta / (1.0 + (1.0 - effective_tax_rate) * debt_to_equity - cash_to_firm_value)
}

// 包装为Polars Expr函数
fn unlevered_beta_expr(
    levered_beta: Expr,
    debt_to_equity: Expr,
    marginal_tax_rate: f64,
    effective_tax_rate: Expr,
    cash_to_firm_value: Expr,
) -> Expr {
    // 将多列打包，逐行应用自定义函数
    levered_beta
        .zip([debt_to_equity, effective_tax_rate, cash_to_firm_value])
        .map(
            |series| {
                // 提取各列的数值迭代器
                let lb_series = series.f64()?;
                let dte_series = series.f64()?;
                let etr_series = series.f64()?;
                let ctfv_series = series.f64()?;

                // 逐行计算并收集结果
                let result = lb_series
                    .into_iter()
                    .zip(dte_series.into_iter())
                    .zip(etr_series.into_iter())
                    .zip(ctfv_series.into_iter())
                    .map(|(((lb, dte), etr), ctfv)| {
                        match (lb, dte, etr, ctfv) {
                            (Some(lb_val), Some(dte_val), Some(etr_val), Some(ctfv_val)) => {
                                Some(unlevered_beta_f(lb_val, dte_val, marginal_tax_rate, etr_val, ctfv_val))
                            }
                            _ => None, // 处理缺失值
                        }
                    })
                    .collect::<Series>();

                Ok(result)
            },
            GetOutput::from_type(DataType::Float64), // 指定输出列类型
        )
}

// 调用方式同场景1
fn industry_total_beta_f(marginal_tax_rate: f64) -> Result<DataFrame> {
    let df = CsvReader::from_path("your_data.csv")?.finish()?;

    let df_with_unlevered_beta = df.with_columns([
        unlevered_beta_expr(
            col("Average of Beta"),
            col("Sum of Total Debt incl leases (in US $)") / col("Sum of Market Cap (in US $)"),
            marginal_tax_rate,
            col("Average of Effective Tax Rate"),
            col("Sum of Cash") / col("Sum of Firm Value (in US $)"),
        ).alias("Average Unlevered Beta")
    ]);

    Ok(df_with_unlevered_beta?)
}

关键注意点 #

• 所有传入with_columns的操作必须是Expr类型，不能直接调用处理单个数值的普通函数。
• 处理缺失值时要在map中显式判断，避免运行时 panic。
• 优先使用场景1的实现方式，因为Polars对内置表达式会做优化，性能远高于逐行处理的map方法。

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