datachain 0.14.2__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,12 +1,15 @@
 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:
     from datachain.catalog import Catalog
     from datachain.data_storage import AbstractMetastore, AbstractWarehouse
+    from datachain.query.udf import AbstractUDFDistributor
 
 METASTORE_SERIALIZED = "DATACHAIN__METASTORE"
 METASTORE_IMPORT_PATH = "DATACHAIN_METASTORE"
@@ -14,13 +17,16 @@ METASTORE_ARG_PREFIX = "DATACHAIN_METASTORE_ARG_"
 WAREHOUSE_SERIALIZED = "DATACHAIN__WAREHOUSE"
 WAREHOUSE_IMPORT_PATH = "DATACHAIN_WAREHOUSE"
 WAREHOUSE_ARG_PREFIX = "DATACHAIN_WAREHOUSE_ARG_"
+DISTRIBUTED_IMPORT_PYTHONPATH = "DATACHAIN_DISTRIBUTED_PYTHONPATH"
 DISTRIBUTED_IMPORT_PATH = "DATACHAIN_DISTRIBUTED"
-DISTRIBUTED_ARG_PREFIX = "DATACHAIN_DISTRIBUTED_ARG_"
+DISTRIBUTED_DISABLED = "DATACHAIN_DISTRIBUTED_DISABLED"
 
 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
 
@@ -61,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
 
@@ -100,31 +108,32 @@ def get_warehouse(in_memory: bool = False) -> "AbstractWarehouse":
     return warehouse_class(**warehouse_args)
 
 
-def get_distributed_class(**kwargs):
-    distributed_import_path = os.environ.get(DISTRIBUTED_IMPORT_PATH)
-    distributed_arg_envs = get_envs_by_prefix(DISTRIBUTED_ARG_PREFIX)
-    # Convert env variable names to keyword argument names by lowercasing them
-    distributed_args = {k.lower(): v for k, v in distributed_arg_envs.items()}
+def get_udf_distributor_class() -> type["AbstractUDFDistributor"] | None:
+    if os.environ.get(DISTRIBUTED_DISABLED) == "True":
+        return None
 
-    if not distributed_import_path:
-        raise RuntimeError(
-            f"{DISTRIBUTED_IMPORT_PATH} import path is required "
-            "for distributed UDF processing."
-        )
-    # Distributed class paths are specified as (for example):
-    # module.classname
+    if not (distributed_import_path := os.environ.get(DISTRIBUTED_IMPORT_PATH)):
+        return None
+
+    # Distributed class paths are specified as (for example): module.classname
     if "." not in distributed_import_path:
         raise RuntimeError(
             f"Invalid {DISTRIBUTED_IMPORT_PATH} import path: {distributed_import_path}"
         )
+
+    # Optional: set the Python path to look for the module
+    distributed_import_pythonpath = os.environ.get(DISTRIBUTED_IMPORT_PYTHONPATH)
+    if distributed_import_pythonpath and distributed_import_pythonpath not in sys.path:
+        sys.path.insert(0, distributed_import_pythonpath)
+
     module_name, _, class_name = distributed_import_path.rpartition(".")
     distributed = import_module(module_name)
-    distributed_class = getattr(distributed, class_name)
-    return distributed_class(**distributed_args | kwargs)
+    return getattr(distributed, class_name)
 
 
 def get_catalog(
-    client_config: Optional[dict[str, Any]] = None, in_memory: bool = False
+    client_config: dict[str, Any] | None = None,
+    in_memory: bool = False,
 ) -> "Catalog":
     """
     Function that creates Catalog instance with appropriate metastore
@@ -139,8 +148,9 @@ def get_catalog(
     """
     from datachain.catalog import Catalog
 
+    metastore = get_metastore(in_memory=in_memory)
     return Catalog(
-        metastore=get_metastore(in_memory=in_memory),
+        metastore=metastore,
         warehouse=get_warehouse(in_memory=in_memory),
         client_config=client_config,
         in_memory=in_memory,

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,7 +3,6 @@ import os
 import sys
 import traceback
 from multiprocessing import freeze_support
-from typing import Optional
 
 from datachain.cli.utils import get_logging_level
 
@@ -16,7 +15,6 @@ from .commands import (
     index,
     list_datasets,
     ls,
-    query,
     rm_dataset,
     show,
 )
@@ -25,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
@@ -34,8 +32,10 @@ def main(argv: Optional[list[str]] = None) -> int:
     datachain_parser = get_parser()
     args = datachain_parser.parse_args(argv)
 
-    if args.command in ("internal-run-udf", "internal-run-udf-worker"):
-        return handle_udf(args.command)
+    if args.command == "internal-run-udf":
+        return handle_udf()
+    if args.command == "internal-run-udf-worker":
+        return handle_udf_runner()
 
     if args.command is None:
         datachain_parser.print_help(sys.stderr)
@@ -59,16 +59,22 @@ 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)
     except BrokenPipeError as exc:
         error, return_code = handle_broken_pipe_error(exc)
         return return_code
