datachain 0.30.5__py3-none-any.whl → 0.39.0__py3-none-any.whl

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/catalog/loader.py

@@ -1,8 +1,9 @@
 import os
 import sys
 from importlib import import_module
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
+from datachain.plugins import ensure_plugins_loaded
 from datachain.utils import get_envs_by_prefix
 
 if TYPE_CHECKING:
@@ -24,6 +25,8 @@ IN_MEMORY_ERROR_MESSAGE = "In-memory is only supported on SQLite"
 
 
 def get_metastore(in_memory: bool = False) -> "AbstractMetastore":
+    ensure_plugins_loaded()
+
     from datachain.data_storage import AbstractMetastore
     from datachain.data_storage.serializer import deserialize
 
@@ -64,6 +67,8 @@ def get_metastore(in_memory: bool = False) -> "AbstractMetastore":
 
 
 def get_warehouse(in_memory: bool = False) -> "AbstractWarehouse":
+    ensure_plugins_loaded()
+
     from datachain.data_storage import AbstractWarehouse
     from datachain.data_storage.serializer import deserialize
 
@@ -103,7 +108,7 @@ def get_warehouse(in_memory: bool = False) -> "AbstractWarehouse":
     return warehouse_class(**warehouse_args)
 
 
-def get_udf_distributor_class() -> Optional[type["AbstractUDFDistributor"]]:
+def get_udf_distributor_class() -> type["AbstractUDFDistributor"] | None:
     if os.environ.get(DISTRIBUTED_DISABLED) == "True":
         return None
 
