Получить имя столбца максимальной даты на полярах
Я пытаюсь получить имя столбца, содержащее максимальное значение даты в моем DataFrame Polars. Я нашел похожий вопрос, на который уже был дан ответ здесь.
Однако в моем случае столбцов много, и добавлять их вручную было бы утомительно. Я хотел бы использовать селекторы столбцов cs.datetime() и попробовал следующее:
import polars as pl
from datetime import datetime
import polars.selectors as cs
data = {
"ID" : [1,2,3],
"Values_A" : [datetime(1,1,2),datetime(1,1,3),datetime(1,1,4)],
"Values_B" : [datetime(1,1,4),datetime(1,1,7),datetime(1,1,2)]
}
df = pl.DataFrame(data)
def arg_max_horizontal(*columns: pl.Expr) -> pl.Expr:
return (
pl.concat_list(columns)
.list.arg_max()
.replace_strict({i: col_name for i, col_name in enumerate(columns)})
)
(
df
.with_columns(
Largest=arg_max_horizontal(pl.select(cs.datetime()))
)
)
Мне нужно имя столбца с максимальной датой в строке
Ответы 1
Вы были в одном шаге от вас, и вам просто нужно было выбрать названия интересующих столбцов, используя df.select(cs.datetime()).columns. Затем мы можем распаковать список при вызове функции.
Примечание. Я соответствующим образом адаптировал подсказку arg_max_horizontal. Более того (спасибо @Cameron Riddell) мы можем упростить преобразование в строковое представление, используя приведение к подходящему pl.Enum.
def arg_max_horizontal(*columns: str) -> pl.Expr:
return (
pl.concat_list(columns)
.list.arg_max()
.cast(pl.Enum(columns))
)
(
df
.with_columns(
Largest=arg_max_horizontal(*df.select(cs.datetime()).columns)
)
)
shape: (3, 4)
┌─────┬─────────────────────┬─────────────────────┬──────────┐
│ ID ┆ Values_A ┆ Values_B ┆ Largest │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ datetime[μs] ┆ datetime[μs] ┆ str │
╞═════╪═════════════════════╪═════════════════════╪══════════╡
│ 1 ┆ 0001-01-02 00:00:00 ┆ 0001-01-04 00:00:00 ┆ Values_B │
│ 2 ┆ 0001-01-03 00:00:00 ┆ 0001-01-07 00:00:00 ┆ Values_B │
│ 3 ┆ 0001-01-04 00:00:00 ┆ 0001-01-02 00:00:00 ┆ Values_A │
└─────┴─────────────────────┴─────────────────────┴──────────┘
Вы можете избежать .replace_strict, используя вместо этого .cast(pl.Enum(columns)), поскольку list.arg_max() вернет вам основные физические значения категориально-подобного ряда.
@CameronRiddell Спасибо! Я отредактировал ответ, включив в него такой состав.
Вам нужна одна строка с именем столбца, содержащим максимальную дату в целом, или имя столбца с максимальной датой для каждой строки?