reorganizte utility functions in parquet and filasystem modules; support pathlib.Path in Filesystem API

python/pyarrow/filesystem.py

@@ -15,10 +15,26 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from os.path import join as pjoin
 import os
+import six
+import inspect
 import posixpath
 
+from os.path import join as pjoin
+from six.moves.urllib.parse import urlparse
+
+try:
+    # pathlib might not be available
+    try:
+        import pathlib
+    except ImportError:
+        # python 2 backport
+        import pathlib2 as pathlib
+    _has_pathlib = True
+except ImportError:
+    _has_pathlib = False
+
+import pyarrow as pa
 from pyarrow.util import implements
 
 
@@ -68,6 +84,7 @@ class FileSystem(object):
         -------
         usage : int
         """
+        path = _ensure_str_path(path)
         path_info = self.stat(path)
         if path_info['kind'] == 'file':
             return path_info['size']
@@ -199,10 +216,12 @@ class LocalFileSystem(FileSystem):
 
     @implements(FileSystem.ls)
     def ls(self, path):
+        path = _ensure_str_path(path)
         return sorted(pjoin(path, x) for x in os.listdir(path))
 
     @implements(FileSystem.mkdir)
     def mkdir(self, path, create_parents=True):
+        path = _ensure_str_path(path)
         if create_parents:
             os.makedirs(path)
         else:
@@ -210,10 +229,12 @@ class LocalFileSystem(FileSystem):
 
     @implements(FileSystem.isdir)
     def isdir(self, path):
+        path = _ensure_str_path(path)
         return os.path.isdir(path)
 
     @implements(FileSystem.isfile)
     def isfile(self, path):
+        path = _ensure_str_path(path)
         return os.path.isfile(path)
 
     @implements(FileSystem._isfilestore)
@@ -222,6 +243,7 @@ class LocalFileSystem(FileSystem):
 
     @implements(FileSystem.exists)
     def exists(self, path):
+        path = _ensure_str_path(path)
         return os.path.exists(path)
 
     @implements(FileSystem.open)
@@ -229,17 +251,19 @@ class LocalFileSystem(FileSystem):
         """
         Open file for reading or writing
         """
+        path = _ensure_str_path(path)
         return open(path, mode=mode)
 
     @property
     def pathsep(self):
         return os.path.sep
 
-    def walk(self, top_dir):
+    def walk(self, path):
         """
         Directory tree generator, see os.walk
         """
-        return os.walk(top_dir)
+        path = _ensure_str_path(path)
+        return os.walk(path)
 
 
 class DaskFileSystem(FileSystem):
@@ -268,14 +292,17 @@ class DaskFileSystem(FileSystem):
 
     @implements(FileSystem.delete)
     def delete(self, path, recursive=False):
+        path = _ensure_str_path(path)
         return self.fs.rm(path, recursive=recursive)
 
     @implements(FileSystem.exists)
     def exists(self, path):
+        path = _ensure_str_path(path)
         return self.fs.exists(path)
 
     @implements(FileSystem.mkdir)
     def mkdir(self, path, create_parents=True):
+        path = _ensure_str_path(path)
         if create_parents:
             return self.fs.mkdirs(path)
         else:
@@ -286,22 +313,26 @@ class DaskFileSystem(FileSystem):
         """
         Open file for reading or writing
         """
+        path = _ensure_str_path(path)
         return self.fs.open(path, mode=mode)
 
     def ls(self, path, detail=False):
+        path = _ensure_str_path(path)
         return self.fs.ls(path, detail=detail)
 
-    def walk(self, top_path):
+    def walk(self, path):
         """
         Directory tree generator, like os.walk
         """
-        return self.fs.walk(top_path)
+        path = _ensure_str_path(path)
+        return self.fs.walk(path)
 
 
 class S3FSWrapper(DaskFileSystem):
 
     @implements(FileSystem.isdir)
     def isdir(self, path):
+        path = _ensure_str_path(path)
         try:
             contents = self.fs.ls(path)
             if len(contents) == 1 and contents[0] == path:
@@ -313,6 +344,7 @@ class S3FSWrapper(DaskFileSystem):
 
     @implements(FileSystem.isfile)
     def isfile(self, path):
+        path = _ensure_str_path(path)
         try:
             contents = self.fs.ls(path)
             return len(contents) == 1 and contents[0] == path
@@ -326,7 +358,7 @@ class S3FSWrapper(DaskFileSystem):
         Generator version of what is in s3fs, which yields a flattened list of
         files
         """
-        path = path.replace('s3://', '')
+        path = _ensure_str_path(path).replace('s3://', '')
         directories = set()
         files = set()
 
@@ -350,3 +382,61 @@ class S3FSWrapper(DaskFileSystem):
         for directory in directories:
             for tup in self.walk(directory, refresh=refresh):
                 yield tup
