introduce dataset dependency (#10164)

* introduce dataset dependency

* support Annotated types

* split

* drop dvc.api.datasets

* drop pydantic and typing-extensions

* drop dvc.datasets

dvc/dependency/__init__.py

@@ -5,6 +5,7 @@ from typing import Any
 from dvc.output import ARTIFACT_SCHEMA, DIR_FILES_SCHEMA, Output
 
 from .base import Dependency
+from .dataset import DatasetDependency
 from .db import DbDependency
 from .param import ParamsDependency
 from .repo import RepoDependency
@@ -33,6 +34,8 @@ def _get(stage, p, info, **kwargs):
         return DbDependency(stage, d)
 
     assert p
+    if DatasetDependency.is_dataset(p):
+        return DatasetDependency(stage, p, info)
     if repo:
         return RepoDependency(repo, stage, p, info)
     return Dependency(stage, p, info, **kwargs)

dvc/dependency/dataset.py

@@ -0,0 +1,69 @@
+from typing import TYPE_CHECKING, Any, ClassVar
+from urllib.parse import urlparse
+
+from funcy import compact, merge
+
+from dvc_data.hashfile.hash_info import HashInfo
+
+from .db import AbstractDependency
+
+if TYPE_CHECKING:
+    from dvc.stage import Stage
+
+
+class DatasetDependency(AbstractDependency):
+    PARAM_DATASET = "dataset"
+    DATASET_SCHEMA: ClassVar[dict] = {PARAM_DATASET: dict}
+
+    def __init__(self, stage: "Stage", p, info, *args, **kwargs):
+        super().__init__(stage, info, *args, **kwargs)
+        self.def_path = p
+        self.name = urlparse(p).netloc
+        dataset_info = info.get(self.PARAM_DATASET) or {}
+        self.hash_info = HashInfo(self.PARAM_DATASET, dataset_info)  # type: ignore[arg-type]
+        self.hash_name = self.PARAM_DATASET
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.def_path!r})"
+
+    def __str__(self):
+        return self.def_path
+
+    @classmethod
+    def is_dataset(cls, p: str):
+        return urlparse(p).scheme == "ds"
+
+    @property
+    def protocol(self):
+        return None
+
+    def dumpd(self, **kwargs):
+        return compact({self.PARAM_PATH: self.def_path, **self.hash_info.to_dict()})
+
+    def fill_values(self, values=None):
+        """Load params values dynamically."""
+        self.hash_info = HashInfo(
+            self.PARAM_DATASET, merge(self.hash_info.value, values or {})
+        )
+
+    def workspace_status(self):
+        registered = self.repo.index.datasets.get(self.name, {})
+        info: dict[str, Any] = self.hash_info.value if self.hash_info else {}  # type: ignore[assignment]
+        if info != registered:
+            return {str(self): "modified"}
+        return {}
+
+    def status(self):
+        return self.workspace_status()
+
+    def get_hash(self):
+        return HashInfo(self.PARAM_DATASET, self.repo.index.datasets.get(self.name, {}))
+
+    def save(self):
+        self.hash_info = self.get_hash()
+
+    def download(self, to, jobs=None):
+        raise NotImplementedError
+
+    def update(self, rev=None):
+        raise NotImplementedError

dvc/dvcfile.py

@@ -179,6 +179,7 @@ class SingleStageFile(FileMixin):
     from dvc.schema import COMPILED_SINGLE_STAGE_SCHEMA as SCHEMA
     from dvc.stage.loader import SingleStageLoader as LOADER  # noqa: N814
 
+    datasets: ClassVar[list[dict[str, Any]]] = []
     metrics: ClassVar[list[str]] = []
     plots: ClassVar[Any] = {}
     params: ClassVar[list[str]] = []
@@ -311,6 +312,10 @@ class ProjectFile(FileMixin):
     def params(self) -> list[str]:
         return self.contents.get("params", [])
 
+    @property
+    def datasets(self) -> list[dict[str, Any]]:
+        return self.contents.get("datasets", [])
+
     @property
     def artifacts(self) -> dict[str, Optional[dict[str, Any]]]:
         return self.contents.get("artifacts", {})

dvc/repo/graph.py

