Polars中合并共享ID列但数据列不同的DataFrame以获取完整ID集合的方法

我最近碰到了一个需求：要合并两个共享id列，但各自拥有不同数据列的Polars DataFrame，最终得到包含所有id的完整结果——不管这个id只在其中一个DataFrame里存在，还是在两个里都有。我是这么写代码实现的，给大家参考：

import polars as pl
import sys

red_data = pl.DataFrame(
    [
        pl.Series("id", [0, 1, 2], dtype=pl.UInt8()),
        pl.Series("red_data", [1, 0, 1], dtype=pl.UInt8()),
    ]
)
blue_data = pl.DataFrame(
    [
        pl.Series("id", [0, 2, 3], dtype=pl.UInt8()),
        pl.Series("blue_data", [0, 1, 1], dtype=pl.UInt8()),
    ]
)

# 获取同时存在于red和blue中的数据
red_and_blue = red_data.join(blue_data, on=["id"])
# 获取只在red中存在、blue中没有的数据，补充blue_data列为null
red_not_blue = red_data.join(blue_data, on=["id"], how="anti").with_columns(
    blue_data=pl.lit(None, dtype=pl.UInt8())
)
# 获取只在blue中存在、red中没有的数据，补充red_data列为null
blue_not_red = blue_data.join(red_data, on=["id"], how="anti").with_columns(
    red_data=pl.lit(None, dtype=pl.UInt8())
)

# 统一列顺序并合并所有结果
columns = ["id", "red_data", "blue_data"]
sys.displayhook(
    pl.concat(
        [
            red_and_blue.select(columns),
            red_not_blue.select(columns),
            blue_not_red.select(columns),
        ]
    )
)

运行这段代码后，得到的结果如下：

shape: (4, 3)
┌─────┬──────────┬───────────┐
│ id  ┆ red_data ┆ blue_data │
│ --- ┆ ---      ┆ ---       │
│ u8  ┆ u8       ┆ u8        │
╞═════╪══════════╪═══════════╡
│ 0   ┆ 1        ┆ 0         │
│ 2   ┆ 1        ┆ 1         │
│ 1   ┆ 0        ┆ null      │
│ 3   ┆ null     ┆ 1         │
└─────┴──────────┴───────────┘

这个方法通过分别处理三种情况（交集、左独有、右独有），再合并结果，完美覆盖了所有id的情况，并且保证了数据类型的一致性。

备注：内容来源于stack exchange，提问作者bzm3r