Repository: iterative/dvc

Issue:
hdfs: migrate to fsspec/pyarrow.fs
To complete the fsspec migration, we need to migrate hdfs to the fsspec. 


Discussion:
It seems like `PyArrowHDFS` doesn't support per-file `info()`, neither it does support the `checksum()` logic we have. 
@isidentical Also it seems that fsspec's is based on CLI hadoop, while we are based on libhdfs. But not 100% sure.

Regarding `checksum`, we definitely could consider dropping it, as it is mostly being misused right now for external outputs scenario. Or could keep it as it is right now - based on raw hadoop CLI for now.
> it seems that fsspec's is based on CLI hadoop, while we are based on libhdfs. But not 100% sure.

I am not sure what is the difference, but from what I can assume they seem to be also using libhdfs:
https://github.com/intake/filesystem_spec/blob/d952ace9db6eea7ccf202e616dc8121a549a4dc9/fsspec/implementations/hdfs.py#L40-L41
prerequisite: https://github.com/intake/filesystem_spec/pull/749
From what I see now is that we can not directly use the one on the fsspec (even with my changed applied) since `pyarrow.hdfs` is deprecated and using `pyarrow.fs.HadoopFileSystem` would require substantial changes on the fsspec side :( There seems to be an open issue about supporting it on the fsspec, and I probably will take it on (there is some starter code over it) https://github.com/intake/filesystem_spec/issues/663
@isidentical Looks like you've posted a link to this same issue in your last comment. You probably meant something else.
Right, i fixed it in the original comment @efiop.
With https://github.com/intake/filesystem_spec/pull/754, now it also works on our test suite and can actually support the mocked one.


Resolving patch:
diff --git a/dvc/fs/hdfs.py b/dvc/fs/hdfs.py
--- a/dvc/fs/hdfs.py
+++ b/dvc/fs/hdfs.py
@@ -1,282 +1,95 @@
-import errno
-import io
-import logging
 import os
 import re
-import shutil
 import subprocess
-from collections import deque
-from contextlib import closing, contextmanager
+import threading
 
-from tqdm.utils import CallbackIOWrapper
+from funcy import cached_property, wrap_prop
 
 from dvc.hash_info import HashInfo
-from dvc.progress import DEFAULT_CALLBACK
 from dvc.scheme import Schemes
-from dvc.utils import fix_env, tmp_fname
+from dvc.utils import fix_env
 
-from .base import BaseFileSystem, RemoteCmdError
-from .pool import get_connection
+from .fsspec_wrapper import CallbackMixin, FSSpecWrapper
 
-logger = logging.getLogger(__name__)
+CHECKSUM_REGEX = re.compile(r".*\t.*\t(?P<checksum>.*)")
 
 
-def _hadoop_fs(cmd, user=None):
-    cmd = "hadoop fs -" + cmd
-    if user:
-        cmd = f"HADOOP_USER_NAME={user} " + cmd
-
-    # NOTE: close_fds doesn't work with redirected stdin/stdout/stderr.
-    # See https://github.com/iterative/dvc/issues/1197.
-    close_fds = os.name != "nt"
-
-    executable = os.getenv("SHELL") if os.name != "nt" else None
-    p = subprocess.Popen(
-        cmd,
-        shell=True,
-        close_fds=close_fds,
-        executable=executable,
-        env=fix_env(os.environ),
-        stdin=subprocess.PIPE,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
-    )
-    out, err = p.communicate()
-    if p.returncode != 0:
-        raise RemoteCmdError("hdfs", cmd, p.returncode, err)
-    return out.decode("utf-8")
-
-
-def _group(regex, s, gname):
-    match = re.match(regex, s)
-    assert match is not None
-    return match.group(gname)
-
-
-def _hadoop_fs_checksum(path_info):
-    # NOTE: pyarrow doesn't support checksum, so we need to use hadoop
-    regex = r".*\t.*\t(?P<checksum>.*)"
-    stdout = _hadoop_fs(f"checksum {path_info.url}", user=path_info.user)
-    return _group(regex, stdout, "checksum")
-
-
-class HDFSFileSystem(BaseFileSystem):
+# pylint: disable=abstract-method
+class HDFSFileSystem(CallbackMixin, FSSpecWrapper):
     scheme = Schemes.HDFS
-    REQUIRES = {"pyarrow": "pyarrow"}
-    REGEX = r"^hdfs://((?P<user>.*)@)?.*$"
+    REQUIRES = {"fsspec": "fsspec", "pyarrow": "pyarrow"}
     PARAM_CHECKSUM = "checksum"
-    TRAVERSE_PREFIX_LEN = 2
-
-    BLOCK_SIZE = 2 ** 26  # 128MB
 
-    def __init__(self, **config):
-        super().__init__(**config)
-
-        self.host = config["host"]
-        self.user = config.get("user")
-        self.port = config.get("port")
-
-        self.kerb_ticket = config.get("kerb_ticket")
+    def _with_bucket(self, path):
+        if isinstance(path, self.PATH_CLS):
+            return path.path
+        return super()._with_bucket(path)
 
     @staticmethod
     def _get_kwargs_from_urls(urlpath):
-        from fsspec.implementations.hdfs import PyArrowHDFS
+        from fsspec.implementations.arrow import HadoopFileSystem
 
         # pylint:disable=protected-access
-        return PyArrowHDFS._get_kwargs_from_urls(urlpath)
-
-    def hdfs(self, path_info):
-        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(
-            HDFSConnection,
-            path_info.host,
-            path_info.port,
-            user=path_info.user,
-            kerb_ticket=self.kerb_ticket,
-        )
-
-    @contextmanager
-    def open(self, path_info, mode="r", encoding=None, **kwargs):
-        assert mode in {"r", "rt", "rb"}
-
-        try:
-            with self.hdfs(path_info) as hdfs, closing(
-                hdfs.open_input_stream(path_info.path)
-            ) as fd:
-                if mode == "rb":
-                    yield fd
-                else:
-                    yield io.TextIOWrapper(fd, encoding=encoding)
-        except OSError as e:
-            # Empty .errno and not specific enough error class in pyarrow,
-            # see https://issues.apache.org/jira/browse/ARROW-6248
-            if "file does not exist" in str(e):
-                raise FileNotFoundError(*e.args)
-            raise
-
-    def exists(self, path_info) -> bool:
-        assert not isinstance(path_info, list)
-        assert path_info.scheme == "hdfs"
-        with self.hdfs(path_info) as hdfs:
-            import pyarrow.fs
-
-            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()
+        return HadoopFileSystem._get_kwargs_from_urls(urlpath)
 
-        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)
+    def _prepare_credentials(self, **config):
+        return config
 
