使用zip和df.apply()遍历Pandas DataFrame运行缓慢的优化求助

Hey there! Let's fix that slow Pandas code—60k rows shouldn't take 40 minutes, that's way too inefficient. The core problem here is your row-wise apply approach, which essentially runs a Python loop under the hood, plus you're repeatedly slicing the entire DataFrame inside each iteration (that's an O(n²) operation, which blows up with 60k rows). Let's break down the fixes step by step:

First, Diagnose the Big Issues #

1. Repeated DataFrame Slicing: Every call to calculate slices the full DataFrame to get df_a, meaning you're doing 60k separate slice operations—each one is expensive.
2. In-Place Modifications: Calling df.drop(row['index'], inplace=True) inside apply is risky (it modifies the DataFrame while iterating over it) and adds unnecessary overhead.
3. Row-Wise Loops: Even with zip, you're still looping through each row in df_a for every row in the original DataFrame—no use of Pandas' vectorized superpowers.

Step 1: Preprocess to Eliminate In-Place Drops #

First, filter out rows that don't have at least 30 subsequent matches before starting calculations. This avoids modifying the DataFrame mid-iteration:

import pandas as pd
import numpy as np

# Ensure your DataFrame is sorted by date first (critical for "subsequent matches" logic)
df = df.sort_values(['tourney_date']).reset_index(drop=True)

# Create a combined list of all matches involving each player (winner or loser)
df_winner = df.rename(columns={'winner_id': 'player_id'})
df_loser = df.rename(columns={'loser_id': 'player_id'})
player_matches = pd.concat([df_winner, df_loser]).sort_values(['player_id', 'tourney_date'])

# For each player, count how many matches come after each row
player_matches['remaining_matches'] = player_matches.groupby('player_id').cumcount(ascending=False) - 1

# Merge this count back to the original DataFrame (for winner_id)
df = df.merge(
    player_matches[['player_id', 'tourney_date', 'remaining_matches']],
    left_on=['winner_id', 'tourney_date'],
    right_on=['player_id', 'tourney_date'],
    how='left'
)

# Filter out rows with fewer than 30 subsequent matches upfront
df = df[df['remaining_matches'] >= 30].reset_index(drop=True)

Step 2: Replace Row-Wise Loops with Vectorized Group Operations #

The key is to process all matches for a player in one go using groupby, instead of handling each row individually. First, update your helper functions to work with NumPy arrays/Series:

def yrs_between(date_series, base_date):
    # Calculate years between a series of dates and a single base date
    return (date_series - base_date).dt.total_seconds() / (365.25 * 24 * 3600)

def time_discount(t):
    # Replace with your actual time discount logic (example exponential decay)
    return np.exp(-0.5 * t)

def opp_weight(rank):
    # Replace with your actual opponent weight logic; handle 0 ranks to avoid division by zero
    return np.where(rank == 0, np.nan, 1 / rank)

Now, process each player's matches to compute the weighted average efficiently:

def compute_player_weighted_avg(group):
    # Sort group by date ascending (process earliest to latest)
    group = group.sort_values('tourney_date').reset_index(drop=True)
    n_matches = len(group)
    
    # Broadcast to compute time differences between all pairs of matches (i < j)
    dates = group['tourney_date'].values
    date_diff = dates[np.newaxis, :] - dates[:, np.newaxis]
    # Keep only j > i (matches after the current row)
    upper_tri_mask = np.triu(np.ones((n_matches, n_matches), dtype=bool), k=1)
    date_diff = date_diff[upper_tri_mask].reshape(n_matches, -1)
    
    # Convert time differences to years and compute time weights
    yrs_diff = date_diff / np.timedelta64(1, 'Y')
    time_weights = time_discount(yrs_diff)
    
    # Get opponent weights and serve percentages for subsequent matches
    opp_weights = opp_weight(group['loser_rank'].values)[np.newaxis, :].repeat(n_matches, axis=0)[upper_tri_mask].reshape(n_matches, -1)
    serve_pcts = group['winner_serve_pts_pct'].values[np.newaxis, :].repeat(n_matches, axis=0)[upper_tri_mask].reshape(n_matches, -1)
    
    # Calculate weighted values and sum them up
    weighted_values = serve_pcts * time_weights * opp_weights
    sum_values = weighted_values.sum(axis=1)
    sum_weights = time_weights.sum(axis=1)
    
    # Compute final average (handle cases where sum of weights is 0)
    group['new'] = np.where(sum_weights == 0, np.nan, sum_values / sum_weights)
    return group

# Apply the function to each player's matches
player_matches_processed = player_matches.groupby('player_id').apply(compute_player_weighted_avg)

# Merge the computed 'new' column back to the original DataFrame
df = df.merge(
    player_matches_processed[['player_id', 'tourney_date', 'new']],
    left_on=['winner_id', 'tourney_date'],
    right_on=['player_id', 'tourney_date'],
    how='left'
)

Why This Is Faster #

• Vectorization: We use NumPy broadcasting to compute all time differences and weights in bulk, avoiding Python loops entirely.
• Grouped Processing: By handling all matches for a player at once, we eliminate repeated slicing of the full DataFrame.
• Early Filtering: We remove rows that don't meet the 30-match requirement upfront, reducing the total number of calculations.

Quick Additional Tips #

• Ensure tourney_date is a datetime64 type (run df['tourney_date'] = pd.to_datetime(df['tourney_date']) if not) to avoid slow type conversions.
• Test with a small subset of your data first (e.g., df.sample(1000)) to verify results match your original code before running the full dataset.

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