Pandas多列取首个非空值及来源方法

问题背景与传统方法局限

在数据处理中，我们经常面临需要从多个潜在来源列中，根据特定条件（例如，选择第一个非空值）提取数据，并同时记录该数据来源列名的情况。numpy.select是处理条件赋值的强大工具，但其设计主要针对单列输出。当尝试使用np.select同时返回多个列时，例如一个值列和一个来源列，通常会遇到类型不匹配或维度错误，迫使开发者采取多次独立的np.select调用，即使这些调用共享相同的条件逻辑。

考虑以下DataFrame示例：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'A': [1.0, 2.0, np.nan],
    'B': [4, 5, 6]
})
print("原始DataFrame:")
print(df)

期望的输出是新增val和val_source两列，其中val优先取A列的值，如果A列为NaN，则取B列的值；val_source则记录该值来自A或B。

传统的np.select方法可能如下所示：

# 传统但繁琐的方法
conds = [df['A'].notna(), True]
choices_val = [df['A'], df['B']]
choices_src = ['A', 'B']

df['val'] = np.select(conds, choices_val, default=np.nan)
df['val_source'] = np.select(conds, choices_src, default=np.nan)
print("\n使用传统np.select方法的结果:")
print(df)

尽管此方法能够达到目的，但其缺点在于需要重复定义条件并进行两次独立的np.select操作，当涉及更多列或更复杂的条件时，代码会显得冗长且效率不高。

高效解决方案：结合Pandas布尔索引与NumPy argmax

为了解决上述问题，我们可以利用Pandas的布尔索引能力和NumPy的argmax函数，以一种更简洁高效的方式实现目标。核心思想是首先确定每行中第一个非空值所在的列位置，然后利用这个位置信息直接从原始DataFrame中提取值和对应的列名。

1. 确定每行第一个非空值的列索引

首先，我们需要识别DataFrame中哪些位置是非空的。df.notna()会返回一个布尔型的DataFrame，其中True表示非空，False表示空（NaN）。

# 步骤1: 获取非空值的布尔DataFrame并转换为NumPy数组
notna_df = df.notna()
print("\ndf.notna()结果:")
print(notna_df)

notna_array = notna_df.to_numpy()
print("\n转换为NumPy数组:")
print(notna_array)

接下来，使用NumPy的argmax函数。argmax(axis=1)会沿着每行（axis=1）查找第一个True值（即第一个非空值）的索引。如果一行中所有值都为False，argmax会返回0（即该行的第一个列索引），这在我们的场景中是可接受的，因为它会按照我们定义的优先级（例如，A优先于B）选择。

# 步骤2: 找出每行第一个非空值的列索引
idx = notna_array.argmax(axis=1)
print("\n每行第一个非空值的列索引 (idx):")
print(idx)

此时，idx数组 [0 0 1] 表示：

• 第0行：第一个非空值在索引0（即'A'列）。
• 第1行：第一个非空值在索引0（即'A'列）。
• 第2行：第一个非空值在索引1（即'B'列）。

2. 提取对应的值 (val)

有了每行第一个非空值的列索引idx，我们可以使用NumPy的“花式索引”（fancy indexing）来高效地从原始DataFrame的NumPy表示中提取对应的值。

# 步骤3: 提取对应的值
df_numpy = df.to_numpy()
val = df_numpy[(df.index, idx)] # df.index提供行索引，idx提供列索引
print("\n提取的val值:")
print(val)

df.to_numpy()[(df.index, idx)] 的工作原理是：df.index 提供了所有行的索引（0, 1, 2），而 idx 提供了每行对应的列索引（0, 0, 1）。NumPy会根据这些配对的索引，从df_numpy中选择出对应的元素。

3. 提取对应的来源列名 (val_source)

同样，利用idx数组，我们可以从DataFrame的列名列表中直接获取对应的列名。

# 步骤4: 提取对应的来源列名
val_source = df.columns[idx]
print("\n提取的val_source列名:")
print(val_source)

df.columns 返回一个包含所有列名的索引对象，通过df.columns[idx]，我们可以根据idx中的列索引，批量获取对应的列名。

4. 将结果添加到原始DataFrame

最后，将计算出的val和val_source添加到原始DataFrame中：

df['val'] = val
df['val_source'] = val_source

print("\n最终DataFrame:")
print(df)

完整示例代码

将上述步骤整合，形成完整的解决方案：

import pandas as pd
import numpy as np

# 原始DataFrame
df = pd.DataFrame({
    'A': [1.0, 2.0, np.nan],
    'B': [4, 5, 6]
})

print("原始DataFrame:")
print(df)

# 确定每行第一个非空值的列索引
# df.notna() -> 布尔型DataFrame
# .to_numpy() -> 转换为NumPy数组
# .argmax(axis=1) -> 找出每行第一个True（非空）的索引
idx = df.notna().to_numpy().argmax(axis=1)

# 提取对应的值
# df.to_numpy() -> 原始DataFrame的NumPy表示
# (df.index, idx) -> 使用花式索引，根据行索引和计算出的列索引提取值
df['val'] = df.to_numpy()[(df.index, idx)]

# 提取对应的来源列名
# df.columns -> DataFrame的列名列表
# df.columns[idx] -> 根据idx中的列索引获取对应的列名
df['val_source'] = df.columns[idx]

print("\n转换后的DataFrame:")
print(df)

优势与注意事项

• 简洁高效： 相比于多次调用np.select，此方法代码量更少，逻辑更集中，且利用了NumPy的向量化操作，性能优越，尤其适用于大型数据集。
• 灵活性： 这种方法能够灵活地处理多列选择，并同时获取值和其来源信息。
• 优先级处理： argmax默认返回第一个最大值的索引，这天然地符合了“优先选择第一个非空列”的需求。
• 适用场景： 主要适用于需要从一组列中，按照从左到右的优先级，选择第一个非空值及其来源的场景。如果需要更复杂的条件逻辑（例如，基于多个列的组合条件），可能仍需结合其他Pandas或NumPy功能。
• 全NaN行处理： 如果某一行所有列均为NaN，argmax会返回0（即第一列的索引）。在这种情况下，val将是NaN，val_source将是第一列的列名。这通常是可接受的行为，但如果需要特殊处理全NaN行，可能需要额外的检查。

总结

通过结合Pandas的布尔索引和NumPy的argmax函数，我们能够以一种优雅且高效的方式，解决从DataFrame多列中条件性地提取值并记录其来源的问题。这种方法不仅提升了代码的可读性和简洁性，也充分利用了底层库的性能优势，是处理此类数据转换任务的推荐实践。