Как получить значение указанного индексного номера в результате сортировки столбца и заполнить его нулевым значением, если оно отсутствует?
import polars as pl
df = pl.DataFrame(
{"name": list("abcdef"), "age": [21, 31, 32, 53, 45, 26], "country": list("AABBBC")}
)
df.group_by("country").agg(
pl.col("name").sort_by("age").first().alias("age_sort_1"),
pl.col("name").sort_by("age").get(2).alias("age_sort_2"), # OutOfBoundsError: index out of bounds
# pl.col("name").sort_by("age").arr.get(2, null_on_oob=True).alias("age_2"),
# SchemaError: invalid series dtype: expected `FixedSizeList`, got `str`
pl.col("name").sort_by("age").last().alias("age_sort_-1")
)
Как показано в приведенном выше коде, я хочу получить имя в каждой стране, возраст которой указан в определенном порядке.
Однако Expr.get не предоставляет параметр null_on_oob. Как автоматически заполнить ноль при возникновении ситуации выхода за пределы поля?
Кроме того, метод .arr.get предоставляет null_on_oob parameter, но сообщает об ошибке SchemaError: invalid series dtype: expected "FixedSizeList", got "str". Я не знаю, к чему относится эта ошибка и как ее решить.
PS: В приведенном выше коде используется повторяющийся код pl.col("name").sort_by("age") много раз. Есть ли более краткий метод?
Ответы 2
import polars as pl
df = pl.DataFrame(
{"name": list("abcdef"), "age": [21, 31, 32, 53, 45, 26], "country": list("AABBBC")}
)
df.group_by("country").agg(
pl.col("name").sort_by("age").first().alias("age_sort_1"),
pl.col("name").sort_by("age").implode().list.get(2, null_on_oob=True).get(0).alias("age_sort_2"),
pl.col("name").sort_by("age").last().alias("age_sort_-1")
)
Приведенный выше код может решить проблему.
Но я не знаю, почему возвращаемый тип pl.col("name").sort_by("age").implode().list.get(2, null_on_oob=True) — list[str] вместо str и должен выполнить .get(0) снова, чтобы получить правильный результат.
Сейчас есть открытая проблема, связанная с вашим вопросом - Polars.Expr.take возвращает значение null, если ComputeError: индекс выходит за пределы.
На данный момент вы можете использовать shift() (maintain_order = True просто, чтобы сделать его более читабельным):
exp = pl.col.name.sort_by("age")
(
df
.group_by("country", maintain_order = True).agg(
exp.first().alias("age_sort_1"),
exp.shift(-2).first().alias("age_sort_2"),
exp.last().alias("age_sort_-1"),
)
)
┌─────────┬────────────┬────────────┬─────────────┐
│ country ┆ age_sort_1 ┆ age_sort_2 ┆ age_sort_-1 │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ str ┆ str │
╞═════════╪════════════╪════════════╪═════════════╡
│ A ┆ a ┆ null ┆ b │
│ B ┆ c ┆ d ┆ d │
│ C ┆ f ┆ null ┆ f │
└─────────┴────────────┴────────────┴─────────────┘
Или просто объедините свои данные в список, а затем используйте .list.get(), который позволяет вам использовать параметр null_on_oob:
(
df
.group_by("country").agg(
pl.col.name.sort_by("age")
)
.with_columns(
pl.col.name
.list.get(i, null_on_oob = True).alias(f"age_sort_{c}")
for i, c in [(0, 1), (2, 2), (-1, -1)]
).drop("name")
)
В качестве альтернативы вы можете использовать list.gather() , чтобы получить список из 3 нужных вам элементов, преобразовать его в struct с помощью метода list.to_struct() , а затем unnest() в столбцы:
(
df
.group_by("country").agg(
pl.col.name.sort_by("age")
)
.with_columns(
pl.col.name
.list.gather([0,2,-1], null_on_oob = True)
.list.to_struct(fields=["age_sort_1","age_sort_2","age_sort_-1"])
).unnest("name")
)