极地：填补每组缺失的月份

Question

我想根据每组中的最小和最大日期填充每组数据框中缺失的月份。

此方法有效，但使用

polars.apply

。

import polars as pl
import numpy as np
from datetime import date

DATA_SIZE = 10000000
raw_df = pl.DataFrame({
    "id": np.random.choice(range(0, 10000), DATA_SIZE),
    "date": pl.date_range(date(1940, 1, 1), date(2020, 1, 1), interval="1mo").sample(DATA_SIZE, with_replacement=True),
    "value": np.random.rand(DATA_SIZE)
})

clean_df = (raw_df
    .groupby("id")
    .agg(
        pl.struct([pl.col("date").min().alias("startDate"), pl.col("date").max().alias("endDate")])
        .apply(lambda row: pl.date_range(row["startDate"], row["endDate"], interval="1mo")).alias("date")
    )
    .explode("date")
    .join(raw_df, how="left", on=["id", "date"])
    .with_column(pl.col("value").fill_null(0))
)

既然

date_range

接受表达式，应该有一个更优雅、更快的解决方案，但我无法让它工作：

ValueError: did not expect value <built-in function id> of type <class 'builtin_function_or_method'>

(raw_df
    .groupby("id")
    .agg([
        pl.col("date").min().alias("startDate"), pl.col("date").max().alias("endDate")
    ])
    .select([
        id, pl.date_range(pl.col("startDate"), pl.col("endDate"), interval="1mo")
    ])
)

如何将

date_range

与表达式一起使用？

Answer 1

您的答案是正确的，但有一点语法错误，您错过了“id”周围的引号。错误消息有点神秘，因为

id

也是一个函数。

(raw_df
    .group_by("id")
    .agg(
        pl.col("date").min().alias("startDate"), pl.col("date").max().alias("endDate")
    )
    .select(
        "id", pl.date_ranges(pl.col("startDate"), pl.col("endDate"), interval="1mo").alias("date")
    )
    .explode("date")
    .join(raw_df, how="left", on=["id", "date"])
    .with_columns(pl.col("value").fill_null(0))
)