使用 Pandas 实现动态查询及处理数据并导出 Excel 文件


需求是这么来的:根据参数执行查询,然后生成 Excel 表供下载,简单说就是查询并导出。

一种思路是根据参数动态生成 select 查询对象,然后执行拿到结果,大概是这样:

...
sql_select = select(columns).select_from(join_obj).where(whereclause).offset(offset).limit(limit).order_by(order_by)

result = session.execute(sql_select)
count = result.rowcount
rows = [dict(row) for row in result.fetchall()]

然后拿到 rows 就可以去该做什么处理就做什么处理了,如果要生存 Excel 表,就可以使用 XlsxWriter 这个工具去写表文件,反正就是对 rows 开个 for 循环就是了。

不过呢,还有没有更好的办法?

假设如果我要对 rows 的数据要做些高级处理,而不是开个 for 循环。毕竟如果是处理数据的话,Pandas 这个工具在思维方式和代码组织上比 for 循环要更利于思考和维护,而同时我们也知道 Pandas 也能直接输出 Excel 文件,那么是否可以用上它?

按常规,查询得到 rows,然后再生成一个 df,对 df 进行一些处理后再去输出 Excel,甚至同时还要把修改后的数据输出为新的 rows,那么这里:查询得到结果、把结果生成 df 这两个步骤涉及到两步数据转化,理论上来说肯定是有一些消耗的,如果是在 API 服务上使用,则需要考虑性能问题,能少一步就少一步。

那么,是否可以一步到位呢?直接查询结果就得到 df,而不是自行再多做一步中转处理。

Pandas 是可以直接 read_sql 的,具体说明如下:

def read_sql(sql, con, index_col=None, coerce_float=True, params=None,
             parse_dates=None, columns=None, chunksize=None):
    """
    Read SQL query or database table into a DataFrame.

    Parameters
    ----------
    sql : string SQL query or SQLAlchemy Selectable (select or text object) to be executed, or database table name.
    con : SQLAlchemy connectable(engine/connection) or database string URI or DBAPI2 connection (fallback mode) Using SQLAlchemy makes it possible to use any DB supported by that library.
    If a DBAPI2 object, only sqlite3 is supported.

所以 Pandas 可以接受 raw sql 也可以接受 SQLAlchemy Selectable 对象。

这就好说了,如果把做好的的 sql_select 不是传给 session.execute 而是传给 pd.read_sql,这事不就实现了。

import pandas as pd

df = pd.read_sql(sql_select, con=self.get_engine())
logger.debug(df)

当然,pd 的 con 参数接受的不是一个 session 实例,而是一个 engine,也就是创建 session 一步的那个 engine 参数。

同时,如果需要调试查看 sql 语句和参数,可以多做点工作:

from sqlalchemy.dialects import mysql
import pandas as pd

sql_raw = sql_select.compile(dialect=mysql.dialect())
params = sql_select.compile(dialect=mysql.dialect()).params
logger.debug(sql_raw)
logger.debug(params)

df = pd.read_sql(sql_select, con=self.get_engine())
logger.debug(df)

当然其实也可以把 sql_raw 和 params 分开传给 pd.read_sql,但貌似也没这个必要性了。

而后就可以对 df 做处理,以及最终可以调用 df.to_excel('./output/test.xlsx', index=False) 生成 Excel 表格文件,这里加了个 index=False 是为了不让默认的行数索引写入表格。

当然也可以继续调用 df.to_dict(orient=‘records’) 生成 rows 来用,或是直接 df.to_json(orient=‘records’) 生成 json 输出了。