在分组的Polars DataFrame中查找日期交集并统一各组日期

原始数据 #

import polars as pl

df = pl.from_repr("""
┌────────┬────────────┐
│ symbol ┆ date       │
│ ---    ┆ ---        │
│ str    ┆ str        │
╞════════╪════════════╡
│ AAPL   ┆ 2023-01-01 │ 
│ AAPL   ┆ 2023-01-02 │
│ AAPL   ┆ 2023-01-03 │
│ AAPL   ┆ 2023-01-04 │
│ AAPL   ┆ 2023-01-05 │ # AAPL has 5 dates
│ GOOGL  ┆ 2023-01-01 │ 
│ GOOGL  ┆ 2023-01-02 │
│ GOOGL  ┆ 2023-01-03 │ # GOOGL has 3 dates
│ MSFT   ┆ 2023-01-01 │ 
│ MSFT   ┆ 2023-01-02 │
│ MSFT   ┆ 2023-01-03 │
│ MSFT   ┆ 2023-01-04 │ # MSFT has 4 dates
└────────┴────────────┘
""")

with pl.Config(tbl_rows=-1):
    print(df)

需求说明 #

过滤数据，只保留所有symbol分组都存在的日期，让每个分组的日期集合完全一致。

解决方案 #

方法一：分组统计+内连接

先统计不同symbol的总数，再找出每个日期对应的symbol数量等于总数的日期，最后通过内连接过滤原数据：

# 获取唯一symbol的数量
symbol_total = df.select(pl.col("symbol").n_unique()).item()

# 筛选所有分组都包含的日期
common_dates = df.group_by("date").agg(
    pl.col("symbol").n_unique().alias("symbol_count")
).filter(pl.col("symbol_count") == symbol_total).select("date")

# 过滤原数据
result = df.join(common_dates, on="date", how="inner")

# 输出结果
with pl.Config(tbl_rows=-1):
    print(result)

方法二：窗口函数过滤

用窗口函数直接计算每个日期对应的symbol数量，再筛选符合条件的记录，步骤更简洁：

symbol_total = df.select(pl.col("symbol").n_unique()).item()

result = df.with_columns(
    # 按日期分组，统计每个日期下的唯一symbol数
    pl.col("symbol").n_unique().over("date").alias("symbol_count")
).filter(pl.col("symbol_count") == symbol_total).drop("symbol_count")

with pl.Config(tbl_rows=-1):
    print(result)

输出结果 #

shape: (9, 2)
┌────────┬────────────┐
│ symbol ┆ date       │
│ ---    ┆ ---        │
│ str    ┆ str        │
╞════════╪════════════╡
│ AAPL   ┆ 2023-01-01 │
│ AAPL   ┆ 2023-01-02 │
│ AAPL   ┆ 2023-01-03 │
│ GOOGL  ┆ 2023-01-01 │
│ GOOGL  ┆ 2023-01-02 │
│ GOOGL  ┆ 2023-01-03 │
│ MSFT   ┆ 2023-01-01 │
│ MSFT   ┆ 2023-01-02 │
│ MSFT   ┆ 2023-01-03 │
└────────┴────────────┘

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