-    except (KeyboardInterrupt, Exception) as exc:
+    except (KeyboardInterrupt, Exception) as exc:  # noqa: BLE001
         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)
@@ -89,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),
@@ -98,8 +103,10 @@ def handle_command(args, catalog, client_config) -> int:
 
     handler = command_handlers.get(args.command)
     if handler:
-        handler()
-        return 0
+        return_code = handler()
+        if return_code is None:
+            return 0
+        return return_code
     print(f"invalid command: {args.command}", file=sys.stderr)
     return 1
 
@@ -149,10 +156,7 @@ def handle_dataset_command(args, catalog):
             args.name,
             new_name=args.new_name,
             description=args.description,
-            labels=args.labels,
-            studio=args.studio,
-            local=args.local,
-            all=args.all,
+            attrs=args.attrs,
             team=args.team,
         ),
         "ls": lambda: list_datasets(
@@ -170,8 +174,6 @@ def handle_dataset_command(args, catalog):
             version=args.version,
             force=args.force,
             studio=args.studio,
-            local=args.local,
-            all=args.all,
             team=args.team,
         ),
         "remove": lambda: rm_dataset(
@@ -180,8 +182,6 @@ def handle_dataset_command(args, catalog):
             version=args.version,
             force=args.force,
             studio=args.studio,
-            local=args.local,
-            all=args.all,
             team=args.team,
         ),
     }
@@ -263,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
@@ -303,13 +294,13 @@ def handle_general_exception(exc, args, logging_level):
     return error, 1
 
 
-def handle_udf(command):
-    if command == "internal-run-udf":
-        from datachain.query.dispatch import udf_entrypoint
+def handle_udf() -> int:
+    from datachain.query.dispatch import udf_entrypoint
+
+    return udf_entrypoint()
 
-        return udf_entrypoint()
 
-    if command == "internal-run-udf-worker":
-        from datachain.query.dispatch import udf_worker_entrypoint
+def handle_udf_runner() -> int:
+    from datachain.query.dispatch import udf_worker_entrypoint
 
-        return udf_worker_entrypoint()
+    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,29 +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 DatasetNotFoundError
+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(
@@ -31,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:
@@ -94,23 +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.name, v.version)
+            yield d.full_name, v.version
 
 
-def list_datasets_local_versions(catalog: "Catalog", name: str):
-    ds = catalog.get_dataset(name)
+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,
     }
@@ -127,49 +147,60 @@ def _datasets_tabulate_row(name, both, local_version, studio_version):
 def rm_dataset(
     catalog: "Catalog",
     name: str,
-    version: Optional[int] = None,
-    force: Optional[bool] = False,
-    studio: bool = False,
-    local: bool = False,
-    all: bool = True,
-    team: Optional[str] = None,
-):
-    from datachain.studio import remove_studio_dataset
-
-    token = Config().read().get("studio", {}).get("token")
-    all, local, studio = determine_flavors(studio, local, all, token)
-
-    if all or local:
+    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:
+        # removing Studio dataset from CLI
+        from datachain.studio import remove_studio_dataset
+
+        if Config().read().get("studio", {}).get("token"):
+            remove_studio_dataset(
+                team, name, namespace_name, project_name, version, force
+            )
+        else:
+            raise DataChainError(
+                "Not logged in to Studio. Log in with 'datachain auth login'."
+            )
+    else:
         try:
-            catalog.remove_dataset(name, version=version, force=force)
+            project = catalog.metastore.get_project(project_name, namespace_name)
+            catalog.remove_dataset(name, project, version=version, force=force)
         except DatasetNotFoundError:
             print("Dataset not found in local", file=sys.stderr)
 
-    if (all or studio) and token:
-        remove_studio_dataset(team, name, version, force)
-
 
 def edit_dataset(
     catalog: "Catalog",
     name: str,
-    new_name: Optional[str] = None,
-    description: Optional[str] = None,
-    labels: Optional[list[str]] = None,
-    studio: bool = False,
-    local: bool = False,
-    all: bool = True,
-    team: Optional[str] = None,
-):
-    from datachain.studio import edit_studio_dataset
+    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
 
-    token = Config().read().get("studio", {}).get("token")
-    all, local, studio = determine_flavors(studio, local, all, token)
+    namespace_name, project_name, name = catalog.get_full_dataset_name(name)
 
-    if all or local:
+    if is_studio() or is_namespace_local(namespace_name):
         try:
-            catalog.edit_dataset(name, new_name, description, labels)
+            catalog.edit_dataset(
+                name, catalog.metastore.default_project, new_name, description, attrs
+            )
         except DatasetNotFoundError:
             print("Dataset not found in local", file=sys.stderr)
-
-    if (all or studio) and token:
-        edit_studio_dataset(team, name, new_name, description, labels)
+    else:
+        from datachain.studio import edit_studio_dataset
+
+        if Config().read().get("studio", {}).get("token"):
+            edit_studio_dataset(
+                team, name, namespace_name, project_name, new_name, description, attrs
+            )
+        else:
+            raise DataChainError(
+                "Not logged in to Studio. Log in with 'datachain auth login'."
+            )