Repository: iterative/dvc

Issue:
S3 version tracking is working, but version switch is not working
# Bug Report

<!--
## Issue name

Issue names must follow the pattern `command: description` where the command is the dvc command that you are trying to run. The description should describe the consequence of the bug.

Example: `repro: doesn't detect input changes`
-->
s3 version tracking with dvc is work, but version switch is not working

## Description

<!--
A clear and concise description of what the bug is.
-->
I want to use DVC to manage data versioning in S3 locally without downloading it.
I tried import-url --no-download --version-aware to make sure that the dvc file is versioned.
But I can't switch data between versions with dvc checkout. Even when I do a dvc pull with a specific version of the dvc file, I get the entire s3 data downloaded, not just the file written in the dvc.

### Reproduce

<!--
Step list of how to reproduce the bug
-->
```
git init
dvc init
dvc remote add -d s3 s3://temp/ds/
dvc remote modify s3 version_aware true
mkdir -p data
# Files in s3://temp/ds/dummy : a.text, b.text
dvc import-url --no-download --version-aware s3://temp/ds/dummy data
git add data/dummy.dvc data/.gitignore .dvc/config
git commit -m "version 1"
git tag -a "v1" -m "version 1 tagging"

# Add File in s3://temp/ds/dummy : a.text, b.text, c.text
dvc update --no-download data/dummy.dvc
git add data/dummy.dvc
git commit -m "version 2"
git tag -a "v2" -m "version 2 tagging"
dvc pull

git checkout v1
dvc checkout

# But a.text, b.text, c.text is existed in data/dummy/{} | version 1 only have a.text, b.text

git clean -fdx
dvc gc -w

dvc pull 
# It also download a.text, b.text, c.text  | version 1 only have to download a.text, b.text
```

<!--
Example:

1. dvc init
2. Copy dataset.zip to the directory
3. dvc add dataset.zip
4. dvc run -d dataset.zip -o model ./train.sh
5. modify dataset.zip
6. dvc repro
-->

### Expected

<!--
A clear and concise description of what you expect to happen.
-->
I want to switch data version with dvc checkout.


### Environment information

<!--
This is required to ensure that we can reproduce the bug.
-->

**Output of `dvc doctor`:**

```console
❯ dvc doctor
DVC version: 3.33.3 (pip)
-------------------------
Platform: Python 3.9.18 on macOS-14.1.1-arm64-arm-64bit
Subprojects:
	dvc_data = 2.22.6
	dvc_objects = 1.4.9
	dvc_render = 1.0.0
	dvc_task = 0.3.0
	scmrepo = 1.5.0
Supports:
	http (aiohttp = 3.9.1, aiohttp-retry = 2.8.3),
	https (aiohttp = 3.9.1, aiohttp-retry = 2.8.3),
	s3 (s3fs = 2023.12.1, boto3 = 1.33.1)
Config:
	Global: /Users/a13080/Library/Application Support/dvc
	System: /Library/Application Support/dvc
```

**Additional Information (if any):**

<!--
Please check https://github.com/iterative/dvc/wiki/Debugging-DVC on ways to gather more information regarding the issue.

If applicable, please also provide a `--verbose` output of the command, eg: `dvc add --verbose`.
If the issue is regarding the performance, please attach the profiling information and the benchmark comparisons.
-->

