如何用Polars以向量化方式实现FIFO销售与采购数据匹配？

我有两个Polars DataFrame，分别存储采购数据和销售数据，希望采用FIFO（先进先出）规则匹配每笔销售对应的采购来源。已知可以通过循环实现该逻辑，但想找到更高效的向量化实现方式。

基础示例数据 #

采购数据 #

代码

from datetime import date
import polars as pl

df_buy = pl.DataFrame(
    {
        "BuyId": [1, 2],
        "Item": ["A", "A"],
        "BuyDate": [date.fromisoformat("2023-01-01"), date.fromisoformat("2024-03-07")],
        "Quantity": [40, 50],
    }
)

数据表格

销售数据 #

代码

df_sell = pl.DataFrame(
    {
        "SellId": [3, 4],
        "Item": ["A", "A"],
        "SellDate": [date.fromisoformat("2024-04-01"), date.fromisoformat("2024-05-01")],
        "Quantity": [10, 45],
    }
)

数据表格

预期FIFO匹配结果 #

补充测试数据 #

采购数据 #

df_buy = pl.DataFrame(
    {
        "BuyId": [5, 1, 2],
        "Item": ["B", "A", "A"],
        "BuyDate": [date.fromisoformat("2023-01-01"), date.fromisoformat("2023-01-01"), date.fromisoformat("2024-03-07")],
        "Quantity": [10, 40, 50],
    }
)

销售数据 #

df_sell = pl.DataFrame(
    {
        "SellId": [6, 3, 4],
        "Item": ["B", "A", "A"],
        "SellDate": [
            date.fromisoformat("2024-04-01"),
            date.fromisoformat("2024-04-01"),
            date.fromisoformat("2024-05-01"),
        ],
        "Quantity": [5, 10, 45],
    }
)

向量化实现方案 #

利用Polars的窗口函数、累积计算和区间匹配可以实现完全向量化的FIFO逻辑，避免循环，充分发挥Polars的性能优势。核心思路是按商品分组后，分别计算采购的累积库存区间和销售的累积销量区间，通过区间重叠关联匹配采购与销售记录。

实现代码 #

import polars as pl
from datetime import date

def fifo_match(df_buy: pl.DataFrame, df_sell: pl.DataFrame) -> pl.DataFrame:
    # 处理采购数据：按商品分组、采购日期排序，计算累积库存区间
    df_buy_processed = df_buy.sort(["Item", "BuyDate"]).with_columns(
        pl.col("Quantity").cumsum().over("Item").alias("cum_buy"),
        (pl.col("Quantity").cumsum().over("Item") - pl.col("Quantity")).alias("prev_cum_buy"),
        pl.col("Quantity").alias("RemainingQuantity")
    )
    
    # 处理销售数据：按商品分组、销售日期排序，计算累积销量区间
    df_sell_processed = df_sell.sort(["Item", "SellDate"]).with_columns(
        pl.col("Quantity").cumsum().over("Item").alias("cum_sell"),
        (pl.col("Quantity").cumsum().over("Item") - pl.col("Quantity")).alias("prev_cum_sell")
    )
    
    # 关联采购与销售数据，筛选区间重叠的记录
    joined = df_buy_processed.join(
        df_sell_processed,
        on="Item",
        how="inner"
    ).filter(
        (pl.col("prev_cum_sell") < pl.col("cum_buy")) & (pl.col("cum_sell") > pl.col("prev_cum_buy"))
    )
    
    # 计算实际销售数量和剩余库存
    result = joined.with_columns(
        pl.min([
            pl.col("cum_sell") - pl.col("prev_cum_buy"),
            pl.col("cum_buy") - pl.col("prev_cum_sell"),
            pl.col("Quantity"),
            pl.col("Quantity_right")
        ]).alias("SellQuantity"),
        (pl.col("RemainingQuantity") - pl.col("SellQuantity")).alias("QuantityAfterSell")
    ).select(
        "BuyId", "Item", "BuyDate", "RemainingQuantity",
        "SellId", "SellDate", "SellQuantity", "QuantityAfterSell"
    ).sort(["Item", "BuyDate", "SellDate"])
    
    return result

# 测试基础示例
base_result = fifo_match(df_buy, df_sell)
print("基础示例结果：")
print(base_result)

# 测试补充示例
ext_result = fifo_match(df_buy_ext, df_sell_ext)
print("\n补充示例结果：")
print(ext_result)

代码说明 #

1. 采购数据处理：按商品分组并按采购日期排序，计算每批次采购的累积库存区间（prev_cum_buy到cum_buy），保留初始剩余库存。
2. 销售数据处理：按商品分组并按销售日期排序，计算每笔销售的累积销量区间（prev_cum_sell到cum_sell）。
3. 区间关联：通过商品字段关联采购和销售数据，筛选出累积销量区间与采购累积库存区间有重叠的记录，这些就是FIFO规则下匹配的采购-销售对。
4. 匹配数量计算：取重叠区间的最小值作为本次销售从该采购批次中消耗的数量，再计算销售后的剩余库存。
5. 结果整理：选择目标字段并按商品、采购日期、销售日期排序，得到符合预期的结构化结果。

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