pandas函数的时空复杂度对比

众所周知，不同的函数实现方式可能会导致执行时间和内存消耗的显著差异。本文以leetcode的题1193为例，比对不同函数实现方式的时空复杂度。

一、问题背景

1193 题要求编写一个sql查询，查找每个月和每个国家/地区的事务数及其总金额、已批准的事务数及其总金额。

表：Transactions

Column Name	Type
id	int
country	varchar
state	enum
amount	int
trans_date	date

id 是这个表的主键。该表包含有关传入事务的信息。state列类型为[“approved”, “declined”]之一。编写一个sql查询来查找每个月和每个国家/地区的事务数及其总金额、已批准的事务数及其总金额。以任意顺序返回结果表。

二、实现方式与性能对比

尝试一

使用apply函数

import pandas as pd

def monthly_transactions(transactions: pd.DataFrame) -> pd.DataFrame:
    transactions['month'] = transactions['trans_date'].apply(lambda x: x.strftime('%Y-%m'))
    transactions['country'].fillna('null', inplace=True)
    data = transactions.groupby(['month', 'country']).agg(
    trans_count = ('month', 'count'),
    approved_count = ('state', lambda x: sum(x == 'approved')),
    trans_total_amount = ('amount', 'sum'),
    approved_total_amount = ('amount', lambda x: sum(x[transactions['state'] == 'approved']))
).reset_index()
    data['country'].replace('null', None, inplace=True)
    return data

• 性能分析
  • 执行用时1953ms
  • 消耗内存68.14MB

尝试二

使用dt.strftime替代apply

def monthly_transactions(transactions: pd.DataFrame) -> pd.DataFrame:
    transactions['month'] = transactions['trans_date'].dt.strftime('%Y-%m')
    data = transactions.groupby(['month', 'country'], dropna=False).agg(
        trans_count=('month', 'count'),
        approved_count=('state', lambda x: sum(x == 'approved')),
        trans_total_amount=('amount', 'sum'),
        approved_total_amount=('amount', lambda x: sum(x[transactions['state'] == 'approved']))
    ).reset_index()
    return data

​​

• 性能分析
  • 执行用时1647ms
  • 消耗内存68.05M

尝试三

前两次尝试都是先用字符串替代country为空，再使用groupby函数，最后再替代回null。可以直接使用groupby函数内置参数dropna，设置dropna为False

def monthly_transactions(transactions: pd.DataFrame) -> pd.DataFrame:
    transactions['month'] = transactions['trans_date'].dt.strftime('%Y-%m')
    data = transactions.groupby(['month', 'country'], dropna=False).agg(
        trans_count=('month', 'count'),
        approved_count=('state', lambda x: sum(x == 'approved')),
        trans_total_amount=('amount', 'sum'),
        approved_total_amount=('amount', lambda x: sum(x[transactions['state'] == 'approved']))
    ).reset_index()
    return data

• 性能分析
  • 执行用时2187ms
  • 消耗内存68.36M

emmmmm…