Использование пользовательской функции в .when().then() [дубликат] ⇐ Python

[Anonymous]

Вот вопрос, в котором я спросил, как удобно создать выражение для длинной цепочки .when().then() без жесткого кодирования значений. Альтернативы для длинных цепочек .when().then().when().then().otherwise()
Я хотел создать новый вопрос, так как мне нужно использовать пользовательскую функцию в цепочке .when().then(), и я хотел бы также спросить об этом.
Имея следующий фрейм данных:

Код: Выделить всё

import polars as pl
import numpy as np
import datetime as dt

df = pl.DataFrame(
{
"Market":["AT", "AT", "DE", "DE", "CA", "DE", "UK", "US"],
"Number of Days":[1, 2, 3, 4, 3, 4, 2, 1],
}
)

Затем в словаре используются следующие параметры пользователя.
Дата начала = некоторая дата начала, определяемая пользователем.
Маска недели = указание дней недели, которые необходимо учитывать, например: «1111100» с понедельника по пятницу, а «1111101» — с понедельника по пятницу и вс.
В приведенном ниже случае все рынки имеют дату начала 01 января 2023 года, за исключением США — 01 июня 2023 года.

Код: Выделить всё

params = {
"AT": {"Start Date":dt.date(2023,1,1), "weekmask":"1111100"},
"DE": {"Start Date":dt.date(2023,1,1), "weekmask":"1111100"},
"CA": {"Start Date":dt.date(2023,1,1), "weekmask":"1111100"},
"UK": {"Start Date":dt.date(2023,1,1), "weekmask":"1111100"},
"US": {"Start Date":dt.date(2023,6,1), "weekmask":"1111111"},
}

Цель состоит в том, чтобы вычислить новую дату на основе вводимых пользователем данных и исторического «количества дней» в кадре данных, но с учетом рабочих дней, специфичных для рынка. Для этого я использую np.busday_offset и приведенную ниже пользовательскую функцию:

Код: Выделить всё

def calculate_target_date(days_offset, start_date, weekmask):
# Calculate the target dates with np.busday_offset
dates_array = np.busday_offset(
dates=start_date,
offsets=days_offset,
weekmask=weekmask
)
return pl.Series(dates_array.tolist())

Создание выражения .when().then() путем перебора словаря параметров.

Код: Выделить всё

# Creating the expression
expr = pl.when(False).then(None)
for market, data in params.items():
condition = pl.col("Market") == market
result = pl.col("Number of Days").map(
lambda s: calculate_target_date(
days_offset=s,
start_date=data["Start Date"],
weekmask=data["weekmask"],
)
)
expr = expr.when(condition).then(result)

expr = expr.alias("New Target Date")

А затем используя выражение в Polars

Код: Выделить всё

df.with_columns([expr])
shape: (8, 3)
┌────────┬────────────────┬─────────────────┐
│ Market ┆ Number of Days ┆ New Target Date │
│ ---    ┆ ---            ┆ ---             │
│ str    ┆ i64            ┆ date            │
╞════════╪════════════════╪═════════════════╡
│ AT     ┆ 1              ┆ 2023-06-02      │
│ AT     ┆ 2              ┆ 2023-06-03      │
│ DE     ┆ 3              ┆ 2023-06-04      │
│ DE     ┆ 4              ┆ 2023-06-05      │
│ CA     ┆ 3              ┆ 2023-06-04      │
│ DE     ┆ 4              ┆ 2023-06-05      │
│ UK     ┆ 2              ┆ 2023-06-03      │
│ US     ┆ 1              ┆ 2023-06-02      │
└────────┴────────────────┴─────────────────┘

Выходные данные рассчитали busday_offset для всех рынков, используя «Дата начала» рынка США (июнь), а не «Дата начала» для конкретного рынка в параметрах. Почему это так и как мне заставить его работать правильно?

Подробнее здесь: https://stackoverflow.com/questions/765 ... -when-then