datachain 0.35.2__tar.gz → 0.36.0__tar.gz

{datachain-0.35.2 → datachain-0.36.0}/.github/workflows/benchmarks.yml

@@ -25,7 +25,7 @@ jobs:
           python-version: '3.13'
 
       - name: Setup uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           enable-cache: true
           cache-suffix: benchmarks

{datachain-0.35.2 → datachain-0.36.0}/.github/workflows/release.yml

@@ -27,7 +27,7 @@ jobs:
           python-version: '3.13'
 
       - name: Setup uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
 
       - name: Install nox
         run: uv pip install nox --system

{datachain-0.35.2 → datachain-0.36.0}/.github/workflows/tests-studio.yml

@@ -97,7 +97,7 @@ jobs:
           python-version: ${{ matrix.pyv }}
 
       - name: Setup uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           enable-cache: true
           cache-suffix: studio

{datachain-0.35.2 → datachain-0.36.0}/.github/workflows/tests.yml

@@ -29,7 +29,7 @@ jobs:
           python-version: '3.10'
 
       - name: Setup uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           enable-cache: true
           cache-suffix: lint
@@ -80,7 +80,7 @@ jobs:
 
       - name: Setup PostgreSQL
         if: runner.os != 'Windows'
-        uses: ikalnytskyi/action-setup-postgres@10ab8a56cc77b4823c2bfa57b1d4dd5605ef0481 # v7
+        uses: ikalnytskyi/action-setup-postgres@c4dda34aae1c821e3a771b68b73b13af3198a7ee # v8
         with:
           username: test
           password: test
@@ -102,7 +102,7 @@ jobs:
           python-version: ${{ matrix.pyv }}
 
       - name: Setup uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           enable-cache: true
           cache-suffix: tests-${{ matrix.pyv }}
@@ -188,7 +188,7 @@ jobs:
           python-version: ${{ matrix.pyv }}
 
       - name: Setup uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           enable-cache: true
           cache-suffix: examples-${{ matrix.pyv }}

{datachain-0.35.2 → datachain-0.36.0}/.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.13.3'
+    rev: 'v0.14.0'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.35.2 → datachain-0.36.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.35.2
+Version: 0.36.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.35.2 → datachain-0.36.0}/src/datachain/catalog/catalog.py

@@ -54,6 +54,7 @@ from datachain.sql.types import DateTime, SQLType
 from datachain.utils import DataChainDir
 
 from .datasource import DataSource
+from .dependency import build_dependency_hierarchy, populate_nested_dependencies
 
 if TYPE_CHECKING:
     from datachain.data_storage import AbstractMetastore, AbstractWarehouse
@@ -1203,6 +1204,38 @@ class Catalog:
         assert isinstance(dataset_info, dict)
         return DatasetRecord.from_dict(dataset_info)
 
+    def get_dataset_dependencies_by_ids(
+        self,
+        dataset_id: int,
+        version_id: int,
+        indirect: bool = True,
+    ) -> list[DatasetDependency | None]:
+        dependency_nodes = self.metastore.get_dataset_dependency_nodes(
+            dataset_id=dataset_id,
+            version_id=version_id,
+        )
+
+        if not dependency_nodes:
+            return []
+
+        dependency_map, children_map = build_dependency_hierarchy(dependency_nodes)
+
+        root_key = (dataset_id, version_id)
+        if root_key not in children_map:
+            return []
+
+        root_dependency_ids = children_map[root_key]
+        root_dependencies = [dependency_map[dep_id] for dep_id in root_dependency_ids]
+
+        if indirect:
+            for dependency in root_dependencies:
+                if dependency is not None:
+                    populate_nested_dependencies(
+                        dependency, dependency_nodes, dependency_map, children_map
+                    )
+
+        return root_dependencies
+
     def get_dataset_dependencies(
         self,
         name: str,
@@ -1216,29 +1249,21 @@ class Catalog:
             namespace_name=namespace_name,
             project_name=project_name,
         )
-
-        direct_dependencies = self.metastore.get_direct_dataset_dependencies(
-            dataset, version
-        )
+        dataset_version = dataset.get_version(version)
+        dataset_id = dataset.id
+        dataset_version_id = dataset_version.id
 
         if not indirect:
-            return direct_dependencies
-
-        for d in direct_dependencies:
-            if not d:
-                # dependency has been removed
-                continue
-            if d.is_dataset:
-                # only datasets can have dependencies
-                d.dependencies = self.get_dataset_dependencies(
-                    d.name,
-                    d.version,
-                    namespace_name=d.namespace,
-                    project_name=d.project,
-                    indirect=indirect,
-                )
+            return self.metastore.get_direct_dataset_dependencies(
+                dataset,
+                version,
+            )
 
