pandas读取excel时获取读取进度的实现

写在前面

QQ群里偶然看到群友问这个问题, pandas读取大文件时怎么才能获取进度? 我第一反应是: 除非pandas的read_excel等函数提供了回调函数的接口, 否则应该没办法做到. 搜索了一下官方文档和网上的帖子, 果然是没有现成的方案, 只能自己动手.

准备工作

确定方案

一开始我就确认了实现方案, 那就是增加回调函数. 这里现学现卖科普一下什么是回调函数. 简单的说就是:

所使用的模块里面, 会调用一个你给定的外部方法/函数, 就是回调函数. 拿本次的尝试作为例子, 我会编写一个"显示进度函数", 通过传参的方式传入pd.read_excel, 这样pd在读取excel时, 会边读取边调用"显示进度函数". 为什么不直接在pd里面增加? 因为pd读取excel文件时是阻塞的, 内部方法在被调用时无法抛出进度信息. (如有谬误请指正)

理解读取方式

先得了解一下pandas是怎么读取excel的. 在pycharm里面按住control点击read_excel, 再浏览一下代码根据关键的函数继续跳转, 还是挺容易得到调用的路径的.

在这里插入图片描述

最后OpenpyxlReader读取excel的方法代码如下. 很明显重点就在其中的for循环里. 调用get_sheet_data时, 已经通过一系列方法获得了目标sheet(这里细节不赘述), 然后在for循环里逐行读取数据并返回data最后生成dataframe.

def get_sheet_data(self, sheet, convert_float: bool) -> List[List[Scalar]]:
        # GH 39001
        # Reading of excel file depends on dimension data being correct but
        # writers sometimes omit or get it wrong
        import openpyxl

        version = LooseVersion(get_version(openpyxl))

        # There is no good way of determining if a sheet is read-only
        # https://foss.heptapod.net/openpyxl/openpyxl/-/issues/1605
        is_readonly = hasattr(sheet, "reset_dimensions")

        if version >= "3.0.0" and is_readonly:
            sheet.reset_dimensions()

        data: List[List[Scalar]] = []
        last_row_with_data = -1
        for row_number, row in enumerate(sheet.rows):
            converted_row = [self._convert_cell(cell, convert_float) for cell in row]
            if not all(cell == "" for cell in converted_row):
                last_row_with_data = row_number
            data.append(converted_row)

        # Trim trailing empty rows
        data = data[: last_row_with_data + 1]

        if version >= "3.0.0" and is_readonly and len(data) > 0:
            # With dimension reset, openpyxl no longer pads rows
            max_width = max(len(data_row) for data_row in data)
            if min(len(data_row) for data_row in data) < max_width:
                empty_cell: List[Scalar] = [""]
                data = [
                    data_row + (max_width - len(data_row)) * empty_cell
                    for data_row in data
                ]

        return data

开始改动

这里直接暴力更改pandas库源文件!(仅用于调试, 注意备份和保护自己的工作环境)

主程序代码

编写main.py, 代码比较简单, 相关功能我都用注释作为解释. 其中show_pd_read_excel_progress就是我编写的回调函数, 通过命令行的方式输出实时的读取进度. 当然你如果编写的是GUI程序比如PYQT5, 也可以在这个回调函数中发送signal给main UI, 做成progress bar或者其他的GUI样式.

import pandas as pd
from datetime import datetime

'''
定义回调函数
cur: 读取时的当前行数
tt: 读取文件的总行数
'''
def show_pd_read_excel_progress(cur, tt):
    # 进度数值
    progress = " {:.2f}%".format(cur/tt*100)
    # 进度条
    bar = " ".join("█" for _ in range(int(cur/tt*100/10)))
    # 显示进度
    print("\r进度:" + bar + progress, end="", flush=True)

# 记录开始时间
t = datetime.now()
# 开始读取excel
print("pd.read_excel: test_4.xlsx...")
xl_data = pd.read_excel("test_4.xlsx", callback=show_pd_read_excel_progress)
# 打印excel头几行
print(xl_data.head())
print("\n")
# 显示花费的时间
print("Time spent:", datetime.now()-t)

修改pandas源码

再自己观察一下, 我在pd.read_excel方法的参数里增加了callback参数, 这个参数是原版read_excel方法里没有的. 所以我们需要处理pandas源码, 这个源码在…/pandas/io/excel/_base.py中, pycharm中按住control点击read_excel可以快速跳转. 这个地方我增加了一个参数callback, 默认值为None. 下方io.parse同样把callback参数传递给ExcelFile类.

def read_excel( io, sheet_name=0, header=0, names=None, index_col=None, usecols=None, squeeze=False, dtype=None, engine=None, converters=None, true_values=None, false_values=None, skiprows=None, nrows=None, ah读取csv或者sql之类的, 也可以照猫画虎.