-        if topdown:
-            yield root, dirs, nondirs
+    @wrap_prop(threading.Lock())
+    @cached_property
+    def fs(self):
+        from fsspec.implementations.arrow import HadoopFileSystem
 
-        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, top, topdown=True, onerror=None, **kwargs):
-        if not self.isdir(top):
-            return
-
-        with self.hdfs(top) as hdfs:
-            for root, dnames, fnames in self._walk(
-                hdfs, top.path, topdown=topdown
-            ):
-                yield top.replace(path=root), dnames, fnames
-
-    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":
-            raise NotImplementedError
-
-        if self.exists(path_info):
-            logger.debug(f"Removing {path_info.path}")
-            with self.hdfs(path_info) as hdfs:
-                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, **kwargs):
-        with self.hdfs(path_info) as hdfs:
-            # 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 closing(hdfs.open_input_stream(from_info.path)) as from_fobj:
-                tmp_info = to_info.parent / tmp_fname(to_info.name)
-                try:
-                    with closing(
-                        hdfs.open_output_stream(tmp_info.path)
-                    ) as tmp_fobj:
-                        shutil.copyfileobj(from_fobj, tmp_fobj)
-                    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:
-            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:
-            import pyarrow.fs
-
-            file_info = hdfs.get_file_info(path_info.path)
-            return file_info.type == pyarrow.fs.FileType.Directory
-
-    def info(self, path_info):
-        with self.hdfs(path_info) as hdfs:
-            from pyarrow.fs import FileType
-
-            file_info = hdfs.get_file_info(path_info.path)
-            if file_info.type is FileType.Directory:
-                kind = "directory"
-            elif file_info.type is FileType.File:
-                kind = "file"
-            else:
-                raise FileNotFoundError(
-                    errno.ENOENT, os.strerror(errno.ENOENT), path_info
-                )
-
-            return {"size": file_info.size, "type": kind}
+        return HadoopFileSystem(**self.fs_args)
 
     def checksum(self, path_info):
         return HashInfo(
             "checksum",
-            _hadoop_fs_checksum(path_info),
+            self._checksum(path_info),
             size=self.getsize(path_info),
         )
 
