主页 / coding / 问题 / 78944749
Accepted
Phil-ZXX
Asked: 2024-09-03 22:07:15 +0800 CST

在多列上分解极坐标行，但逻辑不同

  • 772

我有这段代码，它将一product列拆分为一个列表，然后用它explode来扩展它：

import polars as pl
import datetime as dt
from dateutil.relativedelta import relativedelta

def get_3_month_splits(product: str) -> list[str]:
    front, start_dt, total_m = product.rsplit('.', 2)
    start_dt = dt.datetime.strptime(start_dt, '%Y%m')
    total_m  = int(total_m)
    return [f'{front}.{(start_dt+relativedelta(months=m)).strftime("%Y%m")}.3' for m in range(0, total_m, 3)]

df = pl.DataFrame({
    'product':    ['CHECK.GB.202403.12', 'CHECK.DE.202506.6', 'CASH.US.202509.12'],
    'qty':        [10, -20, 50],
    'price_paid': [1400, -3300, 900],
})

print(df.with_columns(pl.col('product').map_elements(get_3_month_splits, return_dtype=pl.List(str))).explode('product'))

目前这给出

shape: (10, 3)
┌───────────────────┬─────┬────────────┐
│ product           ┆ qty ┆ price_paid │
│ ---               ┆ --- ┆ ---        │
│ str               ┆ i64 ┆ i64        │
╞═══════════════════╪═════╪════════════╡
│ CHECK.GB.202403.3 ┆ 10  ┆ 1400       │
│ CHECK.GB.202406.3 ┆ 10  ┆ 1400       │
│ CHECK.GB.202409.3 ┆ 10  ┆ 1400       │
│ CHECK.GB.202412.3 ┆ 10  ┆ 1400       │
│ CHECK.DE.202506.3 ┆ -20 ┆ -3300      │
│ CHECK.DE.202509.3 ┆ -20 ┆ -3300      │
│ CASH.US.202509.3  ┆ 50  ┆ 900        │
│ CASH.US.202512.3  ┆ 50  ┆ 900        │
│ CASH.US.202603.3  ┆ 50  ┆ 900        │
│ CASH.US.202606.3  ┆ 50  ┆ 900        │
└───────────────────┴─────┴────────────┘

但是，我想保持总数price paid不变。因此，在将行拆分为几个“子类别”后，我想将表格更改为：

shape: (10, 3)
┌───────────────────┬─────┬────────────┐
│ product           ┆ qty ┆ price_paid │
│ ---               ┆ --- ┆ ---        │
│ str               ┆ i64 ┆ i64        │
╞═══════════════════╪═════╪════════════╡
│ CHECK.GB.202403.3 ┆ 10  ┆ 1400       │
│ CHECK.GB.202406.3 ┆ 10  ┆ 0          │
│ CHECK.GB.202409.3 ┆ 10  ┆ 0          │
│ CHECK.GB.202412.3 ┆ 10  ┆ 0          │
│ CHECK.DE.202506.3 ┆ -20 ┆ -3300      │
│ CHECK.DE.202509.3 ┆ -20 ┆ 0          │
│ CASH.US.202509.3  ┆ 50  ┆ 900        │
│ CASH.US.202512.3  ┆ 50  ┆ 0          │
│ CASH.US.202603.3  ┆ 50  ┆ 0          │
│ CASH.US.202606.3  ┆ 50  ┆ 0          │
└───────────────────┴─────┴────────────┘

即只将 保留price_paid在第一个展开行中。因此我支付的总价保持不变。qty保持原样就可以了。

我尝试了例如with_columns(price_arr=pl.col('product').cast(pl.List(pl.Float64)))但无法向列表的第一个元素添加任何内容。或者with_columns(price_arr=pl.col(['product', 'price_paid']).map_elements(price_func))但似乎无法map_elements使用pl.col([...])。

python

2 个回答

  1. Best Answer

    在同时调用和之前，将适当数量的尾随0s 连接到：price_paid.explode()productprice_paid

    print(
    df.with_columns(
        pl.col("product").map_elements(get_3_month_splits, return_dtype=pl.List(str))
    )
    .with_columns(
        pl.concat_list(
            pl.col("price_paid"), pl.lit(0).repeat_by(pl.col("product").list.len() - 1)
        )
    )
    .explode("product", "price_paid")
)

    输出：

    shape: (10, 3)
