如何用简洁高效的Pandas方法按规则填充连续NaN数据块？

场景与示例数据 #

我是Python和Pandas新手，多个传感器以不同间隔（通常每分钟多次）向数据库写入数据。将这些数据读入Pandas DataFrame后，重采样为1分钟粒度并聚合数值，由于col_a和col_b存在无数据的分钟，重采样后的DataFrame中这些列包含NaN，示例代码如下：

import pandas as pd
import numpy as np

d = {'tstamptz': [  '2022-08-09 10:30:00+00:02',
                    '2022-08-09 10:30:01+00:02',
                    '2022-08-09 10:30:02+00:02',
                    '2022-08-09 10:30:03+00:02',
                    '2022-08-09 10:30:04+00:02',
                    '2022-08-09 10:30:05+00:02',
                    '2022-08-09 10:30:06+00:02',
                    '2022-08-09 10:30:07+00:02',
                    '2022-08-09 10:30:08+00:02',
                    '2022-08-09 10:30:09+00:02',
                    '2022-08-09 10:30:10+00:02'],
     'col_a': [413.8, 409.6, np.nan, np.nan, 409.6, 405.1, 417.3, np.nan, np.nan, np.nan, 426.4],
     'col_b': [409.6, 407.9, np.nan, np.nan, 405.1, 417.3, 431.4, np.nan, np.nan, np.nan, 419.9],
     'col_c': [5.1, 5.1, 5.0, 5.0, 5.0, 5.0, 4.9, 4.9, 5.0, 5.0, 5.1]}

df = pd.DataFrame(d)
df.set_index('tstamptz', inplace=True)

缺失值填充规则 #

• 将col_b的最后有效值填入连续NaN块的第一行col_a；
• 将col_a在NaN块后的第一个有效值填入连续NaN块的第一行col_b；
• 后续连续NaN行将col_b的值向前填充到col_a；
• 规则需应用于所有连续NaN块，遇到非NaN行后重新执行。

现有实现方案 #

当前方案可正常运行，但实现繁琐：

result = df.copy()
last_index_position = result.index.get_loc(result.index[-1])

# While there are any NaN rows
while result[['col_a','col_b']].isnull().values.any():

    # Get index position of first NaN row
   first_nan = list(np.where(result[['col_a','col_b']].isna()))[0][0]
   
   try:
      # Check, up to which index position it contains consecutive NaN rows
      next_nonan_index = result.index.get_loc(
                              result[['col_a','col_b']].iloc[first_nan:].first_valid_index()
                         )
   # Catch the KeyError, if there is no valid index left in dataframe (last index is NaN)
   except KeyError:
      next_nonan_index = last_index_position + 1 # add 1 to include last index in ffill method

   # Make sure, that there's a non-NaN row in front of the NaN row(s)
   # If the row on position 0 is NaN, take 0 instead of -1
   if first_nan - 1 >= 0:
      fillblock_start = first_nan -1
   else:
      fillblock_start = 0
   
   # Fill the first NaN row
   result['col_a'].iloc[fillblock_start:next_nonan_index].fillna(
      df['col_b'].ffill(), limit = 1, inplace = True)
   result['col_b'].iloc[fillblock_start:next_nonan_index].fillna(
      df['col_a'].bfill(), limit = 1, inplace = True)
   

   # Fill the rest of the rows
   result['col_a'].iloc[fillblock_start:next_nonan_index].fillna(
      result['col_b'].ffill(), inplace = True)   
   result['col_b'].iloc[fillblock_start:next_nonan_index].fillna(
      result['col_b'].ffill(), inplace = True)

优化需求 #

如何使用更简洁且高效的Pandas方法实现上述需求？

解决方案 #

可以通过识别连续NaN块，结合ffill、bfill和分组处理来简化实现，避免循环操作：

import pandas as pd
import numpy as np

# 复制原数据
result = df.copy()

# 标记连续的NaN块：当col_a或col_b为NaN时标记，然后生成分组键
is_nan_block = result[['col_a', 'col_b']].isna().any(axis=1)
group_key = (is_nan_block != is_nan_block.shift()).cumsum()

# 处理每个分组
for _, group in result.groupby(group_key):
    # 仅处理NaN块
    if is_nan_block.loc[group.index[0]]:
        start_idx = group.index[0]
        end_idx = group.index[-1]
        
        # 填充NaN块第一行的col_a：取col_b的最后有效值
        result.loc[start_idx, 'col_a'] = result['col_b'].ffill().loc[start_idx]
        # 填充NaN块第一行的col_b：取col_a的下一个有效值
        result.loc[start_idx, 'col_b'] = result['col_a'].bfill().loc[start_idx]
        
        # 后续行将col_b向前填充到col_a，同时完成col_b自身的前向填充
        result.loc[start_idx:end_idx, 'col_a'] = result.loc[start_idx:end_idx, 'col_b'].ffill()
        result.loc[start_idx:end_idx, 'col_b'] = result.loc[start_idx:end_idx, 'col_b'].ffill()

代码说明 #

1. 识别连续NaN块：通过is_nan_block标记每行是否属于NaN块，再用cumsum()生成分组键，将连续的NaN行归为同一组。
2. 分组处理：遍历每个分组，仅处理NaN块：
   • 对块的第一行，用ffill()获取col_b的历史有效值填充col_a，用bfill()获取col_a的未来有效值填充col_b；
   • 对块内后续行，直接将col_b的前向填充值赋值给col_a，同时确保col_b自身完成前向填充。
3. 效率提升：避免了原方案的循环判断，利用Pandas的分组和向量化操作，运行效率更高，代码更简洁。

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