fix test for older versions of pyarrow

Author: jorisvandenbossche

pandas/io/parquet.py

@@ -86,9 +86,12 @@ def __init__(self):
                 "\nor via pip\n"
                 "pip install -U pyarrow\n"
             )
+
+        self._pyarrow_lt_060 = (
+            LooseVersion(pyarrow.__version__) < LooseVersion('0.6.0'))
         self._pyarrow_lt_070 = (
-            LooseVersion(pyarrow.__version__) < LooseVersion('0.7.0')
-        )
+            LooseVersion(pyarrow.__version__) < LooseVersion('0.7.0'))
+
         self.api = pyarrow
 
     def write(self, df, path, compression='snappy',
@@ -99,17 +102,23 @@ def write(self, df, path, compression='snappy',
                 df, path, compression, coerce_timestamps, **kwargs
             )
         path, _, _ = get_filepath_or_buffer(path)
-        table = self.api.Table.from_pandas(df)
-        self.api.parquet.write_table(
-            table, path, compression=compression,
-            coerce_timestamps=coerce_timestamps, **kwargs)
+
+        if self._pyarrow_lt_060:
+            table = self.api.Table.from_pandas(df, timestamps_to_ms=True)
+            self.api.parquet.write_table(
+                table, path, compression=compression, **kwargs)
+
+        else:
+            table = self.api.Table.from_pandas(df)
+            self.api.parquet.write_table(
+                table, path, compression=compression,
+                coerce_timestamps=coerce_timestamps, **kwargs)
 
     def read(self, path, columns=None, **kwargs):
         path, _, _ = get_filepath_or_buffer(path)
         parquet_file = self.api.parquet.ParquetFile(path)
         if self._pyarrow_lt_070:
-            parquet_file.path = path
-            return self._read_lt_070(parquet_file, columns, **kwargs)
+            return self._read_lt_070(path, parquet_file, columns, **kwargs)
         kwargs['use_pandas_metadata'] = True
         return parquet_file.read(columns=columns, **kwargs).to_pandas()
 
@@ -143,17 +152,17 @@ def _validate_write_lt_070(self, df, path, compression='snappy',
                 "on a default index"
             )
 
-    def _read_lt_070(self, parquet_file, columns, **kwargs):
+    def _read_lt_070(self, path, parquet_file, columns, **kwargs):
         # Compatibility shim for pyarrow < 0.7.0
         # TODO: Remove in pandas 0.22.0
         from itertools import chain
         import json
         if columns is not None:
-            metadata = json.loads(parquet_file.metadata.metadata[b'pandas'])
+            metadata = json.loads(
+                parquet_file.metadata.metadata[b'pandas'].decode('utf-8'))
             columns = set(chain(columns, metadata['index_columns']))
         kwargs['columns'] = columns
-        kwargs['path'] = parquet_file.path
-        return self.api.parquet.read_table(**kwargs).to_pandas()
+        return self.api.parquet.read_table(path, **kwargs).to_pandas()
 
 
 class FastParquetImpl(BaseImpl):

pandas/tests/io/test_parquet.py

@@ -296,9 +296,14 @@ def test_read_columns(self, engine):
                               write_kwargs={'compression': None},
                               read_kwargs={'columns': ['string']})
 
-    def test_write_with_index(self, engine):
+    def test_write_index(self, engine):
         check_names = engine != 'fastparquet'
 
+        if engine == 'pyarrow':
+            import pyarrow
+            if LooseVersion(pyarrow.__version__) < LooseVersion('0.7.0'):
+                pytest.skip("pyarrow is < 0.7.0")
+
         df = pd.DataFrame({'A': [1, 2, 3]})
         self.check_round_trip(df, engine, write_kwargs={'compression': None})
 
@@ -314,34 +319,31 @@ def test_write_with_index(self, engine):
             self.check_round_trip(
                 df, engine,
                 write_kwargs={'compression': None},
-                check_names=check_names,
-            )
-        if engine != 'fastparquet':
-            # Not suppoprted in fastparquet as of 0.1.3
-            index = pd.MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1)])
-            df.index = index
-            self.check_round_trip(
-                df, engine,
-                write_kwargs={'compression': None},
-            )
+                check_names=check_names)
+
         # index with meta-data
         df.index = [0, 1, 2]
         df.index.name = 'foo'
-        self.check_round_trip(
-            df, engine,
-            write_kwargs={'compression': None}
-        )
+        self.check_round_trip( df, engine, write_kwargs={'compression': None})
+
+    def test_write_multiindex(self, pa_ge_070):
+        # Not suppoprted in fastparquet as of 0.1.3 or older pyarrow version
+        engine = pa_ge_070
+
+        df = pd.DataFrame({'A': [1, 2, 3]})
+        index = pd.MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1)])
+        df.index = index
+        self.check_round_trip(df, engine, write_kwargs={'compression': None})
 
+    def test_write_column_multiindex(self, engine):
         # column multi-index
-        df.index = [0, 1, 2]
-        df.columns = pd.MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1)]),
+        mi_columns = pd.MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1)])
+        df = pd.DataFrame(np.random.randn(4, 3), columns=mi_columns)
         self.check_error_on_write(df, engine, ValueError)
 
-    def test_multiindex_with_columns(self, engine):
-        if engine == 'fastparquet':
-            msg = "fastparquet doesn't support mulit-indexes as of 0.1.3"
-            pytest.xfail(msg)
+    def test_multiindex_with_columns(self, pa_ge_070):
 
+        engine = pa_ge_070
         dates = pd.date_range('01-Jan-2018', '01-Dec-2018', freq='MS')
         df = pd.DataFrame(randn(2 * len(dates), 3), columns=list('ABC'))
         index1 = pd.MultiIndex.from_product(