如何将Python中pandas DataFrame操作转换为Polars实现？

转换思路说明 #

Polars的行级操作逻辑与Pandas存在差异：Pandas中axis=1对应的行级聚合，在Polars中通过水平聚合函数（如max_horizontal、mean_horizontal）实现；同时Polars更推荐使用表达式语法完成逐元素计算，能充分发挥其向量化性能优势。

完整Polars实现代码 #

import polars as pl

# 创建与原Pandas结构一致的Polars DataFrame
df_a = pl.DataFrame({
    "feature1": [1, 2, 3],
    "feature2": [7, 8, 9],
})

df_b = pl.DataFrame({
    "feature1": [3, 8, 2],
    "feature2": [7, 4, 9],
})

selection_mode = "option1"  # 可切换为"option2"

if selection_mode == 'option1':
    # 1. 计算每行最大值（对应Pandas的df_a.max(axis=1)）
    max_values = df_a.select(pl.max_horizontal(pl.all()).alias("max_val"))
    
    # 2. 标记每行中等于最大值的元素（对应Pandas的df_a.eq(max_values, axis=0)）
    selected_features = df_a.with_columns(
        pl.all().eq(pl.col("max_val")).suffix("_selected")
    ).drop("max_val")
    
    # 3. 计算最终结果：选中特征对应df_b值求和 / 选中特征数量
    final_result = (
        selected_features.to_numpy() * df_b.to_numpy()
    ).sum(axis=1) / selected_features.to_numpy().sum(axis=1)
    
    # 更贴合Polars特性的表达式写法（无需转numpy）
    final_result = df_a.with_columns(
        max_val=pl.max_horizontal(pl.all()),
        selected=pl.all().eq(pl.col("max_val")),
        weighted=pl.col("selected") * df_b.to_dict(as_series=True),
    ).select(
        pl.sum_horizontal(pl.col("weighted.*")) / pl.sum_horizontal(pl.col("selected.*"))
    ).to_series()

elif selection_mode == 'option2':
    # 1. 计算每行均值（对应Pandas的df_a.mean(axis=1)）
    row_mean = df_a.select(pl.mean_horizontal(pl.all()).alias("row_mean"))
    
    # 2. 标记每行中大于等于均值的元素（对应Pandas的df_a.ge(df_a.mean(axis=1), axis=0)）
    above_avg = df_a.with_columns(
        pl.all().ge(pl.col("row_mean")).suffix("_above_avg")
    ).drop("row_mean")
    
    # 3. 计算分子：above_avg * df_a * df_b 行求和
    sum_combined = (above_avg.to_numpy() * df_a.to_numpy() * df_b.to_numpy()).sum(axis=1)
    
    # 4. 计算分母：above_avg * df_a 行求和
    sum_above_avg = (above_avg.to_numpy() * df_a.to_numpy()).sum(axis=1)
    
    # 5. 最终结果
    final_result = sum_combined / sum_above_avg
    
    # 更贴合Polars特性的表达式写法
    final_result = df_a.with_columns(
        row_mean=pl.mean_horizontal(pl.all()),
        above_avg=pl.all().ge(pl.col("row_mean")),
        combined=pl.col("above_avg.*") * pl.col("*") * df_b.to_dict(as_series=True),
        sum_above=pl.col("above_avg.*") * pl.col("*"),
    ).select(
        pl.sum_horizontal(pl.col("combined.*")) / pl.sum_horizontal(pl.col("sum_above.*"))
    ).to_series()

关键语法对应说明 #

• Pandas的df.max(axis=1) → Polars的pl.max_horizontal(pl.all())：水平聚合所有列，得到每行最大值
• Pandas的df.eq(max_values, axis=0) → Polars的pl.all().eq(pl.col("max_val"))：逐元素与行最大值比较
• Pandas的df.sum(axis=1) → Polars的pl.sum_horizontal(pl.col("列前缀.*"))：水平聚合指定前缀的列，实现行求和
• 若追求性能，优先使用Polars表达式语法直接关联df_b的列，避免转numpy的额外开销

内容的提问来源于stack exchange，提问作者Míriam