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:

Resolving patch:
diff --git a/dvc/repo/index.py b/dvc/repo/index.py
--- a/dvc/repo/index.py
+++ b/dvc/repo/index.py
@@ -116,6 +116,28 @@ def is_out_or_ignored(root, directory):
         dirs[:] = [d for d in dirs if not is_out_or_ignored(root, d)]
 
 
+def _load_data_from_tree(index, prefix, ws, key, tree):
+    from dvc_data.index import DataIndexEntry, Meta
+
+    parents = set()
+
+    for okey, ometa, ohi in tree:
+        for key_len in range(1, len(okey)):
+            parents.add((*key, *okey[:key_len]))
+
+        fkey = (*key, *okey)
+        index[(*prefix, ws, *fkey)] = DataIndexEntry(
+            key=fkey,
+            meta=ometa,
+            hash_info=ohi,
+        )
+
+    for parent in parents:
+        index[(*prefix, ws, *parent)] = DataIndexEntry(
+            key=parent, meta=Meta(isdir=True), loaded=True
+        )
+
+
 def _load_data_from_outs(index, prefix, outs):
     from dvc_data.index import DataIndexEntry, Meta
 
@@ -129,25 +151,25 @@ def _load_data_from_outs(index, prefix, outs):
         for key_len in range(1, len(key)):
             parents.add((ws, key[:key_len]))
 
-        loaded = None
-        if out.files:
-            loaded = True
-            for okey, ometa, ohi in out.get_obj():
-                for key_len in range(1, len(okey)):
-                    parents.add((ws, (*key, *okey[:key_len])))
-
-                fkey = (*key, *okey)
-                index[(*prefix, ws, *fkey)] = DataIndexEntry(
-                    key=fkey,
-                    meta=ometa,
-                    hash_info=ohi,
-                )
+        tree = None
+        if (
+            out.stage.is_import
+            and not out.stage.is_repo_import
+            and not out.stage.is_db_import
+            and out.stage.deps[0].files
+        ):
+            tree = out.stage.deps[0].get_obj()
+        elif out.files:
+            tree = out.get_obj()
+
+        if tree is not None:
+            _load_data_from_tree(index, prefix, ws, key, tree)
 
         entry = DataIndexEntry(
             key=key,
             meta=out.meta,
             hash_info=out.hash_info,
-            loaded=loaded,
+            loaded=None if tree is None else True,
         )
 
         if (
@@ -194,7 +216,7 @@ def _load_storage_from_import(storage_map, key, out):
             )
         )
 
-    if out.stage.is_repo_import or not out.hash_info:
+    if out.stage.is_repo_import or not out.hash_info or dep.fs.version_aware:
         storage_map.add_remote(FileStorage(key, dep.fs, dep.fs_path, read_only=True))
 
 


Test patch: