Поляры, объединение продаж и покупок, метод ФИФО

1. Рассчитайте общее количество купленных на данный момент и общее количество проданных товаров.

• cum_sum() для расчета промежуточной суммы.

df_buy_total = (
    df_buy
    .with_columns(QuantityTotal = pl.col.Quantity.cum_sum().over("Item"))
)

df_sell_total = (
    df_sell
    .with_columns(QuantityTotal = pl.col.Quantity.cum_sum().over("Item"))
)

┌───────┬──────┬────────────┬──────────┬───────────────┐
│ BuyId ┆ Item ┆ BuyDate    ┆ Quantity ┆ QuantityTotal │
│ ---   ┆ ---  ┆ ---        ┆ ---      ┆ ---           │
│ i64   ┆ str  ┆ date       ┆ i64      ┆ i64           │
╞═══════╪══════╪════════════╪══════════╪═══════════════╡
│ 1     ┆ A    ┆ 2023-01-01 ┆ 40       ┆ 40            │
│ 2     ┆ A    ┆ 2024-03-07 ┆ 50       ┆ 90            │
└───────┴──────┴────────────┴──────────┴───────────────┘

┌────────┬──────┬────────────┬──────────┬───────────────┐
│ SellId ┆ Item ┆ SellDate   ┆ Quantity ┆ QuantityTotal │
│ ---    ┆ ---  ┆ ---        ┆ ---      ┆ ---           │
│ i64    ┆ str  ┆ date       ┆ i64      ┆ i64           │
╞════════╪══════╪════════════╪══════════╪═══════════════╡
│ 3      ┆ A    ┆ 2024-04-01 ┆ 10       ┆ 10            │
│ 4      ┆ A    ┆ 2024-05-01 ┆ 45       ┆ 55            │
└────────┴──────┴────────────┴──────────┴───────────────┘

2. Найдите, какие BuyId и SellId принадлежат друг другу. Для этого мы можем использовать join_asof()

df_sell2buy = (
    df_sell_total
    .join_asof(
        df_buy_total, on = "QuantityTotal", strategy = "forward", by = "Item", coalesce=False, suffix = "Buy"
    )
    .select(
        "BuyId", "SellId", "Item", "BuyDate", "SellDate",
        "QuantityTotalBuy",
        pl.col.QuantityTotal.alias("QuantityTotalSell")
    )
    .filter(pl.col.BuyId.is_not_null(), pl.col.QuantityTotalBuy > pl.col.QuantityTotalSell)
)

df_buy2sell = (
    df_buy_total
    .join_asof(
        df_sell_total, on = "QuantityTotal", strategy = "forward", by = "Item", coalesce=False, suffix = "Sell"
    )
    .select(
        "BuyId", "SellId", "Item", "BuyDate", "SellDate",
        pl.col.QuantityTotal.alias("QuantityTotalBuy"),
        pl.col.QuantityTotalSell.forward_fill().over("Item").fill_null(0)
    )
)

┌───────┬────────┬──────┬────────────┬────────────┬──────────────────┬───────────────────┐
│ BuyId ┆ SellId ┆ Item ┆ BuyDate    ┆ SellDate   ┆ QuantityTotalBuy ┆ QuantityTotalSell │
│ ---   ┆ ---    ┆ ---  ┆ ---        ┆ ---        ┆ ---              ┆ ---               │
│ i64   ┆ i64    ┆ str  ┆ date       ┆ date       ┆ i64              ┆ i64               │
╞═══════╪════════╪══════╪════════════╪════════════╪══════════════════╪═══════════════════╡
│ 1     ┆ 3      ┆ A    ┆ 2023-01-01 ┆ 2024-04-01 ┆ 40               ┆ 10                │
│ 2     ┆ 4      ┆ A    ┆ 2024-03-07 ┆ 2024-05-01 ┆ 90               ┆ 55                │
└───────┴────────┴──────┴────────────┴────────────┴──────────────────┴───────────────────┘

