datachain 0.35.2__tar.gz → 0.36.1__tar.gz

{datachain-0.35.2 → datachain-0.36.1}/.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.1}/.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.1}/.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.1}/.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.1}/.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.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.35.2
+Version: 0.36.1
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -64,7 +64,6 @@ Requires-Dist: torch>=2.1.0; extra == "torch"
 Requires-Dist: torchvision; extra == "torch"
 Requires-Dist: transformers>=4.36.0; extra == "torch"
 Provides-Extra: audio
-Requires-Dist: torchaudio; extra == "audio"
 Requires-Dist: soundfile; extra == "audio"
 Provides-Extra: remote
 Requires-Dist: lz4; extra == "remote"
@@ -76,6 +75,7 @@ Requires-Dist: numba>=0.60.0; extra == "hf"
 Requires-Dist: datasets[vision]>=4.0.0; extra == "hf"
 Requires-Dist: datasets[audio]>=4.0.0; (sys_platform == "linux" or sys_platform == "darwin") and extra == "hf"
 Requires-Dist: fsspec>=2024.12.0; extra == "hf"
+Requires-Dist: torch<2.9.0; extra == "hf"
 Provides-Extra: video
 Requires-Dist: ffmpeg-python; extra == "video"
 Requires-Dist: imageio[ffmpeg,pyav]>=2.37.0; extra == "video"
@@ -117,6 +117,7 @@ Requires-Dist: huggingface_hub[hf_transfer]; extra == "examples"
 Requires-Dist: ultralytics; extra == "examples"
 Requires-Dist: open_clip_torch; extra == "examples"
 Requires-Dist: openai; extra == "examples"
+Requires-Dist: torchaudio<2.9.0; extra == "examples"
 Dynamic: license-file
 
 ================

{datachain-0.35.2 → datachain-0.36.1}/pyproject.toml

@@ -73,7 +73,6 @@ torch = [
   "transformers>=4.36.0"
 ]
 audio = [
-  "torchaudio",
   "soundfile"
 ]
 remote = [
@@ -88,7 +87,11 @@ hf = [
   "datasets[vision]>=4.0.0",
   # https://github.com/pytorch/torchcodec/issues/640
   "datasets[audio]>=4.0.0 ; (sys_platform == 'linux' or sys_platform == 'darwin')",
-  "fsspec>=2024.12.0"
+  "fsspec>=2024.12.0",
+  # Until datasets solve the issue, run test_hf_audio test to see if this can be removed
+  # https://github.com/meta-pytorch/torchcodec/issues/912
+  # https://github.com/huggingface/transformers/pull/41610
+  "torch<2.9.0"
 ]
 video = [
   "ffmpeg-python",
@@ -134,7 +137,9 @@ examples = [
   "huggingface_hub[hf_transfer]",
   "ultralytics",
   "open_clip_torch",
-  "openai"
+  "openai",
+  # Transformers still require it
+  "torchaudio<2.9.0"
 ]
 
 [project.urls]

{datachain-0.35.2 → datachain-0.36.1}/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.1/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.1}/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.1}/src/datachain/data_storage/schema.py

