如何在Polars时间序列DataFrame中获取列的首次、末次非零值出现日期

示例数据 #

import polars as pl

df = pl.from_repr("""
┌────────────┬────────────┬────────────┐
│ date       ┆ column_one ┆ column_two │
│ ---        ┆ ---        ┆ ---        │
│ date       ┆ f64        ┆ i64        │
╞════════════╪════════════╪════════════╡
│ 2024-06-01 ┆ 0.0        ┆ 0          │
│ 2024-06-02 ┆ 0.0        ┆ 1          │
│ 2024-06-03 ┆ 1.0        ┆ 2          │
│ 2024-06-04 ┆ 1.2        ┆ 3          │
└────────────┴────────────┴────────────┘
""")

期望结果 #

shape: (2, 3)
┌────────────┬──────────────────┬─────────────────┐
│ columns    ┆ first_appearance ┆ last_appearance │
│ ---        ┆ ---              ┆ ---             │
│ str        ┆ date             ┆ date            │
╞════════════╪══════════════════╪═════════════════╡
│ column_one ┆ 2024-06-03       ┆ 2024-06-04      │
│ column_two ┆ 2024-06-02       ┆ 2024-06-04      │
└────────────┴──────────────────┴─────────────────┘

解决方案 #

方法1：遍历列处理（直观易懂） #

先筛选出需要处理的数值列，逐个过滤并提取日期：

# 获取除date外的数值列
value_cols = [col for col in df.columns if col != "date"]

result_rows = []
for col in value_cols:
    # 过滤当前列大于0的行
    filtered = df.filter(pl.col(col) > 0)
    # 提取首次和末次日期
    first_date = filtered.select(pl.col("date").first()).item()
    last_date = filtered.select(pl.col("date").last()).item()
    # 收集结果
    result_rows.append({
        "columns": col,
        "first_appearance": first_date,
        "last_appearance": last_date
    })

# 转换为目标DataFrame
result_df = pl.DataFrame(result_rows)
print(result_df)

方法2：Polars链式操作（更优雅高效） #

利用melt转长表，结合分组聚合实现：

value_cols = [col for col in df.columns if col != "date"]

result_df = (
    df.melt(id_vars="date", value_vars=value_cols)
    # 过滤值大于0的行
    .filter(pl.col("value") > 0)
    # 按列名分组
    .group_by("variable")
    # 聚合首次和末次日期
    .agg(
        first_appearance=pl.col("date").first(),
        last_appearance=pl.col("date").last()
    )
    # 重命名列匹配期望格式
    .rename({"variable": "columns"})
)

print(result_df)

两种方法都能得到符合要求的结果，方法2更贴合Polars的向量化操作风格，处理大数据集时效率更高。

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