如何在Polars中提取k标记组的成对j值并按列布局？

测试数据构造 #

先还原示例中的DataFrame（可固定数据方便验证）：

import polars as pl
import numpy as np

# 随机生成数据
df = pl.DataFrame(dict(
    j=np.random.randint(10, 99, 9),
    k=np.tile([1, 2, 2], 3),
))

# 或者使用示例固定数据
# df = pl.DataFrame({
#     "j": [47,22,82,19,85,15,89,74,26],
#     "k": [1,2,2,1,2,2,1,2,2]
# })

核心实现代码 #

result = (
    df
    # 生成组ID：k=1的行标记新组，后续行继承组ID
    .with_columns(
        group_id=pl.when(pl.col("k") == 1)
                  .then(pl.int_range(0, pl.count()))
                  .forward_fill()
    )
    # 按组ID分组，保留原始顺序
    .group_by("group_id", maintain_order=True)
    .agg(
        # 提取组内k=1对应的j值（每组仅一个）
        start_j=pl.col("j").filter(pl.col("k") == 1).first(),
        # 提取组内最后一个k=2对应的j值
        end_j=pl.col("j").filter(pl.col("k") == 2).last()
    )
    # 重命名列匹配需求格式
    .rename({"start_j": "j", "end_j": "k"})
    .drop("group_id")
)

print(result)

代码说明 #

1. 生成组ID：利用when+forward_fill，在k=1的行生成递增的组序号，后续k=2的行自动继承上一组的ID，实现按k=1分隔分组。
2. 分组聚合：
   • start_j筛选组内k=1的j值，因为每组只有一个起始行，用first()确保取到唯一值；
   • end_j筛选组内所有k=2的j值，用last()取该组的最后一个末尾值。
3. 格式调整：将聚合后的列重命名为需求的j和k，删除临时的group_id列。

输出结果 #

shape: (3, 2)
┌─────┬─────┐
│ j   ┆ k   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 47  ┆ 82  │
│ 19  ┆ 15  │
│ 89  ┆ 26  │
└─────┴─────┘

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