使用lambda和pandas优化测试数据处理的Python代码

前置准备 #

首先定义通用常量避免硬编码，后续调整字段无需修改核心逻辑：

import pandas as pd

# 需参与计算的数值列，排除标识类字段
NUM_COLS = ["D1", "D10", "D50", "D90", "D99", "Q3_15", "Q3_10", "l-Q3_63", "RPM"]

第一步：数据清洗（满足规则1、2） #

def clean_data(df: pd.DataFrame) -> pd.DataFrame:
    # 规则1：剔除TEST_NUMBER=11行存在NaN的产品所有数据
    valid_product = df[df["TEST_NUMBER"] == 11]\
        .groupby("PRODUCT_NO")[NUM_COLS]\
        .apply(lambda x: x.notna().all(axis=1).all())
    valid_product = valid_product[valid_product].index
    df_clean = df[df["PRODUCT_NO"].isin(valid_product)].copy()

    # 规则2：仅删除TEST_NUMBER≠11的含NaN单行
    df_clean = df_clean[~((df_clean["TEST_NUMBER"] != 11) & (df_clean[NUM_COLS].isna().any(axis=1)))]
    return df_clean.reset_index(drop=True)

第二步：差值计算（满足规则3） #

def calc_diff(df_clean: pd.DataFrame) -> pd.DataFrame:
    def group_calc(group: pd.DataFrame) -> pd.DataFrame:
        # 2开头的行与TEST_NUMBER=11行做差
        base_11 = group[group["TEST_NUMBER"] == 11][NUM_COLS].iloc[0]
        mask_2 = group["TEST_NUMBER"].astype(str).str.startswith("2")
        group.loc[mask_2, NUM_COLS] = group.loc[mask_2, NUM_COLS].sub(base_11, axis=1)

        # 3开头的行与2开头最大TEST_NUMBER行做差
        max_test_2 = group[mask_2]["TEST_NUMBER"].max()
        if pd.notna(max_test_2):
            base_2 = group[group["TEST_NUMBER"] == max_test_2][NUM_COLS].iloc[0]
            mask_3 = group["TEST_NUMBER"].astype(str).str.startswith("3")
            group.loc[mask_3, NUM_COLS] = group.loc[mask_3, NUM_COLS].sub(base_2, axis=1)
        return group

    return df_clean.groupby("PRODUCT_NO", group_keys=False)\
        .apply(group_calc)\
        .reset_index(drop=True)

完整调用示例 #

if __name__ == "__main__":
    # 替换为你的原始数据读取逻辑
    raw_df = pd.read_csv("your_test_data.csv")
    cleaned_df = clean_data(raw_df)
    result_df = calc_diff(cleaned_df)

优化点说明 #

• 全程采用pandas原生向量化操作，避免逐行遍历，性能比常规循环写法高1~2个数量级
• 数值字段统一用NUM_COLS常量维护，后续增减统计字段无需改动核心逻辑
• 分组计算逻辑复用pandas的groupby.apply特性，天然实现不同产品独立计算的需求
• 所有过滤逻辑用布尔掩码实现，语义清晰无需嵌套判断，可读性和可维护性更高

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