python数据分析一例：使用SQL和pandas对数据进行聚合和diff

对一系列数据聚合后进行diff，是一种常见的数据分析需求。例如，我们可能会需要将每个月的财务支出流水数据进行分类汇总，再对不同月的汇总数据进行比较，看看哪些分类支出变多了，哪些变少了。此次我将使用SQL和pandas来实现上面所述需求，具体来说，使用SQL实现数据聚合功能，使用pandas对聚合后的数据进行diff。

虽说单独使用pandas也足够实现这个需求，但是考虑到类似的流水数据很多都是存在DB里的，检索数据时聚合一下也就顺手的事，因此聚合数据就交给SQL来做了。况且SQL相比pandas性能更好，泛用性也更广，如果不是用SQL来diff数据没有比较直观的方案，完全使用SQL来实现这个需求更加符合我的心意。本文中使用的数据库是python内置的sqlite3。

Talk is cheap，直接上代码，解释和注意事项都在注释里面了：

python 复制代码

#!/usr/bin/python3

from collections import namedtuple
import logging
import os
import sqlite3

import pandas as pd

logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')


DB_NAME = "expenses_aggregate_diff.db"

# 流水记录，总共四个字段：部门、类型、月份和金额，一般来说月份这个字段应该细化为时间戳，这里为了简化就用月份了
Record = namedtuple('Record', ['department', 'type', 'month', 'amount'])


# 插一些样例数据
def insert_example_records():
    records = [
        Record('A', 'R&D', '202406', 1000),
        Record('A', 'Operating', '202406', 2000),
        Record('A', 'R&D', '202406', 4000),
        Record('A', 'R&D', '202406', 5000),
        Record('B', 'Operating', '202406', 2000),
        Record('B', 'Sales', '202406', 8000),
        Record('A', 'Operating', '202407', 7000),
        Record('B', 'Operating', '202407', 4000),
        Record('B', 'Sales', '202407', 2000),
        Record('B', 'R&D', '202407', 1000),
        Record('B', 'R&D', '202407', 9000),
    ]

    logging.info(f'save {len(records)} records to db')
    conn = sqlite3.connect(DB_NAME)
    cursor = conn.cursor()

    cursor.execute('''
        CREATE TABLE IF NOT EXISTS records (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            department TEXT NOT NULL,
            type TEXT NOT NULL,
            month TEXT NOT NULL,
            amount INT NOT NULL
        )
    ''')

    for record in records:
        cursor.execute('''
            INSERT INTO records (department, type, month, amount)
            VALUES (?, ?, ?, ?)
        ''', (record.department, record.type, record.month, record.amount))

    conn.commit()
    conn.close()


# 获得某个月份聚合后的流水数据，以group_by_fields为分组字段，可接受的参数是部门、类型或两者的组合，对金额进行聚合
def get_aggregated_amount_by_month(month: str, group_by_fields: list[str]) -> pd.DataFrame:
    conn = sqlite3.connect(DB_NAME)
    cursor = conn.cursor()

    # 为了简便，就不对group_by_fields做校验了，这种代码绝不能上产线的 XD
    group_by_clause = ", ".join(group_by_fields)
    sql = f"""
        SELECT {group_by_clause}, SUM(amount) FROM records WHERE month = ? GROUP BY {group_by_clause}
    """
    cursor.execute(sql, (month,))
    rows = cursor.fetchall()
    conn.close()
    logging.info(f'get {len(rows)} records for {month}')

    return pd.DataFrame(rows, columns=group_by_fields + ['amount'])


# diff某两个月聚合后的数据
def diff_aggregated_data(month0: str, month1: str, group_by_fields: list[str]) -> pd.DataFrame:
    # 聚合
    df0 = get_aggregated_amount_by_month(month0, group_by_fields)
    df1 = get_aggregated_amount_by_month(month1, group_by_fields)

    # 合并到一个表里面，合并方式为outer join，缺失项会被填充为NaN
    df_diff = pd.merge(df0, df1, on=group_by_fields, how='outer', suffixes=(month0, month1))
    # 将NaN填充为0
    df_diff = df_diff.fillna(0)
    # 对数据进行diff
    df_diff['amount_diff'] = df_diff['amount' + month1] - df_diff['amount' + month0]
    # 只保留diff后的数据
    df_diff = df_diff[group_by_fields + ['amount_diff']]
    return df_diff


if __name__ == '__main__':
    if not os.path.exists(DB_NAME):
        insert_example_records()

    group_by_fields = ['department', 'type']
    month0 = '202406'
    month1 = '202407'

    print(diff_aggregated_data(month0, month1, group_by_fields))

输出为：

  department       type  amount_diff
0          A  Operating       5000.0
1          A        R&D     -10000.0
2          B  Operating       2000.0
3          B        R&D      10000.0
4          B      Sales      -6000.0