From a26f2995ef3dea58f47acd44b76d59cd0305050c Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Wed, 24 Feb 2021 09:15:20 +0200 Subject: [PATCH] dir_info: only build trie when needed (#5512) --- dvc/checkout.py | 6 +- dvc/dir_info.py | 45 ++++++++++----- dvc/fs/dvc.py | 4 +- dvc/objects/__init__.py | 17 +++--- dvc/objects/stage.py | 2 +- dvc/output/base.py | 7 +-- dvc/path_info.py | 3 + tests/unit/test_dir_info.py | 106 +++++++----------------------------- 8 files changed, 69 insertions(+), 121 deletions(-) diff --git a/dvc/checkout.py b/dvc/checkout.py index 04659d5a6..6a0034833 100644 --- a/dvc/checkout.py +++ b/dvc/checkout.py @@ -177,7 +177,7 @@ def _checkout_file( def _remove_redundant_files(path_info, fs, dir_info, cache, force): existing_files = set(fs.walk_files(path_info)) - needed_files = {info for info, _ in dir_info.items(path_info)} + needed_files = {path_info.joinpath(*key) for key, _ in dir_info.items()} redundant_files = existing_files - needed_files for path in redundant_files: _remove(path, fs, cache, force) @@ -197,9 +197,9 @@ def _checkout_dir( logger.debug("Linking directory '%s'.", path_info) - for entry_info, entry_hash_info in obj.hash_info.dir_info.items(path_info): + for entry_key, entry_hash_info in obj.hash_info.dir_info.items(): entry_modified = _checkout_file( - entry_info, + path_info.joinpath(*entry_key), fs, cache.get(entry_hash_info), cache, diff --git a/dvc/dir_info.py b/dvc/dir_info.py index 55fab135c..beea89a20 100644 --- a/dvc/dir_info.py +++ b/dvc/dir_info.py @@ -1,6 +1,7 @@ import posixpath from operator import itemgetter +from funcy import cached_property from pygtrie import Trie from .hash_info import HashInfo @@ -48,28 +49,36 @@ class DirInfo: PARAM_RELPATH = "relpath" def __init__(self): - self.trie = Trie() + self._dict = {} + + def as_dict(self): + return self._dict.copy() + + def add(self, key, hash_info): + self.__dict__.pop("trie", None) + self._dict[key] = hash_info + + def get(self, key, default=None): + return self._dict.get(key, default) + + @cached_property + def trie(self): + return Trie(self._dict) @property def size(self): try: - return sum( - hash_info.size - for _, hash_info in self.trie.iteritems() # noqa: B301 - ) + return sum(hash_info.size for _, hash_info in self._dict.items()) except TypeError: return None @property def nfiles(self): - return len(self.trie) + return len(self._dict) - def items(self, path_info=None): - for key, hash_info in self.trie.iteritems(): # noqa: B301 - path = posixpath.sep.join(key) - if path_info is not None: - path = path_info / path - yield path, hash_info + def items(self): + for key, hash_info in self._dict.items(): + yield key, hash_info @classmethod def from_list(cls, lst): @@ -78,7 +87,8 @@ class DirInfo: entry = _entry.copy() relpath = entry.pop(cls.PARAM_RELPATH) parts = tuple(relpath.split(posixpath.sep)) - ret.trie[parts] = HashInfo.from_dict(entry) + hash_info = HashInfo.from_dict(entry) + ret.add(parts, hash_info) return ret def to_list(self): @@ -91,12 +101,17 @@ class DirInfo: hash_info.name: hash_info.value, self.PARAM_RELPATH: posixpath.sep.join(parts), } - for parts, hash_info in self.trie.iteritems() # noqa: B301 + for parts, hash_info in self._dict.items() # noqa: B301 ), key=itemgetter(self.PARAM_RELPATH), ) def merge(self, ancestor, their): + merged_dict = _merge( + ancestor.as_dict(), self.as_dict(), their.as_dict(), + ) + merged = DirInfo() - merged.trie = _merge(ancestor.trie, self.trie, their.trie) + for key, hi in merged_dict.items(): + merged.add(key, hi) return merged diff --git a/dvc/fs/dvc.py b/dvc/fs/dvc.py index fd50b7faf..ea07d7b53 100644 --- a/dvc/fs/dvc.py +++ b/dvc/fs/dvc.py @@ -48,7 +48,7 @@ class DvcFileSystem(BaseFileSystem): # pylint:disable=abstract-method # NOTE: use string paths here for performance reasons key = tuple(relpath(path_info, out.path_info).split(os.sep)) out.get_dir_cache(remote=remote) - file_hash = out.hash_info.dir_info.trie.get(key) + file_hash = out.hash_info.dir_info.get(key) if file_hash: return file_hash raise FileNotFoundError @@ -145,7 +145,7 @@ class DvcFileSystem(BaseFileSystem): # pylint:disable=abstract-method self._fetch_dir(out, **kwargs) base = out.path_info.parts - for key in out.dir_cache.trie.iterkeys(): # noqa: B301 + for key, _ in out.dir_cache.items(): # noqa: B301 trie[base + key] = None def _walk(self, root, trie, topdown=True, **kwargs): diff --git a/dvc/objects/__init__.py b/dvc/objects/__init__.py index 02689064b..eb2656b93 100644 --- a/dvc/objects/__init__.py +++ b/dvc/objects/__init__.py @@ -87,11 +87,7 @@ class Tree(HashFile): return self.hash_info.dir_info.nfiles def __iter__(self): - for ( - key, - hash_info, - ) in self.hash_info.dir_info.trie.iteritems(): # noqa: B301 - yield (key, hash_info) + yield from self.hash_info.dir_info.items() @classmethod def save_dir_info(cls, odb, dir_info, hash_info=None): @@ -168,17 +164,18 @@ class Tree(HashFile): hi = self.save_dir_info( odb, self.src_hash_info.dir_info, self.hash_info ) - for entry_info, entry_hash in Tqdm( - hi.dir_info.items(self.src_path_info), + for entry_key, entry_hash in Tqdm( + hi.dir_info.items(), desc="Saving " + self.src_path_info.name, unit="file", ): + entry_info = self.src_path_info.joinpath(*entry_key) entry_obj = HashFile(entry_info, self.src_fs, entry_hash) entry_obj.save(odb, **kwargs) self.src_fs.repo.state.save(self.src_path_info, self.src_fs, hi) def filter(self, odb, prefix): - hash_info = self.hash_info.dir_info.trie.get(prefix) + hash_info = self.hash_info.dir_info.get(prefix) if hash_info: return load(odb, hash_info) @@ -186,7 +183,7 @@ class Tree(HashFile): dir_info = DirInfo() try: for key, value in self.hash_info.dir_info.trie.items(prefix): - dir_info.trie[key[depth:]] = value + dir_info.add(key[depth:], value) except KeyError: return None @@ -297,7 +294,7 @@ def _transfer_directory(odb, from_fs, from_info, jobs, no_progress_bar=False): (entry_tmp_info, entry_hash), ) in _transfer_directory_contents(odb, from_fs, from_info, jobs, pbar): odb.add(entry_tmp_info, odb.fs, entry_hash) - dir_info.trie[entry_info.parts] = entry_hash + dir_info.add(entry_info.parts, entry_hash) return Tree.save_dir_info(odb, dir_info) diff --git a/dvc/objects/stage.py b/dvc/objects/stage.py index b339b6558..c1106f4fd 100644 --- a/dvc/objects/stage.py +++ b/dvc/objects/stage.py @@ -76,7 +76,7 @@ def _collect_dir(path_info, fs, name, state, **kwargs): # # Yes, this is a BUG, as long as we permit "/" in # filenames on Windows and "\" on Unix - dir_info.trie[fi.relative_to(path_info).parts] = hi + dir_info.add(fi.relative_to(path_info).parts, hi) return dir_info diff --git a/dvc/output/base.py b/dvc/output/base.py index c74543639..be250f112 100644 --- a/dvc/output/base.py +++ b/dvc/output/base.py @@ -551,11 +551,8 @@ class BaseOutput: path = str(self.path_info) filter_path = str(filter_info) if filter_info else None - is_win = os.name == "nt" - for entry_relpath, entry_hash_info in self.dir_cache.items(): - if is_win: - entry_relpath = entry_relpath.replace("/", os.sep) - entry_path = os.path.join(path, entry_relpath) + for entry_key, entry_hash_info in self.dir_cache.items(): + entry_path = os.path.join(path, *entry_key) if ( not filter_path or entry_path == filter_path diff --git a/dvc/path_info.py b/dvc/path_info.py index 6f6aa57d8..ff45dd994 100644 --- a/dvc/path_info.py +++ b/dvc/path_info.py @@ -191,6 +191,9 @@ class URLInfo(_BasePath): def __div__(self, other): return self.replace(path=posixpath.join(self._spath, other)) + def joinpath(self, *args): + return self.replace(path=posixpath.join(self._spath, *args)) + __truediv__ = __div__ @property diff --git a/tests/unit/test_dir_info.py b/tests/unit/test_dir_info.py index 2222a5e13..53a0f920c 100644 --- a/tests/unit/test_dir_info.py +++ b/tests/unit/test_dir_info.py @@ -1,11 +1,9 @@ from operator import itemgetter import pytest -from pygtrie import Trie from dvc.dir_info import DirInfo, _merge from dvc.hash_info import HashInfo -from dvc.path_info import PosixPathInfo, WindowsPathInfo @pytest.mark.parametrize( @@ -54,7 +52,7 @@ from dvc.path_info import PosixPathInfo, WindowsPathInfo ) def test_list(lst, trie_dict): dir_info = DirInfo.from_list(lst) - assert dir_info.trie == Trie(trie_dict) + assert dir_info._dict == trie_dict assert dir_info.to_list() == sorted(lst, key=itemgetter("relpath")) @@ -84,7 +82,7 @@ def test_list(lst, trie_dict): ) def test_size(trie_dict, size): dir_info = DirInfo() - dir_info.trie = Trie(trie_dict) + dir_info._dict = trie_dict assert dir_info.size == size @@ -114,92 +112,31 @@ def test_size(trie_dict, size): ) def test_nfiles(trie_dict, nfiles): dir_info = DirInfo() - dir_info.trie = Trie(trie_dict) + dir_info._dict = trie_dict assert dir_info.nfiles == nfiles @pytest.mark.parametrize( - "trie_dict, items", + "trie_dict", [ - ({}, []), - ( - { - ("a",): HashInfo("md5", "abc"), - ("b",): HashInfo("md5", "def"), - ("c",): HashInfo("md5", "ghi"), - ("dir", "foo"): HashInfo("md5", "jkl"), - ("dir", "bar"): HashInfo("md5", "mno"), - ("dir", "baz"): HashInfo("md5", "pqr"), - ("dir", "subdir", "1"): HashInfo("md5", "stu"), - ("dir", "subdir", "2"): HashInfo("md5", "vwx"), - ("dir", "subdir", "3"): HashInfo("md5", "yz"), - }, - [ - ("a", HashInfo("md5", "abc")), - ("b", HashInfo("md5", "def")), - ("c", HashInfo("md5", "ghi")), - ("dir/foo", HashInfo("md5", "jkl")), - ("dir/bar", HashInfo("md5", "mno")), - ("dir/baz", HashInfo("md5", "pqr")), - ("dir/subdir/1", HashInfo("md5", "stu")), - ("dir/subdir/2", HashInfo("md5", "vwx")), - ("dir/subdir/3", HashInfo("md5", "yz")), - ], - ), + {}, + { + ("a",): HashInfo("md5", "abc"), + ("b",): HashInfo("md5", "def"), + ("c",): HashInfo("md5", "ghi"), + ("dir", "foo"): HashInfo("md5", "jkl"), + ("dir", "bar"): HashInfo("md5", "mno"), + ("dir", "baz"): HashInfo("md5", "pqr"), + ("dir", "subdir", "1"): HashInfo("md5", "stu"), + ("dir", "subdir", "2"): HashInfo("md5", "vwx"), + ("dir", "subdir", "3"): HashInfo("md5", "yz"), + }, ], ) -def test_items(trie_dict, items): +def test_items(trie_dict): dir_info = DirInfo() - dir_info.trie = Trie(trie_dict) - assert list(dir_info.items()) == items - - -@pytest.mark.parametrize( - "path_info, trie_dict, items", - [ - (PosixPathInfo(), {}, []), - (WindowsPathInfo(), {}, []), - ( - PosixPathInfo("/some/path"), - { - ("a",): HashInfo("md5", "abc"), - ("dir", "foo"): HashInfo("md5", "jkl"), - ("dir", "sub", "1"): HashInfo("md5", "stu"), - }, - [ - (PosixPathInfo("/some/path/a"), HashInfo("md5", "abc")), - (PosixPathInfo("/some/path/dir/foo"), HashInfo("md5", "jkl")), - ( - PosixPathInfo("/some/path/dir/sub/1"), - HashInfo("md5", "stu"), - ), - ], - ), - ( - WindowsPathInfo("C:\\some\\path"), - { - ("a",): HashInfo("md5", "abc"), - ("dir", "foo"): HashInfo("md5", "jkl"), - ("dir", "sub", "1"): HashInfo("md5", "stu"), - }, - [ - (WindowsPathInfo("C:\\some\\path\\a"), HashInfo("md5", "abc")), - ( - WindowsPathInfo("C:\\some\\path\\dir\\foo"), - HashInfo("md5", "jkl"), - ), - ( - WindowsPathInfo("C:\\some\\path\\dir\\sub\\1"), - HashInfo("md5", "stu"), - ), - ], - ), - ], -) -def test_items_with_path(path_info, trie_dict, items): - dir_info = DirInfo() - dir_info.trie = Trie(trie_dict) - assert list(dir_info.items(path_info)) == items + dir_info._dict = trie_dict + assert list(dir_info.items()) == list(trie_dict.items()) @pytest.mark.parametrize( @@ -268,6 +205,5 @@ def test_items_with_path(path_info, trie_dict, items): ], ) def test_merge(ancestor_dict, our_dict, their_dict, merged_dict): - actual = _merge(Trie(ancestor_dict), Trie(our_dict), Trie(their_dict)) - expected = Trie(merged_dict) - assert actual == expected + actual = _merge(ancestor_dict, our_dict, their_dict) + assert actual == merged_dict