修改Pandas interpolate逻辑：仅填充长度≤limit的连续NaN段

问题说明 #

pandas原生interpolate方法的limit参数逻辑为：对长连续NaN段，最多填充前limit个值，剩余位置保留NaN，无法满足以下需求：
仅当单列内某段连续NaN的总长度≤设定limit阈值时，才对该段所有NaN执行插值；若连续NaN总长度超过limit，整段全部保留NaN不填充

测试数据 #

import pandas as pd
import numpy as np

df = pd.DataFrame({'col1': [0, np.nan, np.nan, np.nan, 3, 4],
                   'col2': [np.nan, 1, 2, np.nan, 4, np.nan],
                   'col3': [4, np.nan, np.nan, 7, 10, 11]})
print(df)

初始数据输出：

col1  col2  col3
0   0.0   NaN   4.0
1   NaN   1.0   NaN
2   NaN   2.0   NaN
3   NaN   NaN   7.0
4   3.0   4.0  10.0
5   4.0   NaN  11.0

原生方法调用df.interpolate(method="linear", limit=2, limit_area="inside")时，col1索引1-3的3个连续NaN会被填充前2个、保留第3个，不符合预期。

预期结果（limit=2，线性插值，仅填充内部NaN段） #

col1  col2  col3
0   0.0   NaN   4.0
1    NaN   1.0   5.0
2    NaN   2.0   6.0
3    NaN   3.0   7.0
4   3.0   4.0  10.0
5   4.0   NaN  11.0

结果校验逻辑：

• col1索引1-3共3个连续NaN，长度超过limit=2，整段不填充
• col2索引3位置共1个连续内部NaN，长度符合阈值，插值得到3.0；首尾边缘NaN不填充
• col3索引1-2共2个连续内部NaN，长度符合阈值，线性插值得到5.0、6.0

实现代码 #

核心逻辑：逐列识别连续NaN段，标记长度超过阈值的段位置，先执行原生插值，再将超长NaN段的位置还原为NaN。

def custom_interpolate(df, method="linear", limit=2, limit_area="inside"):
    result = df.copy()
    for col in result.columns:
        col_series = result[col]
        is_na = col_series.isna()
        # 为每段连续NaN分配唯一分组ID
        na_group_ids = (~is_na).cumsum()[is_na]
        # 统计每个NaN段的长度
        group_length = na_group_ids.map(na_group_ids.value_counts())
        # 标记长度超过阈值的NaN位置
        long_na_mask = group_length > limit
        # 执行原生插值
        interpolated_col = col_series.interpolate(method=method, limit_area=limit_area)
        # 超长NaN段还原为空值
        interpolated_col[long_na_mask] = np.nan
        result[col] = interpolated_col
    return result

验证调用 #

output = custom_interpolate(df, method="linear", limit=2, limit_area="inside")
print(output)

运行输出和预期结果完全一致。

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