hdfs: Replace deprecated pyarrow.hdfs.connect (#4973)

* hdfs: Replace deprecated pyarrow.hdfs.connect

Replace with pyarrow.fs.HadoopFileSystem, which has a different API.

Fixes #4905

* hdfs: fix copy/walk

* hdfs: Remove setting CLASSPATH

* Revert "hdfs: Remove setting CLASSPATH"

This reverts commit 8925a5d1d16ddbfc5e80c9b7974d96abdd4620a5.

Co-authored-by: Ruslan Kuprieiev <[email protected]>

dvc/tree/hdfs.py

@@ -46,10 +46,17 @@ class HDFSTree(BaseTree):
 
     @staticmethod
     def hdfs(path_info):
-        import pyarrow
+        import pyarrow.fs
+
+        # NOTE: HadoopFileSystem is not meant to be closed by us and doesn't
+        # have close() method or alternative, so we add a noop one to satisfy
+        # our connection pool.
+        class HDFSConnection(pyarrow.fs.HadoopFileSystem):
+            def close(self):
+                pass
 
         return get_connection(
-            pyarrow.hdfs.connect,
+            HDFSConnection,
             path_info.host,
             path_info.port,
             user=path_info.user,
@@ -61,7 +68,7 @@ class HDFSTree(BaseTree):
 
         try:
             with self.hdfs(path_info) as hdfs, closing(
-                hdfs.open(path_info.path, mode="rb")
+                hdfs.open_input_stream(path_info.path)
             ) as fd:
                 if mode == "rb":
                     yield fd
@@ -78,31 +85,51 @@ class HDFSTree(BaseTree):
         assert not isinstance(path_info, list)
         assert path_info.scheme == "hdfs"
         with self.hdfs(path_info) as hdfs:
-            return hdfs.exists(path_info.path)
+            import pyarrow.fs
 
-    def walk_files(self, path_info, **kwargs):
-        if not self.exists(path_info):
+            file_info = hdfs.get_file_info(path_info.path)
+            return file_info.type != pyarrow.fs.FileType.NotFound
+
+    def _walk(self, hdfs, root, topdown=True):
+        import posixpath
+
+        from pyarrow.fs import FileSelector, FileType
+
+        dirs = deque()
+        nondirs = deque()
+
+        selector = FileSelector(root)
+        for entry in hdfs.get_file_info(selector):
+            if entry.type == FileType.Directory:
+                dirs.append(entry.base_name)
+            else:
+                nondirs.append(entry.base_name)
+
+        if topdown:
+            yield root, dirs, nondirs
+
+        for dname in dirs:
+            # NOTE: posixpath.join() is slower
+            yield from self._walk(
+                hdfs, f"{root}{posixpath.sep}{dname}", topdown=topdown
+            )
+
+        if not topdown:
+            yield root, dirs, nondirs
+
+    def walk(self, path_info, **kwargs):
+        if not self.isdir(path_info):
             return
 
-        root = path_info.path
-        dirs = deque([root])
+        with self.hdfs(path_info) as hdfs:
+            for root, dnames, fnames in self._walk(
+                hdfs, path_info.path, **kwargs
+            ):
+                yield path_info.replace(path=root), dnames, fnames
 
-        with self.hdfs(self.path_info) as hdfs:
-            if not hdfs.exists(root):
-                return
-            while dirs:
-                try:
-                    entries = hdfs.ls(dirs.pop(), detail=True)
-                    for entry in entries:
-                        if entry["kind"] == "directory":
-                            dirs.append(urlparse(entry["name"]).path)
-                        elif entry["kind"] == "file":
-                            path = urlparse(entry["name"]).path
-                            yield path_info.replace(path=path)
-                except OSError:
-                    # When searching for a specific prefix pyarrow raises an
-                    # exception if the specified cache dir does not exist
-                    pass
+    def walk_files(self, path_info, **kwargs):
+        for root, _, fnames in self.walk(path_info):
+            yield from (root / fname for fname in fnames)
 
     def remove(self, path_info):
         if path_info.scheme != "hdfs":
@@ -111,34 +138,49 @@ class HDFSTree(BaseTree):
         if self.exists(path_info):
             logger.debug(f"Removing {path_info.path}")
             with self.hdfs(path_info) as hdfs:
-                hdfs.rm(path_info.path)
+                import pyarrow.fs
+
+                file_info = hdfs.get_file_info(path_info.path)
+                if file_info.type == pyarrow.fs.FileType.Directory:
+                    hdfs.delete_dir(path_info.path)
+                else:
+                    hdfs.delete_file(path_info.path)
 
     def makedirs(self, path_info):
         with self.hdfs(path_info) as hdfs:
-            # NOTE: hdfs.mkdir creates parents by default
-            hdfs.mkdir(path_info.path)
+            # NOTE: fs.create_dir creates parents by default
+            hdfs.create_dir(path_info.path)
 
     def copy(self, from_info, to_info, **_kwargs):
+        # NOTE: hdfs.copy_file is not supported yet in pyarrow
         with self.hdfs(to_info) as hdfs:
             # NOTE: this is how `hadoop fs -cp` works too: it copies through
             # your local machine.
-            with hdfs.open(from_info.path, "rb") as from_fobj:
+            with closing(hdfs.open_input_stream(from_info.path)) as from_fobj:
                 tmp_info = to_info.parent / tmp_fname(to_info.name)
                 try:
-                    with hdfs.open(tmp_info.path, "wb") as tmp_fobj:
-                        tmp_fobj.upload(from_fobj)
-                    hdfs.rename(tmp_info.path, to_info.path)
+                    with closing(
+                        hdfs.open_output_stream(tmp_info.path)
+                    ) as tmp_fobj:
+                        tmp_fobj.write(from_fobj.read())
+                    hdfs.move(tmp_info.path, to_info.path)
                 except Exception:
                     self.remove(tmp_info)
                     raise
 
     def isfile(self, path_info):
         with self.hdfs(path_info) as hdfs:
-            return hdfs.isfile(path_info.path)
+            import pyarrow.fs
+
+            file_info = hdfs.get_file_info(path_info.path)
+            return file_info.type == pyarrow.fs.FileType.File
 
     def isdir(self, path_info):
         with self.hdfs(path_info) as hdfs:
-            return hdfs.isdir(path_info.path)
+            import pyarrow.fs
+
+            file_info = hdfs.get_file_info(path_info.path)
+            return file_info.type == pyarrow.fs.FileType.Directory
 
     def hadoop_fs(self, cmd, user=None):
         cmd = "hadoop fs -" + cmd
@@ -182,7 +224,8 @@ class HDFSTree(BaseTree):
         )
 
         with self.hdfs(path_info) as hdfs:
-            hash_info.size = hdfs.info(path_info.path)["size"]
+            file_info = hdfs.get_file_info(path_info.path)
+            hash_info.size = file_info.size
 
         return hash_info
 
@@ -200,14 +243,16 @@ class HDFSTree(BaseTree):
                     total=total,
                     disable=no_progress_bar,
                 ) as wrapped:
