From f50f12ad6148b7e0777c315d2e2841353adfac5a Mon Sep 17 00:00:00 2001 From: jwokaty Date: Wed, 2 Dec 2020 14:46:35 -0500 Subject: [PATCH] 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 --- dvc/tree/hdfs.py | 122 +++++++++++++++++++++++++++++------------- tests/remotes/hdfs.py | 55 +++++++++++-------- 2 files changed, 117 insertions(+), 60 deletions(-) diff --git a/dvc/tree/hdfs.py b/dvc/tree/hdfs.py index 4ea629e9e..f564f4b9b 100644 --- a/dvc/tree/hdfs.py +++ b/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()) diff --git a/tests/remotes/hdfs.py b/tests/remotes/hdfs.py index 953ba0c68..35963a205 100644 --- a/tests/remotes/hdfs.py +++ b/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