@@ -117,6 +117,7 @@ def build_graph(stages, outs_trie=None):
     """
     import networkx as nx
 
+    from dvc.dependency import DatasetDependency
     from dvc.exceptions import StagePathAsOutputError
 
     from .trie import build_outs_trie
@@ -140,6 +141,8 @@ def build_graph(stages, outs_trie=None):
             continue
 
         for dep in stage.deps:
+            if isinstance(dep, DatasetDependency):
+                continue
             dep_key = dep.fs.parts(dep.fs_path)
             overlapping = [n.value for n in outs_trie.prefixes(dep_key)]
             if outs_trie.has_subtrie(dep_key):
@@ -157,6 +160,8 @@ def build_graph(stages, outs_trie=None):
 def build_outs_graph(graph, outs_trie):
     import networkx as nx
 
+    from dvc.dependency import DatasetDependency
+
     outs_graph = nx.DiGraph()
 
     outs_graph.add_nodes_from(outs_trie.values())
@@ -166,6 +171,8 @@ def build_outs_graph(graph, outs_trie):
         if stage.is_db_import:
             continue
         for dep in stage.deps:
+            if isinstance(dep, DatasetDependency):
+                continue
             dep_key = dep.fs.parts(dep.fs_path)
             overlapping = [n.value for n in outs_trie.prefixes(dep_key)]
             if outs_trie.has_subtrie(dep_key):

dvc/repo/index.py

@@ -288,6 +288,7 @@ class Index:
         plots: Optional[dict[str, list[str]]] = None,
         params: Optional[dict[str, Any]] = None,
         artifacts: Optional[dict[str, Any]] = None,
+        datasets: Optional[dict[str, list[dict[str, Any]]]] = None,
     ) -> None:
         self.repo = repo
         self.stages = stages or []
@@ -295,6 +296,7 @@ class Index:
         self._plots = plots or {}
         self._params = params or {}
         self._artifacts = artifacts or {}
+        self._datasets: dict[str, list[dict[str, Any]]] = datasets or {}
         self._collected_targets: dict[int, list["StageInfo"]] = {}
 
     @cached_property
@@ -318,6 +320,7 @@ class Index:
         plots = {}
         params = {}
         artifacts = {}
+        datasets = {}
 
         onerror = onerror or repo.stage_collection_error_handler
         for _, idx in collect_files(repo, onerror=onerror):
@@ -326,6 +329,7 @@ class Index:
             plots.update(idx._plots)
             params.update(idx._params)
             artifacts.update(idx._artifacts)
+            datasets.update(idx._datasets)
         return cls(
             repo,
             stages=stages,
@@ -333,6 +337,7 @@ class Index:
             plots=plots,
             params=params,
             artifacts=artifacts,
+            datasets=datasets,
         )
 
     @classmethod
@@ -347,6 +352,7 @@ class Index:
             plots={path: dvcfile.plots} if dvcfile.plots else {},
             params={path: dvcfile.params} if dvcfile.params else {},
             artifacts={path: dvcfile.artifacts} if dvcfile.artifacts else {},
+            datasets={path: dvcfile.datasets} if dvcfile.datasets else {},
         )
 
     def update(self, stages: Iterable["Stage"]) -> "Index":
@@ -361,6 +367,7 @@ class Index:
             plots=self._plots,
             params=self._params,
             artifacts=self._artifacts,
+            datasets=self._datasets,
         )
 
     @cached_property
@@ -398,6 +405,11 @@ class Index:
         for stage in self.stages:
             yield from stage.outs
 
+    @cached_property
+    def datasets(self) -> dict[str, dict[str, Any]]:
+        datasets = chain.from_iterable(self._datasets.values())
+        return {dataset["name"]: dataset for dataset in datasets}
+
     @cached_property
     def out_data_keys(self) -> dict[str, set["DataIndexKey"]]:
         by_workspace: dict[str, set["DataIndexKey"]] = defaultdict(set)

dvc/schema.py

@@ -35,6 +35,7 @@ DATA_SCHEMA: dict[Any, Any] = {
     Output.PARAM_CLOUD: CLOUD_SCHEMA,
     Output.PARAM_FILES: [DIR_FILES_SCHEMA],
     Output.PARAM_HASH: str,
+    **dependency.DatasetDependency.DATASET_SCHEMA,
 }
 LOCK_FILE_STAGE_SCHEMA = {
     vol.Required(StageParams.PARAM_CMD): vol.Any(str, list),
@@ -122,7 +123,12 @@ FOREACH_IN = {
 SINGLE_PIPELINE_STAGE_SCHEMA = {
     str: either_or(STAGE_DEFINITION, FOREACH_IN, [FOREACH_KWD, DO_KWD])
 }
+
+DATASET_SCHEMA = vol.Schema(
+    {vol.Required("type"): str, vol.Required("name"): str}, extra=vol.ALLOW_EXTRA
+)
 MULTI_STAGE_SCHEMA = {
+    "datasets": [DATASET_SCHEMA],
     PLOTS: [vol.Any(str, SINGLE_PLOT_SCHEMA)],
     STAGES: SINGLE_PIPELINE_STAGE_SCHEMA,
     VARS_KWD: VARS_SCHEMA,

dvc/stage/loader.py

@@ -73,6 +73,10 @@ class StageLoader(Mapping):
             info = info.copy()
             info.pop("path", None)
 
+            if isinstance(item, dependency.DatasetDependency):
+                item.fill_values(info.get(dependency.DatasetDependency.PARAM_DATASET))
+                continue
+
             hash_name = info.pop(Output.PARAM_HASH, None)
             item.meta = Meta.from_dict(merge_file_meta_from_cloud(info))
             item.hash_name, item.hash_info = item._compute_hash_info_from_meta(

dvc/stage/serialize.py

@@ -148,6 +148,7 @@ def to_pipeline_file(stage: "PipelineStage"):
 
 def to_single_stage_lockfile(stage: "Stage", **kwargs) -> dict:
     from dvc.cachemgr import LEGACY_HASH_NAMES
+    from dvc.dependency import DatasetDependency
     from dvc.output import (
         _serialize_hi_to_dict,
         _serialize_tree_obj_to_files,
@@ -158,6 +159,9 @@ def to_single_stage_lockfile(stage: "Stage", **kwargs) -> dict:
     assert stage.cmd
 
     def _dumpd(item: "Output"):
+        if isinstance(item, DatasetDependency):
+            return item.dumpd()
+
         ret: dict[str, Any] = {item.PARAM_PATH: item.def_path}
         if item.hash_name not in LEGACY_HASH_NAMES:
             ret[item.PARAM_HASH] = "md5"