Polars中从多列取非空值生成指定DataFrame并校验唯一性

需求说明 #

给定Polars DataFrame，需按照指定映射规则（remapping）将多列合并为单列，要求每行的源列组合仅能存在一个非空值，若出现多非空值则抛出异常；同时保留无关列，并支持批量处理多组合并场景。

实现代码 #

import polars as pl

# 原始DataFrame
df = pl.DataFrame(
    {
        "a": [1, None, None],
        "b": [None, 2, None],
        "c": [None, None, 3],
        "non_relevant": [1, 2, 3],
    }
)

# 映射规则
remapping = {"a": ("a", "b", "c"), "other_key": ("aN", "bN", "cN")}

def merge_columns_with_validation(df: pl.DataFrame, remapping: dict) -> pl.DataFrame:
    # 收集所有待处理的源列
    all_source_cols = set()
    for cols in remapping.values():
        all_source_cols.update(cols)
    
    # 保留无需处理的无关列
    non_source_cols = [col for col in df.columns if col not in all_source_cols]
    result_df = df.select(non_source_cols)
    
    # 批量处理每个映射规则
    for target_col, source_cols in remapping.items():
        # 计算每行源列的非空值数量
        non_null_count = df.select(pl.sum_horizontal(pl.col(source_cols).is_not_null()).alias("count"))
        
        # 校验是否存在多非空值的行
        invalid_rows = non_null_count.filter(pl.col("count") > 1)
        if len(invalid_rows) > 0:
            raise ValueError(f"列组合 {source_cols} 中存在 {len(invalid_rows)} 行有多个非空值，不符合要求")
        
        # 提取每行唯一非空值作为目标列
        merged_col = df.select(pl.coalesce(pl.col(source_cols)).alias(target_col))
        result_df = result_df.hstack([merged_col])
    
    return result_df

# 执行处理
try:
    result = merge_columns_with_validation(df, remapping)
    print(result)
except ValueError as e:
    print(e)

代码解析 #

1. 列分类处理：先区分待合并的源列和无关列，保留无关列作为结果基础，避免重复操作。
2. 非空唯一性校验：用sum_horizontal统计每行源列的非空值数量，若存在数量>1的行，直接抛出异常并提示问题行数。
3. 列合并逻辑：借助Polars内置的coalesce函数，自动提取每行源列中的唯一非空值，生成目标列。
4. 批量适配：通过遍历remapping字典，自动处理所有列合并需求，无需为每个组合单独编写代码。

测试结果 #

运行代码后，输出的result与预期一致：

shape: (3, 2)
┌─────┬──────────────┐
│ a   ┆ non_relevant │
│ --- ┆ ---          │
│ i64 ┆ i64          │
╞═════╪══════════════╡
│ 1   ┆ 1            │
│ 2   ┆ 2            │
│ 3   ┆ 3            │
└─────┴──────────────┘

异常场景验证 #

若修改原始DataFrame，让某一行存在多个非空值：

df_invalid = pl.DataFrame(
    {
        "a": [1, 5, None],
        "b": [None, 2, None],
        "c": [None, None, 3],
        "non_relevant": [1, 2, 3],
    }
)

执行处理时会抛出异常：列组合 ('a', 'b', 'c') 中存在 1 行有多个非空值，不符合要求

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