┌───────┬────────┬──────┬────────────┬────────────┬──────────────────┬───────────────────┐
│ BuyId ┆ SellId ┆ Item ┆ BuyDate    ┆ SellDate   ┆ QuantityTotalBuy ┆ QuantityTotalSell │
│ ---   ┆ ---    ┆ ---  ┆ ---        ┆ ---        ┆ ---              ┆ ---               │
│ i64   ┆ i64    ┆ str  ┆ date       ┆ date       ┆ i64              ┆ i64               │
╞═══════╪════════╪══════╪════════════╪════════════╪══════════════════╪═══════════════════╡
│ 1     ┆ 4      ┆ A    ┆ 2023-01-01 ┆ 2024-05-01 ┆ 40               ┆ 55                │
│ 2     ┆ null   ┆ A    ┆ 2024-03-07 ┆ null       ┆ 90               ┆ 55                │
└───────┴────────┴──────┴────────────┴────────────┴──────────────────┴───────────────────┘

3. Объедините оба связанных DataFrame

df_result = (
    pl.concat([df_sell2buy, df_buy2sell])
    .filter(~(
        pl.col.SellId.is_null() &
        pl.col.SellId.is_not_null().any().over("BuyId")
    ))
    .sort("Item", "BuyId", "SellId", nulls_last=True)
)

┌───────┬────────┬──────┬────────────┬────────────┬──────────────────┬───────────────────┐
│ BuyId ┆ SellId ┆ Item ┆ BuyDate    ┆ SellDate   ┆ QuantityTotalBuy ┆ QuantityTotalSell │
│ ---   ┆ ---    ┆ ---  ┆ ---        ┆ ---        ┆ ---              ┆ ---               │
│ i64   ┆ i64    ┆ str  ┆ date       ┆ date       ┆ i64              ┆ i64               │
╞═══════╪════════╪══════╪════════════╪════════════╪══════════════════╪═══════════════════╡
│ 1     ┆ 3      ┆ A    ┆ 2023-01-01 ┆ 2024-04-01 ┆ 40               ┆ 10                │
│ 1     ┆ 4      ┆ A    ┆ 2023-01-01 ┆ 2024-05-01 ┆ 40               ┆ 55                │
│ 2     ┆ 4      ┆ A    ┆ 2024-03-07 ┆ 2024-05-01 ┆ 90               ┆ 55                │
└───────┴────────┴──────┴────────────┴────────────┴──────────────────┴───────────────────┘

4. Подсчитать конечный результат

(
    df_result
    .with_columns(
        QuantityRunning = pl.min_horizontal("QuantityTotalBuy", "QuantityTotalSell"),
    )    
    .with_columns(
        QuantityRunningPrev = pl.col.QuantityRunning.shift(1).over("Item").fill_null(0)
    )
    .select(
        "BuyId", "Item", "BuyDate",
        (pl.col.QuantityTotalBuy - pl.col.QuantityRunningPrev).alias("RemainingQuantity"),
        "SellId", "SellDate", 
        (pl.col.QuantityRunning - pl.col.QuantityRunningPrev).alias("SellQuantity"),
        (pl.col.QuantityTotalBuy - pl.col.QuantityRunning).alias("QuantityAfterSell")
    )
)

┌───────┬──────┬────────────┬───────────────────┬────────┬────────────┬──────────────┬───────────────────┐
│ BuyId ┆ Item ┆ BuyDate    ┆ RemainingQuantity ┆ SellId ┆ SellDate   ┆ SellQuantity ┆ QuantityAfterSell │
│ ---   ┆ ---  ┆ ---        ┆ ---               ┆ ---    ┆ ---        ┆ ---          ┆ ---               │
│ i64   ┆ str  ┆ date       ┆ i64               ┆ i64    ┆ date       ┆ i64          ┆ i64               │
╞═══════╪══════╪════════════╪═══════════════════╪════════╪════════════╪══════════════╪═══════════════════╡
│ 1     ┆ A    ┆ 2023-01-01 ┆ 40                ┆ 3      ┆ 2024-04-01 ┆ 10           ┆ 30                │
│ 1     ┆ A    ┆ 2023-01-01 ┆ 30                ┆ 4      ┆ 2024-05-01 ┆ 30           ┆ 0                 │
│ 2     ┆ A    ┆ 2024-03-07 ┆ 50                ┆ 4      ┆ 2024-05-01 ┆ 15           ┆ 35                │
└───────┴──────┴────────────┴───────────────────┴────────┴────────────┴──────────────┴───────────────────┘