psycopg2経由でAmazon Redshiftからpandasのデータフレームへ取り込んでみた

こんにちは。データインテグレーション部改めデータアナリティクス事業本部の岩澤です。

最近はAWSの学習の合間にpythonもライブラリ仕様とかを調べつつ動かして癒されているのですが、

• postgresqlからpsycopg2で取得したデータをpandasのデータフレームに取り込める
• Redshiftからpsycopg2でデータを取得できる

と知りまして。
「つまりAmazon Redshiftからpandasのデータフレームにも取り込めるよね？」と思い、試してみました。
（AWSもpythonもまだまだ修行中ですので、なっていない箇所があるかもしれませんがご容赦を）

環境

OS:Mac OS Mojave (version 10.14.5)
Shell:fish 2.7.1 python:3.7.3
pythonのライブラリ:
numpy 1.16.4
pandas 0.24.2
pip 19.1.1
psycopg2 2.8.3
python-dateutil 2.8.0
pytz 2019.1
setuptools 40.8.0
six 1.12.0

テスト環境構築(psycopg2とpandasのインストール)

仮想環境を構築して実行しています。 ※fish shellなので仮想環境を有効化するコマンドが違いますが、基本bashとかと同じです。

$ python  -m venv test_psycopg2
$ source test_psycopg2/bin/activate.fish
$ pip install --upgrade pip
$ pip install psycopg2
$ pip install pandas

テーブルとデータ

今回取得してみるテーブルとデータです。

=> \d test_date;
 id       | integer               |           |          |
 value    | character varying(10) |           |          |
 reg_date | date                  |           |          |

=> SELECT * FROM test_date;
 10 | hige  | 2019-07-04
 20 | hoge  | 2019-07-05
 30 | mige  | 2019-07-06
 40 | moge  | 2019-07-07
 50 |       | 2019-07-06
 60 | moge  |
 70 |       |
 80 | None  | 2019-07-06

例1:psycopg2の出力をpandasに渡す

psycopg2はSELECTの結果をタプル型のデータとして取得します。
最初はそれをpandas.DataFrame()に渡して取り込んでみます。

このとき引数columnsに列名を渡す必要があるので注意が必要です。
列名の一覧は[x.name for x in cur.description]で取得できます。

import psycopg2
import pandas as pd

def get_connection():
    dsn = {
            "host":"",      # ホスト名をここに記載
            "port":"",      # ポート番号をここに記載
            "database":"",  # DB名をここに記載
            "user":"",      # ユーザをここに記載
            "password":"",  # パスワードをここに記載
            }
    con = psycopg2.connect(**dsn)
    return con

def sql_test1(sql):
    with get_connection() as conn:
        with conn.cursor() as cur:
            cur.execute(sql)
            rows = cur.fetchall()
            return pd.DataFrame(rows, columns=[x.name for x in cur.description])

sql = "SELECT id, value, TO_CHAR(reg_date, 'YYYY-MM-DD') AS reg_date FROM test_date"
df = sql_test1(sql)
print(df)
print(df.isnull())

実行結果です。

   id value    reg_date
0  10  hige  2019-07-04
1  20  hoge  2019-07-05
2  30  mige  2019-07-06
3  40  moge  2019-07-07
4  50  None  2019-07-06
5  60  moge        None
6  70  None        None
7  80  None  2019-07-06
      id  value  reg_date
0  False  False     False
1  False  False     False
2  False  False     False
3  False  False     False
4  False   True     False
5  False  False      True
6  False   True      True
7  False  False     False

あっさり。
ちゃんとNullの判定もできています。

例2:pandas.read_sql()を使ってみる

pandasにはSQLの結果を直接データフレームに取り込むpandas.read_sql()関数があります。
次はこれを利用してみます。

import psycopg2
import pandas as pd

def get_connection():
    dsn = {
            "host":"",     # ホスト名
            "port":"",     # ポート番号
            "database":"", # DB名
            "user":"",     # ユーザ
            "password":"", # パスワード
            }
    con = psycopg2.connect(**dsn)
    return con

def sql_test2(sql):
    with get_connection() as conn:
        return pd.read_sql(sql=sql, con=conn)

sql = "SELECT id, value, TO_CHAR(reg_date, 'YYYY-MM-DD') AS reg_date FROM test_date"
df = sql_test2(sql)
print(df)
print(df.isnull())
print(f"reg_date:dtype is:{df['reg_date'].dtype}")

   id value    reg_date
0  10  hige  2019-07-04
1  20  hoge  2019-07-05
2  30  mige  2019-07-06
3  40  moge  2019-07-07
4  50  None  2019-07-06
5  60  moge        None
6  70  None        None
7  80  None  2019-07-06
      id  value  reg_date
0  False  False     False
1  False  False     False
2  False  False     False
3  False  False     False
4  False   True     False
5  False  False      True
6  False   True      True
7  False  False     False
reg_date:dtype is:object

こっちのが、列名も自動的に設定してくれるので、さらに簡単ですね。
Nullもバッチリです。

日時として取り込む

例2にてreg_dateのdtypeを出力していますが、このときははobjectでした。
実は、pandas.read_sql()の引数parse_datesを利用すると日時として取り込むことができます。

parse_dates引数にはディクショナリ型({"列名1":"書式1","列名2":"書式2", ...})を渡します。

例2のsql_test2()関数を以下のように書き換えます。

def sql_test2(sql):
    with get_connection() as conn:
        # parse_dates引数に設定すれば日付型として取り込む
        return pd.read_sql(sql=sql, con=conn, parse_dates = {"reg_date":"%Y-%m-%d",})

実行結果はこうなります。

   id value   reg_date
0  10  hige 2019-07-04
1  20  hoge 2019-07-05
2  30  mige 2019-07-06
3  40  moge 2019-07-07
4  50  None 2019-07-06
5  60  moge        NaT
6  70  None        NaT
7  80  None 2019-07-06
      id  value  reg_date
0  False  False     False
1  False  False     False
2  False  False     False
3  False  False     False
4  False   True     False
5  False  False      True
6  False   True      True
7  False  False     False
reg_date:dtype is:datetime64[ns]

dtypeをdatetime64[ns]としてreg_Dateを取り込めました。

まとめ

今回はRedshiftからデータフレームに取り込めるかやってみました。 ODBCとか触っていた身としては、DBからのデータ取得すら簡単に書けるあたりpythonって便利だなーと思いました。

参考

O'Reilly Japan - Pythonによるデータ分析入門 第2版
psycopg2によるPython2.7からのAmazon Redshiftアクセスサンプル ｜ DevelopersIO
pandas.read_sql — pandas 0.24.2 documentation
Python PostgreSQLのテーブルをPandasのDataFrameへ読み込む - け日記