Polars分组后按年龄排序取指定索引值，越界填充null及代码优化

先看修正后的可运行代码：

import polars as pl

df = pl.DataFrame(
    {"name": list("abcdef"), "age": [21, 31, 32, 53, 45, 26], "country": list("AABBBC")}
)

result = df.group_by("country").agg(
    # 定义分组内按age排序后的name列表，避免重复代码
    sorted_names=pl.col("name").sort_by("age").list(),
    # 取排序后的第一个元素
    age_sort_1=pl.col("sorted_names").arr.get(0),
    # 取索引2的元素，越界时自动填充null
    age_sort_2=pl.col("sorted_names").arr.get(2, null_on_oob=True),
    # 取排序后的最后一个元素
    age_sort_-1=pl.col("sorted_names").arr.last()
).drop("sorted_names")  # 可选：删除中间变量列

print(result)

输出结果：

shape: (3, 4)
┌─────────┬──────────┬──────────┬────────────┐
│ country ┆ age_sort_1 ┆ age_sort_2 ┆ age_sort_-1 │
│ ---     ┆ ---      ┆ ---      ┆ ---        │
│ str     ┆ str      ┆ str      ┆ str        │
╞═════════╪══════════╪══════════╪════════════╡
│ A       ┆ a        ┆ null     ┆ b          │
│ B       ┆ f        ┆ d        ┆ d          │
│ C       ┆ e        ┆ null     ┆ e          │
└─────────┴──────────┴──────────┴────────────┘

逐个问题解析 #

问题3：简化重复代码 #

在agg中先定义中间列sorted_names，存储每个分组内按age排序后的name列表，后续所有依赖排序结果的操作直接引用该列即可，避免重复编写pl.col("name").sort_by("age")。不需要保留中间列时，用drop("sorted_names")删除即可。

问题2：arr.get报错原因及解决 #

报错SchemaError: invalid series dtype: expected FixedSizeList, got str的核心原因是：分组后直接调用pl.col("name").sort_by("age")返回的仍是str类型的Series（每个组的元素为单个字符串），而arr.get是列表/数组类型专属的表达式方法。

解决方式是在排序后追加.list()，将每个分组的name列转换为List[str]类型，这样就能正常调用arr.get方法。

问题1：索引越界自动填充null #

将分组后的name转为列表后，使用arr.get(index, null_on_oob=True)参数，即可在索引超出列表长度时自动返回null，而非抛出OutOfBoundsError。比如国家A的组仅2个元素，索引2越界，返回null；国家B的组有3个元素，索引2有效，返回对应值。

内容的提问来源于stack exchange，提问作者PierXuY