如何高效过滤DataFrame？现有代码运行过慢求优化方案

我正在做一个项目，需要识别DataFrame中特定列值相同的行组，之后检查这些行的某一KPI列均值是否超过阈值。这需要通过多列值反复把DataFrame筛选成更小的子集。

当前筛选函数（循环实现） #

def filter_df(df, cols, vals):
    """This function takes a dataframe, a list of columns, and a list of values.
     It returns a dataframe with only the rows that match the values for the columns."""
    temp_df = df
    for i in range(len(cols)):
        temp_df = temp_df[(temp_df[cols[i]] == vals[i])]
    return temp_df

固定列数的优化版（n=6） #

我针对固定n=6列写了另一种实现，运行速度大概减半：

def filter_df(df, cols, vals):
    """This function takes a dataframe, a list of columns, and a list of values.
     It returns a dataframe with only the rows that match the values for the columns."""
    temp_df = df
    temp_df = temp_df[(temp_df[cols[0]] == vals[0]) & (temp_df[cols[1]] == vals[1]) & (temp_df[cols[2]] == vals[2]) & (temp_df[cols[3]] == vals[3]) & (temp_df[cols[4]] == vals[4]) & (temp_df[cols[5]] == vals[5])]
    return temp_df

主逻辑与性能瓶颈 #

我的数据集大概40行，但程序需要遍历每一种n列组合的所有可能值组合。主函数如下：

def identify_clusters(df, KPI, n, KPI_threshhold):
    """This function takes in a dataframe, a KPI, a number of conditions, and a KPI threshhold.
    It returns a list of lists with each list of groups, the values list, and the ad_ids in the cluster."""
    groups = set_groups(df, n)
    output = []
    for i in range(len(groups)):
        vals_lst = set_values(groups[i], df)
        for comb in vals_lst:
            temp_df = filter_df(df, groups[i], comb)
            if temp_df[KPI].mean() > KPI_threshhold:
                output.append([groups[i], comb, temp_df['ad_id'].values])
    return output

目前外层循环有大约300万个“组”，每个组的内层循环有240-1960个值组合。观察到外层单次迭代耗时0.5到2秒，整体运行时间过长。

可能低效的set_values函数 #

我怀疑set_values也是性能瓶颈，但不知道怎么提速：

def set_values(group, df):
    """This function takes a list of columns and a dataframe. It returns a list of lists of all possible
     combinations of values for each column in the group in the dataframe."""
    vals_lst = []
    for col in group:
        vals_lst.append(list(df[col].dropna().unique()))
    output = list(itertools.product(*vals_lst))
    return output

优化需求 #

我不确定低效来自filter_df、identify_clusters还是其他地方。我知道DataFrame的for循环过滤通常可以用向量化操作替代，但不知道当filter_df接收列列表参数时怎么实现，求相关优化建议。

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