Фильтрация внутри групп в полярах
Я новичок в Polars и мне нужен совет экспертов. У меня есть рабочий код, но я верю, что есть более быстрый и/или более элегантный способ сделать это. У меня есть большой фрейм данных с релевантными столбцами cik(int), form(string) и period(date). Форма может иметь значение «10-Q» или «10-K». Каждый цикл будет иметь множество строк двух типов форм с разными периодами. В итоге я хочу, чтобы для каждой группы cik остались только самые последние 10-Q и только самые последние 10 10-K. Конечно, если форм 10-К меньше 10, все должны остаться. Вот что я делаю сейчас (работает):
def filter_sub_for_11_rows_per_cik(df_):
df = df_.sort('cik')
# Keep only the last 10-Q
q_filtered_df = df.group_by('cik').map_groups(
lambda g:
g.sort('period', descending=True).filter(pl.col('form').eq('10-Q')).head(1))
# Keep the last up to 10 10-Ks
k_filtered_df = df.group_by('cik').map_groups(
lambda g:
g.sort('period', descending=True)
.filter(pl.col('form').eq('10-K'))
.slice(0, min(10, g.filter(pl.col('form').eq('10-K')).shape[0]))
)
return pl.concat([q_filtered_df, k_filtered_df])
@ouroboros1 Я не думаю, что это дубликат, поскольку в зависимости от групповых ключей (формы) следует применять разные фильтры.
Ответы 1
Чтобы упростить пример, я рассматриваю фрейм данных с 3 записями 10-Q и 2 10-K для каждого из двух значений cik. Я отфильтрую две самые последние строки 10-K и самую последнюю строку 10-Q для каждой группы, определенной cik.
import polars as pl
import datetime
df = pl.DataFrame({
"cik": [0] * 5 + [1] * 5,
"form": (["10-Q"] * 2 + ["10-K"] * 3) * 2,
"period": [datetime.date(2021, 1, 1+day) for day in range(10)],
})
shape: (10, 3)
┌─────┬──────┬────────────┐
│ cik ┆ form ┆ period │
│ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ date │
╞═════╪══════╪════════════╡
│ 0 ┆ 10-Q ┆ 2021-01-01 │
│ 0 ┆ 10-Q ┆ 2021-01-02 │
│ 0 ┆ 10-K ┆ 2021-01-03 │
│ 0 ┆ 10-K ┆ 2021-01-04 │
│ 0 ┆ 10-K ┆ 2021-01-05 │
│ 1 ┆ 10-Q ┆ 2021-01-06 │
│ 1 ┆ 10-Q ┆ 2021-01-07 │
│ 1 ┆ 10-K ┆ 2021-01-08 │
│ 1 ┆ 10-K ┆ 2021-01-09 │
│ 1 ┆ 10-K ┆ 2021-01-10 │
└─────┴──────┴────────────┘
Чтобы отфильтровать фрейм данных для каждой группы, определенной cik, мы можем просто использовать pl.DataFrame.filter вместе с pl.Expr.over (для определения групп) следующим образом.
(
df
.sort(by=["cik", "form", "period"], descending=[False, False, True])
.filter(
(
((pl.col("form") == "10-Q") & (pl.int_range(pl.len()) == 0)) |
((pl.col("form") == "10-K") & (pl.int_range(pl.len()) < 2))
)
.over("cik", "form")
)
)
shape: (6, 3)
┌─────┬──────┬────────────┐
│ cik ┆ form ┆ period │
│ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ date │
╞═════╪══════╪════════════╡
│ 0 ┆ 10-K ┆ 2021-01-05 │
│ 0 ┆ 10-K ┆ 2021-01-04 │
│ 0 ┆ 10-Q ┆ 2021-01-02 │
│ 1 ┆ 10-K ┆ 2021-01-10 │
│ 1 ┆ 10-K ┆ 2021-01-09 │
│ 1 ┆ 10-Q ┆ 2021-01-07 │
└─────┴──────┴────────────┘
Объяснение.
- Мы сортируем DataFrame в порядке убывания по дате для каждой группы, определенной cik и form.
- Мы фильтруем фрейм данных для строк с формой 10-K и индексом строки меньше 2 (0 или 1 - с вашими данными вы будете фильтровать строки с индексом строки меньше 10) или с формой 10-K и строкой индекс равен 0, то есть самая последняя запись. Мы используем
pl.Expr.over, чтобы выполнить эту фильтрацию отдельно для каждой группы, определенной cik и формы (чтобы гарантировать правильный сброс индекса для каждой формы).
Я убедился, что это дает тот же результат в моем большом фрейме данных, что и мой предыдущий код. Мне потребуется немного больше изучения, чтобы понять, как это работает, но это очень аккуратно.
Похоже, что есть способ использовать это для решения другой связанной проблемы — добавления столбца, содержащего номер элемента 10-K, начиная с 1 для самого последнего. (Все 10-Q можно установить в 0, предварительно инициализировав весь столбец как 0). Мне нужно лучше понять, что здесь происходит, чтобы добавить это
@MikeP Рад, что помогло! Я расскажу немного больше, как только вернусь к своему ноутбуку. Вы определенно можете получить желаемый столбец, используя выражения, аналогичные приведенным в фильтре выше.
Спасибо за помощь. Мне удалось найти способ получить числа, добавив в ваш код следующее: .with_columns(pl.col('form').cum_count().alias('fy_idx').over([' cik', 'form'])) .with_columns(pl.when(pl.col('form') == '10-Q') .then(0) .otherwise(pl.col('fy_idx')). псевдоним('fy_idx'))
@MikeP Это правильный подход - молодец!
Можете ли вы привести пример фрейма данных вместе с выводом данной функции?