Расчленение строк Polars по нескольким столбцам, но с другой логикой

У меня есть код, который разбивает столбец 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 python-polars

03.09.2024 16:07

Почему в Python есть оператор "pass"?

Оператор pass в Python - это простая концепция, которую могут быстро освоить даже новички без опыта программирования.

Некоторые методы, о которых вы не знали, что они существуют в Python

Python - самый известный и самый простой в изучении язык в наши дни. Имея широкий спектр применения в области машинного обучения, Data Science,...

Основы Python Часть I

Вы когда-нибудь задумывались, почему в программах на Python вы видите приведенный ниже код?

LeetCode - 1579. Удаление максимального числа ребер для сохранения полной проходимости графа

Алиса и Боб имеют неориентированный граф из n узлов и трех типов ребер:

Оптимизация кода с помощью тернарного оператора Python

И последнее, что мы хотели бы показать вам, прежде чем двигаться дальше, это

Советы по эффективной веб-разработке с помощью Python

Как веб-разработчик, Python может стать мощным инструментом для создания эффективных и масштабируемых веб-приложений.

Перейти к ответу Данный вопрос помечен как решенный

Ответы 2

Ответ принят как подходящий

Соедините соответствующее количество конечных 0 с price_paid перед вызовом .explode() на product и price_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          │
└───────────────────┴─────┴────────────┘

03.09.2024 16:53

вы можете использовать is_first_distinct() и когда/то:

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

03.09.2024 19:17

Другие вопросы по теме

Вычисление перекрестных рейтингов с использованием аккуратного полярного фрейма данных

Как перевести операции DataFrame pandas в Polars в Python?

Примените Scaler() к каждому идентификатору в фрейме данных Polars

Добавить X дней в серию Datetime

Продолжайте обучение модели Pytorch на новых данных

Условная дедупликация в полярах

Анализ числовых данных с разделителем тысяч в полярах

Как я могу объединить все строки с одним и тем же идентификатором в Polars?

Используйте поляры .when() вместо объединения

Polars – неожиданное поведение при использовании drop_nans() для всех столбцов

Расчленение строк Polars по нескольким столбцам, но с другой логикой

Ответы 2

Другие вопросы по теме

Похожие вопросы