Repository: iterative/dvc
Pull request: checkout: fix partial file checkouts in the same directory
Resolves #5158

Problem statement:
pull/checkout: pulling multiple files consecutively from a folder checks out only first
# Bug Report

## Description

When attempting to pull individual files from dvc tracked folder, only the first pull creates file links. Similar behavior with `checkout` unless `--relink` is used.

### Reproduce
0. (git + dvc init)

1. mkdir db
2. touch db/a.txt && touch db/b.txt
3. dvc add db
4. git add db.dvc && git commit -m "dvc"
5. dvc remote add -d tmp /tmp/remote
6. dvc push
7. rm db/*
8. dvc pull db/a.txt
9. dvc pull db/b.txt

### Expected

db/a.txt & db/b.txt exist (When in fact b.txt doesn't). 

Getting messages:
dvc pull db/a.txt --> 1 file modified
dvc pull db/b.txt --> Everything is up to date. 

### Environment information

DVC version: 1.11.9 (pip)
---------------------------------
Platform: Python 3.6.10 on Linux-4.18.0-1024-azure-x86_64-with-debian-buster-sid
Supports: http, https, s3
Cache types: hardlink, symlink
Cache directory: ext4 on /dev/sda1
Caches: local
Remotes: local
Workspace directory: ext4 on /dev/sda1
Repo: dvc, git

Hints from issue discussion:
I'd like to propose a patch for this if no one else already prepared something (cc: @efiop) [have a working demo, no tests yet though]

Commit message: Update dvc/cache/base.py

Co-authored-by: Ruslan Kuprieiev <kupruser@gmail.com>

Patch to review:
diff --git a/dvc/cache/base.py b/dvc/cache/base.py
index 1f8ce26b91..eae344cd84 100644
--- a/dvc/cache/base.py
+++ b/dvc/cache/base.py
@@ -89,7 +89,26 @@ def load_dir_cache(self, hash_info):
 
         return DirInfo.from_list(d)
 
-    def changed(self, path_info, hash_info):
+    def _get_filtered_hash_info(self, filter_info, hash_info, path_info):
+        dir_info = self.get_dir_cache(hash_info)
+        hash_key = filter_info.relative_to(path_info).parts
+
+        # Check whether it is a file that exists on the trie
+        if hash_key in dir_info.trie:
+            return dir_info.trie[hash_key]
+
+        depth = len(hash_key)
+        filtered_dir_info = DirInfo()
+        try:
+            for key, value in dir_info.trie.items(hash_key):
+                filtered_dir_info.trie[key[depth:]] = value
+        except KeyError:
+            return None
+
+        filtered_hash_info, _ = self._get_dir_info_hash(filtered_dir_info)
+        return filtered_hash_info
+
+    def changed(self, path_info, hash_info, filter_info=None):
         """Checks if data has changed.
 
         A file is considered changed if:
@@ -101,23 +120,33 @@ def changed(self, path_info, hash_info):
         Args:
             path_info: dict with path information.
             hash: expected hash value for this data.
+            filter_info: an optional argument to target a specific path.
 
         Returns:
             bool: True if data has changed, False otherwise.
         """
 
+        path = filter_info or path_info
         logger.trace(
             "checking if '%s'('%s') has changed.", path_info, hash_info
         )
 
-        if not self.tree.exists(path_info):
-            logger.debug("'%s' doesn't exist.", path_info)
+        if not self.tree.exists(path):
+            logger.debug("'%s' doesn't exist.", path)
             return True
 
         if not hash_info:
             logger.debug("hash value for '%s' is missing.", path_info)
             return True
 
+        if path != path_info:
+            filtered_hash_info = self._get_filtered_hash_info(
+                path, hash_info, path_info
+            )
+            if filtered_hash_info is not None:
+                path_info = filter_info
+                hash_info = filtered_hash_info
+
         if self.changed_cache(hash_info):
             logger.debug(
                 "cache for '%s'('%s') has changed.", path_info, hash_info
@@ -134,7 +163,7 @@ def changed(self, path_info, hash_info):
             )
             return True
 
-        logger.trace("'%s' hasn't changed.", path_info)
+        logger.trace("'%s' hasn't changed.", path)
         return False
 
     def link(self, from_info, to_info):
@@ -520,7 +549,9 @@ def checkout(
             self.safe_remove(path_info, force=force)
             failed = path_info
 
-        elif not relink and not self.changed(path_info, hash_info):
+        elif not relink and not self.changed(
+            path_info, hash_info, filter_info=filter_info
+        ):
             logger.trace("Data '%s' didn't change.", path_info)
             skip = True
 
diff --git a/tests/func/test_checkout.py b/tests/func/test_checkout.py
index 712b26df60..5cd3a10701 100644
--- a/tests/func/test_checkout.py
+++ b/tests/func/test_checkout.py
@@ -880,3 +880,101 @@ def test_checkout_executable(tmp_dir, dvc):
         assert not isexec
     else:
         assert isexec
+
+
+def test_checkout_partial(tmp_dir, dvc):
+    tmp_dir.dvc_gen(
+        {"db": {"a.txt": "a", "b.txt": "b", "sub_dir": {"c.txt": "c"}}}
+    )
+
+    db_dir = tmp_dir / "db"
+    file_1, file_2, sub_dir_file = [
+        db_dir / filename for filename in ("a.txt", "b.txt", "sub_dir/c.txt")
+    ]
+
+    shutil.rmtree(db_dir)
+
+    dvc.checkout(targets=str(file_1))
+    assert db_dir.read_text() == {"a.txt": "a"}
+
+    dvc.checkout(targets=str(sub_dir_file))
+    assert db_dir.read_text() == {"a.txt": "a", "sub_dir": {"c.txt": "c"}}
+
+    dvc.checkout(targets=str(file_2))
+    assert db_dir.read_text() == {
+        "a.txt": "a",
+        "b.txt": "b",
+        "sub_dir": {"c.txt": "c"},
+    }
+
+
+def test_checkout_partial_unchanged(tmp_dir, dvc):
+    original_dir_shape = {
+        "a.txt": "a",
+        "b.txt": "b",
+        "sub_dir": {"c.txt": "c"},
+        "empty_sub_dir": {},
+    }
+    tmp_dir.dvc_gen({"db": original_dir_shape})
+
+    db_dir = tmp_dir / "db"
+    sub_dir = db_dir / "sub_dir"
+    file_1, file_2, sub_dir_file = [
+        db_dir / filename for filename in ("a.txt", "b.txt", "sub_dir/c.txt")
+    ]
+
+    # Nothing changed, nothing added/deleted/modified
+    stats = dvc.checkout(targets=str(file_2))
+    assert not any(stats.values())
+
+    # Irrelevant file changed, still nothing added/deleted/modified
+    file_1.unlink()
+    stats = dvc.checkout(targets=str(file_2))
+    assert not any(stats.values())
+
+    # Relevant change, one modified
+    file_2.unlink()
+    stats = dvc.checkout(targets=str(file_2))
+    assert len(stats["modified"]) == 1
+
+    # No changes inside db/sub
+    stats = dvc.checkout(targets=str(sub_dir))
+    assert not any(stats.values())
+
+    # Relevant change, one modified
+    sub_dir_file.unlink()
+    stats = dvc.checkout(targets=str(sub_dir))
+    assert len(stats["modified"]) == 1
+
+    stats = dvc.checkout(targets=str(db_dir / "empty_sub_dir"))
+    assert not any(stats.values())
+
+    dvc.checkout(targets=str(db_dir))
+
+    # Everything is in place, no action taken
+    stats = dvc.checkout(targets=str(db_dir))
+    assert not any(stats.values())
+
+
+def test_checkout_partial_subdir(tmp_dir, dvc):
+    tmp_dir.dvc_gen(
+        {"db": {"a.txt": "a", "sub_dir": {"b.txt": "b", "c.txt": "c"}}}
+    )
+
+    db_dir = tmp_dir / "db"
+    sub_dir = db_dir / "sub_dir"
+    sub_dir_file_2 = sub_dir / "c.txt"
+
+    shutil.rmtree(sub_dir)
+    dvc.checkout(targets=str(sub_dir))
+    assert db_dir.read_text() == {
+        "a.txt": "a",
+        "sub_dir": {"b.txt": "b", "c.txt": "c"},
+    }
+
+    sub_dir_file_2.unlink()
+    dvc.checkout(targets=str(sub_dir_file_2))
+    assert db_dir.read_text() == {
+        "a.txt": "a",
+        "sub_dir": {"b.txt": "b", "c.txt": "c"},
+    }


Code review comments:
diff_hunk: @@ -89,7 +89,26 @@ def load_dir_cache(self, hash_info):
 
         return DirInfo.from_list(d)
 
-    def changed(self, path_info, hash_info):
+    def _get_filtered_hash_info(self, filter_info, hash_info, path_info):
+        dir_info = self.get_dir_cache(hash_info)
+        hash_key = filter_info.relative_to(path_info).parts
+
+        # Check whether it is a file that exists on the trie
+        if hash_key in dir_info.trie:
+            return dir_info.trie[hash_key]
+
+        depth = len(hash_key)
+        filtered_dir_info = DirInfo()
+        try:
+            for key, value in dir_info.trie.items(hash_key):
+                filtered_dir_info.trie[key[depth:]] = value
+        except KeyError:
+            return None
+
+        filtered_hash_info, _ = self._get_dir_info_hash(filtered_dir_info)
+        return filtered_hash_info
path: dvc/cache/base.py
text: @user1:
```suggestion
        return self._get_dir_info_hash(filtered_dir_info)[0]
```

diff_hunk: @@ -101,23 +120,33 @@ def changed(self, path_info, hash_info):
         Args:
             path_info: dict with path information.
             hash: expected hash value for this data.
+            filter_info: an optional argument to target a specific path.
 
         Returns:
             bool: True if data has changed, False otherwise.
         """
 
+        path = filter_info or path_info
         logger.trace(
             "checking if '%s'('%s') has changed.", path_info, hash_info
         )
 
-        if not self.tree.exists(path_info):
-            logger.debug("'%s' doesn't exist.", path_info)
+        if not self.tree.exists(path):
+            logger.debug("'%s' doesn't exist.", path)
             return True
 
         if not hash_info:
             logger.debug("hash value for '%s' is missing.", path_info)
             return True
 
+        if path != path_info:
path: dvc/cache/base.py
text: @user1:
```suggestion
        if filter_info:
```

@author:
This is not always true, since filter_info might equal to the path_info sometimes. 

@user1:
ah, good point!