如何用Polars避免双重循环实现多场景产品营收预测？

一、核心需求：为每个产品匹配多场景价格列 #

你需要将pixies_df中的多场景价格（每行对应一个场景）匹配到main_df的对应产品行，生成场景列。无需循环的实现方式如下：

示例数据 #

import polars as pl

main_df = pl.from_repr("""
┌─────┬───────┐
│ xxx ┆ price │
│ --- ┆ ---   │
│ str ┆ str   │
╞═════╪═══════╡
│ A   ┆ 100   │
│ B   ┆ 150   │
│ C   ┆ 200   │
│ D   ┆ 250   │
│ A   ┆ 230   │
└─────┴───────┘
""")

pixies_df = pl.from_repr("""
┌─────┬─────┬─────┬─────┐
│ A   ┆ B   ┆ C   ┆ D   │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╪═════╡
│ 110 ┆ 160 ┆ 210 ┆ 260 │
│ 120 ┆ 170 ┆ 220 ┆ 270 │
│ 130 ┆ 180 ┆ 230 ┆ 280 │
└─────┴─────┴─────┴─────┘
""")

实现步骤 #

1. 将pixies_df重塑为长格式，添加场景索引（行号）
2. 与main_df按产品名称左连接
3. 将连接结果重塑回宽格式，得到每个场景的价格列

# 步骤1：重塑pixies_df为长格式，添加场景id
pixies_long = pixies_df.with_row_index("scene_id").melt(
    id_vars="scene_id",
    variable_name="xxx",
    value_name="scene_price"
)

# 步骤2：与main_df左连接
joined = main_df.join(pixies_long, on="xxx", how="left")

# 步骤3：重塑为宽格式，生成场景列
result = joined.pivot(
    index=["xxx", "price"],
    columns="scene_id",
    values="scene_price"
).sort("xxx")

print(result)

执行后会得到你期望的输出格式，全程依赖Polars的向量化操作，效率远高于循环实现。

二、大规模数据下的高效营收计算 #

针对SKU超10万、场景超千的情况，循环处理每个场景子集会导致内存崩溃，需改用长格式批量处理的方式，利用Polars的向量化和内存优化特性：

优化后的代码 #

import polars as pl
import numpy as np
from scipy.stats import norm

# 假设aa是包含产品信息和多场景价格列（simprice_开头）的主数据框
# 第一步：将所有场景列转成长格式，避免重复复制主数据
aa_long = aa.melt(
    id_vars=[col for col in aa.columns if not col.startswith("simprice_")],
    variable_name="scene_name",
    value_name="new_price"
)

# 定义向量化营收计算函数，直接在Polars DataFrame上操作
def revenue_function(df):
    prices = df["price"].to_numpy()
    new_prices = df["new_price"].to_numpy()
    
    # 价格调整
    price_adjusted = new_prices * np.random.uniform(0.9, 1.1, size=len(new_prices))
    # 正态分布因子计算
    mean_price = np.mean(prices)
    norm_dist_factor = norm.cdf(price_adjusted / mean_price)
    # 营收比例计算
    price_to_revenue_ratio = 1 + norm_dist_factor * np.random.uniform(0.8, 1.2, size=len(new_prices))
    # 基础营收计算
    revenues = price_adjusted * price_to_revenue_ratio
    # 添加随机噪声
    noise = np.random.normal(0, 0.05, size=len(new_prices))
    revenues = revenues * (1 + noise)
    revenues = np.round(revenues)
    
    # 返回带营收列的新DataFrame
    return df.with_columns(revenues=pl.Series(revenues))

# 批量计算所有场景的营收
aa_with_revenue = revenue_function(aa_long)

# 按场景聚合营收总和
scene_revenue_totals = aa_with_revenue.group_by("scene_name").agg(
    total_revenue=pl.col("revenues").sum()
)

print(scene_revenue_totals)

关键优化点 #

1. 长格式转换：将多场景列转成「场景标识+单价格列」的长格式，避免重复复制主数据，大幅降低内存占用
2. 向量化计算：利用NumPy的向量化操作替代逐行循环，Polars与NumPy兼容良好，性能拉满
3. 内存高效：Polars的延迟加载和内存优化机制，处理百万级数据时内存占用远低于Pandas

注意事项 #

• 原代码中prices_df.with_columns(revenues = pl.Series(revenues))未返回新对象，需改为return df.with_columns(...)，否则修改不会生效
• 若场景数极多，可改用Polars的lazy API进一步优化：将aa_long改为aa.lazy().melt(...)，计算时用.collect()触发执行

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