python - 优化 Polars 中具有不同 'weekmask' 逻辑的两个日期之间的 timedelta 计算

Question

Polars 库给我留下了深刻的印象，并试图更好地学习它。 :)

现在，我正在尝试计算 Polars 中数百万行的两个日期之间的天数，但在某些情况下，对于某些行，我需要排除某些工作日。在 Pandas/Numpy 中，我使用了 np.busday_count ，我可以在其中定义一个周掩码，其中计算每个条件的工作日，并在需要时排除假期。

我在快速计算有条件的天数时遇到困难，因为我无法弄清楚如何在表达式中执行此操作。

示例数据框:

df = (pl
    .DataFrame({"Market": ["AT", "DE", "AT", "CZ", "GB", "CZ"],
              "Service": ["Standard", "Express", "Standard", "Standard", "Standard", "Standard"],
              "Day1": ["2022-01-02","2022-01-03", "2022-01-04", "2022-01-05", "2022-01-06", "2022-01-07"],
              "Day2": ["2022-01-03","2022-01-04", "2022-01-05", "2022-01-06", "2022-01-07", "2022-01-08"]
             }
            )
    .with_columns(pl.col(["Day1", "Day2"]).str.strptime(pl.Date, "%Y-%m-%d"))
)

我能够通过 struct 和 apply 方法将数据传递给 np.busday_function 。然而，与 Pandas 分配(262 毫秒)相比，真实数据集(34.4 秒)的执行速度要慢得多。

下面是我在 Polars 中想出的代码。我正在寻找一种更快的优化方法。

(df
    .with_column(
        pl.struct([pl.col("Day1"), pl.col("Day2")])
        .apply(lambda x: np.busday_count(x["Day1"], x["Day2"], weekmask='1110000'))
        .alias("Result"))
)

编辑，预期输出:

┌────────┬──────────┬────────────┬────────────┬────────┐
│ Market ┆ Service  ┆ Day1       ┆ Day2       ┆ Result │
│ ---    ┆ ---      ┆ ---        ┆ ---        ┆ ---    │
│ str    ┆ str      ┆ date       ┆ date       ┆ i64    │
╞════════╪══════════╪════════════╪════════════╪════════╡
│ AT     ┆ Standard ┆ 2022-01-02 ┆ 2022-01-03 ┆ 0      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ DE     ┆ Express  ┆ 2022-01-03 ┆ 2022-01-04 ┆ 1      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ AT     ┆ Standard ┆ 2022-01-04 ┆ 2022-01-05 ┆ 1      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ CZ     ┆ Standard ┆ 2022-01-05 ┆ 2022-01-06 ┆ 1      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ GB     ┆ Standard ┆ 2022-01-06 ┆ 2022-01-07 ┆ 0      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ CZ     ┆ Standard ┆ 2022-01-07 ┆ 2022-01-08 ┆ 0      │
└────────┴──────────┴────────────┴────────────┴────────┘

Answer 1

当您在 select 上下文中使用 apply 时，您将创建一个 Python 字典，并将其提供给列表中每个元素的 lambda。这很贵。

您可以通过使用map而不是apply来利用矢量化。这样我们就可以一次将整列发送到 numpys busday_count。

(df
    .with_column(
        pl.struct([pl.col("Day1"), pl.col("Day2")])
        .map(lambda x: np.busday_count(x.struct["Day1"], x.struct["Day2"], weekmask='1110000'))
        .alias("Result"))
)