Pandas批量生成多时段Shift特征的性能优化问题

Hey there! I totally get the pain of slow feature generation when dealing with massive datasets—10 million rows is no joke, and those 84 extra columns can quickly turn into a performance bottleneck if you're not using the right tools. Let's break down how to fix this:

Core Principles to Boost Performance #

First, forget about any row-wise or naive column-wise loops—those are going to crawl on 10M rows. We need lean, vectorized operations that leverage pandas' optimized C-backed internals.

Step 1: Prep Your Data #

Before generating features, make sure your data is sorted correctly (this is critical for shift to work as intended) and optimize data types to save memory:

import pandas as pd

# Assume your timestamp column is named 'timestamp'
df = df.sort_values('timestamp').reset_index(drop=True)

# Optimize dtypes to cut memory usage (if precision allows)
# Convert float64 to float32 halves memory footprint
feature_cols = df.columns.tolist()  # Your 14 continuous features
df[feature_cols] = df[feature_cols].astype('float32')

Step 2: Generate Shifted Features Efficiently #

We'll use pd.DataFrame.shift() combined with pd.concat() to create all your shifted columns in bulk. This avoids repeated copies of your data and keeps operations vectorized:

# Define the range of hours to shift (1 to 6 in your case)
shift_hours = range(1, 7)

# Generate all shifted features for each column
shifted_feature_dfs = []
for col in feature_cols:
    # Create a DataFrame with all shifted versions of the current column
    shifted_cols = pd.concat(
        [df[col].shift(hour).rename(f"{col}_{hour}") for hour in shift_hours],
        axis=1
    )
    shifted_feature_dfs.append(shifted_cols)

# Merge all shifted features back to the original DataFrame
final_df = pd.concat([df] + shifted_feature_dfs, axis=1)

Step 3: Handle Memory Constraints (If Needed) #

If even optimized pandas can't fit the full dataset in memory, use Dask to process the data in chunks. Dask mimics pandas' API but works with out-of-core data:

import dask.dataframe as dd

# Convert pandas DataFrame to Dask DataFrame (adjust partitions based on your CPU cores)
dask_df = dd.from_pandas(df, npartitions=4)

shifted_feature_dfs = []
for col in feature_cols:
    shifted_cols = dd.concat(
        [dask_df[col].shift(hour).rename(f"{col}_{hour}") for hour in shift_hours],
        axis=1
    )
    shifted_feature_dfs.append(shifted_cols)

# Combine and compute the result (this will process chunks in parallel)
final_dask_df = dd.concat([dask_df] + shifted_feature_dfs, axis=1)
final_df = final_dask_df.compute()

Why This Works Better Than Your Previous Method #

• Vectorized Operations: shift and concat are implemented in optimized C code, so they avoid the overhead of Python loops.
• Minimized Data Copies: By generating all shifted columns for a feature at once, we reduce the number of intermediate DataFrames created.
• Memory Optimization: Using smaller dtypes like float32 cuts down on memory usage, which prevents slowdowns from swapping data to disk.

Just double-check that your time granularity matches the shift amount (e.g., each row represents exactly one hour of data)—otherwise, you'll need to adjust the shift parameter to match your actual time intervals.

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