ENH/TST SQL: add multi-index support to to_sql

Author: jorisvandenbossche

pandas/io/sql.py

@@ -439,19 +439,25 @@ def maybe_asscalar(self, i):
     def insert(self):
         ins = self.insert_statement()
         data_list = []
-        # to avoid if check for every row
-        keys = self.frame.columns
+
         if self.index is not None:
-            for t in self.frame.itertuples():
-                data = dict((k, self.maybe_asscalar(v))
-                            for k, v in zip(keys, t[1:]))
-                data[self.index] = self.maybe_asscalar(t[0])
-                data_list.append(data)
+            temp = self.frame.copy()
+            temp.index.names = self.index
+            try:
+                temp.reset_index(inplace=True)
+            except ValueError as err:
+                raise ValueError(
+                    "duplicate name in index/columns: {0}".format(err))
         else:
-            for t in self.frame.itertuples():
-                data = dict((k, self.maybe_asscalar(v))
-                            for k, v in zip(keys, t[1:]))
-                data_list.append(data)
+            temp = self.frame
+
+        keys = temp.columns
+
+        for t in temp.itertuples():
+            data = dict((k, self.maybe_asscalar(v))
+                        for k, v in zip(keys, t[1:]))
+            data_list.append(data)
+
         self.pd_sql.execute(ins, data_list)
 
     def read(self, coerce_float=True, parse_dates=None, columns=None):
@@ -486,12 +492,24 @@ def read(self, coerce_float=True, parse_dates=None, columns=None):
 
     def _index_name(self, index, index_label):
         if index is True:
+            nlevels = self.frame.index.nlevels
+            # if index_label is specified, set this as index name(s)
             if index_label is not None:
-                return _safe_col_name(index_label)
-            elif self.frame.index.name is not None:
-                return _safe_col_name(self.frame.index.name)
+                if not isinstance(index_label, list):
+                    index_label = [index_label]
+                if len(index_label) != nlevels:
+                    raise ValueError(
+                        "Length of 'index_label' should match number of "
+                        "levels, which is {0}".format(nlevels))
+                else:
+                    return index_label
+            # return the used column labels for the index columns
+            if nlevels == 1 and 'index' not in self.frame.columns and self.frame.index.name is None:
+                return ['index']
             else:
-                return self.prefix + '_index'
+                return [l if l is not None else "level_{0}".format(i)
+                        for i, l in enumerate(self.frame.index.names)]
+
         elif isinstance(index, string_types):
             return index
         else:
@@ -507,10 +525,10 @@ def _create_table_statement(self):
                    for name, typ in zip(safe_columns, column_types)]
 
         if self.index is not None:
-            columns.insert(0, Column(self.index,
-                                     self._sqlalchemy_type(
-                                         self.frame.index),
-                                     index=True))
+            for i, idx_label in enumerate(self.index[::-1]):
+                idx_type = self._sqlalchemy_type(
+                    self.frame.index.get_level_values(i))
+                columns.insert(0, Column(idx_label, idx_type, index=True))
 
         return Table(self.name, self.pd_sql.meta, *columns)
 
@@ -788,6 +806,17 @@ def insert(self):
         cur.close()
         self.pd_sql.con.commit()
 
+    def _index_name(self, index, index_label):
+        if index is True:
+            if self.frame.index.name is not None:
+                return _safe_col_name(self.frame.index.name)
+            else:
+                return 'pandas_index'
+        elif isinstance(index, string_types):
+            return index
+        else:
+            return None
+
     def _create_table_statement(self):
         "Return a CREATE TABLE statement to suit the contents of a DataFrame."
 

pandas/io/tests/test_sql.py

@@ -7,7 +7,7 @@
 import nose
 import numpy as np
 
-from pandas import DataFrame, Series
+from pandas import DataFrame, Series, MultiIndex
 from pandas.compat import range, lrange, iteritems
 #from pandas.core.datetools import format as date_format
 
@@ -266,7 +266,7 @@ def _roundtrip(self):
         self.pandasSQL.to_sql(self.test_frame1, 'test_frame_roundtrip')
         result = self.pandasSQL.read_sql('SELECT * FROM test_frame_roundtrip')
 
-        result.set_index('pandas_index', inplace=True)
+        result.set_index('level_0', inplace=True)
         # result.index.astype(int)
 
         result.index.name = None
@@ -391,7 +391,7 @@ def test_roundtrip(self):
 
         # HACK!
         result.index = self.test_frame1.index
-        result.set_index('pandas_index', inplace=True)
+        result.set_index('level_0', inplace=True)
         result.index.astype(int)
         result.index.name = None
         tm.assert_frame_equal(result, self.test_frame1)
@@ -476,10 +476,10 @@ def connect(self):
     def test_to_sql_index_label(self):
         temp_frame = DataFrame({'col1': range(4)})
 
