pandasででかいテーブルをMySQLからHDF5に落としこむ

pd.read_sqlしてto_hdfしたい所だが、メモリ上にのりきらないようなサイズだと動かない。

pd.read_sqlにはchunksizeという便利なものがあり、この引数を指定するとchunkに切られたDataFrameのiteratorを返してくれるので、どうにかなりそうに見えるが、read_sqlのソースを読むと、SQLのクエリは依然一度にexecuteされるのでSELECT *した結果が全部メモリに乗らないようなサイズだとchunksizeを使っても依然どうにもならない。

そこでちゃんとSQLクエリでLIMITとOFFSETを使って少しずつとってくる。

import pandas as pd

engine = 'mysql://(mysql host)/(your database)'

def convert():
    limit = 50000
    offset = 0
    while True:
        print('offset: %d' % offset)
        df = pd.read_sql_query('SELECT * FROM ticks LIMIT %d OFFSET %d' % (limit, offset), engine)
        df = df.set_index('id')
        if len(df) < limit:
            break
        offset += limit
        yield df
    yield df

store = pd.HDFStore('ticks.h5', complib='zlib')
key = 'ticks'

first_time = True

for chunk in convert():
    if first_time:
        store.put(key, chunk, format='table')
    else:
        store.append(key, chunk)
    first_time = False

store.close()