Замените ячейку в столбце на основе ячейки в другом столбце в полярном DataFrame
Обратите внимание на следующее pl.DataFrame:
import polars as pl
df = pl.DataFrame(
{
"symbol": ["s1", "s1", "s2", "s2"],
"signal": [0, 1, 2, 0],
"trade": [None, 1, None, -1],
}
)
shape: (4, 3)
┌────────┬────────┬───────┐
│ symbol ┆ signal ┆ trade │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 │
╞════════╪════════╪═══════╡
│ s1 ┆ 0 ┆ null │
│ s1 ┆ 1 ┆ 1 │
│ s2 ┆ 2 ┆ null │
│ s2 ┆ 0 ┆ -1 │
└────────┴────────┴───────┘
Теперь мне нужно сгруппировать фрейм данных по symbol и проверить, не равна ли первая строка в каждой группе в столбце signal 0 (нулю). Если оно равно True, мне нужно заменить соответствующую ячейку в столбце trade значением в ячейке signal.
Вот что я на самом деле ищу:
shape: (4, 3)
┌────────┬────────┬───────┐
│ symbol ┆ signal ┆ trade │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 │
╞════════╪════════╪═══════╡
│ s1 ┆ 0 ┆ null │
│ s1 ┆ 1 ┆ 1 │
│ s2 ┆ 2 ┆ 2 │ <- copy value from the ``signal`` column
│ s2 ┆ 0 ┆ -1 │
└────────┴────────┴───────┘
Не обязательно.
Хорошо, хорошо. Я просто спросил, потому что иначе можно было бы упростить с помощью pl.Expr.fill_null.
Ответы 4
Для этого можно использовать конструкцию когда-то-иначе.
- Мы создаем условие, которое оценивается как
Trueровно для первых строк (создавайте индекс на лету, используя pl.int_range) в каждой группе сsignal, не равным 0. - В зависимости от этого условия мы выбираем значение в столбце
signalилиtrade.
df.with_columns(
trade=pl.when(
pl.col("signal") != 0,
pl.int_range(pl.len()) == 0,
).then("signal").otherwise("trade").over("symbol")
)
shape: (4, 3)
┌────────┬────────┬───────┐
│ symbol ┆ signal ┆ trade │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 │
╞════════╪════════╪═══════╡
│ s1 ┆ 0 ┆ null │
│ s1 ┆ 1 ┆ 1 │
│ s2 ┆ 2 ┆ 2 │
│ s2 ┆ 0 ┆ -1 │
└────────┴────────┴───────┘
Это умное решение! Мне было интересно, можем ли мы работать с first() вместо создания индекса на лету?
@ Энди Нет, насколько я знаю. Мы могли бы использовать pl.Expr.first, если бы хотели указать значение первой строки. Однако мы хотели бы поставить условие для текущей обработки первой строки (например, «нахождение» в первой строке).
@Andi Если у вас есть индексный столбец или любой столбец, содержащий уникальные значения в каждой группе symbol, вы можете заменить условие на pl.col("index") == pl.col("index").first(). В приведенном выше решении мы сами создаем индекс и сравниваем его с первым значением нашего индекса, которое всегда равно 0.
Это звучит разумно. Спасибо!
df.with_columns(
trade = pl.when(
pl.int_range(pl.len()) == 0,
pl.col.signal != 0
).then(
pl.col.signal
).otherwise(
pl.col.trade
).over("symbol")
)
┌────────┬────────┬───────┐
│ symbol ┆ signal ┆ trade │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 │
╞════════╪════════╪═══════╡
│ s1 ┆ 0 ┆ null │
│ s1 ┆ 1 ┆ 1 │
│ s2 ┆ 2 ┆ 2 │
│ s2 ┆ 0 ┆ -1 │
└────────┴────────┴───────┘
Использование cum_count над символом и , когда:
df.with_columns(
trade=pl.when(
pl.cum_count('symbol').over('symbol').eq(1),
pl.col('signal') != 0,
)
.then('signal')
.otherwise('trade')
)
Выход:
┌────────┬────────┬───────┐
│ symbol ┆ signal ┆ trade │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 │
╞════════╪════════╪═══════╡
│ s1 ┆ 0 ┆ null │
│ s1 ┆ 1 ┆ 1 │
│ s2 ┆ 2 ┆ 2 │
│ s2 ┆ 2 ┆ -2 │
└────────┴────────┴───────┘
Использование pl.cum_count работает только в том случае, если trade гарантированно отсутствует в первой строке каждой группы symbol.
@Hericks, действительно, он отличается от панд, тогда pl.cum_count('symbol').over('symbol').eq(1) все будет в порядке?
Да, в пандах это будет работать, поскольку cum_count создает индекс строки. В полярах он подсчитывает ненулевые элементы. Поскольку мы не можем быть уверены, что первый элемент в каждой группе (не) равен нулю, его нельзя использовать здесь. Вместо этого pl.int_range(pl.len()) можно использовать для создания индекса строки.
Думаю, эквивалент groupby.cumcount — это int_range.over ;)
Вы можете использовать .is_first_distinct() для идентификации нужных строк.
df.with_columns(pl.col.symbol.is_first_distinct().alias("checkme"))
shape: (4, 4)
┌────────┬────────┬───────┬─────────┐
│ symbol ┆ signal ┆ trade ┆ checkme │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 ┆ bool │
╞════════╪════════╪═══════╪═════════╡
│ s1 ┆ 0 ┆ null ┆ true │
│ s1 ┆ 1 ┆ 1 ┆ false │
│ s2 ┆ 2 ┆ null ┆ true │
│ s2 ┆ 0 ┆ -1 ┆ false │
└────────┴────────┴───────┴─────────┘
И используйте как часть логики «когда/то»:
df.with_columns(
pl.when(
pl.col.symbol.is_first_distinct(),
pl.col.signal != 0
)
.then("signal")
.otherwise("trade")
.alias("trade")
)
shape: (4, 3)
┌────────┬────────┬───────┐
│ symbol ┆ signal ┆ trade │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 │
╞════════╪════════╪═══════╡
│ s1 ┆ 0 ┆ null │
│ s1 ┆ 1 ┆ 1 │
│ s2 ┆ 2 ┆ 2 │
│ s2 ┆ 0 ┆ -1 │
└────────┴────────┴───────┘
это хорошо.
Будет ли здесь всегда отсутствовать значение столбца
trade?