checkout/fetch: finer control over steps

See https://github.com/iterative/dvc-data/pull/390 and https://github.com/iterative/dvc-data/pull/391

dvc/repo/checkout.py

@@ -75,8 +75,7 @@ def checkout(  # noqa: C901
     from dvc.repo.index import build_data_index
     from dvc.stage.exceptions import StageFileBadNameError, StageFileDoesNotExistError
     from dvc_data.hashfile.checkout import CheckoutError as IndexCheckoutError
-    from dvc_data.index.checkout import ADD, DELETE, MODIFY
-    from dvc_data.index.checkout import checkout as icheckout
+    from dvc_data.index.checkout import ADD, DELETE, MODIFY, apply, compare
 
     stats: Dict[str, List[str]] = {
         "added": [],
@@ -119,21 +118,24 @@ def checkout(  # noqa: C901
 
     new = view.data["repo"]
 
+    with Callback.as_tqdm_callback(
+        desc="Comparing indexes",
+        unit="entry",
+    ) as cb:
+        diff = compare(old, new, relink=relink, delete=True, callback=cb)
+
     with Callback.as_tqdm_callback(
         unit="file",
         desc="Checkout",
     ) as cb:
         try:
-            changes = icheckout(
-                new,
+            changes = apply(
+                diff,
                 self.root_dir,
                 self.fs,
-                old=old,
                 callback=cb,
-                delete=True,
                 prompt=prompt.confirm,
                 update_meta=False,
-                relink=relink,
                 force=force,
                 allow_missing=allow_missing,
                 state=self.state,

dvc/repo/experiments/executor/base.py

@@ -515,7 +515,6 @@ class BaseExecutor(ABC):
                     targets=targets,
                     with_deps=targets is not None,
                     force=True,
-                    quiet=True,
                     allow_missing=True,
                     recursive=kwargs.get("recursive", False),
                 )

dvc/repo/fetch.py

@@ -36,6 +36,7 @@ def fetch(  # noqa: C901, PLR0913
     from fsspec.utils import tokenize
 
     from dvc.fs.callbacks import Callback
+    from dvc_data.index.fetch import collect
     from dvc_data.index.fetch import fetch as ifetch
 
     if isinstance(targets, str):
@@ -76,17 +77,27 @@ def fetch(  # noqa: C901, PLR0913
 
     cache_key = ("fetch", tokenize(sorted(index_keys)))
 
+    with Callback.as_tqdm_callback(
+        desc="Collecting",
+        unit="entry",
+    ) as cb:
+        data = collect(
+            indexes, cache_index=self.data_index, cache_key=cache_key, callback=cb
+        )
+
     with Callback.as_tqdm_callback(
         desc="Fetching",
         unit="file",
     ) as cb:
-        fetch_transferred, fetch_failed = ifetch(
-            indexes,
-            jobs=jobs,
-            callback=cb,
-            cache_index=self.data_index,
-            cache_key=cache_key,
-        )  # pylint: disable=assignment-from-no-return
+        try:
+            fetch_transferred, fetch_failed = ifetch(
+                data,
+                jobs=jobs,
+                callback=cb,
+            )  # pylint: disable=assignment-from-no-return
+        finally:
+            for fs_index in data.values():
+                fs_index.close()
 
     if fetch_transferred:
         # NOTE: dropping cached index to force reloading from newly saved cache

dvc/repo/imports.py

@@ -100,7 +100,8 @@ def save_imports(
         Objects which were downloaded from source location.
     """
     from dvc.stage.exceptions import DataSourceChanged
-    from dvc_data.index import checkout, md5, save
+    from dvc_data.index import md5, save
+    from dvc_data.index.checkout import apply, compare
     from dvc_objects.fs.callbacks import Callback
 
     downloaded: Set["HashInfo"] = set()
@@ -124,7 +125,8 @@ def save_imports(
                     desc="Downloading imports from source",
                     unit="files",
                 ) as cb:
-                    checkout(data_view, tmpdir, cache.fs, callback=cb, storage="remote")
+                    diff = compare(None, data_view)
+                    apply(diff, tmpdir, cache.fs, callback=cb, storage="remote")
                 md5(data_view, name=cache.hash_name)
                 save(data_view, odb=cache, hardlink=True)
 

dvc/repo/worktree.py

@@ -112,8 +112,7 @@ def push_worktree(
     **kwargs: Any,
 ) -> int:
     from dvc.repo.index import build_data_index
-    from dvc_data.index import checkout
-    from dvc_data.index.checkout import VersioningNotSupported
+    from dvc_data.index.checkout import VersioningNotSupported, apply, compare
 
     pushed = 0
     stages: Set["Stage"] = set()
@@ -139,6 +138,18 @@ def push_worktree(
         else:
             diff_kwargs = {}
 
+        with Callback.as_tqdm_callback(
+            unit="entry",
+            desc=f"Comparing indexes for remote {remote_obj.name!r}",
+        ) as cb:
+            diff = compare(
+                old_index,
+                new_index,
+                callback=cb,
+                delete=remote_obj.worktree,
+                **diff_kwargs,
+            )
+
         total = len(new_index)
         with Callback.as_tqdm_callback(
             unit="file",
@@ -147,16 +158,13 @@ def push_worktree(
         ) as cb:
             cb.set_size(total)
             try:
-                stats = checkout(
-                    new_index,
+                stats = apply(
+                    diff,
                     remote_obj.path,
                     remote_obj.fs,
-                    old=old_index,
-                    delete=remote_obj.worktree,
                     callback=cb,
                     latest_only=remote_obj.worktree,
                     jobs=jobs,
-                    **diff_kwargs,
                 )
                 pushed += sum(len(changes) for changes in stats.values())
             except VersioningNotSupported:
@@ -316,23 +324,33 @@ def _fetch_out_changes(
     remote_index: Union["DataIndex", "DataIndexView"],
     remote: "Remote",
 ):
-    from dvc_data.index import checkout
+    from dvc_data.index.checkout import apply, compare
 
     old, new = _get_diff_indexes(out, local_index, remote_index)
+
+    with Callback.as_tqdm_callback(
+        unit="entry",
+        desc="Comparing indexes",
+    ) as cb:
+        diff = compare(
+            old,
+            new,
+            delete=True,
+            meta_only=True,
+            meta_cmp_key=partial(_meta_checksum, remote.fs),
+            callback=cb,
+        )
+
     total = len(new)
     with Callback.as_tqdm_callback(
         unit="file", desc=f"Updating '{out}'", disable=total == 0
     ) as cb:
         cb.set_size(total)
-        checkout(
-            new,
+        apply(
+            diff,
             out.repo.root_dir,
             out.fs,
-            old=old,
-            delete=True,
             update_meta=False,
-            meta_only=True,
-            meta_cmp_key=partial(_meta_checksum, remote.fs),
             storage="data",
             callback=cb,
         )

pyproject.toml

@@ -34,7 +34,7 @@ dependencies = [
     "configobj>=5.0.6",
     "distro>=1.3",
     "dpath<3,>=2.1.0",
-    "dvc-data>=1.11.0,<1.12.0",
+    "dvc-data>=2.0.1,<2.1.0",
     "dvc-http>=2.29.0",
     "dvc-render>=0.3.1,<1",
     "dvc-studio-client>=0.9.2,<1",