Python Pandas to_sql, how to create a table with a primary key?

后端 未结 3 628
悲哀的现实
悲哀的现实 2020-12-01 08:01

I would like to create a MySQL table with Pandas\' to_sql function which has a primary key (it is usually kind of good to have a primary key in a mysql table) as so:

3条回答
  •  生来不讨喜
    2020-12-01 08:28

    Disclaimer: this answer is more experimental then practical, but maybe worth mention.

    I found that class pandas.io.sql.SQLTable has named argument key and if you assign it the name of the field then this field becomes the primary key:

    Unfortunately you can't just transfer this argument from DataFrame.to_sql() function. To use it you should:

    1. create pandas.io.SQLDatabase instance

      engine = sa.create_engine('postgresql:///somedb')
      pandas_sql = pd.io.sql.pandasSQL_builder(engine, schema=None, flavor=None)
      
    2. define function analoguous to pandas.io.SQLDatabase.to_sql() but with additional *kwargs argument which is passed to pandas.io.SQLTable object created inside it (i've just copied original to_sql() method and added *kwargs):

      def to_sql_k(self, frame, name, if_exists='fail', index=True,
                 index_label=None, schema=None, chunksize=None, dtype=None, **kwargs):
          if dtype is not None:
              from sqlalchemy.types import to_instance, TypeEngine
              for col, my_type in dtype.items():
                  if not isinstance(to_instance(my_type), TypeEngine):
                      raise ValueError('The type of %s is not a SQLAlchemy '
                                       'type ' % col)
      
          table = pd.io.sql.SQLTable(name, self, frame=frame, index=index,
                           if_exists=if_exists, index_label=index_label,
                           schema=schema, dtype=dtype, **kwargs)
          table.create()
          table.insert(chunksize)
      
    3. call this function with your SQLDatabase instance and the dataframe you want to save

      to_sql_k(pandas_sql, df2save, 'tmp',
              index=True, index_label='id', keys='id', if_exists='replace')
      

    And we get something like

    CREATE TABLE public.tmp
    (
      id bigint NOT NULL DEFAULT nextval('tmp_id_seq'::regclass),
    ...
    )
    

    in the database.

    PS You can of course monkey-patch DataFrame, io.SQLDatabase and io.to_sql() functions to use this workaround with convenience.

提交回复
热议问题