-                    hdfs.upload(tmp_file, wrapped)
-            hdfs.rename(tmp_file, to_info.path)
+                    with hdfs.open_output_stream(tmp_file) as sobj:
+                        sobj.write(wrapped.read())
+            hdfs.move(tmp_file, to_info.path)
 
     def _download(
         self, from_info, to_file, name=None, no_progress_bar=False, **_kwargs
     ):
         with self.hdfs(from_info) as hdfs:
-            total = hdfs.info(from_info.path)["size"]
+            file_info = hdfs.get_file_info(from_info.path)
+            total = file_info.size
             with open(to_file, "wb+") as fobj:
                 with Tqdm.wrapattr(
                     fobj,
@@ -216,4 +261,5 @@ class HDFSTree(BaseTree):
                     total=total,
                     disable=no_progress_bar,
                 ) as wrapped:
-                    hdfs.download(from_info.path, wrapped)
+                    with hdfs.open_input_stream(from_info.path) as sobj:
+                        wrapped.write(sobj.read())

tests/remotes/hdfs.py

@@ -1,6 +1,7 @@
 import locale
 import os
 import platform
+import subprocess
 import uuid
 from contextlib import contextmanager
 
@@ -15,25 +16,31 @@ from .base import Base
 class HDFS(Base, URLInfo):  # pylint: disable=abstract-method
     @contextmanager
     def _hdfs(self):