@@ -127,7 +132,7 @@ def get_udf_distributor_class() -> Optional[type["AbstractUDFDistributor"]]:
 
 
 def get_catalog(
-    client_config: Optional[dict[str, Any]] = None,
+    client_config: dict[str, Any] | None = None,
     in_memory: bool = False,
 ) -> "Catalog":
     """

datachain/checkpoint.py

@@ -0,0 +1,43 @@
+import uuid
+from dataclasses import dataclass
+from datetime import datetime
+
+
+@dataclass
+class Checkpoint:
+    """
+    Represents a checkpoint within a job run.
+
+    A checkpoint marks a successfully completed stage of execution. In the event
+    of a failure, the job can resume from the most recent checkpoint rather than
+    starting over from the beginning.
+
+    Checkpoints can also be created in a "partial" mode, which indicates that the
+    work at this stage was only partially completed. For example, if a failure
+    occurs halfway through running a UDF, already computed results can still be
+    saved, allowing the job to resume from that partially completed state on
+    restart.
+    """
+
+    id: str
+    job_id: str
+    hash: str
+    partial: bool
+    created_at: datetime
+
+    @classmethod
+    def parse(
+        cls,
+        id: str | uuid.UUID,
+        job_id: str,
+        _hash: str,
+        partial: bool,
+        created_at: datetime,
+    ) -> "Checkpoint":
+        return cls(
+            str(id),
+            job_id,
+            _hash,
+            bool(partial),
+            created_at,
+        )

datachain/cli/__init__.py

@@ -3,10 +3,8 @@ import os
 import sys
 import traceback
 from multiprocessing import freeze_support
-from typing import Optional
 
 from datachain.cli.utils import get_logging_level
-from datachain.error import DataChainError as DataChainError
 
 from .commands import (
     clear_cache,
@@ -17,7 +15,6 @@ from .commands import (
     index,
     list_datasets,
     ls,
-    query,
     rm_dataset,
     show,
 )
@@ -26,7 +23,7 @@ from .parser import get_parser
 logger = logging.getLogger("datachain")
 
 
-def main(argv: Optional[list[str]] = None) -> int:
+def main(argv: list[str] | None = None) -> int:
     from datachain.catalog import get_catalog
 
     # Required for Windows multiprocessing support
@@ -38,7 +35,7 @@ def main(argv: Optional[list[str]] = None) -> int:
     if args.command == "internal-run-udf":
         return handle_udf()
     if args.command == "internal-run-udf-worker":
-        return handle_udf_runner(args.fd)
+        return handle_udf_runner()
 
     if args.command is None:
         datachain_parser.print_help(sys.stderr)
@@ -62,6 +59,7 @@ def main(argv: Optional[list[str]] = None) -> int:
 
     error = None
 
+    catalog = None
     try:
         catalog = get_catalog(client_config=client_config)
         return handle_command(args, catalog, client_config)
@@ -72,6 +70,11 @@ def main(argv: Optional[list[str]] = None) -> int:
         error, return_code = handle_general_exception(exc, args, logging_level)
         return return_code
     finally:
+        if catalog is not None:
+            try:
+                catalog.close()
+            except Exception:
+                logger.exception("Failed to close catalog")
         from datachain.telemetry import telemetry
 
         telemetry.send_cli_call(args.command, error=error)
@@ -92,7 +95,6 @@ def handle_command(args, catalog, client_config) -> int:
         "find": lambda: handle_find_command(args, catalog),
         "index": lambda: handle_index_command(args, catalog),
         "completion": lambda: handle_completion_command(args),
-        "query": lambda: handle_query_command(args, catalog),
         "clear-cache": lambda: clear_cache(catalog),
         "gc": lambda: garbage_collect(catalog),
         "auth": lambda: process_auth_cli_args(args),
@@ -261,15 +263,6 @@ def handle_completion_command(args):
     print(completion(args.shell))
 
 
-def handle_query_command(args, catalog):
-    query(
-        catalog,
-        args.script,
-        parallel=args.parallel,
-        params=args.param,
-    )
-
-
 def handle_broken_pipe_error(exc):
     # Python flushes standard streams on exit; redirect remaining output
     # to devnull to avoid another BrokenPipeError at shutdown
@@ -307,7 +300,7 @@ def handle_udf() -> int:
     return udf_entrypoint()
 
 
-def handle_udf_runner(fd: Optional[int] = None) -> int:
+def handle_udf_runner() -> int:
     from datachain.query.dispatch import udf_worker_entrypoint
 
-    return udf_worker_entrypoint(fd)
+    return udf_worker_entrypoint()

datachain/cli/commands/__init__.py

@@ -1,14 +1,8 @@
-from .datasets import (
-    edit_dataset,
-    list_datasets,
-    list_datasets_local,
-    rm_dataset,
-)
+from .datasets import edit_dataset, list_datasets, list_datasets_local, rm_dataset
 from .du import du
 from .index import index
 from .ls import ls
 from .misc import clear_cache, completion, garbage_collect
-from .query import query
 from .show import show
 
 __all__ = [
@@ -21,7 +15,6 @@ __all__ = [
     "list_datasets",
     "list_datasets_local",
     "ls",
-    "query",
     "rm_dataset",
     "show",
 ]

datachain/cli/commands/datasets.py

@@ -1,30 +1,41 @@
 import sys
-from typing import TYPE_CHECKING, Optional
+from collections.abc import Iterable, Iterator
+from typing import TYPE_CHECKING
 
 from tabulate import tabulate
 
-if TYPE_CHECKING:
-    from datachain.catalog import Catalog
-
+from datachain import semver
 from datachain.catalog import is_namespace_local
 from datachain.cli.utils import determine_flavors
 from datachain.config import Config
 from datachain.error import DataChainError, DatasetNotFoundError
 from datachain.studio import list_datasets as list_datasets_studio
 
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+
+def group_dataset_versions(
+    datasets: Iterable[tuple[str, str]], latest_only=True
+) -> dict[str, str | list[str]]:
+    grouped: dict[str, list[tuple[int, int, int]]] = {}
 
-def group_dataset_versions(datasets, latest_only=True):
-    grouped = {}
     # Sort to ensure groupby works as expected
     # (groupby expects consecutive items with the same key)
     for name, version in sorted(datasets):
-        grouped.setdefault(name, []).append(version)
+        grouped.setdefault(name, []).append(semver.parse(version))
 
     if latest_only:
         # For each dataset name, pick the highest version.
-        return {name: max(versions) for name, versions in grouped.items()}
+        return {
+            name: semver.create(*(max(versions))) for name, versions in grouped.items()
+        }
+
     # For each dataset name, return a sorted list of unique versions.
-    return {name: sorted(set(versions)) for name, versions in grouped.items()}
+    return {
+        name: [semver.create(*v) for v in sorted(set(versions))]
+        for name, versions in grouped.items()
+    }
 
 
 def list_datasets(
@@ -32,10 +43,10 @@ def list_datasets(
     studio: bool = False,
     local: bool = False,
     all: bool = True,
-    team: Optional[str] = None,
+    team: str | None = None,
     latest_only: bool = True,
-    name: Optional[str] = None,
-):
+    name: str | None = None,
+) -> None:
     token = Config().read().get("studio", {}).get("token")
     all, local, studio = determine_flavors(studio, local, all, token)
     if name:
@@ -95,27 +106,31 @@ def list_datasets(
     print(tabulate(rows, headers="keys"))
 
 
-def list_datasets_local(catalog: "Catalog", name: Optional[str] = None):
+def list_datasets_local(
+    catalog: "Catalog", name: str | None = None
+) -> Iterator[tuple[str, str]]:
     if name:
         yield from list_datasets_local_versions(catalog, name)
         return
 
     for d in catalog.ls_datasets():
         for v in d.versions:
-            yield (d.full_name, v.version)
+            yield d.full_name, v.version
 
 
-def list_datasets_local_versions(catalog: "Catalog", name: str):
+def list_datasets_local_versions(
+    catalog: "Catalog", name: str
+) -> Iterator[tuple[str, str]]:
     namespace_name, project_name, name = catalog.get_full_dataset_name(name)
 
     ds = catalog.get_dataset(
         name, namespace_name=namespace_name, project_name=project_name
     )
     for v in ds.versions:
-        yield (name, v.version)
+        yield name, v.version
 
 
-def _datasets_tabulate_row(name, both, local_version, studio_version):
+def _datasets_tabulate_row(name, both, local_version, studio_version) -> dict[str, str]:
     row = {
         "Name": name,
     }
@@ -132,11 +147,11 @@ def _datasets_tabulate_row(name, both, local_version, studio_version):
 def rm_dataset(
     catalog: "Catalog",
     name: str,
-    version: Optional[str] = None,
-    force: Optional[bool] = False,
-    studio: Optional[bool] = False,
-    team: Optional[str] = None,
-):
+    version: str | None = None,
+    force: bool | None = False,
+    studio: bool | None = False,
+    team: str | None = None,
+) -> None:
     namespace_name, project_name, name = catalog.get_full_dataset_name(name)
 
     if studio:
@@ -162,11 +177,11 @@ def rm_dataset(
 def edit_dataset(
     catalog: "Catalog",
     name: str,
-    new_name: Optional[str] = None,
-    description: Optional[str] = None,
-    attrs: Optional[list[str]] = None,
-    team: Optional[str] = None,
-):
+    new_name: str | None = None,
+    description: str | None = None,
+    attrs: list[str] | None = None,
+    team: str | None = None,
+) -> None:
     from datachain.lib.dc.utils import is_studio
 
     namespace_name, project_name, name = catalog.get_full_dataset_name(name)

datachain/cli/commands/ls.py

@@ -1,13 +1,14 @@
 import shlex
 from collections.abc import Iterable, Iterator
 from itertools import chain
-from typing import TYPE_CHECKING, Optional
-
-if TYPE_CHECKING:
-    from datachain.catalog import Catalog
+from typing import TYPE_CHECKING
 
 from datachain.cli.utils import determine_flavors
 from datachain.config import Config
+from datachain.query.session import Session
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
 
 
 def ls(
@@ -16,7 +17,7 @@ def ls(
     studio: bool = False,
     local: bool = False,
     all: bool = True,
-    team: Optional[str] = None,
+    team: str | None = None,
     **kwargs,
 ):
     token = Config().read().get("studio", {}).get("token")
@@ -32,18 +33,15 @@ def ls(
 def ls_local(
     sources,
     long: bool = False,
-    catalog: Optional["Catalog"] = None,
+    catalog=None,
     client_config=None,
     **kwargs,
 ):
     from datachain import listings
 
     if sources:
-        if catalog is None:
-            from datachain.catalog import get_catalog
-
-            catalog = get_catalog(client_config=client_config)
-
+        session = Session.get(catalog=catalog, client_config=client_config)
+        catalog = session.catalog
         actual_sources = list(ls_urls(sources, catalog=catalog, long=long, **kwargs))
         if len(actual_sources) == 1:
             for _, entries in actual_sources:
@@ -78,7 +76,7 @@ def format_ls_entry(entry: str) -> str:
 def ls_remote(
     paths: Iterable[str],
     long: bool = False,
-    team: Optional[str] = None,
+    team: str | None = None,
 ):
     from datachain.node import long_line_str
     from datachain.remote.studio import StudioClient
@@ -145,7 +143,7 @@ def _ls_urls_flat(
     long: bool,
     catalog: "Catalog",
     **kwargs,
-) -> Iterator[tuple[str, Iterator[str]]]:
+) -> Iterator[tuple[str, Iterable[str]]]:
     from datachain.client import Client
     from datachain.node import long_line_str
 
@@ -154,7 +152,9 @@ def _ls_urls_flat(
         if client_cls.is_root_url(source):
             buckets = client_cls.ls_buckets(**catalog.client_config)
             if long:
-                values = (long_line_str(b.name, b.created) for b in buckets)
+                values: Iterable[str] = (
+                    long_line_str(b.name, b.created) for b in buckets
+                )
             else:
                 values = (b.name for b in buckets)
             yield source, values
@@ -164,7 +164,7 @@ def _ls_urls_flat(
             if long:
                 fields.append("last_modified")
             for data_source, results in catalog.ls([source], fields=fields, **kwargs):
-                values = (_node_data_to_ls_values(r, long) for r in results)
+                values = [_node_data_to_ls_values(r, long) for r in results]
                 found = True
                 yield data_source.dirname(), values
             if not found:

datachain/cli/commands/show.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 from datachain.lib.signal_schema import SignalSchema
 
@@ -10,7 +10,7 @@ if TYPE_CHECKING:
 def show(
     catalog: "Catalog",
     name: str,
-    version: Optional[str] = None,
+    version: str | None = None,
     limit: int = 10,
     offset: int = 0,
     columns: Sequence[str] = (),

datachain/cli/parser/__init__.py

@@ -3,7 +3,7 @@ from importlib.metadata import PackageNotFoundError, version
 
 import shtab
 
-from datachain.cli.utils import BooleanOptionalAction, KeyValueArgs
+from datachain.cli.utils import BooleanOptionalAction
 
 from .job import add_jobs_parser
 from .studio import add_auth_parser
@@ -16,9 +16,7 @@ from .utils import (
     add_update_arg,
     find_columns_type,
 )
-from .utils import (
-    CustomArgumentParser as ArgumentParser,
-)
+from .utils import CustomArgumentParser as ArgumentParser
 
 
 def get_parser() -> ArgumentParser:  # noqa: PLR0915
@@ -467,37 +465,6 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     show_parser.add_argument("--schema", action="store_true", help="Show schema")
     add_show_args(show_parser)
 
-    query_parser = subp.add_parser(
-        "query",
-        parents=[parent_parser],
-        description="Create a new dataset with a query script.",
-        formatter_class=CustomHelpFormatter,
-    )
-    add_anon_arg(query_parser)
-    query_parser.add_argument(
-        "script", metavar="<script.py>", type=str, help="Filepath for script"
-    )
-    query_parser.add_argument(
-        "--parallel",
-        nargs="?",
-        type=int,
-        const=-1,
-        default=None,
-        metavar="N",
-        help=(
-            "Use multiprocessing to run any query script UDFs with N worker processes. "
-            "N defaults to the CPU count"
-        ),
-    )
-    query_parser.add_argument(
-        "-p",
-        "--param",
-        metavar="param=value",
-        nargs=1,
-        action=KeyValueArgs,
-        help="Query parameters",
-    )
-
     parse_clear_cache = subp.add_parser(
         "clear-cache",
         parents=[parent_parser],
@@ -515,14 +482,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     add_anon_arg(parse_gc)
 
     subp.add_parser("internal-run-udf", parents=[parent_parser])
-    run_udf_worker = subp.add_parser("internal-run-udf-worker", parents=[parent_parser])
-    run_udf_worker.add_argument(
-        "--fd",
-        type=int,
-        action="store",
-        default=None,
-        help="File descriptor to write results to",
-    )
+    subp.add_parser("internal-run-udf-worker", parents=[parent_parser])
 
     add_completion_parser(subp, [parent_parser])
     return parser

datachain/cli/parser/job.py

@@ -83,7 +83,7 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
     studio_run_parser.add_argument(
         "--python-version",
         action="store",
-        help="Python version for the job (e.g., 3.9, 3.10, 3.11)",
+        help="Python version for the job (e.g., 3.10, 3.11, 3.12, 3.13)",
     )
     studio_run_parser.add_argument(
         "--repository",