Я написал специальную функцию в Polars для создания списка выражений с горизонтальным заполнением вперед/назад. Функция принимает итерацию выражений (или имен столбцов) для определения порядка заполнения.
Я хочу использовать все столбцы через pl.all() по умолчанию. Проблема в том, что pl.all() возвращает одно выражение, а не итерацию, поэтому попытка перевернуть или перебрать его приводит к ошибке типа.
Есть ли способ преобразования между отдельными выражениями и итерациями выражений?
Любые предложения или обходные пути приветствуются!
Вот функция:
from typing import Iterable
from polars._typing import IntoExpr
import polars as pl
def fill_horizontal(exprs: Iterable[IntoExpr], forward: bool = True) -> list[pl.Expr]:
"""Generate a horizontal forward/backward fill list of expressions."""
# exprs = exprs or pl.all() # use all columns as default
cols = [col for col in reversed(exprs)] if forward else exprs
return [pl.coalesce(cols[i:]) for i in range(0, len(cols) - 1)]
Вот пример:
df = pl.DataFrame({
"col1": [1, None, 2],
"col2": [1, 2, None],
"col3": [None, None, 3]})
print(df)
# shape: (3, 3)
# ┌──────┬──────┬──────┐
# │ col1 ┆ col2 ┆ col3 │
# │ --- ┆ --- ┆ --- │
# │ i64 ┆ i64 ┆ i64 │
# ╞══════╪══════╪══════╡
# │ 1 ┆ 1 ┆ null │
# │ null ┆ 2 ┆ null │
# │ 2 ┆ null ┆ 3 │
# └──────┴──────┴──────┘
print('forward_fill')
print(df.with_columns(fill_horizontal(df.columns, forward=True)))
# shape: (3, 3)
# ┌──────┬──────┬──────┐
# │ col1 ┆ col2 ┆ col3 │
# │ --- ┆ --- ┆ --- │
# │ i64 ┆ i64 ┆ i64 │
# ╞══════╪══════╪══════╡
# │ 1 ┆ 1 ┆ 1 │
# │ null ┆ 2 ┆ 2 │
# │ 2 ┆ 2 ┆ 3 │
# └──────┴──────┴──────┘
print('backward_fill')
print(df.with_columns(fill_horizontal(df.columns, forward=False)))
# shape: (3, 3)
# ┌──────┬──────┬──────┐
# │ col1 ┆ col2 ┆ col3 │
# │ --- ┆ --- ┆ --- │
# │ i64 ┆ i64 ┆ i64 │
# ╞══════╪══════╪══════╡
# │ 1 ┆ 1 ┆ null │
# │ 2 ┆ 2 ┆ null │
# │ 2 ┆ 3 ┆ 3 │
# └──────┴──────┴──────┘
Изменить: объединение ответа @Henry Harbeck и комментария @jqurious кажется не идеальным, но на данный момент достаточным решением.
def fill_horizontal(
exprs: Iterable[IntoExpr] | None = None,
*,
forward: bool = True,
ncols: int = 1000) -> pl.Expr:
"""Generate a horizontal forward/backward fill expression."""
if exprs is None:
# if forward is false, ncols has to be defined with the present number of cols or more
cols = pl.all() if forward else pl.nth(range(ncols, -1, -1))
else:
cols = exprs if forward else reversed(exprs)
return pl.cum_reduce(lambda s1, s2: pl.coalesce(s2, s1), cols).struct.unnest()
Подробнее здесь: https://stackoverflow.com/questions/794 ... -in-polars
