优化大型数据集上Pandas groupby()聚合二进制列的性能

First, let's confirm your core need: you want to aggregate binary columns by the template group, where if any row in the group has a 1 in an is_* column, the result should stay 1 (which is exactly what max() does for 0/1 data). The standard groupby.max() works, but it can get slow on large datasets—here are several optimized approaches to speed this up:

1. Use np.any() Instead of max() for Vectorized Efficiency #

Since your columns are binary, checking if any value in the group is 1 is equivalent to taking the max, but NumPy's vectorized any() operation is faster than pandas' built-in max() for large data:

import numpy as np

# Aggregate with np.any, then convert back to integer (since any returns bool)
result = df.groupby('template', as_index=False)[['is_a', 'is_b', 'is_c', 'is_d', 'is_e']].agg(np.any).astype(int)

2. Switch to pivot_table for Optimized Grouping #

Pandas' pivot_table uses different under-the-hood optimizations that can outperform standard groupby in some cases, especially with structured grouping and aggregation:

result = pd.pivot_table(
    df,
    index='template',
    values=['is_a', 'is_b', 'is_c', 'is_d', 'is_e'],
    aggfunc='max'
).reset_index()

3. Parallel Processing with Dask (For Very Large Datasets) #

If your dataset is too big to fit comfortably in memory, use Dask DataFrame to split the work across multiple CPU cores:

import dask.dataframe as dd

# Convert pandas DataFrame to Dask DataFrame (adjust npartitions to match your CPU cores)
ddf = dd.from_pandas(df, npartitions=4)
# Run grouped max in parallel, then compute to get a pandas DataFrame
result = ddf.groupby('template')[['is_a', 'is_b', 'is_c', 'is_d', 'is_e']].max().compute().reset_index()

4. Reduce Memory Footprint First #

Binary columns don't need to use pandas' default int64 dtype—switching to int8 or bool cuts down memory usage drastically, which directly speeds up all operations:

# Convert binary columns to int8 (uses 1 byte per value instead of 8)
for col in ['is_a', 'is_b', 'is_c', 'is_d', 'is_e']:
    df[col] = df[col].astype('int8')

# Now run your original groupby (it'll be faster with less memory overhead)
result = df.groupby('template', as_index=False)[['is_a', 'is_b', 'is_c', 'is_d', 'is_e']].max()

Bonus: Remove Redundant Rows First #

Since duplicate rows don't affect the max() result (if a group already has a 1 in a column, adding more rows with 0/1 won't change the outcome), you can drop duplicates before grouping to reduce the data size:

# Keep only unique combinations of template and is_* columns
df_clean = df.drop_duplicates(subset=['template'] + [col for col in df.columns if col.startswith('is_')])
# Group on the cleaned data
result = df_clean.groupby('template', as_index=False)[['is_a', 'is_b', 'is_c', 'is_d', 'is_e']].max()

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