-        return direct_dependencies
+        return self.get_dataset_dependencies_by_ids(
+            dataset_id,
+            dataset_version_id,
+            indirect,
+        )
 
     def ls_datasets(
         self,

datachain-0.36.0/src/datachain/catalog/dependency.py

@@ -0,0 +1,164 @@
+import builtins
+from dataclasses import dataclass
+from datetime import datetime
+from typing import TypeVar
+
+from datachain.dataset import DatasetDependency
+
+DDN = TypeVar("DDN", bound="DatasetDependencyNode")
+
+
+@dataclass
+class DatasetDependencyNode:
+    namespace: str
+    project: str
+    id: int
+    dataset_id: int | None
+    dataset_version_id: int | None
+    dataset_name: str | None
+    dataset_version: str | None
+    created_at: datetime
+    source_dataset_id: int
+    source_dataset_version_id: int | None
+    depth: int
+
+    @classmethod
+    def parse(
+        cls: builtins.type[DDN],
+        namespace: str,
+        project: str,
+        id: int,
+        dataset_id: int | None,
+        dataset_version_id: int | None,
+        dataset_name: str | None,
+        dataset_version: str | None,
+        created_at: datetime,
+        source_dataset_id: int,
+        source_dataset_version_id: int | None,
+        depth: int,
+    ) -> "DatasetDependencyNode | None":
+        return cls(
+            namespace,
+            project,
+            id,
+            dataset_id,
+            dataset_version_id,
+            dataset_name,
+            dataset_version,
+            created_at,
+            source_dataset_id,
+            source_dataset_version_id,
+            depth,
+        )
+
+    def to_dependency(self) -> "DatasetDependency | None":
+        return DatasetDependency.parse(
+            namespace_name=self.namespace,
+            project_name=self.project,
+            id=self.id,
+            dataset_id=self.dataset_id,
+            dataset_version_id=self.dataset_version_id,
+            dataset_name=self.dataset_name,
+            dataset_version=self.dataset_version,
+            dataset_version_created_at=self.created_at,
+        )
+
+
+def build_dependency_hierarchy(
+    dependency_nodes: list[DatasetDependencyNode | None],
+) -> tuple[
+    dict[int, DatasetDependency | None], dict[tuple[int, int | None], list[int]]
+]:
+    """
+    Build dependency hierarchy from dependency nodes.
+
+    Args:
+        dependency_nodes: List of DatasetDependencyNode objects from the database
+
+    Returns:
+        Tuple of (dependency_map, children_map) where:
+        - dependency_map: Maps dependency_id -> DatasetDependency
+        - children_map: Maps (source_dataset_id, source_version_id) ->
+          list of dependency_ids
+    """
+    dependency_map: dict[int, DatasetDependency | None] = {}
+    children_map: dict[tuple[int, int | None], list[int]] = {}
+
+    for node in dependency_nodes:
+        if node is None:
+            continue
+        dependency = node.to_dependency()
+        parent_key = (node.source_dataset_id, node.source_dataset_version_id)
+
+        if dependency is not None:
+            dependency_map[dependency.id] = dependency
+            children_map.setdefault(parent_key, []).append(dependency.id)
+        else:
+            # Handle case where dependency creation failed (e.g., deleted dependency)
+            dependency_map[node.id] = None
+            children_map.setdefault(parent_key, []).append(node.id)
+
+    return dependency_map, children_map
+
+
+def populate_nested_dependencies(
+    dependency: DatasetDependency,
+    dependency_nodes: list[DatasetDependencyNode | None],
+    dependency_map: dict[int, DatasetDependency | None],
+    children_map: dict[tuple[int, int | None], list[int]],
+) -> None:
+    """
+    Recursively populate nested dependencies for a given dependency.
+
+    Args:
+        dependency: The dependency to populate nested dependencies for
+        dependency_nodes: All dependency nodes from the database
+        dependency_map: Maps dependency_id -> DatasetDependency
+        children_map: Maps (source_dataset_id, source_version_id) ->
+        list of dependency_ids
+    """
+    # Find the target dataset and version for this dependency
+    target_dataset_id, target_version_id = find_target_dataset_version(
+        dependency, dependency_nodes
+    )
+
+    if target_dataset_id is None or target_version_id is None:
+        return
+
+    # Get children for this target
+    target_key = (target_dataset_id, target_version_id)
+    if target_key not in children_map:
+        dependency.dependencies = []
+        return
+
+    child_dependency_ids = children_map[target_key]
+    child_dependencies = [dependency_map[child_id] for child_id in child_dependency_ids]
+
+    dependency.dependencies = child_dependencies
+
+    # Recursively populate children
+    for child_dependency in child_dependencies:
+        if child_dependency is not None:
+            populate_nested_dependencies(
+                child_dependency, dependency_nodes, dependency_map, children_map
+            )
+
+
+def find_target_dataset_version(
+    dependency: DatasetDependency,
+    dependency_nodes: list[DatasetDependencyNode | None],
+) -> tuple[int | None, int | None]:
+    """
+    Find the target dataset ID and version ID for a given dependency.
+
+    Args:
+        dependency: The dependency to find target for
+        dependency_nodes: All dependency nodes from the database
+
+    Returns:
+        Tuple of (target_dataset_id, target_version_id) or (None, None) if not found
+    """
+    for node in dependency_nodes:
+        if node is not None and node.id == dependency.id:
+            return node.dataset_id, node.dataset_version_id
+    return None, None

{datachain-0.35.2 → datachain-0.36.0}/src/datachain/data_storage/metastore.py

@@ -22,10 +22,12 @@ from sqlalchemy import (
     Text,
     UniqueConstraint,
     desc,
+    literal,
     select,
 )
 from sqlalchemy.sql import func as f
 
+from datachain.catalog.dependency import DatasetDependencyNode
 from datachain.checkpoint import Checkpoint
 from datachain.data_storage import JobQueryType, JobStatus
 from datachain.data_storage.serializer import Serializable
@@ -78,6 +80,7 @@ class AbstractMetastore(ABC, Serializable):
     dataset_list_class: type[DatasetListRecord] = DatasetListRecord
     dataset_list_version_class: type[DatasetListVersion] = DatasetListVersion
     dependency_class: type[DatasetDependency] = DatasetDependency
+    dependency_node_class: type[DatasetDependencyNode] = DatasetDependencyNode
     job_class: type[Job] = Job
     checkpoint_class: type[Checkpoint] = Checkpoint
 
@@ -366,6 +369,12 @@ class AbstractMetastore(ABC, Serializable):
     ) -> list[DatasetDependency | None]:
         """Gets direct dataset dependencies."""
 
