У меня есть код, который разбивает столбец 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 имеет богатую стандартную библиотеку, включая модули для работы с текстом, файлами и сетями.
Соедините соответствующее количество конечных 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 │
└───────────────────┴─────┴────────────┘
вы можете использовать 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 │
└───────────────────┴─────┴────────────┘