如何为Polars DataFrame连接添加基数验证？

Polars目前没有内置对应Pandas merge方法中validate参数的功能，但可以通过检查连接键的唯一性来实现完全相同的校验逻辑，比添加行号的方案简洁高效得多。

核心校验逻辑 #

对应Pandas的四种校验类型，我们可以通过Polars的is_unique()方法直接验证连接键的唯一性：

• 1:1（one_to_one）：左右DataFrame的连接键必须都唯一
• 1:m（one_to_many）：仅左表的连接键需要唯一
• m:1（many_to_one）：仅右表的连接键需要唯一
• m:m（many_to_many）：不做任何校验

实现代码示例 #

可以封装成一个工具函数，先执行校验再执行连接：

import polars as pl

def validated_join(left: pl.DataFrame, right: pl.DataFrame, on: str | list[str], validate: str):
    # 处理校验逻辑
    if validate in ("one_to_one", "1:1"):
        # 检查左表连接键唯一性
        if not left.select(on).is_unique().all():
            raise ValueError("左表连接键存在重复，不符合1:1连接要求")
        # 检查右表连接键唯一性
        if not right.select(on).is_unique().all():
            raise ValueError("右表连接键存在重复，不符合1:1连接要求")
    elif validate in ("one_to_many", "1:m"):
        if not left.select(on).is_unique().all():
            raise ValueError("左表连接键存在重复，不符合1:m连接要求")
    elif validate in ("many_to_one", "m:1"):
        if not right.select(on).is_unique().all():
            raise ValueError("右表连接键存在重复，不符合m:1连接要求")
    elif validate in ("many_to_many", "m:m"):
        # 多对多连接不校验
        pass
    else:
        raise ValueError(f"不支持的校验类型：{validate}")
    
    # 校验通过后执行连接
    return left.join(right, on=on)

使用示例 #

# 创建测试数据
left_11 = pl.DataFrame({"id": [1,2,3], "val_left": ["a","b","c"]})
right_11 = pl.DataFrame({"id": [1,2,3], "val_right": ["x","y","z"]})

# 执行1:1连接校验
validated_join(left_11, right_11, on="id", validate="1:1")

# 如果左表有重复键，会抛出错误
left_1m = pl.DataFrame({"id": [1,1,2], "val_left": ["a","a1","b"]})
validated_join(left_1m, right_11, on="id", validate="1:1")  # 触发ValueError

补充：连接后验证（可选） #

如果需要在连接完成后验证基数，也可以通过行数对比判断：

• 1:1连接：连接后行数 = 左表行数 = 右表行数（键无缺失的情况下）
• 1:m连接：连接后行数 ≥ 左表行数，且左表连接键唯一
• m:1连接：连接后行数 ≥ 右表行数，且右表连接键唯一

不过连接前校验的效率更高，可以避免不必要的大表连接操作。

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