PandasのDataFrameでPostgreSQLに読み書きする方法

こんにちは、データサイエンティストのたぬ(@tanuhack)です!

僕は普段、Python(PandasのDataFrame)でデータ分析しているので、SQL文を書かずに全部Pythonでやりたいなー思っていました。シームレスだし。

PostgreSQLからDataFrameに格納して、DataFrameからPostgreSQLにまとめて追加、書き換え出来たら最高じゃないですか...!

ということで今回は、PostgreSQLからDataFrameに格納したり、DataFrameからPostgreSQLに追加、書き換えしたりする方法を具体例を元に、実践における注意点を踏まえて紹介します。

PostgreSQLの準備

次のコマンドは、記事の理解を深めるために用意した具体例です。

データベースの読み込みと書き出しだけサクッと知りたい方は、この章を読み飛ばしてPostgreSQLを読み込むから読んでくださいね。

STEP.1
データベースの起動状況を確認する
ターミナル
$ brew services list
STEP.2
postgresqlを起動する
ターミナル
$ brew services start postgresql
STEP.3
データベースを新規作成する

『df-postgres』という名前にデータベースを作成します。

ターミナル
$ createdb df-postgres
STEP.4
データベースに接続する

STEP.3で作成したデータベースに接続します。

ターミナル
$ psql df-postgres
STEP.5
テーブルを作成する

『test』という名前のテーブルを作成します。

  • 1列目:id(int)
  • 2列目:name(text)
ターミナル
create table test (id int, name text);
STEP.6
testテーブルに2行、値を挿入する
ターミナル
insert into test (id, name) values (1,'tanu');
insert into test (id, name) values (2,'hack');

pandas-postgres-readto1

MEMO
上の画像は、PG CommanderというPostgreSQLのGUIツールの画面を表示しています。便利なので、よかったら使ってみてください。

さっそくこのテーブルをPandasのDataFrameに読み込んでみましょう。

PostgreSQLを読み込む

PythonでPostgreSQLを操作するためには、psycopg2モジュールを使用します。

ターミナル
$ pip install psycopg2

PostgreSQLからデータベースをPandasのDataFrameに格納するためには、pandas.DataFrameのメソッドread_sqlを使用します。

import pandas as pd
import psycopg2

# データベースの接続情報
connection_config = {
    'user': 'ユーザー名',
    'password': '',
    'host': 'localhost',
    'port': 'ポート番号', # なくてもOK
    'database': 'データベース名'
}

# PostgreSQLに接続する
connection = psycopg2.connect(**connection_config)

# 変数dfにテーブルの全レコードを格納する
df = pd.read_sql(sql='SELECT * FROM テーブル名;', con=connection)

pandas-postgres-readto3

データベースの接続情報がわからなければ、psqlコマンドの\conninfoを叩けば、データベース名・ユーザー名・ポート番号が返り値として得られます。

データベースの情報を表示する|psqlコマンド
\conninfo

pandas-postgres-readto2

PostgreSQLに書き込む

DataFrameの中身をPostgreSQLに書き込むには、pandas.DataFrameのメソッドto_sqlを使用します。

しかし、to_sqlのデフォルトのデータベースはSQLiteになっているので、これをPostgreSQLに変更しないといけません。

そのために、sqlalchemyモジュールからcreate_engineを使用します。

import pandas as pd
from sqlalchemy import create_engine

# データベースの接続情報
connection_config = {
    'user': 'ユーザー名',
    'password': '',
    'host': 'localhost',
    'port': 'ポート番号', # なくてもOK
    'database': 'データベース名'
}

engine = create_engine('postgresql://{user}:{password}@{host}:{port}/{database}'.format(**connection_config))

# PostgreSQLに書き込む
df.to_sql('テーブル名', con=engine, if_exists='appendまたはreplace', index=False)

テーブルに追加する

to_sqlのオプションでif_exists='append'と設定すると、元あるテーブルにDataFrameをPostgreSQLに追加することができます。

pandas-postgres-readto4

# PostgreSQLのテーブルにDataFrameを追加する
df.to_sql('テーブル名', con=engine, if_exists='append', index=False)

pandas-postgres-readto5

[aside]補足
appendするDataFrameのカラムはデータベースと同じ名前、列数にしないとエラーが発生します。
[/aside]

テーブルを上書きする

to_sqlのオプションでif_exists='replace'と設定すると、元あるテーブルに上書きする形でDataFrameをPostgreSQLに格納することができます。

ただ、公式ドキュメントにも記載してあるように、テーブルを上書きする前に元から存在するテーブルを削除しないとエラーが発生します。

if_exists : {‘fail’, ‘replace’, ‘append’}, default ‘fail’. How to behave if the table already exists.

  • fail: Raise a ValueError.
  • replace: Drop the table before inserting new values.
  • append: Insert new values to the existing table.

pandas.DataFrame.to_sql公式ドキュメント|pandas 0.24.2

pandas-postgres-readto6

# テーブルを削除するSQLを実行する
with connection as conn:
    with conn.cursor() as cur:
        cur.execute('DROP TABLE テーブル名')

# PostgreSQLのテーブルにDataFrameを上書きする
df.to_sql('テーブル名', con=engine, if_exists='replace', index=False)

pandas-postgres-readto7

ご覧の通り、上書きされていることが確認できますね。