Кадр данных Polars: добавление столбцов при условии, что другой столбец дает разную длину ⇐ Python

[Anonymous]

У меня есть следующий кадр данных в полярах.

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

df = pl.DataFrame({
"Buy_Signal": [1, 0, 1, 0, 0],
"Returns": np.random.normal(0, 0.1, 5),
})

В конечном счете, я хочу выполнять агрегирование по столбцам Returns с учетом разных интервалов, которые задаются столбцом Buy_Signal. В приведенном выше случае длина равна единице до конца кадра данных.
Каков наиболее полярный способ сделать это?< /p>
Мое «глупое» решение (прежде чем я смогу применить агрегирование) выглядит следующим образом:

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

df = (df
.with_columns(pl.col("Buy_Signal").cum_sum().alias("Holdings"))

# (1) Determine returns for each holding period (>=1, >=2) and unpivot into one column
.with_columns(
pl.when(pl.col("Holdings") >= 1).then(pl.col("Returns")).alias("Port_1"),
pl.when(pl.col("Holdings") >= 2).then(pl.col("Returns")).alias("Port_2"),
)
)

Очевидно, что решение не работает для многих «сигналов на покупку». Поэтому я написал отдельную функцию Calc_port_returns, которая включает в себя цикл for, который затем передается в функцию Polar Pipe. См. здесь:

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

def calc_port_returns(_df: pl.DataFrame) -> pl.DataFrame:
n_portfolios = _df["Buy_Signal"].sum()
data = pl.DataFrame()
_df = _df.with_columns(pl.col("Buy_Signal").cum_sum().alias("Holdings"))
for i in range(1, n_portfolios + 1):
tmp = (
_df.with_columns(
pl.when(pl.col("Holdings") >= i).then(pl.col("Returns")).alias(f"Port_{i}"),
)
.select(pl.col(f"Port_{i}"))
)
data = pl.concat([data, tmp], how="horizontal")
_df = pl.concat([_df, data], how="horizontal")
return _df

df = pl.DataFrame({
"Buy_Signal": [1, 0, 1, 0, 0],
"Returns": np.random.normal(0, 0.1, 5),
})

df.pipe(calc_port_returns)

Каким «полярным» способом это сделать? В pandas я мог бы представить решение этой проблемы с помощью df.assign({f"Port_{i}": ... for i in range(1, ...)}) с несколькими предварительными дополнительными столбцами/побочными вычислениями .
Спасибо за любые предложения.

Подробнее здесь: https://stackoverflow.com/questions/778 ... ferent-len