PostgreSQLのデータをPandasのデータフレームとして読み書きする

最近、JupyterやPandasを使ってデータを処理する機会が増えてきました。

とは言え、手元のデータはPostgreSQLに溜まっていたり、あるいはSQLで処理したい、ということがよくあります。

というわけで、Jupyterを使っている時に、「PostgreSQLからデータを取り出して、Pandasやら何やらでいろいろ処理した後、結果をPostgreSQLを書き出す」というユースケースを想定して、その方法を調べてみました。

■やりたいこと

やりたいことは、PostgreSQLのデータをJupyter上でPandasのデータフレームとして読み込み、集計やデータ分析をした結果をPostgreSQLに書き戻す、ということです。

データの加工や整形など（データ前処理）はPostgreSQLの方が高速に行えるのでSQLで、複雑なアルゴリズムの適用はPythonで行いたい、そしてその結果をPostgreSQLに書き戻して利用したい、というケースを想定しています。

あるいはPostgreSQLのデータをmatplotlibを使って可視化したい、といった場合にも使えるでしょう。（この場合は書き戻しは必要ありませんが）

■必要なもの

必要なものは以下の通りです。

• PostgreSQL
• Jupyter
• pscycopg2
• SQLAlchemy

PostgreSQLからDataFrameにデータを読み込むだけならpsycopg2だけで実現可能なのですが、DataFrameのデータをPostgreSQLに書き戻すにはSQLAlchemyが必要になります。

つまり、最初からSQLAlchemyを用意しておいた方がいいでしょう。

■PostgreSQLに接続する

まず、PostgreSQLに接続する必要があります。

Python（つまりJupyter上で）必要なモジュールのインポートは以下の通りです。

from sqlalchemy import create_engine
import pandas as pd

そして、SQLAlchemyのcreate_engine()関数を使ってPostgreSQLのエンジンオブジェクトを作成します。

engine = create_engine('postgresql://postgres:postgres@localhost:5432/postgres')

■PostgreSQLのデータをPandasに読み込む

PostgreSQLのデータをPandasのDataFrameに読み込むには、read_sql()メソッドを使います。

以下の例では、pg_databaseテーブルの内容をDataFrame dfに読み込んでいます。

df = pd.read_sql("SELECT * FROM pg_database", engine)

このメソッドを一回呼び出すだけで、PostgreSQLのデータをDataFrameに取り込んで扱うことができます。

■PandasのデータをPostgreSQLに書き出す

DataFrameを使っていろいろ処理した結果をPostgreSQLに書き出すには、DataFrameのto_sql()メソッドを使います。

• pandas.DataFrame.to_sql — pandas 0.22.0 documentation
  https://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.to_sql.html

Pandasのドキュメントにも書かれていますが、SQLAlchemyを使う場合には、SQLAlchemyによってサポートされているDBMSで書き出しを行うことができます。そうでない場合には、SQLiteのみがサポートされます。よって、PostgreSQLに書き出したい場合には、SQLAlchemyを使う必要があります。

以下の例では、DataFrame dfの内容をt1テーブルに書き出しています。その際、既にt1テーブルが存在した場合には、そのテーブルを置き換えるように指定しています。（if_exists引数）

df.to_sql('t1', engine, if_exists='replace')

その結果は以下の通りです。

to_sql()メソッドにはif_exists以外にもいろいろなオプションがありますので、ぜひ確認してみていただければと思います。

■まとめ

以上、非常に簡単ではありましたが、Python（というかJupyter）からPandas（DataFrame）を介してPostgreSQLのデータを読み書きする方法をご紹介しました。

データベースエンジニアの目の前には、さまざまなデータ、多くのデータが横たわっています。

私自身は、PL/PythonやMADlibを始めとするIn-Database処理が好みではありますが、一方でPandasで扱うデータの前処理をPostgreSQLで実現できたら便利だなぁ、と感じることも多々あるのが現実だったりします。

ぜひ、さまざまなツールを知ることで、それらのデータを活用していただければと思います。

では、また。