-    def upload_fobj(self, fobj, to_info, **kwargs):
-        with self.hdfs(to_info) as hdfs:
-            with hdfs.open_output_stream(to_info.path) as fdest:
-                shutil.copyfileobj(fobj, fdest, self.BLOCK_SIZE)
-
-    def put_file(
-        self, from_file, to_info, callback=DEFAULT_CALLBACK, **kwargs
-    ):
-        with self.hdfs(to_info) as hdfs:
-            hdfs.create_dir(to_info.parent.path)
+    def _checksum(self, path_info, **kwargs):
+        # PyArrow doesn't natively support retrieving the
+        # checksum, so we have to use hadoop fs
 
-            tmp_file = tmp_fname(to_info.path)
-            total: int = os.path.getsize(from_file)
-            callback.set_size(total)
-
-            with open(from_file, "rb") as fobj:
-                wrapped = CallbackIOWrapper(callback.relative_update, fobj)
-                with hdfs.open_output_stream(tmp_file) as sobj:
-                    shutil.copyfileobj(wrapped, sobj, self.BLOCK_SIZE)
-            hdfs.move(tmp_file, to_info.path)
-
-    def get_file(
-        self, from_info, to_file, callback=DEFAULT_CALLBACK, **kwargs
-    ):
-        with self.hdfs(from_info) as hdfs:
-            file_info = hdfs.get_file_info(from_info.path)
-            total = file_info.size
-            if total:
-                callback.set_size(total)
+        result = self._run_command(
+            f"checksum {path_info.url}",
+            env=fix_env(os.environ),
+            user=path_info.user,
+        )
+        if result is None:
+            return None
+
+        match = CHECKSUM_REGEX.match(result)
+        if match is None:
+            return None
+
+        return match.group("checksum")
+
+    def _run_command(self, cmd, env=None, user=None):
+        cmd = "hadoop fs -" + cmd
+        if user:
+            cmd = f"HADOOP_USER_NAME={user} " + cmd
+
+        # NOTE: close_fds doesn't work with redirected stdin/stdout/stderr.
+        # See https://github.com/iterative/dvc/issues/1197.
+        close_fds = os.name != "nt"
+
+        executable = os.getenv("SHELL") if os.name != "nt" else None
+        p = subprocess.Popen(
+            cmd,
+            shell=True,
+            close_fds=close_fds,
+            executable=executable,
+            env=env or os.environ,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+        out, err = p.communicate()
 
-            with hdfs.open_input_stream(from_info.path) as sobj:
-                with open(to_file, "wb+") as fobj:
-                    wrapped = CallbackIOWrapper(callback.relative_update, sobj)
-                    shutil.copyfileobj(wrapped, fobj, self.BLOCK_SIZE)
+        if p.returncode != 0:
+            raise subprocess.CalledProcessError(p.returncode, cmd, out, err)
+        else:
+            return out.decode("utf-8")


Test patch:
diff --git a/tests/remotes/hdfs.py b/tests/remotes/hdfs.py
--- a/tests/remotes/hdfs.py
+++ b/tests/remotes/hdfs.py
@@ -195,8 +195,7 @@ def md5md5crc32c(path):
     return "000002000000000000000000" + md5md5.hexdigest()
 
 
-def hadoop_fs_checksum(path_info):
-
+def hadoop_fs_checksum(_, path_info):
     return md5md5crc32c(Path(_hdfs_root.name) / path_info.path.lstrip("/"))
 
 
@@ -216,25 +215,48 @@ def _path(self, path):
                 path.allow_not_found,
                 path.recursive,
             )
