Python手遊び(CSVとPostgreSQL間での相互運用) [アップしたコード間違ってたので差し替えました版]

この記事、何?

ずっと前から考えていたこと。
プログラム上では変数に持つけど、ShutDownするからCSVとかにしたい。っていうか、横断検索考えると、DBに置きたい。でも、相互変換をいちいち書くのめんどい・・・

列情報から最小限でテーブル作ってそこに値突っ込んで・・・とか考えてできたのがこれ。

追記 [2020/07/06]

駄目だ。疲れてたんだけどさ。
タイトルのメインのコード、1個前の記事のために書いたコードを再掲しちゃってた。。。
というわけで、今日(07/06)の夜差し替えました。

つまり何作ったの?

まず、CSVファイルから列名を拾う。
さすがに各列の型を自動判別にしたくないので指定する。
まあ、目的は機械学習なので列の規定形式は倍精度浮動小数。なので、整数型とテキスト型の列を指定する。
で、それをもとにさっきとった列名からCREATE TABLE, INSERT INTOで突っ込む。
(まあ、INSERTはPandas使ってもいいけど、この流れを作っておいたほうがほかでいろいろやりやすいからね。)

コード

# ExecuteSQLPostgreSQL.py
import ReadCSVColumn
import CreateSQL


import psycopg2
import psycopg2.extras


# 接続用文字列を返す
# O : 接続用文字列
def get_connection_string():
    hostname = '192.168.52.109'
    port = '5432'
    database = 'db01'
    username = 'user01'
    password = 'manager'
    return f'postgresql://{username}:{password}@{hostname}:{port}/{database}'


# PostgreSQL接続用のカーソルを返す
# O : connection
#     cursor
def get_cur():
    conn = psycopg2.connect(get_connection_string())
    cur = conn.cursor()
    return conn, cur


# PostgreSQL接続用のカーソルを返す(Ex版:列名取得可能)
# O : connection
#     cursor
def get_cur_ex():
    conn = psycopg2.connect(get_connection_string())
    cur_ex = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
    return conn, cur_ex


# 各種テスト用処理
def execute_sample_sqls():

    # connection
    conn, cur = get_cur()

    table_name = 'RDKit'

    # drop
    sql = CreateSQL.get_sql_drop_table(table_name)
    cur.execute(sql)

    # create
    csvpath = 'solubility.test.csv'
    columns = ReadCSVColumn.get_columns_list(csvpath)
    columns_int = ['SampleID', 'Atoms', 'Bonds']
    columns_text = ['SampleName', 'Structure', 'ID', 'NAME', 'SOL', 'SMILES', 'SOL_classification']
    sql = CreateSQL.get_sql_create_table(table_name, columns, columns_int, columns_text)
    cur.execute(sql)

    # read_csv
    import pandas as pd
    df = pd.read_csv('solubility.test.csv')
    # pd.DataFrameの参照もろもろ
    # print(type(df))
    # print('------------------------')
    # print(df.iloc[0, 0])
    # print(len(df))
    # print(len(df.columns))
    # print('------------------------')

    # insert
    sql = CreateSQL.get_sql_insert_table(table_name, columns)
    for rowid in range(len(df)):
        paras = list()
        for columnid in range(len(df.columns)):
            paras.append(str(df.iloc[rowid, columnid]))
        cur.execute(sql, paras)

    # select
    sql = CreateSQL.get_sql_select(table_name, columns)
    print(sql)
    df = pd.read_sql(sql=sql, con=conn)
    print(df)
    df.to_csv('solubility.test2.csv', index=False)
    df.columns = columns
    df.to_csv('solubility.test3.csv', index=False)

    # commit/close
    conn.commit()
    conn.close()


# 主処理
def main():
    execute_sample_sqls()


# 起動処理
if __name__ == '__main__':
    main()

も1個。

# CreateSQL.py
import ReadCSVColumn


# SQLを返す：DROP
# I : テーブル名
# O : SQL
def get_sql_drop_table(table_name):
    sql = f'DROP TABLE IF EXISTS {table_name}'
    return sql


# SQLを返す：SELECT
# I : テーブル名
#     列名リスト
# O : SQL
def get_sql_select(table_name, columns):
    sql = f'''
        SELECT
                {','.join(columns)}
        FROM
                {table_name}
    '''
    return sql


# SQLを返す：INSERT
# I : テーブル名
#     列名リスト
# O : SQL
def get_sql_insert_table(table_name, columns):
    sql = f'''
        INSERT INTO
                {table_name}
                (
                    {','.join(columns)}
                )
        VALUES
                (
                    {','.join(['%s' for column in columns])}
                )
    '''
    return sql


# SQLを返す：CREATE TABLE
# I : テーブル名
#     列名リスト
#     列名リスト(再掲：整数型)
#     列名リスト(再掲：テキスト型)
# O : SQL
def get_sql_create_table(table_name, columns, columns_int, columns_text):
    sql_columns = list()
    for column in columns:
        if column in columns_int:
            sql_columns.append(column + ' ' + 'integer')
        elif column in columns_text:
            sql_columns.append(column + ' ' + 'text')
        else:
            sql_columns.append(column + ' ' + 'double precision')
    sql = f'''
        CREATE TABLE {table_name}
        (
            {','.join(sql_columns)}
        )
    '''
    return sql


# 主処理
def main():
    csvpath = 'solubility.test.csv'
    columns = ReadCSVColumn.get_columns_list(csvpath)
    columns_int = ['SampleID', 'Atoms', 'Bonds']
    columns_text = ['SampleName', 'Structure', 'ID', 'NAME', 'SOL', 'SMILES', 'SOL_classification']
    sql = get_sql_create_table('RDKit', columns, columns_int, columns_text)
    print(sql)


# 起動処理
if __name__ == '__main__':
    main()

あっ、も1個あった。

# ReadCSVColumn.py


# テキストファイルを行のリスト形式で読み取る
# I : ファイルパス
# O : テキストリスト
def get_lines_from_textfile(filepath):
    with open(filepath, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    return lines


# CSVパスリストから1行目を読み取り項目リストとして返す
# I : CSVパス
# O : 項目リスト
def get_columns_list(csvpath):
    # 最終項目に改行がついていたので外す
    line_1st = get_lines_from_textfile(csvpath)[0]
    # 1行目の末尾に改行がつくので外す
    line_1st = line_1st.replace('\n', '')
    return line_1st.split(',')


# 主処理
def main():
    csvpath = 'solubility.test.csv'
    columns = get_columns_list(csvpath)
    for column in columns:
        print(column)


# 起動処理
if __name__ == '__main__':
    main()

これ全部動かすと初出のCSVと同型のテーブルをPostgreSQL上に作ってINSERTする、はず。
あとはDB上でマスタとJOINしてGROUP BYするなりご自由に。

まあ・・・PostgreSQL上でPythonが動くならそれが一番いいと思うんだけどねぇ・・・
まあ、将来の課題ということで。

感想

ようやくまともな枠ができかけたような気がする。
さて、次、もうちょっと本題に近いことをするよ。