pandasからSQLiteのデータの読み書き

概要

pandasで一時データを保存しておきたい時に、pickleやcsvで出すのも良いけどDBに出しておいた方が後で楽というケースがたまにある。そんな時に使いそうなSQLiteとの連携。

バージョン情報

Python 3.8.0
pandas==0.25.1

PythonでのSQLiteの利用

SQLiteはPythonに同梱されているので特にインストール作業等をしなくても使える。DBもファイルに書き出されるだけなので気軽に作って気楽に捨てられる。

まずはpandasは使わずにテーブルのcreateとデータのinsertなどを実行してみる。

import sqlite3

with sqlite3.connect("example.db") as conn:
    # テーブルの生成
    conn.execute("create table if not exists example(name text, age integer)")

    # insert
    conn.execute("insert into example values( 'hoge', 10 )")

    # ?でパラメータを設定する場合
    conn.execute("insert into example values( ?, ? )", [ 'fuga', 20 ])

    # executemanyで複数SQLの実行
    conn.executemany("insert into example values( ?, ? )",
               [ ( 'foo', 30 ), ( 'baz', 40 ) ])

    conn.commit()

上記はconnect時に指定した example.db というファイル名でDBファイルが生成される。

ファイルを作成せずにインメモリでも動かせる。

sqlite3.connect(":memory:")

でもpandas使う時はメモリに乗るならわざわざSQLiteに移したりしないか。

インデックスも貼っておこう。

import sqlite3

with sqlite3.connect("example.db") as conn:
    conn.execute("create unique index example_name_index on example (name)")

続いてselect。

import sqlite3

with sqlite3.connect("example.db") as conn:
    cur = conn.cursor()

    # 1行ずつ読み込む
    cur.execute( "select * from example" )
    for row in cur:
        print( '{}, {}'.format(row[0], row[1]))
        #=> hoge, 10
        #=> fuga, 20
        #=> foo, 30
        #=> baz, 40

     # 全行読み込む
    cur.execute( "select * from example" )
    rows = cur.fetchall()
    print(rows)
        #=> [('hoge', 10), ('fuga', 20), ('foo', 30), ('baz', 40)]

    cur.close()

pandasからのSQLiteの利用

DBに入れたデータをpandasで読み込む。

selectの結果を自動でDataFrameにしてくれるので便利。

import pandas as pd 

with sqlite3.connect("example.db") as conn:
    df = pd.read_sql('select * from example', con=conn)

DataFrameにカラムを追加して、結果を別のテーブルに保存してみる。

# カラム追加
df['new_column'] = 1

# データの投入
with sqlite3.connect("example.db") as conn:
    df.to_sql('example2', con=conn)

with sqlite3.connect("example.db") as conn:
    cur = conn.cursor()
    cur.execute( "select * from example2" )
    print(cur.fetchall())

    #=> [(0, 'hoge', 10, 1), (1, 'fuga', 20, 1), (2, 'foo', 30, 1), (3, 'baz', 40, 1)]

追加したカラム以外に、先頭に0, 1, 2, 3のような数値が入っている。これはDataFrameのindex。

indexを入れる必要がない場合は index=False を指定する。

    df.to_sql('example2', con=conn, index=False)

if_existsの挙動

to_sqlを実行するテーブルが既に存在する場合、デフォルトでは例外が投げられる。

# 再度to_sqlを実行
with sqlite3.connect("example.db") as conn:
    df.to_sql('example2', con=conn)

    #=> ValueError: Table 'example2' already exists.

テーブルを置き換える場合は、 if_existsにreplaceを指定する。

# データの投入
with sqlite3.connect("example.db") as conn:
    df.to_sql('example2', con=conn, if_exists='replace')

これでテーブルが再作成されてデータが投入される。

旧データはすべて削除されるけど、カラムの形式が変わってもスキーマを再作成する等の手間は発生しないので便利。

また既に存在するテーブルにデータを追記していきたい場合は、appendを指定する。

with sqlite3.connect("example.db") as conn:
    df.to_sql('example2', con=conn, if_exists='append')

    cur = conn.cursor()
    cur.execute( "select * from example2" )
    print(cur.fetchall())

        # 2回実行するとデータが2回入る
        #=> [(0, 'hoge', 10, 1), (1, 'fuga', 20, 1), (2, 'foo', 30, 1), (3, 'baz', 40, 1), (0, 'hoge', 10, 1), (1, 'fuga', 20, 1), (2, 'foo', 30, 1), (3, 'baz', 40, 1)]

テーブルのindexについて

コマンドラインから生成されたテーブルのスキーマを見てみる。

$ sqlite3 example.db
sqlite> .schema example2

CREATE TABLE IF NOT EXISTS "example2" (
"index" INTEGER,
  "name" TEXT,
  "age" INTEGER,
  "new_column" INTEGER
);
CREATE INDEX "ix_example2_index"ON "example2" ("index");

DataFrameのindexカラムに対して自動でDB側にindexが張られている。

事前にDataFrameにset_indexすることで好きなカラムにindexを貼れる。

with sqlite3.connect("example.db") as conn:
    df.set_index('name') \
        .to_sql('example2', con=conn, if_exists='replace')

CREATE TABLE IF NOT EXISTS "example2" (
"name" TEXT,
  "age" INTEGER,
  "new_column" INTEGER
);
CREATE INDEX "ix_example2_name"ON "example2" ("name");

read/writeの速度

100万行 × 100カラムのfloatが入ったDataFrameを作ってread/writeの速度を見てみる。

実行したマシンのCPUはRyzen2700X。

まずはwriteから。

import numpy as np

# サンプルデータの作成
df = pd.DataFrame(
    np.random.random([1000000, 100]))

# write
with sqlite3.connect("example2.db") as conn:
    %time df.to_sql('example', con=conn)

    #=> CPU times: user 25.2 s, sys: 1.17 s, total: 26.3 s
    #=> Wall time: 33.9 s

33.9秒かかった。ファイルサイズは993M。

ちなみにto_pickleした場合は下記。

%time df.to_pickle('foo.pickle')                                                                                               

    #=> CPU times: user 21.5 ms, sys: 583 ms, total: 604 ms
    #=> Wall time: 7.42 s

コスト低い出力と比べるとちょっと時間かかる。

でもgzipでto_csvするよりは速い。

%time df.to_csv('foo.csv.gz', compression='gzip')

    #=> CPU times: user 5min 44s, sys: 526 ms, total: 5min 45s
    #=> Wall time: 5min 45s

続いてread。

with sqlite3.connect("example2.db") as conn:
    %time df = pd.read_sql('select * from example', con=conn)

    #=> CPU times: user 17.5 s, sys: 1.07 s, total: 18.6 s
    #=> Wall time: 18.6 s

18.6秒。

試した限りでは速度は概ね線形。急ぎの処理でなければ十分使い物になりそう。

改定履歴

Author: Masato Watanabe, Date: 2019-12-28, 記事投稿

iMind Developers Blog

iMind開発者ブログ