SELECT文の結果をDataFrameで受け取る

#困ったこと

psycopg2を使用してPostgreSQLからデータを取得すると、
SQLのデータ型がnumeric,decimal,dateの場合は、object型になってしまうため、
何らかの方法でデータを変換する必要がある。

#データを取得してから変換

pandas.astypeを使用して適切な型にキャストする。

with conn.cursor() as cursor:
    cursor.execute(select_sql)
    column_list = [d.name for d in cursor.description]
    df = pd.DataFrame(cursor.fetchall(), columns=column_list)
    dtype_dict = {}
    for d in cursor.description:
        if d.type_code==1700: dtype_dict[d.name] = 'float64'
        if d.type_code==1082: dtype_dict[d.name] = 'datetime64'
    if len(dtype_dict)> 0: df = df.astype(dtype_dict)

SQLのデータ型とpandasのデータ型の対応関係は以下の通り。

var	sql type	type_code	pd.dtype	astype
x1	smallint	21	int64
x2	integer	23	int64
x3	bigint	20	int64
x4	decimal	1700	object	float64
x5	numeric	1700	object	float64
x6	real	700	float64
x7	double precision	701	float64
x8	timestamp	1114	datetime64[ns]
x9	date	1082	object	datetime64

#データ取得時に変換

PostgreSQLからデータを取得する際に適用する変換関数を登録する方法もある。

from dateutil.parser import parse

def register_typecasters(conn):
    def cast_date(value, cursor):
        return parse(value)

    def cast_numeric(value, cursor):
        return float(value)

    cursor = conn.cursor()
    cursor.execute("SELECT NULL::date")
    date_oid = cursor.description[0][1]
    cursor.execute("SELECT NULL::timestamp")
    timestamp_oid = cursor.description[0][1]
    cursor.execute("SELECT NULL::timestamp with time zone")
    timestamptz_oid = cursor.description[0][1]
    oids = (date_oid, timestamp_oid, timestamptz_oid)
    new_type = psycopg2.extensions.new_type(oids, "DATE", cast_date)
    psycopg2.extensions.register_type(new_type)

    cursor.execute("SELECT NULL::numeric")
    numeric_oid = cursor.description[0][1]
    cursor.execute("SELECT NULL::decimal")
    decimal_oid = cursor.description[0][1]
    oids = (numeric_oid, decimal_oid)
    new_type = psycopg2.extensions.new_type(oids, "NUMERIC", cast_numeric)
    psycopg2.extensions.register_type(new_type)

#参考資料

8.1. 数値データ型
https://www.postgresql.jp/document/12/html/datatype-numeric.html
8.5. 日付/時刻データ型
https://www.postgresql.jp/document/12/html/datatype-datetime.html

Python psycopg2.extensions.register_type() Examples
https://www.programcreek.com/python/example/79078/psycopg2.extensions.register_type