Início / coding / Perguntas / 78944749
Accepted
Phil-ZXX
Asked: 2024-09-03 22:07:15 +0800 CST

Explodir linhas polares em várias colunas, mas com lógica diferente

  • 772

Eu tenho esse código, que divide uma productcoluna em uma lista e depois a usa explodepara expandi-la:

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'))

Isso atualmente dá

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        │
└───────────────────┴─────┴────────────┘

No entanto, quero manter o total price paido mesmo. Então, depois de dividir as linhas em várias "subcategorias", quero mudar a tabela para isto:

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          │
└───────────────────┴─────┴────────────┘

ou seja, mantendo apenas o price_paidna primeira linha expandida. Então meu preço total pago permanece o mesmo. O qtyestá tudo bem em ficar do jeito que está.

Eu tentei eg with_columns(price_arr=pl.col('product').cast(pl.List(pl.Float64)))mas não consegui adicionar nada ao primeiro elemento da lista. Ou with_columns(price_arr=pl.col(['product', 'price_paid']).map_elements(price_func))mas não pareceu possível usar map_elementson pl.col([...]).

python

2 respostas

  1. Best Answer

    Concatene o número apropriado de 0s finais price_paidantes de chamar .explode()ambos producte price_paidao mesmo tempo:

    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")
)

    Saída:

    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. você pode usar is_first_distinct()e 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          │
└───────────────────┴─────┴────────────┘

    Você também pode evitar o uso map_elements()e usar expressões polares puras (é um pouco difícil, mas você entendeu a ideia):

    (
    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

relate perguntas