在 pandas 中为 read_sql 指定 dtypes

Posted 2023-02-24

【中文标题】在 pandas 中为 read_sql 指定 dtypes

【英文标题】：Specifying dtypes for read_sql in pandas

【发布时间】：2016-12-24 08:07:41

【问题描述】：

我想指定执行 pandas.read_sql 时返回的 dtypes。特别是我对节省内存和将浮点值返回为 np.float32 而不是 np.float64 感兴趣。我知道我可以在之后使用 astype(np.float32) 进行转换，但这并不能解决初始查询中需要大量内存的问题。在我的实际代码中，我将提取 8400 万行，而不是此处显示的 5 行。 pandas.read_csv 允许将 dtypes 指定为 dict，但我认为使用 read_sql 无法做到这一点。

我正在使用 mysqldb 和 Python 2.7。

顺便说一句，read_sql 在运行时使用的内存似乎比最终 DataFrame 存储所需的内存要多得多（大约 2 倍）。

In [70]: df=pd.read_sql('select ARP, ACP from train where seq < 5', connection)

In [71]: df
Out[71]: 
   ARP      ACP
0  1.17915  1.42595
1  1.10578  1.21369
2  1.35629  1.12693
3  1.56740  1.61847
4  1.28060  1.05935


In [72]: df.dtypes
Out[72]: 
ARP    float64
ACP    float64
dtype: object

【问题讨论】：

这也可以避免在你有 NaN 时将整数列转换为浮点数。

【参考方案1】：

cast() 和 convert() 呢？

'SELECT cast(ARP as float32()), cast (ACP as float32()) from train where seq < 5'

或类似的东西。

http://www.smallsql.de/doc/sql-functions/system/convert.html

【讨论】：

这是一个非常好的建议，但我发现了 2 个问题：1. cast 不允许 FLOAT (dev.mysql.com/doc/refman/5.6/en/cast-functions.html) 2. 即使尝试 DECIMAL(10,2)，例如， read_sql 仍然返回 float64 值。这有多烦人？

【参考方案2】：

看看thisgithub issue，看起来他们倾向于添加选项。

【参考方案3】：

您可以使用 pandas read_sql_query，它允许您指定返回的 dtypes（仅从 pandas 1.3 开始支持）。

pd.read_sql_query('select ARP, ACP from train where seq < 5', connection,
                  dtype='ARP': np.float32, 'ACP': np.float32)

【讨论】：

终于！完美的解决方案，5 年后;)

【参考方案4】：

顺便说一句，read_sql 在运行时使用的内存似乎比最终 DataFrame 存储所需的内存多得多（大约 2 倍）。

也许你可以试试我们的工具ConnectorX (pip install -U connectorx)，它是在 Rust 中实现的，旨在提高pandas.read_sql 在时间和内存使用方面的性能，并提供类似的接口。要切换到它，您只需：

import connectorx as cx
conn_url = "mysql://username:password@server:port/database"
query = "select ARP, ACP from train where seq < 5"
df = cx.read_sql(conn_url, query)

pandas.read_sql在运行过程中使用大量内存的原因是因为它的中间python对象很大，在ConnectorX我们使用Rust和流处理来解决这个问题。

这是一些基准测试结果：

PostgreSQL：

MySQL：

【讨论】：

看起来不错的工具！