I already reported this issue to [discord](https://discord.com/channels/485586884165107732/563406153334128681/1184403305221328906).
Manager told me to report issue on github


Discussion:
@justHungryMan Could you give an upstream dvc a try, please? https://dvc.org/doc/install/pre-release
@efiop Is it fixed? I already have a same problem.

Data is not changed

```
V1: a.text, b.text
V2: a.text, b.text, c.text
```

If I pull from V2, and `git checkout + dvc checkout to V1`, but c.text is yet alived.
or if I pull from V1 from clean status (first pull), but c.text is also pulled
@justHungryMan Could you show `dvc doctor`, please? Just to make sure.


@efiop Here you are.

```
DVC version: 3.38.0
-------------------
Platform: Python 3.9.18 on macOS-14.1.1-arm64-arm-64bit
Subprojects:
        dvc_data = 3.6.0
        dvc_objects = 3.0.3
        dvc_render = 1.0.0
        dvc_task = 0.3.0
        scmrepo = 2.0.2
Supports:
        http (aiohttp = 3.9.1, aiohttp-retry = 2.8.3),
        https (aiohttp = 3.9.1, aiohttp-retry = 2.8.3),
        s3 (s3fs = 2023.12.2, boto3 = 1.33.13)
Config:
        Global: /Users/a13080/Library/Application Support/dvc
        System: /Library/Application Support/dvc
Cache types: <https://error.dvc.org/no-dvc-cache>
Caches: local
Remotes: None
Workspace directory: apfs on /dev/disk3s1s1
Repo: dvc, git
Repo.site_cache_dir: /Library/Caches/dvc/repo/da5e613fce8e68e2c5daaa19cffb1a83
```


Can confirm, taking a look...

Resolving patch:
diff --git a/dvc/repo/index.py b/dvc/repo/index.py
--- a/dvc/repo/index.py
+++ b/dvc/repo/index.py
@@ -255,6 +255,33 @@ def _load_storage_from_out(storage_map, key, out):
         _load_storage_from_import(storage_map, key, out)
 
 
+def _build_tree_from_outs(outs):
+    from dvc_data.hashfile.tree import Tree
+
+    tree = Tree()
+    for out in outs:
+        if not out.use_cache:
+            continue
+
+        ws, key = out.index_key
+
+        if not out.stage.is_partial_import:
+            tree.add((ws, *key), out.meta, out.hash_info)
+            continue
+
+        dep = out.stage.deps[0]
+        if not dep.files:
+            tree.add((ws, *key), dep.meta, dep.hash_info)
+            continue
+
+        for okey, ometa, ohi in dep.get_obj():
+            tree.add((ws, *key, *okey), ometa, ohi)
+
+    tree.digest()
+
+    return tree
+
+
 class Index:
     def __init__(
         self,
@@ -504,20 +531,7 @@ def plot_keys(self) -> Dict[str, Set["DataIndexKey"]]:
 
     @cached_property
     def data_tree(self):
-        from dvc_data.hashfile.tree import Tree
-
-        tree = Tree()
-        for out in self.outs:
-            if not out.use_cache:
-                continue
-
-            ws, key = out.index_key
-
-            tree.add((ws, *key), out.meta, out.hash_info)
-
-        tree.digest()
-
-        return tree
+        return _build_tree_from_outs(self.outs)
 
     @cached_property
     def data(self) -> "Dict[str, DataIndex]":
@@ -772,20 +786,7 @@ def data_keys(self) -> Dict[str, Set["DataIndexKey"]]:
 
     @cached_property
     def data_tree(self):
-        from dvc_data.hashfile.tree import Tree
-
-        tree = Tree()
-        for out in self.outs:
-            if not out.use_cache:
-                continue
-
-            ws, key = out.index_key
-
-            tree.add((ws, *key), out.meta, out.hash_info)
-
-        tree.digest()
-
-        return tree
+        return _build_tree_from_outs(self.outs)
 
     @cached_property
     def data(self) -> Dict[str, Union["DataIndex", "DataIndexView"]]:


Test patch:
diff --git a/dvc/testing/workspace_tests.py b/dvc/testing/workspace_tests.py
--- a/dvc/testing/workspace_tests.py
+++ b/dvc/testing/workspace_tests.py
@@ -149,14 +149,33 @@ def test_import_dir(self, tmp_dir, dvc, remote_version_aware):
         assert (tmp_dir / "data_dir" / "subdir" / "file").read_text() == "modified"
         assert (tmp_dir / "data_dir" / "new_file").read_text() == "new"
 
-    def test_import_no_download(self, tmp_dir, dvc, remote_version_aware):
+    def test_import_no_download(self, tmp_dir, dvc, remote_version_aware, scm):
         remote_version_aware.gen({"data_dir": {"subdir": {"file": "file"}}})
         dvc.imp_url("remote://upstream/data_dir", version_aware=True, no_download=True)
+        scm.add(["data_dir.dvc", ".gitignore"])
+        scm.commit("v1")
+        scm.tag("v1")
+
         stage = first(dvc.index.stages)
         assert not stage.outs[0].can_push
 
-        dvc.pull()
-        assert (tmp_dir / "data_dir" / "subdir" / "file").read_text() == "file"
+        (remote_version_aware / "data_dir" / "foo").write_text("foo")
+        dvc.update(no_download=True)
+        assert dvc.pull()["fetched"] == 2
+        assert (tmp_dir / "data_dir").read_text() == {
+            "foo": "foo",
+            "subdir": {"file": "file"},
+        }
+        scm.add(["data_dir.dvc", ".gitignore"])
+        scm.commit("update")
+
+        scm.checkout("v1")
+        dvc.cache.local.clear()
+        remove(tmp_dir / "data_dir")
+        assert dvc.pull()["fetched"] == 1
+        assert (tmp_dir / "data_dir").read_text() == {
+            "subdir": {"file": "file"},
+        }
 
         dvc.commit(force=True)
         assert dvc.status() == {}
diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py
--- a/tests/func/test_data_cloud.py
+++ b/tests/func/test_data_cloud.py
@@ -271,7 +271,7 @@ def test_pull_partial_import(tmp_dir, dvc, local_workspace):
     stage = dvc.imp_url("remote://workspace/file", os.fspath(dst), no_download=True)
 
     result = dvc.pull("file")
-    assert result["fetched"] == 0
+    assert result["fetched"] == 1
     assert dst.exists()
 
     assert stage.outs[0].get_hash().value == "d10b4c3ff123b26dc068d43a8bef2d23"