利用Pandas进行CSV文件转存并写入Mysql

最近在做文件写入数据表。纪录下遇到的问题

背景

由于业务变更，需要导入大约两个月的数据。考虑到第三方平台提供的 API 存在速率和其他限制，我们决定使用文件导出数据的方式。为了进行测试，我选择了一个大小约为 3G 的 CSV 文件。这个文件在电脑上的操作不太方便，因为其中只有部分列是我们需要的，其他无用的信息我们不想要。然而，手动逐个处理是不现实的，所以我打算使用 Pandas 来实现处理。

转存数据

去除不需要的列，将需要的列转存为另外一份文件。

import pandas as pd

def dumpfile(csv_file: str, col: list):
    """
    读取csv文件，并选择转存的列，另存为一个新的文件
    :return:
    """
    data = pd.read_csv(csv_file)
    if len(list) > 0:
        selected_columns = data.loc[:, col]
        selected_columns.to_csv('output.csv', index=False)
    else:
        data.to_csv('output.csv', index=False)

        
dumpfile('test.csv', ['a','b','c'])

写入数据到mysql

实际上，在这一步，如果不希望手动编写代码，可以考虑使用第三方工具来导入数据，这也是一个不错的选择。我这边采用的是代码来实现。

from sqlalchemy import create_engine
from sqlalchemy import Column, Integer, String, Date


def write_sql_by_csv(csv_file):
    # 创建数据库连接
    engine = create_engine(f'mysql+pymysql://root:[email protected]:3306/xxxx')
    chunksize = 1000  # 每次读取的行数
    for chunk in pd.read_csv(csv_file, chunksize=chunksize):
        # 在这里处理每个数据块的逻辑
        selected_columns = chunk.loc[:, ['a','b','c']]

        # 指定列的数据类型。如果使用字典，键应该是列名，值应该是 SQLAlchemy 类型或 sqlite3 传统模式的字符串。如果提供标量，它将应用于所有列。
        dtypes = {
            'a': String(128),
            'b': Date,
            'c': Integer,
        }

        # 将数据块插入数据库
        # if_exists: append-将新值插入到现有表中，replace-删除表之前插入数据，fail-抛出异常
        selected_columns.to_sql('player_gems3', con=engine, if_exists='append', index=False, dtype=dtypes)

可以看到，我在代码中指定了dtypes，但也可以选择不指定。如果不指定，pandas会根据数据自动判断字段类型。然而，我个人认为自己定义dtypes会更好，因为这样更符合我的需求。除此之外，代码非常简单，也附有注释，其他的就不再多说了。

然而，需要注意的是，chunksize 并不完全符合我们的预期，它仍然会将所有数据加载到内存中，只是将数据分块加载。并没有有效降低内存使用。这一点我从下面的代码中感受到了，当数据表数量增加时，程序运行速度变慢，有时甚至需要很长时间等待。

df = pd.read_sql_query(sql_query, engine, chunksize=chunk_size)

这点我发现还是有不少网友提出来到的，下面是相关的讨论。

Reading table with chunksize still pumps the memory

Add doc note on memory usage of read_sql with chunksize