Pandas使用tqdm显示to_csv进度条

Question

正如标题所示，我正在尝试在执行 pandas.to_csv 时显示进度条。
我有以下脚本：

def filter_pileup(pileup, output, lists):
    tqdm.pandas(desc='Reading, filtering, exporting', bar_format=BAR_DEFAULT_VIEW)
    # Reading files
    pileup_df = pd.read_csv(pileup, '\t', header=None).progress_apply(lambda x: x)
    lists_df = pd.read_csv(lists, '\t', header=None).progress_apply(lambda x: x)
    # Filtering pileup
    intersection = pd.merge(pileup_df, lists_df, on=[0, 1]).progress_apply(lambda x: x)
    intersection.columns = [i for i in range(len(intersection.columns))]
    intersection = intersection.loc[:, 0:5]
    # Exporting filtered pileup
    intersection.to_csv(output, header=None, index=None, sep='\t')

在前几行中，我找到了一种集成进度条的方法，但是这种方法对于最后一行无效。我如何实现呢？

- Eliran Turgeman

---

以一种非常不专业的方式，您可以创建一个继承自“io.TextIOBase”的包装类，将“output”传递给.write调用，同时更新进度条。虽然我不建议这样做，所以我不会将其发布为答案。 - orlp

你找到解决方案了吗？如果您能将其发布为答案，我会非常感激 :) - Sierox

@Sierox 我还没有找到这个具体问题的解决方法，但最终我使用了dask模块，该模块在dask.diagnostics下自带有进度条。 - Eliran Turgeman

Answer 1

您可以将数据框按n行划分成块，并使用mode='w'保存第一行和mode="a"保存其余的行，逐块将数据框保存为csv文件：

示例：

import numpy as np
import pandas as pd
from tqdm import tqdm

df = pd.DataFrame(data=[i for i in range(0, 10000000)], columns = ["integer"])

print(df.head(10))

chunks = np.array_split(df.index, 100) # split into 100 chunks

for chunck, subset in enumerate(tqdm(chunks)):
    if chunck == 0: # first row
        df.loc[subset].to_csv('data.csv', mode='w', index=True)
    else:
        df.loc[subset].to_csv('data.csv', header=None, mode='a', index=True)

输出：

    integer
0        0
1        1
2        2
3        3
4        4
5        5
6        6
7        7
8        8
9        9

100%|██████████| 100/100 [00:12<00:00,  8.12it/s]

- Chicodelarose

---

array_split函数接受要生成的数组数。注释应该是 # 分成100块，更多信息请参见https://numpy.org/doc/stable/reference/generated/numpy.array_split.html。 - LoneWanderer