Polars中join_asof含null的by列时保留左表空值行的方法

更新说明 #

该问题已解决，直接执行以下代码即可获得预期结果：

df.join_asof(df2, on="time", by=["a", "b"])

原问题场景 #

使用Polars的join_asof方法时，当by参数指定的列存在null值会触发报错，但业务需求需要保留左DataFrame中所有包含null值的行，同时完成asof连接。

示例数据 #

import polars as pl

df = pl.DataFrame(
    {"a": [1, 2, 3, 4, 5, None, 8], "b": [2, 3, 4, 5, 6, 7, None], "time": [1, 2, 3, 4, 5, 6, 7]}
)
df2 = pl.DataFrame({"a": [1, 3, 4, None], "b": [2, 4, 5, 8], "c": [2, 3, 4, 5], "time": [0, 2, 4, 6]})

直接执行的报错情况 #

运行以下代码会触发panic错误：

df.join_asof(df2, on="time", by=["a", "b"])

错误日志：

thread '<unnamed>' panicked at 'called `Result::unwrap()` on an `Err` value: ComputeError(Borrowed("cannot take slice"))', /home/runner/work/polars/polars/polars/polars-core/src/frame/asof_join/groups.rs:253:35

临时规避方案（丢失null行） #

删除by列中的null值后代码可正常运行，但会丢失左表中含null的行：

df.drop_nulls(["a", "b"]).join_asof(df2.drop_nulls(["a", "b"]), on="time", by=["a", "b"])

返回结果：

shape: (5, 4)
┌─────┬─────┬──────┬──────┐
│ a   ┆ b   ┆ time ┆ c    │
│ --- ┆ --- ┆ ---  ┆ ---  │
│ i64 ┆ i64 ┆ i64  ┆ i64  │
╞═════╪═════╪══════╪══════╡
│ 1   ┆ 2   ┆ 1    ┆ 2    │
│ 2   ┆ 3   ┆ 2    ┆ null │
│ 3   ┆ 4   ┆ 3    ┆ 3    │
│ 4   ┆ 5   ┆ 4    ┆ 4    │
│ 5   ┆ 6   ┆ 5    ┆ null │
└─────┴─────┴──────┴──────┘

预期结果 #

需要保留左表所有行，包括含null的行，最终预期输出如下：

┌─────┬─────┬──────┬──────┐
│ a   ┆ b   ┆ time ┆ c    │
│ --- ┆ --- ┆ ---  ┆ ---  │
│ i64 ┆ i64 ┆ i64  ┆ i64  │
╞═════╪═════╪══════╪══════╡
│ 1   ┆ 2   ┆ 1    ┆ 2    │
│ 2   ┆ 3   ┆ 2    ┆ null │
│ 3   ┆ 4   ┆ 3    ┆ 3    │
│ 4   ┆ 5   ┆ 4    ┆ 4    │
│ 5   ┆ 6   ┆ 5    ┆ null │
│ null┆ 7   ┆ 6    ┆ null │
│ 8   ┆ null┆ 7    ┆ null │
└─────┴─────┴──────┴──────┘

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