diff --git a/doc/source/whatsnew/v0.24.1.rst b/doc/source/whatsnew/v0.24.1.rst index 222963a7ff71a..0923b05d41479 100644 --- a/doc/source/whatsnew/v0.24.1.rst +++ b/doc/source/whatsnew/v0.24.1.rst @@ -31,7 +31,6 @@ Fixed Regressions Enhancements ^^^^^^^^^^^^ - .. _whatsnew_0241.bug_fixes: Bug Fixes diff --git a/doc/source/whatsnew/v0.25.0.rst b/doc/source/whatsnew/v0.25.0.rst index 939fb8b9415bd..052f052420e41 100644 --- a/doc/source/whatsnew/v0.25.0.rst +++ b/doc/source/whatsnew/v0.25.0.rst @@ -165,7 +165,7 @@ MultiIndex I/O ^^^ -- +- :func:`get_schema` now accepts an `index` parameter (default: `False`) that includes the index in the generated schema. (:issue:`9084`) - - diff --git a/pandas/io/sql.py b/pandas/io/sql.py index aaface5415384..7e4cefddc2746 100644 --- a/pandas/io/sql.py +++ b/pandas/io/sql.py @@ -1223,8 +1223,9 @@ def drop_table(self, table_name, schema=None): self.get_table(table_name, schema).drop() self.meta.clear() - def _create_sql_schema(self, frame, table_name, keys=None, dtype=None): - table = SQLTable(table_name, self, frame=frame, index=False, keys=keys, + def _create_sql_schema(self, frame, table_name, keys=None, dtype=None, + index=False): + table = SQLTable(table_name, self, frame=frame, index=index, keys=keys, dtype=dtype) return str(table.sql_schema()) @@ -1565,13 +1566,14 @@ def drop_table(self, name, schema=None): name=_get_valid_sqlite_name(name)) self.execute(drop_sql) - def _create_sql_schema(self, frame, table_name, keys=None, dtype=None): - table = SQLiteTable(table_name, self, frame=frame, index=False, + def _create_sql_schema(self, frame, table_name, keys=None, dtype=None, + index=False): + table = SQLiteTable(table_name, self, frame=frame, index=index, keys=keys, dtype=dtype) return str(table.sql_schema()) -def get_schema(frame, name, keys=None, con=None, dtype=None): +def get_schema(frame, name, keys=None, con=None, dtype=None, index=False): """ Get the SQL db table schema for the given frame. @@ -1589,8 +1591,11 @@ def get_schema(frame, name, keys=None, con=None, dtype=None): dtype : dict of column name to SQL type, default None Optional specifying the datatype for columns. The SQL type should be a SQLAlchemy type, or a string for sqlite3 fallback connection. + index : boolean, default False + Whether to include DataFrame index as a column """ pandas_sql = pandasSQL_builder(con=con) - return pandas_sql._create_sql_schema(frame, name, keys=keys, dtype=dtype) + return pandas_sql._create_sql_schema( + frame, name, keys=keys, dtype=dtype, index=index) diff --git a/pandas/tests/io/test_sql.py b/pandas/tests/io/test_sql.py index 75a6d8d009083..e37921441596b 100644 --- a/pandas/tests/io/test_sql.py +++ b/pandas/tests/io/test_sql.py @@ -823,6 +823,21 @@ def test_get_schema_keys(self): constraint_sentence = 'CONSTRAINT test_pk PRIMARY KEY ("A", "B")' assert constraint_sentence in create_sql + @pytest.mark.parametrize("index_arg, expected", [ + ({}, False), + ({"index": False}, False), + ({"index": True}, True), + ]) + def test_get_schema_with_index(self, index_arg, expected): + frame = DataFrame({ + 'one': pd.Series([1, 2, 3], index=['a', 'b', 'c']), + 'two': pd.Series([1, 2, 3], index=['a', 'b', 'c']) + }) + frame.index.name = 'alphabet' + + create_sql = sql.get_schema(frame, 'test', con=self.conn, **index_arg) + assert ('alphabet' in create_sql) == expected + def test_chunksize_read(self): df = DataFrame(np.random.randn(22, 5), columns=list('abcde')) df.to_sql('test_chunksize', self.conn, index=False)