dir_info: only build trie when needed (#5512)

This commit is contained in:
Ruslan Kuprieiev
2021-02-24 09:15:20 +02:00
committed by GitHub
parent 502f0eecc9
commit a26f2995ef
8 changed files with 69 additions and 121 deletions
+3 -3
View File
@@ -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,
+30 -15
View File
@@ -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
+2 -2
View File
@@ -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):
+7 -10
View File
@@ -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)
+1 -1
View File
@@ -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
+2 -5
View File
@@ -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
+3
View File
@@ -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
+21 -85
View File
@@ -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