通过pandas.to_sql插入枚举数据类型

时间:2018-09-18 18:59:59

标签: python postgresql pandas sqlalchemy

我有一个熊猫数据框,如下所示:

In [93]: df = pd.DataFrame([[1, 'a'], [2, 'b'], [3, 'a']], columns=['val', 'types'])

In [94]: df
Out[94]:
   val types
0    1    a
1    2    b
2    3    a

In [95]: df.dtypes
Out[95]:
val       int64
types    object
dtype: object

# convert the 'types' column to category
In [96]: df.types = pd.Categorical(df.types)

假设我已经有一个列types ENUM 的postgres表,如何使用熊猫to_sql方法通过sqlachemy类型将数据插入到数据库中?

我尝试了这个,但是没用:

In [101]: class myEnum(enum.Enum):
     ...:     a = 1
     ...:     b = 2
     ...:

In [102]: from sqlalchemy.types import *

In [103]: sql_dtypes = {'val' : SmallInteger(), 'types': Enum(myEnum)}

In [104]: df.to_sql('tablename', conn, if_exists='replace', index=False,
                    dtype=sql_dtypes)

正在运行,这给我以下错误:

sqlalchemy.exc.ProgrammingError: (psycopg2.ProgrammingError) type "myEnum" already 
exists

[SQL: "CREATE TYPE myEnum AS ENUM ('a', 'b')"] (Background on this 
error at: http://sqlalche.me/e/f405)

使用schema中的pd.io.json.build_table_schema选项进行指定也无济于事。

1 个答案:

答案 0 :(得分:0)

我现在从here发现的一种解决方法是,先将数据类型写为varchar,然后再将其转换为enum

In [104]: df.to_sql('tablename', conn, if_exists='replace', index=False,
                    dtype=sql_dtypes)

In [105]: conn.execute("""ALTER TABLE tablename 
                          ALTER COLUMN types TYPE myEnum using types::myEnum;""")

但我希望有更好的解决方案。