+        if isinstance(path, list):
+            return [self._path(sub_path) for sub_path in path]
 
         return os.fspath(self._root / path.lstrip("/"))
 
-    def create_dir(self, path):
-        return self._fs.create_dir(self._path(path))
+    def create_dir(self, path, **kwargs):
+        return self._fs.create_dir(self._path(path), **kwargs)
 
-    def open_input_stream(self, path):
-        return self._fs.open_input_stream(self._path(path))
+    def open_input_stream(self, path, **kwargs):
+        return self._fs.open_input_stream(self._path(path), **kwargs)
 
-    def open_output_stream(self, path):
+    def open_output_stream(self, path, **kwargs):
         import posixpath
 
         # NOTE: HadoopFileSystem.open_output_stream creates directories
         # automatically.
         self.create_dir(posixpath.dirname(path))
-        return self._fs.open_output_stream(self._path(path))
+        return self._fs.open_output_stream(self._path(path), **kwargs)
+
+    def get_file_info(self, path, **kwargs):
+        from pyarrow.fs import FileInfo
+
+        entries = self._fs.get_file_info(self._path(path), **kwargs)
+        if isinstance(entries, FileInfo):
+            return self._adjust_entry(entries)
+
+        assert isinstance(entries, list)
+        return list(map(self._adjust_entry, entries))
+
+    def _adjust_entry(self, entry):
+        import posixpath
+
+        from pyarrow.fs import FileInfo
 
-    def get_file_info(self, path):
-        return self._fs.get_file_info(self._path(path))
+        mocked_path = os.path.relpath(entry.path, self._root)
+        mocked_parts = mocked_path.split(os.path.sep)
+        return FileInfo(
+            path=posixpath.join(*mocked_parts),
+            type=entry.type,
+            mtime=entry.mtime,
+            size=entry.size,
+        )
 
     def move(self, from_path, to_path):
         self._fs.move(self._path(from_path), self._path(to_path))
@@ -259,7 +281,7 @@ def hdfs(test_config, mocker):
         "pyarrow.fs.HadoopFileSystem", FakeHadoopFileSystem, create=True
     )
 
-    mocker.patch("dvc.fs.hdfs._hadoop_fs_checksum", hadoop_fs_checksum)
+    mocker.patch("dvc.fs.hdfs.HDFSFileSystem._checksum", hadoop_fs_checksum)
 
     url = f"hdfs://example.com:12345/{uuid.uuid4()}"
     yield HDFS(url)
diff --git a/tests/unit/fs/test_hdfs.py b/tests/unit/fs/test_hdfs.py
deleted file mode 100644
--- a/tests/unit/fs/test_hdfs.py
+++ /dev/null
@@ -1,31 +0,0 @@
-import os
-
-from dvc.fs.hdfs import _hadoop_fs_checksum
-from dvc.path_info import URLInfo
-from dvc.utils import fix_env
-
-
-def test_hadoop_fs_checksum(mocker):
-    mock_proc = mocker.Mock()
-
-    out = b"/path/to/file\tMD5-of-0MD5-of-512CRC32C\t123456789"
-    err = b""
-    mock_proc.configure_mock(
-        **{"communicate.return_value": (out, err), "returncode": 0}
-    )
-    mock_popen = mocker.patch("subprocess.Popen", return_value=mock_proc)
-
-    path_info = URLInfo("hdfs://example.com:1234/path/to/file")
-
-    assert _hadoop_fs_checksum(path_info) == "123456789"
-    mock_popen.assert_called_once_with(
-        "hadoop fs -checksum hdfs://example.com:1234/path/to/file",
-        shell=True,
-        close_fds=os.name != "nt",
-        executable=os.getenv("SHELL") if os.name != "nt" else None,
-        env=fix_env(os.environ),
-        stdin=-1,
-        stdout=-1,
-        stderr=-1,
-    )
-    assert mock_proc.communicate.called