@@ -11,7 +11,6 @@ from datachain.sql.types import (
     JSON,
     Boolean,
     DateTime,
-    Int,
     Int64,
     SQLType,
     String,
@@ -269,7 +268,7 @@ class DataTable:
     @classmethod
     def sys_columns(cls):
         return [
-            sa.Column("sys__id", Int, primary_key=True),
+            sa.Column("sys__id", UInt64, primary_key=True),
             sa.Column(
                 "sys__rand", UInt64, nullable=False, server_default=f.abs(f.random())
             ),

{datachain-0.35.2 → datachain-0.36.1}/src/datachain/data_storage/sqlite.py

@@ -868,11 +868,8 @@ class SQLiteWarehouse(AbstractWarehouse):
                 if isinstance(c, BinaryExpression):
                     right_left_join = add_left_rows_filter(c)
 
-        # Use CTE instead of subquery to force SQLite to materialize the result
-        # This breaks deep nesting and prevents parser stack overflow.
         union_cte = sqlalchemy.union(left_right_join, right_left_join).cte()
-
-        return self._regenerate_system_columns(union_cte)
+        return sqlalchemy.select(*union_cte.c).select_from(union_cte)
 
     def _system_row_number_expr(self):
         return func.row_number().over()
@@ -884,11 +881,7 @@ class SQLiteWarehouse(AbstractWarehouse):
         """
         Create a temporary table from a query for use in a UDF.
         """
-        columns = [
-            sqlalchemy.Column(c.name, c.type)
-            for c in query.selected_columns
-            if c.name != "sys__id"
-        ]
+        columns = [sqlalchemy.Column(c.name, c.type) for c in query.selected_columns]
         table = self.create_udf_table(columns)
 
         with tqdm(desc="Preparing", unit=" rows", leave=False) as pbar:

{datachain-0.35.2 → datachain-0.36.1}/src/datachain/data_storage/warehouse.py

@@ -5,7 +5,7 @@ import random
 import string
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Generator, Iterable, Iterator, Sequence
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, Union, cast
 from urllib.parse import urlparse
 
 import attrs
@@ -23,7 +23,7 @@ from datachain.node import DirType, DirTypeGroup, Node, NodeWithPath, get_path
 from datachain.query.batch import RowsOutput
 from datachain.query.schema import ColumnMeta
 from datachain.sql.functions import path as pathfunc
-from datachain.sql.types import Int, SQLType
+from datachain.sql.types import SQLType
 from datachain.utils import sql_escape_like
 
 if TYPE_CHECKING:
@@ -32,6 +32,7 @@ if TYPE_CHECKING:
         _FromClauseArgument,
         _OnClauseArgument,
     )
+    from sqlalchemy.sql.selectable import FromClause
     from sqlalchemy.types import TypeEngine
 
     from datachain.data_storage import schema
@@ -248,45 +249,56 @@ class AbstractWarehouse(ABC, Serializable):
 
     def _regenerate_system_columns(
         self,
-        selectable: sa.Select | sa.CTE,
+        selectable: sa.Select,
         keep_existing_columns: bool = False,
+        regenerate_columns: Iterable[str] | None = None,
     ) -> sa.Select:
         """
-        Return a SELECT that regenerates sys__id and sys__rand deterministically.
+        Return a SELECT that regenerates system columns deterministically.
 
-        If keep_existing_columns is True, existing sys__id and sys__rand columns
-        will be kept as-is if they exist in the input selectable.
-        """
-        base = selectable.subquery() if hasattr(selectable, "subquery") else selectable
-
-        result_columns: dict[str, sa.ColumnElement] = {}
-        for col in base.c:
-            if col.name in result_columns:
-                raise ValueError(f"Duplicate column name {col.name} in SELECT")
-            if col.name in ("sys__id", "sys__rand"):
-                if keep_existing_columns:
-                    result_columns[col.name] = col
-            else:
-                result_columns[col.name] = col
+        If keep_existing_columns is True, existing system columns will be kept as-is
+        even when they are listed in ``regenerate_columns``.
 
-        system_types: dict[str, sa.types.TypeEngine] = {
+        Args:
+            selectable: Base SELECT
+            keep_existing_columns: When True, reuse existing system columns even if
+                they are part of the regeneration set.
+            regenerate_columns: Names of system columns to regenerate. Defaults to
+                {"sys__id", "sys__rand"}. Columns not listed are left untouched.
+        """
+        system_columns = {
             sys_col.name: sys_col.type
             for sys_col in self.schema.dataset_row_cls.sys_columns()
         }
+        regenerate = set(regenerate_columns or system_columns)
+        generators = {
+            "sys__id": self._system_row_number_expr,
+            "sys__rand": self._system_random_expr,
+        }
+
+        base = cast("FromClause", selectable.subquery())
+
+        def build(name: str) -> sa.ColumnElement:
+            expr = generators[name]()
+            return sa.cast(expr, system_columns[name]).label(name)
+
+        columns: list[sa.ColumnElement] = []
+        present: set[str] = set()
+        changed = False
+
+        for col in base.c:
+            present.add(col.name)
+            regen = col.name in regenerate and not keep_existing_columns
+            columns.append(build(col.name) if regen else col)
+            changed |= regen
+
+        for name in regenerate - present:
+            columns.append(build(name))
+            changed = True
+
+        if not changed:
+            return selectable
 
-        # Add missing system columns if needed
-        if "sys__id" not in result_columns:
-            expr = self._system_row_number_expr()
-            expr = sa.cast(expr, system_types["sys__id"])
-            result_columns["sys__id"] = expr.label("sys__id")
-        if "sys__rand" not in result_columns:
-            expr = self._system_random_expr()
-            expr = sa.cast(expr, system_types["sys__rand"])
-            result_columns["sys__rand"] = expr.label("sys__rand")
-
-        # Wrap in subquery to materialize window functions, then wrap again in SELECT
-        # This ensures window functions are computed before INSERT...FROM SELECT
-        columns = list(result_columns.values())
         inner = sa.select(*columns).select_from(base).subquery()
         return sa.select(*inner.c).select_from(inner)
 
@@ -950,10 +962,15 @@ class AbstractWarehouse(ABC, Serializable):
         SQLite TEMPORARY tables cannot be directly used as they are process-specific,
         and UDFs are run in other processes when run in parallel.
         """
+        columns = [
+            c
+            for c in columns
+            if c.name not in [col.name for col in self.dataset_row_cls.sys_columns()]
+        ]
         tbl = sa.Table(
             name or self.udf_table_name(),
             sa.MetaData(),
-            sa.Column("sys__id", Int, primary_key=True),
+            *self.dataset_row_cls.sys_columns(),
             *columns,
         )
         self.db.create_table(tbl, if_not_exists=True)