如何向量化处理：将时间戳在指定区间的DataFrame行置为NaN

问题背景 #

现有两个Pandas DataFrame：

• data：包含时间戳列ts以及数据列val1、val2；
• intervals：包含start和end列，记录data中数据损坏的时间区间。

需求明确：若data某行的ts落在任意一个损坏区间内，就将该行所有值替换为NaN。目前采用for循环的朴素实现，想知道能否通过向量化优化提升效率？

现有朴素实现 #

def _remove_corruptions(corruption_filename: str, df_to_clean: pd.DataFrame):
    corruption_timestamps = pd.read_csv(corruption_filename)
    for j, row in tqdm(corruption_timestamps.iterrows()):
        start, end = row['start'], row['end']
        df_to_clean[(start <= df_to_clean['ts']) & (df_to_clean['ts'] <= end)] = np.nan
    return df_to_clean

示例数据 #

intervals = pd.DataFrame({'start': ['2019-06-23', '2020-01-10'], 'end': ['2019-06-24', '2020-01-12']})
# 输出：
#         start         end
# 0  2019-06-23  2019-06-24
# 1  2020-01-10  2020-01-12

data = pd.DataFrame({'ts': ['2019-06-23', '2019-10-24', '2020-01-11'], 'val1': [1, 1, 1], 'val2': [2, 2, 2]})
# 输出：
#            ts  val1  val2
# 0  2019-06-23     1     2
# 1  2019-10-24     1     2
# 2  2020-01-11     1     2

期望输出 #

out = pd.DataFrame({'ts': [np.nan, '2019-10-24', np.nan], 'val1': [np.nan, 1, np.nan], 'val2': [np.nan, 2, np.nan]})
# 输出：
#            ts  val1  val2
# 0         NaN   NaN   NaN
# 1  2019-10-24   1.0   2.0
# 2         NaN   NaN   NaN

向量化优化方案 #

当然可以！下面提供两种高效的向量化实现，彻底摆脱循环：

方案一：利用Numpy广播机制 #

这种方法直接通过数组运算生成掩码，适配任意数量的损坏区间，效率极高：

import pandas as pd
import numpy as np

def _remove_corruptions(corruption_filename: str, df_to_clean: pd.DataFrame):
    # 读取损坏区间并转换为datetime类型，避免字符串比较的潜在问题
    corruption_timestamps = pd.read_csv(corruption_filename, parse_dates=['start', 'end'])
    # 确保data的ts列也是datetime类型
    df_to_clean['ts'] = pd.to_datetime(df_to_clean['ts'])
    
    # 用广播机制生成布尔掩码：每个ts是否落在任意损坏区间内
    # 将data的ts转为二维数组（每行对应一个ts），和区间的start/end做逐元素比较
    mask = np.any(
        (corruption_timestamps['start'].values <= df_to_clean['ts'].values[:, np.newaxis]) & 
        (df_to_clean['ts'].values[:, np.newaxis] <= corruption_timestamps['end'].values),
        axis=1
    )
    
    # 一次性将符合条件的行全部设为NaN
    df_to_clean.loc[mask, :] = np.nan
    return df_to_clean

优化逻辑：

• 统一时间格式为datetime，避免字符串比较的bug；
• 通过Numpy广播，把data['ts']的一维数组扩展为二维，和所有损坏区间做比较，得到数据行数 × 区间数的布尔矩阵；
• 用np.any沿列取或，得到每个行是否需要标记的掩码；
• 一次性批量修改，避免循环中反复修改DataFrame，大幅减少IO开销。

方案二：使用merge_asof（适合区间不重叠场景） #

如果损坏区间不存在重叠，用merge_asof会更高效，它是Pandas专门为时间序列对齐设计的方法：

import pandas as pd

def _remove_corruptions(corruption_filename: str, df_to_clean: pd.DataFrame):
    corruption_timestamps = pd.read_csv(corruption_filename, parse_dates=['start', 'end'])
    df_to_clean['ts'] = pd.to_datetime(df_to_clean['ts'])
    
    # merge_asof要求右表按连接键排序
    corruption_timestamps = corruption_timestamps.sort_values('start')
    # 按ts左连接，找到每个ts对应的最近的"start <= ts"的区间
    merged = pd.merge_asof(df_to_clean.sort_values('ts'), 
                           corruption_timestamps, 
                           left_on='ts', 
                           right_on='start', 
                           direction='backward')
    # 判断当前ts是否<=对应区间的end，即是否在损坏区间内
    mask = merged['ts'] <= merged['end']
    # 还原原数据的顺序
    df_to_clean = df_to_clean.loc[merged.index]
    # 标记损坏行
    df_to_clean.loc[mask, :] = np.nan
    return df_to_clean

优化逻辑：

• 通过merge_asof快速为每个ts匹配对应的损坏区间；
• 只需一次比较就能生成掩码，时间复杂度为O(n log n)，比循环的O(n*m)高效得多。

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