Add additional tests

python/pyarrow/pandas_compat.py

@@ -530,7 +530,8 @@ def _add_any_metadata(table, pandas_metadata):
 
     schema = table.schema
 
-    n_index_levels = len(pandas_metadata['index_columns'])
+    index_columns = pandas_metadata['index_columns']
+    n_index_levels = len(index_columns)
     n_columns = len(pandas_metadata['columns']) - n_index_levels
 
     # Add time zones
@@ -539,7 +540,7 @@ def _add_any_metadata(table, pandas_metadata):
         raw_name = col_meta['name']
         if i >= n_columns:
             # index columns
-            raw_name = pandas_metadata['index_columns'][i - n_columns]
+            raw_name = index_columns[i - n_columns]
         if raw_name is None:
             raw_name = 'None'
 

python/pyarrow/tests/test_convert_pandas.py

@@ -1232,10 +1232,26 @@ class TestPandasConversion(object):
         result = table_subset.to_pandas()
         tm.assert_frame_equal(result, df[['a']])
 
-        table_subset2 = table.remove_column(1)
+        table_subset2 = table_subset.remove_column(1)
         result = table_subset2.to_pandas()
         tm.assert_frame_equal(result, df[['a']])
 
+        # non-default index
+        for index in [
+                pd.Index(['a', 'b', 'c'], name='index'),
+                pd.date_range("2017-01-01", periods=3, tz='Europe/Brussels')]:
+            df = pd.DataFrame({'a': [1, 2, 3],
+                               'b': [.1, .2, .3]}, index=index)
+            table = pa.Table.from_pandas(df)
+
+            table_subset = table.remove_column(1)
+            result = table_subset.to_pandas()
+            tm.assert_frame_equal(result, df[['a']])
+
+            table_subset2 = table_subset.remove_column(1)
+            result = table_subset2.to_pandas()
+            tm.assert_frame_equal(result, df[['a']].reset_index(drop=True))
+
 
 def _pytime_from_micros(val):
     microseconds = val % 1000000