datachain 0.35.1__tar.gz → 0.36.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.35.1
+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.1 → 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
@@ -133,19 +134,26 @@ def shutdown_process(
             return proc.wait()
 
 
-def _process_stream(stream: "IO[bytes]", callback: Callable[[str], None]) -> None:
+def process_output(stream: IO[bytes], callback: Callable[[str], None]) -> None:
     buffer = b""
-    while byt := stream.read(1):  # Read one byte at a time
-        buffer += byt
 
-        if byt in (b"\n", b"\r"):  # Check for newline or carriage return
-            line = buffer.decode("utf-8")
-            callback(line)
-            buffer = b""  # Clear buffer for next line
+    try:
+        while byt := stream.read(1):  # Read one byte at a time
+            buffer += byt
 
-    if buffer:  # Handle any remaining data in the buffer
-        line = buffer.decode("utf-8")
-        callback(line)
+            if byt in (b"\n", b"\r"):  # Check for newline or carriage return
+                line = buffer.decode("utf-8", errors="replace")
+                callback(line)
+                buffer = b""  # Clear buffer for the next line
+
+        if buffer:  # Handle any remaining data in the buffer
+            line = buffer.decode("utf-8", errors="replace")
+            callback(line)
+    finally:
+        try:
+            stream.close()  # Ensure output is closed
+        except Exception:  # noqa: BLE001, S110
+            pass
 
 
 class DatasetRowsFetcher(NodesThreadPool):
@@ -1196,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,
@@ -1209,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,
@@ -1747,13 +1779,13 @@ class Catalog:
             recursive=recursive,
         )
 
+    @staticmethod
     def query(
-        self,
         query_script: str,
         env: Mapping[str, str] | None = None,
         python_executable: str = sys.executable,
-        capture_output: bool = False,
-        output_hook: Callable[[str], None] = noop,
+        stdout_callback: Callable[[str], None] | None = None,
+        stderr_callback: Callable[[str], None] | None = None,
         params: dict[str, str] | None = None,
         job_id: str | None = None,
         reset: bool = False,
@@ -1773,13 +1805,18 @@ class Catalog:
             },
         )
         popen_kwargs: dict[str, Any] = {}
-        if capture_output:
-            popen_kwargs = {"stdout": subprocess.PIPE, "stderr": subprocess.STDOUT}
+
+        if stdout_callback is not None:
+            popen_kwargs = {"stdout": subprocess.PIPE}
+        if stderr_callback is not None:
+            popen_kwargs["stderr"] = subprocess.PIPE
 
         def raise_termination_signal(sig: int, _: Any) -> NoReturn:
             raise TerminationSignal(sig)
 
-        thread: Thread | None = None
+        stdout_thread: Thread | None = None
+        stderr_thread: Thread | None = None
+
         with subprocess.Popen(cmd, env=env, **popen_kwargs) as proc:  # noqa: S603
             logger.info("Starting process %s", proc.pid)
 
@@ -1793,10 +1830,20 @@ class Catalog:
             orig_sigterm_handler = signal.getsignal(signal.SIGTERM)
             signal.signal(signal.SIGTERM, raise_termination_signal)
             try:
-                if capture_output:
-                    args = (proc.stdout, output_hook)
-                    thread = Thread(target=_process_stream, args=args, daemon=True)
-                    thread.start()
+                if stdout_callback is not None:
+                    stdout_thread = Thread(
+                        target=process_output,
+                        args=(proc.stdout, stdout_callback),
+                        daemon=True,
+                    )
+                    stdout_thread.start()
+                if stderr_callback is not None:
+                    stderr_thread = Thread(
+                        target=process_output,
+                        args=(proc.stderr, stderr_callback),
+                        daemon=True,
+                    )
+                    stderr_thread.start()
 
                 proc.wait()
             except TerminationSignal as exc:
@@ -1814,8 +1861,22 @@ class Catalog:
             finally:
                 signal.signal(signal.SIGTERM, orig_sigterm_handler)
                 signal.signal(signal.SIGINT, orig_sigint_handler)
-                if thread:
-                    thread.join()  # wait for the reader thread
+                # wait for the reader thread
+                thread_join_timeout_seconds = 30
+                if stdout_thread is not None:
+                    stdout_thread.join(timeout=thread_join_timeout_seconds)
+                    if stdout_thread.is_alive():
+                        logger.warning(
+                            "stdout thread is still alive after %s seconds",
+                            thread_join_timeout_seconds,
+                        )
+                if stderr_thread is not None:
+                    stderr_thread.join(timeout=thread_join_timeout_seconds)
+                    if stderr_thread.is_alive():
+                        logger.warning(
+                            "stderr thread is still alive after %s seconds",
+                            thread_join_timeout_seconds,
+                        )
 
         logger.info("Process %s exited with return code %s", proc.pid, proc.returncode)
         if proc.returncode in (

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.1 → 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.1 → datachain-0.36.0}/src/datachain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.35.1
+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.1 → 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

{datachain-0.35.1 → datachain-0.36.0}/tests/unit/test_query.py

@@ -53,12 +53,31 @@ def test_args(catalog, mock_popen, reset):
     mock_popen.assert_called_once_with(["mypython", "-c", "pass"], env=expected_env)
 
 
+def test_capture_stdout(catalog, mock_popen):
+    mock_popen.stdout = io.BytesIO(b"Hello, World!\rLorem Ipsum\nDolor Sit Amet\nconse")
+    stdout = []
+
+    catalog.query("pass", stdout_callback=stdout.append)
+    assert stdout == ["Hello, World!\r", "Lorem Ipsum\n", "Dolor Sit Amet\n", "conse"]
+
+
+def test_capture_stderr(catalog, mock_popen):
+    mock_popen.stderr = io.BytesIO(b"Hello, World!\rLorem Ipsum\nDolor Sit Amet\nconse")
+    stderr = []
+
+    catalog.query("pass", stderr_callback=stderr.append)
+    assert stderr == ["Hello, World!\r", "Lorem Ipsum\n", "Dolor Sit Amet\n", "conse"]
+
+
 def test_capture_output(catalog, mock_popen):
     mock_popen.stdout = io.BytesIO(b"Hello, World!\rLorem Ipsum\nDolor Sit Amet\nconse")
-    lines = []
+    mock_popen.stderr = io.BytesIO(b"foo\nbar")
+    stdout = []
+    stderr = []
 
-    catalog.query("pass", capture_output=True, output_hook=lines.append)
-    assert lines == ["Hello, World!\r", "Lorem Ipsum\n", "Dolor Sit Amet\n", "conse"]
+    catalog.query("pass", stdout_callback=stdout.append, stderr_callback=stderr.append)
+    assert stdout == ["Hello, World!\r", "Lorem Ipsum\n", "Dolor Sit Amet\n", "conse"]
+    assert stderr == ["foo\n", "bar"]
 
 
 def test_canceled_by_user(catalog, mock_popen):