在pandas中为read_sql指定dtypes

Question

我想指定在执行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在运行时(大约2x)似乎比最终的DataFrame存储需要更多的内存.

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

Answer 1

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

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

Answer 2

顺便说一句，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：