python - 从远程服务器中提取大量数据到 DataFrame 中

Question

为了提供尽可能多的上下文，我正在尝试使用 psycopg2 进行连接，将存储在远程 postgres 服务器 (heroku) 上的一些数据提取到 pandas DataFrame 中。

我对两个特定的表感兴趣，users 和 events，并且连接工作正常，因为在下拉用户数据时

import pandas.io.sql as sql 
# [...]
users = sql.read_sql("SELECT * FROM users", conn)

等待几秒钟后，DataFrame 按预期返回。

<class 'pandas.core.frame.DataFrame'>
Int64Index: 67458 entries, 0 to 67457
Data columns (total 35 columns): [...]

然而，当试图直接从 ipython 中提取更大、更重的事件数据时，很长一段时间后，它就崩溃了:

In [11]: events = sql.read_sql("SELECT * FROM events", conn)
vagrant@data-science-toolbox:~$

当从 iPython notebook 尝试时，我得到了Dead kernel 错误

The kernel has died, would you like to restart it? If you do not restart the kernel, you will be able to save the notebook, but running code will not work until the notebook is reopened.

---

更新 #1:

为了更好地了解我尝试引入的事件表的大小，下面是记录数和每个记录的属性数:

In [11]: sql.read_sql("SELECT count(*) FROM events", conn)
Out[11]:
     count
0  2711453

In [12]: len(sql.read_sql("SELECT * FROM events LIMIT 1", conn).columns)
Out[12]: 18

---

更新#2:

内存绝对是当前 read_sql 实现的瓶颈:当拉下 events 并尝试运行 iPython 的另一个实例时，结果是

vagrant@data-science-toolbox:~$ sudo ipython
-bash: fork: Cannot allocate memory

---

更新 #3:

我首先尝试了一个只返回部分数据帧数组的read_sql_chunked实现:

def read_sql_chunked(query, conn, nrows, chunksize=1000):
    start = 0
    dfs = []
    while start < nrows:
        df = pd.read_sql("%s LIMIT %s OFFSET %s" % (query, chunksize, start), conn)
        start += chunksize
        dfs.append(df)
        print "Events added: %s to %s of %s" % (start-chunksize, start, nrows)
    # print "concatenating dfs"
    return dfs

event_dfs = read_sql_chunked("SELECT * FROM events", conn, events_count, 100000)

这很好用，但是当尝试连接 DataFrame 时，内核再次死机。
这是在为 VM 提供 2GB RAM 之后。

根据 Andy 对 read_sql 与 read_csv 在实现和性能方面的差异的解释，我接下来尝试的是将记录附加到 CSV 中，然后全部读取进入数据框:

event_dfs[0].to_csv(path+'new_events.csv', encoding='utf-8')

for df in event_dfs[1:]:
    df.to_csv(path+'new_events.csv', mode='a', header=False, encoding='utf-8')

同样，写入 CSV 成功完成 – 一个 657MB 的文件 – 但从 CSV 读取从未完成。

既然 2GB 似乎还不够，那么如何估算出多少 RAM 足以读取一个 657MB 的 CSV 文件？

---

感觉我缺少对 DataFrames 或 psycopg2 的一些基本理解，但我被卡住了，我什至无法查明瓶颈或优化的地方。

从远程 (postgres) 服务器提取大量数据的正确策略是什么？

Answer 1

我怀疑这里有一些(相关的)事情在起作用导致缓慢:

1. read_sql 是用 python 编写的，所以它有点慢(特别是与 read_csv 相比，它是用 cython 编写的 - 并且为了速度而仔细实现!)并且它依赖于sqlalchemy 而不是一些(可能快得多)C-DBAPI。 迁移到 sqlalchmey 的动力是让迁移在未来变得更容易(以及跨 sql 平台支持)。
2. 您可能会用完内存，因为内存中有太多 python 对象(这与不使用 C-DBAPI 有关)，但可能会被解决...

我认为直接的解决方案是基于 block 的方法(并且有一个 feature request 可以在 pandas read_sql 和 read_sql_table 中本地工作)。

编辑:从 Pandas v0.16.2 开始，这种基于 block 的方法在 read_sql 中原生实现。

---

由于您使用的是 postgres，因此您可以访问 LIMIT and OFFSET queries ，这使得分 block 非常容易。 (我是否认为这些并非在所有 sql 语言中都可用？)

首先，获取表中的行数(或 estimate):

nrows = con.execute('SELECT count(*) FROM users').fetchone()[0]  # also works with an sqlalchemy engine

使用它来遍历表(为了调试，你可以添加一些打印语句来确认它正在工作/没有崩溃!)然后合并结果:

def read_sql_chunked(query, con, nrows, chunksize=1000):
    start = 1
    dfs = []  # Note: could probably make this neater with a generator/for loop
    while start < nrows:
        df = pd.read_sql("%s LIMIT %s OFFSET %s" % (query, chunksize, start), con)
        dfs.append(df)
    return pd.concat(dfs, ignore_index=True)

注意:这假设数据库适合内存!如果不是这样，您将需要处理每个 block (mapreduce 样式)...或投资更多内存!