+    @abstractmethod
+    def get_dataset_dependency_nodes(
+        self, dataset_id: int, version_id: int
+    ) -> list[DatasetDependencyNode | None]:
+        """Gets dataset dependency node from database."""
+
     @abstractmethod
     def remove_dataset_dependencies(
         self, dataset: DatasetRecord, version: str | None = None
@@ -1483,6 +1492,77 @@ class AbstractDBMetastore(AbstractMetastore):
 
         return [self.dependency_class.parse(*r) for r in self.db.execute(query)]
 
+    def get_dataset_dependency_nodes(
+        self, dataset_id: int, version_id: int
+    ) -> list[DatasetDependencyNode | None]:
+        n = self._namespaces_select().subquery()
+        p = self._projects
+        d = self._datasets_select().subquery()
+        dd = self._datasets_dependencies
+        dv = self._datasets_versions
+
+        # Common dependency fields for CTE
+        dep_fields = [
+            dd.c.id,
+            dd.c.source_dataset_id,
+            dd.c.source_dataset_version_id,
+            dd.c.dataset_id,
+            dd.c.dataset_version_id,
+        ]
+
+        # Base case: direct dependencies
+        base_query = select(
+            *dep_fields,
+            literal(0).label("depth"),
+        ).where(
+            (dd.c.source_dataset_id == dataset_id)
+            & (dd.c.source_dataset_version_id == version_id)
+        )
+
+        cte = base_query.cte(name="dependency_tree", recursive=True)
+
+        # Recursive case: dependencies of dependencies
+        recursive_query = select(
+            *dep_fields,
+            (cte.c.depth + 1).label("depth"),
+        ).select_from(
+            cte.join(
+                dd,
+                (cte.c.dataset_id == dd.c.source_dataset_id)
+                & (cte.c.dataset_version_id == dd.c.source_dataset_version_id),
+            )
+        )
+
+        cte = cte.union(recursive_query)
+
+        # Fetch all with full details
+        final_query = select(
+            n.c.name,
+            p.c.name,
+            cte.c.id,
+            cte.c.dataset_id,
+            cte.c.dataset_version_id,
+            d.c.name,
+            dv.c.version,
+            dv.c.created_at,
+            cte.c.source_dataset_id,
+            cte.c.source_dataset_version_id,
+            cte.c.depth,
+        ).select_from(
+            # Use outer joins to handle cases where dependent datasets have been
+            # physically deleted. This allows us to return dependency records with
+            # None values instead of silently omitting them, making broken
+            # dependencies visible to callers.
+            cte.join(d, cte.c.dataset_id == d.c.id, isouter=True)
+            .join(dv, cte.c.dataset_version_id == dv.c.id, isouter=True)
+            .join(p, d.c.project_id == p.c.id, isouter=True)
+            .join(n, p.c.namespace_id == n.c.id, isouter=True)
+        )
+
+        return [
+            self.dependency_node_class.parse(*r) for r in self.db.execute(final_query)
+        ]
+
     def remove_dataset_dependencies(
         self, dataset: DatasetRecord, version: str | None = None
     ) -> None:

{datachain-0.35.2 → datachain-0.36.0}/src/datachain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.35.2
+Version: 0.36.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.35.2 → datachain-0.36.0}/src/datachain.egg-info/SOURCES.txt

@@ -136,6 +136,7 @@ src/datachain.egg-info/top_level.txt
 src/datachain/catalog/__init__.py
 src/datachain/catalog/catalog.py
 src/datachain/catalog/datasource.py
+src/datachain/catalog/dependency.py
 src/datachain/catalog/loader.py
 src/datachain/cli/__init__.py
 src/datachain/cli/utils.py