如何加速百万行数据的自定义无损groupby聚合？

我编写了名为aggregate_nondestructive的自定义聚合函数，可在不删除任何值的前提下，将同一时间点的测量数据合并为尽可能少的数据点。函数代码如下：

import pandas as pd
from pandas import Series, DataFrame

def aggregate_nondestructive(df: Series | DataFrame) -> Series | DataFrame:
    """Aggregate multiple simultaneous measurements in a non-destructive way.
    
    Given a DataFrame of size m×k, 
    this will construct a new DataFrame of size m'×k, where m' = max(df.notna().sum()) is the 
    maximal number of measured not-null values in a column.
    
    Example
    -------

                              Acetate  Base   DOT  Fluo_GFP   Glucose  OD600  Probe_Volume    pH
        measurement_time                                                                        
        2020-12-09 09:48:38      <NA>  <NA>  <NA>      <NA>  4.578233   <NA>          <NA>  <NA>
        2020-12-09 09:48:38      <NA>  <NA>  <NA>      <NA>      <NA>  0.445          <NA>  <NA>
        2020-12-09 09:48:38  0.116585  <NA>  <NA>      <NA>      <NA>   <NA>          <NA>  <NA>
        2020-12-09 09:48:38  0.114842  <NA>  <NA>      <NA>      <NA>   <NA>          <NA>  <NA>
        2020-12-09 09:48:38      <NA>  <NA>  <NA>      <NA>      <NA>  0.485          <NA>  <NA>
        2020-12-09 09:48:38      <NA>  <NA>  <NA>      <NA>      <NA>   <NA>           200  <NA>
        2020-12-09 09:48:38      <NA>  <NA>  <NA>    1112.5      <NA>   <NA>          <NA>  <NA>
        2020-12-09 09:48:38      <NA>  <NA>  <NA>     912.5      <NA>   <NA>          <NA>  <NA>
        2020-12-09 09:48:38      <NA>  <NA>  <NA>      <NA>  4.554859   <NA>          <NA>  <NA>
        
    Returns
    
                              Acetate  Base   DOT  Fluo_GFP   Glucose  OD600  Probe_Volume    pH
        measurement_time                                                                        
        2020-12-09 09:48:38  0.116585  <NA>  <NA>    1112.5  4.578233  0.445           200  <NA>
        2020-12-09 09:48:38  0.114842  <NA>  <NA>     912.5  4.554859  0.485          <NA>  <NA>    
    """
    if isinstance(df, Series):
        return df

    mask = df.notna()
    nitems = mask.sum()
    nrows = nitems.max()
    result = DataFrame(index=df.index[:nrows], columns=df.columns).astype(df.dtypes)
    
    for col in result:
        result[col].iloc[:nitems[col]] = df.loc[mask[col], col]
    return result

现在需要将该函数用于约100万行的表格，通过groupby(["measurement_time"]).apply(aggregate_nondestructive)聚合同时测量的数据，但操作速度极慢。以下是除Numba重写之外的简单提速方法：

1. 抛弃逐组apply，用分组行号映射重构数据 #

apply在处理大量分组时会产生巨大的调用开销，直接通过构造分组内的行号来重组数据是更高效的方案：

import pandas as pd
import numpy as np

# 重置索引，方便后续处理
df_process = df.reset_index().copy()

# 为每个分组内的非空值分配组内行号
for col in df_process.columns.drop("measurement_time"):
    non_null_mask = df_process[col].notna()
    # 仅对非空行生成序号
    df_process.loc[non_null_mask, "row_idx"] = df_process[non_null_mask].groupby("measurement_time").cumcount()

# 按measurement_time和row_idx分组，取每组内的第一个非空值（自动对齐列）
result = df_process.groupby(["measurement_time", "row_idx"]).first().reset_index(level=1, drop=True)

这个方法完全利用pandas的分组内置逻辑，避免了自定义函数的逐组调用，速度提升显著。

2. 优化自定义函数的向量化逻辑 #

原函数中的for循环和iloc链式索引是效率瓶颈，修改为直接操作数组的向量化版本：

def aggregate_nondestructive_fast(df: Series | DataFrame) -> Series | DataFrame:
    if isinstance(df, Series):
        return df

    # 一次性提取所有列的非空值数组
    col_non_null = {col: df[col].dropna().values for col in df.columns}
    max_rows = max(len(vals) for vals in col_non_null.values())
    
    # 预分配结果DataFrame
    result = pd.DataFrame(
        index=df.index[:max_rows],
        columns=df.columns,
        dtype=df.dtypes
    )
    
    # 直接赋值数组，避免索引开销
    for col, vals in col_non_null.items():
        result.loc[:len(vals)-1, col] = vals
    
    return result

这里用dropna().values直接获取非空值的numpy数组，赋值时使用loc而非iloc，减少pandas的索引计算开销。

3. 提前过滤全空行，减少分组处理量 #

如果原数据中存在大量全空的行（即所有列都是NaN），先过滤这些行能大幅减少每个分组的处理规模：

# 过滤掉所有列均为NaN的行
df_clean = df.dropna(how="all")
# 再进行分组聚合
result = df_clean.groupby("measurement_time").apply(aggregate_nondestructive_fast)

4. 使用agg对列单独聚合后合并 #

利用groupby.agg对每个列单独处理非空值，再将结果组合成目标格式：

def agg_single_col(series):
    # 获取列的非空值，不足最大行数的补NaN
    vals = series.dropna().values
    group_max_len = series.notna().sum()
    return np.pad(vals, (0, group_max_len - len(vals)), constant_values=np.nan)

# 对每个列应用聚合函数，转置后得到结果
agg_result = df.groupby("measurement_time").agg(agg_single_col).T
# 重置索引层级，恢复原格式
result = agg_result.reset_index(level=0, drop=True)

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