如何使用 Pandas DataFrame 替换 lambda 和分组以提高性能

Question

也许我的问题看起来很复杂，但本质上很简单。我是 Python 的新手，现在我面临代码太慢的问题。下面是代码的优化版本。我会很感激一个小的代码审查和关于如何加快它的建议。我认为最慢的操作是.apply(lambda和分组，但我不知道如何替换它们。

...
for raw_file in raw_files:
    reader = pd.read_csv(raw_file, chunksize=100000)
    for chunk in reader:
        processed_data = task(chunk)
        for name, data in processed_data:
            save_data(name, data) # some method which saves DataFrame correctly
...


def task(data):
    data = data[data['Quantity'] != 0] # remove zero items
    # add date parts as columns
    data[['dt_year', 'dt_month', 'dt_day', 'dt_day_of_year', 'dt_day_of_week', 'dt_hour']] = \
                data.apply(lambda df: to_date_parts(df['SalesDate']), axis=1)
    # group by location-item to aggregate in different files
    grouped = data.groupby(['LocationID','ItemID'])
    result = []
    for name, group in grouped:
        result += [(name, group)]
    return result



def to_date_parts(str_date):
    date = dt.datetime.strptime(str_date.split(".")[0], '%Y-%m-%d %H:%M:%S')
    dt_year = date.year
    dt_month = date.month
    dt_day = date.day
    dt_day_of_year = date.toordinal() - dt.datetime(date.year, 1, 1).toordinal() + 1
    dt_day_of_week = date.weekday()
    dt_hour = date.hour
    return pd.Series([dt_year, dt_month, dt_day, dt_day_of_year, dt_day_of_week, dt_hour])

Answer 1

Python datetime vs Pandas datetime

您发现性能不佳有两个相互关联的原因：

1. 您使用 Python 内置的 datetime 对象，而不是高效的 Pandas datetime 系列来存储日期。
2. 您使用 Python 级别的 for 循环，而不是 Pandas datetime 系列支持的矢量化操作。

所以首先将您的系列转换为 Pandas datetime 系列：

date_format = '%Y-%m-%d %H:%M:%S'
df['SalesDate'] = pd.to_datetime(df['SalesDate'], format=date_format, errors='coerce')

然后直接从您的系列中提取属性：

from operator import attrgetter

# list attributes
fields = ['year', 'month', 'day', 'dayofyear', 'dayofweek', 'hour']

# extract attributes
attributes = pd.concat(attrgetter(*fields)(df['SalesDate'].dt), axis=1, keys=fields)

# join attributes to dataframe
df = df.join(attributes)

熊猫GroupBy对象

这种将项目串联到list 是不必要的：

grouped = data.groupby(['LocationID','ItemID'])
result = []
for name, group in grouped:
    result += [(name, group)]
return result

由于data.groupby(...) 是一个可迭代对象，你可以只使用return GroupBy 对象：

return data.groupby(['LocationID','ItemID'])