【Pandas】データベーステーブル読込み時のメモリエラーを回避する

5196 ワード
PostgreSQL Python pandas Python テキストリンク

※"データベース" とはPostgreSQL（またはAmazon RedShift）を指しています。

TLDR

pyscopg2のサーバーサイドカーソルを併用してデータを読み込んでいく必要がある。
（※PostgreSQL限定）

サーバーサイドカーソル、クライアントサイドカーソルについては下記記事が大変参考になる。
PythonとDB: DBIのcursorを理解する

事象

pandasでは、一度にメモリに乗り切らない巨大なデータを一定量ごとに読み込んで逐次処理するためのchunksizeオプションが用意されている。

しかし、巨大なDBテーブルの読み込みにchunksizeを指定してもメモリエラーが発生するケースがあった。

>>> import pandas as pd
>>> engine = get_engine()

# 10,000行テーブルの読込み(成功)
>>> small_df = pd.read_sql_table('small-table', con=engine)
>>> len(small_df)
10000

# chunksize を指定して10,000,000行のテーブル読込み(失敗)
>>> large_df = pd.read_sql_table('large-table', con=engine, chunksize=10000)
$ Killed  # メモリエラーで強制終了

解決方法

psycopg2の名前付きカーソル(= サーバーサイドカーソル)を使用する。

import psycopg2
import pandas as pd

def get_psycopg2_connection():
    ...

con = get_psycopg2_connection()

# サーバーサイドカーソルの定義
cursor = con.cursor('large-table-cursor')
cursor.execute('SELECT * FROM "larget-table" ;')

while True:
    rows = cursor.fetchmany(10000)
    if len(rows) == 0:
        break

    # カラム一覧を取得
    columns = [desc.name for desc in cursor.description]

    # データフレームに格納
    df = pd.DataFrame(rows, columns=columns)

con.close()

備考

read_sql_table/read_sql_query関数ではchunksizeを指定してもクライアントサイドカーソルが使われていると思われる（ソースコードレベルでの確証なし）。

Amazon RedShiftのドキュメントによると、巨大なテーブルに対してカーソルを使用することは推奨されていない。
※結果セットを一時的にリーダーノードに保持するため

参考: カーソルを使用するときのパフォーマンスに関する考慮事項

c++の反復器iterator
python自動掃雷の共有