Added a test for non-default indexes

Author: dhirschfeld

pandas/tests/io/test_parquet.py

@@ -247,33 +247,6 @@ def test_columns_dtypes_invalid(self, engine):
                       datetime.datetime(2011, 1, 1, 1, 1)]
         self.check_error_on_write(df, engine, ValueError)
 
-    def test_write_with_index(self, engine):
-
-        df = pd.DataFrame({'A': [1, 2, 3]})
-        self.check_round_trip(df, engine, write_kwargs={'compression': None})
-
-        # non-default index
-        for index in [[2, 3, 4],
-                      pd.date_range('20130101', periods=3),
-                      list('abc'),
-                      [1, 3, 4],
-                      pd.MultiIndex.from_tuples([('a', 1), ('a', 2),
-                                                 ('b', 1)]),
-                      ]:
-
-            df.index = index
-            self.check_error_on_write(df, engine, ValueError)
-
-        # index with meta-data
-        df.index = [0, 1, 2]
-        df.index.name = 'foo'
-        self.check_error_on_write(df, engine, ValueError)
-
-        # column multi-index
-        df.index = [0, 1, 2]
-        df.columns = pd.MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1)]),
-        self.check_error_on_write(df, engine, ValueError)
-
     @pytest.mark.parametrize('compression', [None, 'gzip', 'snappy', 'brotli'])
     def test_compression(self, engine, compression):
 
@@ -360,6 +333,61 @@ def test_categorical_unsupported(self, pa_lt_070):
         df = pd.DataFrame({'a': pd.Categorical(list('abc'))})
         self.check_error_on_write(df, pa, NotImplementedError)
 
+    def test_write_with_index(self, pa_lt_070):
+        engine = pa_lt_070
+        df = pd.DataFrame({'A': [1, 2, 3]})
+        self.check_round_trip(df, engine, write_kwargs={'compression': None})
+
+        # non-default index
+        for index in [[2, 3, 4],
+                      pd.date_range('20130101', periods=3),
+                      list('abc'),
+                      [1, 3, 4],
+                      pd.MultiIndex.from_tuples([('a', 1), ('a', 2),
+                                                 ('b', 1)]),
+                      ]:
+
+            df.index = index
+            self.check_error_on_write(df, engine, ValueError)
+
+        # index with meta-data
+        df.index = [0, 1, 2]
+        df.index.name = 'foo'
+        self.check_error_on_write(df, engine, ValueError)
+
+        # column multi-index
+        df.index = [0, 1, 2]
+        df.columns = pd.MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1)]),
+        self.check_error_on_write(df, engine, ValueError)
+
+    def test_write_with_non_default_index(self, pa_ge_070):
+        engine = pa_ge_070
+
+        df = pd.DataFrame({'A': [1, 2, 3]})
+        self.check_round_trip(df, engine, write_kwargs={'compression': None})
+
+        # non-default index
+        for index in [[2, 3, 4],
+                      pd.date_range('20130101', periods=3),
+                      list('abc'),
+                      [1, 3, 4],
+                      pd.MultiIndex.from_tuples([('a', 1), ('a', 2),
+                                                 ('b', 1)]),
+                      ]:
+
+            df.index = index
+            self.check_round_trip(df, engine)
+
+        # index with meta-data
+        df.index = [0, 1, 2]
+        df.index.name = 'foo'
+        self.check_round_trip(df, engine)
+
+        # column multi-index
+        df.index = [0, 1, 2]
+        df.columns = pd.MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1)]),
+        self.check_error_on_write(df, engine, ValueError)
+
 
 class TestParquetFastParquet(Base):
 
@@ -385,6 +413,34 @@ def test_basic(self, fp):
 
         self.check_round_trip(df, fp, write_kwargs={'compression': None})
 
+    def test_write_with_index(self, fp):
+        engine = fp
+
+        df = pd.DataFrame({'A': [1, 2, 3]})
+        self.check_round_trip(df, engine, write_kwargs={'compression': None})
+
+        # non-default index
+        for index in [[2, 3, 4],
+                      pd.date_range('20130101', periods=3),
+                      list('abc'),
+                      [1, 3, 4],
+                      pd.MultiIndex.from_tuples([('a', 1), ('a', 2),
+                                                 ('b', 1)]),
+                      ]:
+
+            df.index = index
+            self.check_error_on_write(df, engine, ValueError)
+
+        # index with meta-data
+        df.index = [0, 1, 2]
+        df.index.name = 'foo'
+        self.check_error_on_write(df, engine, ValueError)
+
+        # column multi-index
+        df.index = [0, 1, 2]
+        df.columns = pd.MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1)]),
+        self.check_error_on_write(df, engine, ValueError)
+
     @pytest.mark.skip(reason="not supported")
     def test_duplicate_columns(self, fp):