-        # no index name, defaults to 'pandas_index'
+        # no index name, defaults to 'index'
         sql.to_sql(temp_frame, 'test_index_label', self.conn)
         frame = sql.read_table('test_index_label', self.conn)
-        self.assertEqual(frame.columns[0], 'pandas_index')
+        self.assertEqual(frame.columns[0], 'index')
 
         # specifying index_label
         sql.to_sql(temp_frame, 'test_index_label', self.conn,
@@ -489,11 +489,11 @@ def test_to_sql_index_label(self):
                          "Specified index_label not written to database")
 
         # using the index name
-        temp_frame.index.name = 'index'
+        temp_frame.index.name = 'index_name'
         sql.to_sql(temp_frame, 'test_index_label', self.conn,
                    if_exists='replace')
         frame = sql.read_table('test_index_label', self.conn)
-        self.assertEqual(frame.columns[0], 'index',
+        self.assertEqual(frame.columns[0], 'index_name',
                          "Index name not written to database")
 
         # has index name, but specifying index_label
@@ -503,6 +503,43 @@ def test_to_sql_index_label(self):
         self.assertEqual(frame.columns[0], 'other_label',
                          "Specified index_label not written to database")
 
+    def test_to_sql_index_label_multiindex(self):
+        temp_frame = DataFrame({'col1': range(4)},
+            index=MultiIndex.from_product([('A0', 'A1'), ('B0', 'B1')]))
+        
+        # no index name, defaults to 'level_0' and 'level_1'
+        sql.to_sql(temp_frame, 'test_index_label', self.conn)
+        frame = sql.read_table('test_index_label', self.conn)
+        self.assertEqual(frame.columns[0], 'level_0')
+        self.assertEqual(frame.columns[1], 'level_1')
+
+        # specifying index_label
+        sql.to_sql(temp_frame, 'test_index_label', self.conn,
+                   if_exists='replace', index_label=['A', 'B'])
+        frame = sql.read_table('test_index_label', self.conn)
+        self.assertEqual(frame.columns[:2].tolist(), ['A', 'B'],
+                         "Specified index_labels not written to database")
+
+        # using the index name
+        temp_frame.index.names = ['A', 'B']
+        sql.to_sql(temp_frame, 'test_index_label', self.conn,
+                   if_exists='replace')
+        frame = sql.read_table('test_index_label', self.conn)
+        self.assertEqual(frame.columns[:2].tolist(), ['A', 'B'],
+                         "Index names not written to database")
+
+        # has index name, but specifying index_label
+        sql.to_sql(temp_frame, 'test_index_label', self.conn,
+                   if_exists='replace', index_label=['C', 'D'])
+        frame = sql.read_table('test_index_label', self.conn)
+        self.assertEqual(frame.columns[:2].tolist(), ['C', 'D'],
+                         "Specified index_labels not written to database")
+
+        # wrong length of index_label
+        self.assertRaises(ValueError, sql.to_sql, temp_frame,
+                          'test_index_label', self.conn, if_exists='replace',
+                          index_label='C')
+
     def test_read_table_columns(self):
         # test columns argument in read_table
         sql.to_sql(self.test_frame1, 'test_frame', self.conn)
@@ -566,6 +603,23 @@ def test_sql_open_close(self):
 
         tm.assert_frame_equal(self.test_frame2, result)
 
+    def test_roundtrip(self):
+        # this test otherwise fails, Legacy mode still uses 'pandas_index'
+        # as default index column label
+        sql.to_sql(self.test_frame1, 'test_frame_roundtrip',
+                   con=self.conn, flavor='sqlite')
+        result = sql.read_sql(
+            'SELECT * FROM test_frame_roundtrip',
+            con=self.conn,
+            flavor='sqlite')
+
+        # HACK!
+        result.index = self.test_frame1.index
+        result.set_index('pandas_index', inplace=True)
+        result.index.astype(int)
+        result.index.name = None
+        tm.assert_frame_equal(result, self.test_frame1)
+
 
 class _TestSQLAlchemy(PandasSQLTest):
     """
@@ -788,6 +842,16 @@ def setUp(self):
 
         self._load_test1_data()
 
+    def _roundtrip(self):
+        # overwrite parent function (level_0 -> pandas_index in legacy mode)
+        self.drop_table('test_frame_roundtrip')
+        self.pandasSQL.to_sql(self.test_frame1, 'test_frame_roundtrip')
+        result = self.pandasSQL.read_sql('SELECT * FROM test_frame_roundtrip')
+        result.set_index('pandas_index', inplace=True)
+        result.index.name = None
+
+        tm.assert_frame_equal(result, self.test_frame1)
+
     def test_invalid_flavor(self):
         self.assertRaises(
             NotImplementedError, sql.PandasSQLLegacy, self.conn, 'oracle')