首页 > 解决方案 > 使用 pandas 分块将唯一行写入 CSV

问题描述

我有一个 1 亿行的 csv,我必须用 pandas 分块读取,如下所示:

df_chunks = pandas.read_csv(
    'my-file.csv.gz',
    sep='\t',
    chunksize=100000,
    compression='gzip')

for df in df_chunks:
    # here I filter some rows and columns and after that
    # I write to a new csv
    filtered_df.to_csv(
        'my_filtered.csv.gz',
        sep=',',
        columns=['id', 'date'],
        compression='gzip',
        mode='a')

我要写的数据看起来像这样,只有 2 列

id,date
42517544,2019-06-30
42517544,2019-06-30
42517544,2019-07-01
...

现在我可以使用类似的东西,df.drop_duplicates()但由于我是分块写作,我最终可能会出现重复。注意文件很大,大约 10G,所以我需要分块读写。

我想找到一种使用 pandas 的方法,也许还有一个不会消耗太多内存的内存集,因为这也是一个约束。

有什么好的方法呢?

标签: pythonpandas

解决方案


设置

100 万行

np.random.seed([3, 1415])
n = 1_000_000
dfout = pd.DataFrame({
    'id': np.random.randint(1000, size=n),
    'date': np.random.choice(pd.date_range('2019-01-01', periods=1000), size=n)
})

dfout.to_csv('my-file.csv.gz', compression='gzip', sep='\t', index=False)

解决方案

像你一样切块

df_chunks = pd.read_csv(
    'my-file.csv.gz',
    sep='\t',
    chunksize=100000,
    compression='gzip')

每个唯一文件写入单个文件date

for i, df in enumerate(df_chunks):
    for date, d in df.groupby('date'):
        date = pd.Timestamp(date)
        d.drop_duplicates().to_csv(
            f'{date:%Y%m%d}.csv.gz',
            compression='gzip',
            mode='a',
            index=False,
            header=False
        )
    print(f'\r{i}', end='')

读入每个单独的日期文件drop_duplicates,然后写回

from pathlib import Path

path = Path('.')

for i, fh in enumerate(path.glob('[0-9]' * 8 + '.csv.gz')):
    df = pd.read_csv(fh, header=None)
    df.drop_duplicates().to_csv(
        'my_filtered.csv.gz',
        compression='gzip',
        mode='a',
        index=False,
        header=False
    )
    print(f'\r{i}: {fh}', end='')

df = pd.read_csv(
    'my_filtered.csv.gz',
    compression='gzip',
    header=None,
    names=['id', 'date']
)

验证

assert len(df) == len(dfout) - dfout.duplicated().sum()
assert df.duplicated().sum() == 0

推荐阅读