-        import pyarrow
+        import pyarrow.fs
 
-        conn = pyarrow.hdfs.connect(self.host, self.port)
-        try:
-            yield conn
-        finally:
-            conn.close()
+        conn = pyarrow.fs.HadoopFileSystem(self.host, self.port)
+        yield conn
 
     def is_file(self):
         with self._hdfs() as _hdfs:
-            return _hdfs.isfile(self.path)
+            import pyarrow.fs
+
+            file_info = _hdfs.get_file_info(self.path)
+            return file_info.type == pyarrow.fs.FileType.File
 
     def is_dir(self):
         with self._hdfs() as _hdfs:
-            return _hdfs.isfile(self.path)
+            import pyarrow.fs
+
+            file_info = _hdfs.get_file_info(self.path)
+            return file_info.type == pyarrow.fs.FileType.Directory
 
     def exists(self):
         with self._hdfs() as _hdfs:
-            return _hdfs.exists(self.path)
+            import pyarrow.fs
+
+            file_info = _hdfs.get_file_info(self.path)
+            return file_info.type != pyarrow.fs.FileType.NotFound
 
     def mkdir(self, mode=0o777, parents=False, exist_ok=False):
         assert mode == 0o777
@@ -41,13 +48,12 @@ class HDFS(Base, URLInfo):  # pylint: disable=abstract-method
         assert not exist_ok
 
         with self._hdfs() as _hdfs:
-            # NOTE: hdfs.mkdir always creates parents
-            _hdfs.mkdir(self.path)
+            # NOTE: fs.create_dir creates parents by default
+            _hdfs.create_dir(self.path)
 
     def write_bytes(self, contents):
         with self._hdfs() as _hdfs:
-            # NOTE: hdfs.open only supports 'rb', 'wb' or 'ab'
-            with _hdfs.open(self.path, "wb") as fobj:
+            with _hdfs.open_output_stream(self.path) as fobj:
                 fobj.write(contents)
 
     def write_text(self, contents, encoding=None, errors=None):
@@ -58,8 +64,7 @@ class HDFS(Base, URLInfo):  # pylint: disable=abstract-method
 
     def read_bytes(self):
         with self._hdfs() as _hdfs:
-            # NOTE: hdfs.open only supports 'rb', 'wb' or 'ab'
-            with _hdfs.open(self.path, "rb") as fobj:
+            with _hdfs.open_input_stream(self.path) as fobj:
                 return fobj.read()
 
     def read_text(self, encoding=None, errors=None):
@@ -115,10 +120,19 @@ def hadoop(test_config):
     os.environ["HADOOP_HOME"] = hadoop_home
     os.environ["PATH"] += f":{hadoop_home}/bin:{hadoop_home}/sbin"
 
+    # NOTE: must set CLASSPATH to connect using pyarrow.fs.HadoopFileSystem
+    result = subprocess.run(
+        [f"{hadoop_home}/bin/hdfs", "classpath", "--glob"],
+        universal_newlines=True,
+        stdout=subprocess.PIPE,
+        check=False,
+    )
+    os.environ["CLASSPATH"] = result.stdout
+
 
 @pytest.fixture(scope="session")
 def hdfs_server(hadoop, docker_compose, docker_services):
-    import pyarrow
+    import pyarrow.fs
 
     port = docker_services.port_for("hdfs", 8020)
     web_port = docker_services.port_for("hdfs", 50070)
@@ -127,12 +141,9 @@ def hdfs_server(hadoop, docker_compose, docker_services):
         try:
             # NOTE: just connecting or even opening something is not enough,
             # we need to make sure that we are able to write something.
-            conn = pyarrow.hdfs.connect("127.0.0.1", port)
-            try:
-                with conn.open(str(uuid.uuid4()), "wb") as fobj:
-                    fobj.write(b"test")
-            finally:
-                conn.close()
+            conn = pyarrow.fs.HadoopFileSystem("hdfs://127.0.0.1", port)
+            with conn.open_output_stream(str(uuid.uuid4())) as fobj:
+                fobj.write(b"test")
             return True
         except (pyarrow.ArrowException, OSError):
             return False