+
+
+def is_path(x):
+    return (isinstance(x, six.string_types)
+            or (_has_pathlib and isinstance(x, pathlib.Path)))
+
+
+def _ensure_str_path(path):
+    if isinstance(path, six.string_types):
+        return path
+    elif _has_pathlib and isinstance(path, pathlib.Path):
+        return str(path)
+    else:
+        raise TypeError('Path must be a string or an instance of pathlib.Path,'
+                        ' got object with type: {}'.format(type(path)))
+
+
+def _ensure_filesystem(fs):
+    fs_type = type(fs)
+
+    # If the arrow filesystem was subclassed, assume it supports the full
+    # interface and return it
+    if not issubclass(fs_type, FileSystem):
+        for mro in inspect.getmro(fs_type):
+            if mro.__name__ is 'S3FileSystem':
+                return S3FSWrapper(fs)
+            # In case its a simple LocalFileSystem (e.g. dask) use native arrow
+            # FS
+            elif mro.__name__ is 'LocalFileSystem':
+                return LocalFileSystem.get_instance()
+
+        raise IOError('Unrecognized filesystem: {0}'.format(fs_type))
+    else:
+        return fs
+
+
+def _get_fs_from_path(path):
+    """
+    return filesystem from path which could be an HDFS URI
+    """
+    # input can be hdfs URI such as hdfs://host:port/myfile.parquet
+    path = _ensure_str_path(path)
+    # if _has_pathlib and isinstance(path, pathlib.Path):
+    #     path = str(path)
+    parsed_uri = urlparse(path)
+    if parsed_uri.scheme == 'hdfs':
+        netloc_split = parsed_uri.netloc.split(':')
+        host = netloc_split[0]
+        if host == '':
+            host = 'default'
+        port = 0
+        if len(netloc_split) == 2 and netloc_split[1].isnumeric():
+            port = int(netloc_split[1])
+        fs = pa.hdfs.connect(host=host, port=port)
+    else:
+        fs = LocalFileSystem.get_instance()
+
+    return fs

python/pyarrow/parquet.py

@@ -18,21 +18,13 @@
 from collections import defaultdict
 from concurrent import futures
 import os
-import inspect
 import json
 import re
-import six
-from six.moves.urllib.parse import urlparse
-# pathlib might not be available in Python 2
-try:
-    import pathlib
-    _has_pathlib = True
-except ImportError:
-    _has_pathlib = False
 
 import numpy as np
 
-from pyarrow.filesystem import FileSystem, LocalFileSystem, S3FSWrapper
+from pyarrow.filesystem import (LocalFileSystem, is_path,  _ensure_str_path,
+                                _ensure_filesystem, _get_fs_from_path)
 from pyarrow._parquet import (ParquetReader, RowGroupStatistics,  # noqa
                               FileMetaData, RowGroupMetaData,
                               ColumnChunkMetaData,
@@ -634,11 +626,6 @@ class ParquetPartitions(object):
                              filter[1])
 
 
-def is_path(x):
-    return (isinstance(x, six.string_types)
-            or (_has_pathlib and isinstance(x, pathlib.Path)))
-
-
 class ParquetManifest(object):
     """
 
@@ -957,34 +944,6 @@ class ParquetDataset(object):
         self.pieces = [p for p in self.pieces if all_filters_accept(p)]
 
 
-def _ensure_filesystem(fs):
-    fs_type = type(fs)
-
-    # If the arrow filesystem was subclassed, assume it supports the full
-    # interface and return it
-    if not issubclass(fs_type, FileSystem):
-        for mro in inspect.getmro(fs_type):
-            if mro.__name__ is 'S3FileSystem':
-                return S3FSWrapper(fs)
-            # In case its a simple LocalFileSystem (e.g. dask) use native arrow
-            # FS
-            elif mro.__name__ is 'LocalFileSystem':
-                return LocalFileSystem.get_instance()
-
-        raise IOError('Unrecognized filesystem: {0}'.format(fs_type))
-    else:
-        return fs
-
-
-def _ensure_str_path(path):
-    if isinstance(path, six.string_types):
-        return path
-    elif _has_pathlib and isinstance(path, pathlib.Path):
-        return str(path)
-    else:
-        raise TypeError('Path must be a string or an instance of pathlib.Path')
-
-
 def _make_manifest(path_or_paths, fs, pathsep='/', metadata_nthreads=1):
     partitions = None
     common_metadata_path = None
@@ -1103,7 +1062,7 @@ def write_table(table, where, row_group_size=None, version='1.0',
     except Exception:
         if is_path(where):
             try:
-                os.remove(where)
+                os.remove(_ensure_str_path(where))
             except os.error:
                 pass
         raise
@@ -1267,26 +1226,3 @@ def read_schema(where):
     schema : pyarrow.Schema
     """
     return ParquetFile(where).schema.to_arrow_schema()
