如何基于类型-对应id映射列表高效过滤Pandas DataFrame行？

核心思路 #

针对大数据量表、多规则的过滤场景，最优方案要满足两个核心要求：

• 性能优先：尽可能使用Pandas底层向量化操作、哈希匹配，避免逐行Python循环带来的性能损耗
• 可维护性优先：将所有业务规则从核心过滤逻辑中抽离为独立配置，后续新增规则无需改动主体代码

分步实现 #

1. 初始化测试数据 #

import pandas as pd
import numpy as np
# 构造示例DataFrame
df = pd.DataFrame({
    "type": ["type1","type1","type1","type2","type2","type2","type1/type2","type1/type2","type1/type2"],
    "id": [1,2,3,1,2,3,1,2,3],
    "values": [1,np.nan,7,np.nan,5,np.nan,1,1,1]
})

2. 规则配置层（完全独立于过滤逻辑） #

将所有业务规则拆成三类配置，后续新增规则仅需修改配置即可：

# 配置1：各type对应的合法id规则，支持两种形式：连续区间、枚举列表
legal_id_rule = {
    # type1 合法id为[1,3]闭区间（含1、2、3）
    "type1": {"range": [1, 3]},
    # type2 合法id为[2,3]闭区间（含2、3）
    "type2": {"range": [2, 3]},
    # type1/type2 合法id为枚举列表[1,2,3]
    "type1/type2": {"enum": [1,2,3]}
}
# 配置2：特殊NaN排除规则：满足type+id组合、且values为NaN时，需要过滤
nan_exclude_rule = [
    {"type": "type1", "id": 2}
]
# 配置3：特殊NaN保留规则：满足type+id组合时，即使values为NaN也保留
nan_keep_rule = [
    {"type": "type2", "id": 3}
]

3. 向量化标记计算 #

不使用逐行循环，全部通过布尔运算计算每行的规则匹配标记：

# 计算id合法性标记
df["is_legal_id"] = False
for t, rule in legal_id_rule.items():
    # 按type筛选当前行，匹配对应规则
    type_mask = df["type"] == t
    if "range" in rule:
        # 连续区间用向量化between判断，速度远快于逐行
        id_mask = df["id"].between(rule["range"][0], rule["range"][1])
    else:
        # 枚举值用isin判断
        id_mask = df["id"].isin(rule["enum"])
    df.loc[type_mask & id_mask, "is_legal_id"] = True
# 计算NaN排除标记
df["is_nan_exclude"] = False
for rule in nan_exclude_rule:
    mask = (df["type"] == rule["type"]) & (df["id"] == rule["id"])
    df.loc[mask, "is_nan_exclude"] = True
# 计算NaN保留标记
df["is_nan_keep"] = False
for rule in nan_keep_rule:
    mask = (df["type"] == rule["type"]) & (df["id"] == rule["id"])
    df.loc[mask, "is_nan_keep"] = True

4. 最终过滤 #

通过布尔掩码一次性完成过滤，全程无Python层逐行循环：

# 构造过滤掩码：
# 1. 基础要求：id属于合法范围
# 2. 排除条件：values为NaN + 属于需排除的组合 + 不属于需保留的组合
final_mask = (
    df["is_legal_id"]
    & ~(
        df["values"].isna()
        & df["is_nan_exclude"]
        & ~df["is_nan_keep"]
    )
)
# 应用掩码，清理中间标记列
final_df = df[final_mask].drop(columns=["is_legal_id", "is_nan_exclude", "is_nan_keep"])

方案优势 #

• 性能表现：所有判断操作都是Pandas底层C实现的向量化计算，对比逐行apply的实现方式，千万级数据量下处理速度可提升10-100倍
• 可维护性：业务规则与核心过滤逻辑完全解耦，新增规则时仅需在配置字典中添加条目，无需修改过滤主体代码，避免硬编码带来的维护成本
• 调试友好：每个规则的匹配结果都存储为独立的标记列，定位数据问题时可单独核查某类规则的匹配情况，不会出现逻辑混杂难以排查的问题

大数据量优化技巧 #

如果合法id的枚举列表量级很大（单type对应上万个合法id），可以将向量化判断替换为哈希merge的方式，进一步降低计算开销：

1. 将合法id配置转成两列（type、id）的查找小表
2. 通过df.merge(legal_id_lookup, how="left", on=["type", "id"])完成哈希匹配，比逐类型循环判断的速度更快

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