┌───────────────────┬─────┬────────────┐
│ product           ┆ qty ┆ price_paid │
│ ---               ┆ --- ┆ ---        │
│ str               ┆ i64 ┆ i64        │
╞═══════════════════╪═════╪════════════╡
│ CHECK.GB.202403.3 ┆ 10  ┆ 1400       │
│ CHECK.GB.202406.3 ┆ 10  ┆ 0          │
│ CHECK.GB.202409.3 ┆ 10  ┆ 0          │
│ CHECK.GB.202412.3 ┆ 10  ┆ 0          │
│ CHECK.DE.202506.3 ┆ -20 ┆ -3300      │
│ CHECK.DE.202509.3 ┆ -20 ┆ 0          │
│ CASH.US.202509.3  ┆ 50  ┆ 900        │
│ CASH.US.202512.3  ┆ 50  ┆ 0          │
│ CASH.US.202603.3  ┆ 50  ┆ 0          │
│ CASH.US.202606.3  ┆ 50  ┆ 0          │
└───────────────────┴─────┴────────────┘
    • 1

  2. 您可以使用is_first_distinct()和when/then：

    (
    df.with_columns(
        pl.col.product.alias("group"),
        pl.col.product.map_elements(get_3_month_splits, return_dtype=pl.List(str))
    ).explode("product")
    .with_columns(
        pl.when(pl.col.group.is_first_distinct()).then(pl.col.price_paid).otherwise(0)
    ).drop("group")
)

┌───────────────────┬─────┬────────────┐
│ product           ┆ qty ┆ price_paid │
│ ---               ┆ --- ┆ ---        │
│ str               ┆ i64 ┆ i64        │
╞═══════════════════╪═════╪════════════╡
│ CHECK.GB.202403.3 ┆ 10  ┆ 1400       │
│ CHECK.GB.202406.3 ┆ 10  ┆ 0          │
│ CHECK.GB.202409.3 ┆ 10  ┆ 0          │
│ CHECK.GB.202412.3 ┆ 10  ┆ 0          │
│ CHECK.DE.202506.3 ┆ -20 ┆ -3300      │
│ CHECK.DE.202509.3 ┆ -20 ┆ 0          │
│ CASH.US.202509.3  ┆ 50  ┆ 900        │
│ CASH.US.202512.3  ┆ 50  ┆ 0          │
│ CASH.US.202603.3  ┆ 50  ┆ 0          │
│ CASH.US.202606.3  ┆ 50  ┆ 0          │
└───────────────────┴─────┴────────────┘

    您还可以避免使用map_elements()而使用纯极坐标表达式（这有点粗糙，但您明白了）：

    (
    df.with_columns(
        pl.col.product.str.split('.')
    ).with_columns(
        start_dt = pl.col.product.list.get(2).str.to_datetime("%Y%m"),
        offset = pl.col.product.list.get(3).cast(pl.Int64) - 3
    ).with_columns(
        dates = pl.date_ranges(
            pl.col.start_dt,
            pl.col.start_dt.dt.offset_by(pl.col.offset.cast(pl.String) + "mo"),
            interval="3mo"
        )
    ).explode("dates")
    .with_columns(
        pl.concat_list(
            pl.col.product.list.head(2),
            pl.col.start_dt.dt.strftime("%Y%m"),
            pl.lit(3)
        ).list.join(".")
    )
    .with_columns(
        pl.when(pl.col.product.is_first_distinct().over("start_dt", "offset")).then(pl.col.price_paid).otherwise(0)
    ).drop("start_dt", "offset", "dates")
)

┌───────────────────┬─────┬────────────┐
│ product           ┆ qty ┆ price_paid │
│ ---               ┆ --- ┆ ---        │
│ str               ┆ i64 ┆ i64        │
╞═══════════════════╪═════╪════════════╡
│ CHECK.GB.202403.3 ┆ 10  ┆ 1400       │
│ CHECK.GB.202403.3 ┆ 10  ┆ 0          │
│ CHECK.GB.202403.3 ┆ 10  ┆ 0          │
│ CHECK.GB.202403.3 ┆ 10  ┆ 0          │
│ CHECK.DE.202506.3 ┆ -20 ┆ -3300      │
│ CHECK.DE.202506.3 ┆ -20 ┆ 0          │
│ CASH.US.202509.3  ┆ 50  ┆ 900        │
│ CASH.US.202509.3  ┆ 50  ┆ 0          │
│ CASH.US.202509.3  ┆ 50  ┆ 0          │
│ CASH.US.202509.3  ┆ 50  ┆ 0          │
└───────────────────┴─────┴────────────┘
    • 0

相关问题