-
-
-def _get_fs_from_path(path):
-    """
-    return filesystem from path which could be an HDFS URI
-    """
-    # input can be hdfs URI such as hdfs://host:port/myfile.parquet
-    if _has_pathlib and isinstance(path, pathlib.Path):
-        path = str(path)
-    parsed_uri = urlparse(path)
-    if parsed_uri.scheme == 'hdfs':
-        netloc_split = parsed_uri.netloc.split(':')
-        host = netloc_split[0]
-        if host == '':
-            host = 'default'
-        port = 0
-        if len(netloc_split) == 2 and netloc_split[1].isnumeric():
-            port = int(netloc_split[1])
-        fs = pa.hdfs.connect(host=host, port=port)
-    else:
-        fs = LocalFileSystem.get_instance()
-
-    return fs

python/pyarrow/tests/test_parquet.py

@@ -457,10 +457,10 @@ def test_pandas_parquet_pyfile_roundtrip(tempdir):
 
     arrow_table = pa.Table.from_pandas(df)
 
-    with open(filename, 'wb') as f:
+    with filename.open('wb') as f:
         _write_table(arrow_table, f, version="1.0")
 
-    data = io.BytesIO(open(filename, 'rb').read())
+    data = io.BytesIO(filename.read_bytes())
 
     table_read = _read_table(data)
     df_read = table_read.to_pandas()
@@ -1425,10 +1425,10 @@ def _generate_partition_directories(fs, base_dir, partition_spec, df):
                     _write_table(part_table, f)
                 assert fs.exists(file_path)
 
-                _touch(level_dir / '_SUCCESS')
+                (level_dir / '_SUCCESS').touch()
             else:
                 _visit_level(level_dir, level + 1, this_part_keys)
-                _touch(level_dir / '_SUCCESS')
+                (level_dir / '_SUCCESS').touch()
 
     _visit_level(base_dir, 0, [])
 
@@ -1528,11 +1528,6 @@ def _filter_partition(df, part_keys):
     return df[predicate].drop(to_drop, axis=1)
 
 
-def _touch(path):
-    with open(path, 'wb'):
-        pass
-
-
 def test_read_multiple_files(tempdir):
     import pyarrow.parquet as pq
 
@@ -1540,7 +1535,7 @@ def test_read_multiple_files(tempdir):
     size = 5
 
     dirpath = tempdir / guid()
-    os.mkdir(dirpath)
+    dirpath.mkdir()
 
     test_data = []
     paths = []
@@ -1559,7 +1554,7 @@ def test_read_multiple_files(tempdir):
         paths.append(path)
 
     # Write a _SUCCESS.crc file
-    _touch(dirpath / '_SUCCESS.crc')
+    (dirpath / '_SUCCESS.crc').touch()
 
     def read_multiple_files(paths, columns=None, nthreads=None, **kwargs):
         dataset = pq.ParquetDataset(paths, **kwargs)
@@ -1621,7 +1616,7 @@ def test_dataset_read_pandas(tempdir):
     size = 5
 
     dirpath = tempdir / guid()
-    os.mkdir(dirpath)
+    dirpath.mkdir()
 
     test_data = []
     frames = []
@@ -1656,7 +1651,7 @@ def test_dataset_read_pandas_common_metadata(tempdir, preserve_index):
     size = 5
 
     dirpath = tempdir / guid()
-    os.mkdir(dirpath)
+    dirpath.mkdir()
 
     test_data = []
     frames = []
@@ -1708,13 +1703,13 @@ def test_ignore_private_directories(tempdir):
     import pyarrow.parquet as pq
 
     dirpath = tempdir / guid()
-    os.mkdir(dirpath)
+    dirpath.mkdir()
 
     paths = _make_example_multifile_dataset(dirpath, nfiles=10,
                                             file_nrows=5)
 
     # private directory
-    os.mkdir(dirpath / '_impala_staging')
+    (dirpath / '_impala_staging').mkdir()
 
     dataset = pq.ParquetDataset(dirpath)
     assert set(map(str, paths)) == set(x.path for x in dataset.pieces)
@@ -1724,15 +1719,15 @@ def test_ignore_hidden_files(tempdir):
     import pyarrow.parquet as pq
 
     dirpath = tempdir / guid()
-    os.mkdir(dirpath)
+    dirpath.mkdir()
 
     paths = _make_example_multifile_dataset(dirpath, nfiles=10,
                                             file_nrows=5)
 
-    with open(dirpath / '.DS_Store', 'wb') as f:
+    with (dirpath / '.DS_Store').open('wb') as f:
         f.write(b'gibberish')
 
-    with open(dirpath / '.private', 'wb') as f:
+    with (dirpath / '.private').open('wb') as f:
         f.write(b'gibberish')
 
     dataset = pq.ParquetDataset(dirpath)
@@ -1781,7 +1776,7 @@ def test_write_error_deletes_incomplete_file(tempdir):
     except pa.ArrowException:
         pass
 
-    assert not os.path.exists(filename)
+    assert not filename.exists()
 
 
 def test_read_non_existent_file(tempdir):