datachain 0.2.10__py3-none-any.whl → 0.2.12__py3-none-any.whl

datachain/__init__.py

@@ -1,4 +1,4 @@
-from datachain.lib.data_model import DataModel, DataType, FileBasic, is_chain_type
+from datachain.lib.data_model import DataModel, DataType, is_chain_type
 from datachain.lib.dc import C, Column, DataChain, Sys
 from datachain.lib.file import (
     File,
@@ -8,15 +8,14 @@ from datachain.lib.file import (
     TarVFile,
     TextFile,
 )
+from datachain.lib.model_store import ModelStore
 from datachain.lib.udf import Aggregator, Generator, Mapper
 from datachain.lib.utils import AbstractUDF, DataChainError
-from datachain.query.dataset import UDF as BaseUDF  # noqa: N811
 from datachain.query.session import Session
 
 __all__ = [
     "AbstractUDF",
     "Aggregator",
-    "BaseUDF",
     "C",
     "Column",
     "DataChain",
@@ -24,12 +23,12 @@ __all__ = [
     "DataModel",
     "DataType",
     "File",
-    "FileBasic",
     "FileError",
     "Generator",
     "ImageFile",
     "IndexedFile",
     "Mapper",
+    "ModelStore",
     "Session",
     "Sys",
     "TarVFile",

datachain/cache.py

@@ -1,6 +1,7 @@
 import hashlib
 import json
 import os
+from datetime import datetime
 from functools import partial
 from typing import TYPE_CHECKING, Optional
 
@@ -9,6 +10,8 @@ from dvc_data.hashfile.db.local import LocalHashFileDB
 from dvc_objects.fs.local import LocalFileSystem
 from fsspec.callbacks import Callback, TqdmCallback
 
+from datachain.utils import TIME_ZERO
+
 from .progress import Tqdm
 
 if TYPE_CHECKING:
@@ -23,10 +26,13 @@ class UniqueId:
     storage: "StorageURI"
     parent: str
     name: str
-    etag: str
     size: int
-    vtype: str
-    location: Optional[str]
+    etag: str
+    version: str = ""
+    is_latest: bool = True
+    vtype: str = ""
+    location: Optional[str] = None
+    last_modified: datetime = TIME_ZERO
 
     @property
     def path(self) -> str:
@@ -49,7 +55,7 @@ class UniqueId:
     def get_hash(self) -> str:
         etag = f"{self.vtype}{self.location}" if self.vtype else self.etag
         return sha256(
-            f"{self.storage}/{self.parent}/{self.name}/{etag}".encode()
+            f"{self.storage}/{self.parent}/{self.name}/{self.version}/{etag}".encode()
         ).hexdigest()
 
 

datachain/catalog/catalog.py

@@ -84,12 +84,14 @@ if TYPE_CHECKING:
         AbstractMetastore,
         AbstractWarehouse,
     )
+    from datachain.dataset import DatasetVersion
+    from datachain.job import Job
 
 logger = logging.getLogger("datachain")
 
 DEFAULT_DATASET_DIR = "dataset"
 DATASET_FILE_SUFFIX = ".edatachain"
-FEATURE_CLASSES = ["Feature"]
+FEATURE_CLASSES = ["DataModel"]
 
 TTL_INT = 4 * 60 * 60
 
@@ -948,13 +950,9 @@ class Catalog:
                     ms = self.metastore.clone(uri, None)
                     st = self.warehouse.clone()
                     listing = Listing(None, ms, st, client, None)
-                    rows = (
-                        DatasetQuery(
-                            name=dataset.name, version=ds_version, catalog=self
-                        )
-                        .select()
-                        .to_records()
-                    )
+                    rows = DatasetQuery(
+                        name=dataset.name, version=ds_version, catalog=self
+                    ).to_db_records()
                     indexed_sources.append(
                         (
                             listing,
@@ -1160,9 +1158,8 @@ class Catalog:
         if not dataset_version.preview:
             values["preview"] = (
                 DatasetQuery(name=dataset.name, version=version, catalog=self)
-                .select()
                 .limit(20)
-                .to_records()
+                .to_db_records()
             )
 
         if not values:
@@ -1420,6 +1417,25 @@ class Catalog:
             if not d.is_bucket_listing:
                 yield d
 
+    def list_datasets_versions(
+        self,
+    ) -> Iterator[tuple[DatasetRecord, "DatasetVersion", Optional["Job"]]]:
+        """Iterate over all dataset versions with related jobs."""
+        datasets = list(self.ls_datasets())
+
+        # preselect dataset versions jobs from db to avoid multiple queries
+        jobs_ids: set[str] = {
+            v.job_id for ds in datasets for v in ds.versions if v.job_id
+        }
+        jobs: dict[str, Job] = {}
+        if jobs_ids:
+            jobs = {j.id: j for j in self.metastore.list_jobs_by_ids(list(jobs_ids))}
+
+        for d in datasets:
+            yield from (
+                (d, v, jobs.get(v.job_id) if v.job_id else None) for v in d.versions
+            )
+
     def ls_dataset_rows(
         self, name: str, version: int, offset=None, limit=None
     ) -> list[dict]:
@@ -1427,7 +1443,7 @@ class Catalog:
 
         dataset = self.get_dataset(name)
 
-        q = DatasetQuery(name=dataset.name, version=version, catalog=self).select()
+        q = DatasetQuery(name=dataset.name, version=version, catalog=self)
         if limit:
             q = q.limit(limit)
         if offset:
@@ -1435,7 +1451,7 @@ class Catalog:
 
         q = q.order_by("sys__id")
 
-        return q.to_records()
+        return q.to_db_records()
 
     def signed_url(self, source: str, path: str, client_config=None) -> str:
         client_config = client_config or self.client_config
@@ -1609,6 +1625,7 @@ class Catalog:
                 ...
             }
         """
+        from datachain.lib.file import File
         from datachain.lib.signal_schema import DEFAULT_DELIMITER, SignalSchema
 
         version = self.get_dataset(dataset_name).get_version(dataset_version)
@@ -1616,7 +1633,7 @@ class Catalog:
         file_signals_values = {}
 
         schema = SignalSchema.deserialize(version.feature_schema)
-        for file_signals in schema.get_file_signals():
+        for file_signals in schema.get_signals(File):
             prefix = file_signals.replace(".", DEFAULT_DELIMITER) + DEFAULT_DELIMITER
             file_signals_values[file_signals] = {
                 c_name.removeprefix(prefix): c_value
@@ -1657,10 +1674,13 @@ class Catalog:
             row["source"],
             row["parent"],
             row["name"],
-            row["etag"],
             row["size"],
+            row["etag"],
+            row["version"],
+            row["is_latest"],
             row["vtype"],
             row["location"],
+            row["last_modified"],
         )
 
     def ls(
@@ -1992,7 +2012,7 @@ class Catalog:
                 )
             if proc.returncode == QUERY_SCRIPT_INVALID_LAST_STATEMENT_EXIT_CODE:
                 raise QueryScriptRunError(
-                    "Last line in a script was not an instance of DatasetQuery",
+                    "Last line in a script was not an instance of DataChain",
                     return_code=proc.returncode,
                     output=output,
                 )

datachain/cli.py

@@ -3,7 +3,7 @@ import os
 import shlex
 import sys
 import traceback
-from argparse import SUPPRESS, Action, ArgumentParser, ArgumentTypeError, Namespace
+from argparse import Action, ArgumentParser, ArgumentTypeError, Namespace
 from collections.abc import Iterable, Iterator, Mapping, Sequence
 from importlib.metadata import PackageNotFoundError, version
 from itertools import chain
@@ -106,10 +106,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     parser = ArgumentParser(
         description="DataChain: Wrangle unstructured AI data at scale", prog="datachain"
     )
-
     parser.add_argument("-V", "--version", action="version", version=__version__)
-    parser.add_argument("--internal-run-udf", action="store_true", help=SUPPRESS)
-    parser.add_argument("--internal-run-udf-worker", action="store_true", help=SUPPRESS)
 
     parent_parser = ArgumentParser(add_help=False)
     parent_parser.add_argument(
@@ -150,9 +147,15 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help="Drop into the pdb debugger on fatal exception",
     )
 
-    subp = parser.add_subparsers(help="Sub-command help", dest="command")
+    subp = parser.add_subparsers(
+        title="Available Commands",
+        metavar="command",
+        dest="command",
+        help=f"Use `{parser.prog} command --help` for command-specific help.",
+        required=True,
+    )
     parse_cp = subp.add_parser(
-        "cp", parents=[parent_parser], help="Copy data files from the cloud"
+        "cp", parents=[parent_parser], description="Copy data files from the cloud"
     )
     add_sources_arg(parse_cp).complete = shtab.DIR  # type: ignore[attr-defined]
     parse_cp.add_argument("output", type=str, help="Output")
@@ -179,7 +182,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     parse_clone = subp.add_parser(
-        "clone", parents=[parent_parser], help="Copy data files from the cloud"
+        "clone", parents=[parent_parser], description="Copy data files from the cloud"
     )
     add_sources_arg(parse_clone).complete = shtab.DIR  # type: ignore[attr-defined]
     parse_clone.add_argument("output", type=str, help="Output")
@@ -222,7 +225,9 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     parse_pull = subp.add_parser(
-        "pull", parents=[parent_parser], help="Pull specific dataset version from SaaS"
+        "pull",
+        parents=[parent_parser],
+        description="Pull specific dataset version from SaaS",
     )
     parse_pull.add_argument(
         "dataset",
@@ -263,7 +268,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     parse_edit_dataset = subp.add_parser(
-        "edit-dataset", parents=[parent_parser], help="Edit dataset metadata"
+        "edit-dataset", parents=[parent_parser], description="Edit dataset metadata"
     )
     parse_edit_dataset.add_argument("name", type=str, help="Dataset name")
     parse_edit_dataset.add_argument(
@@ -285,9 +290,9 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help="Dataset labels",
     )
 
-    subp.add_parser("ls-datasets", parents=[parent_parser], help="List datasets")
+    subp.add_parser("ls-datasets", parents=[parent_parser], description="List datasets")
     rm_dataset_parser = subp.add_parser(
-        "rm-dataset", parents=[parent_parser], help="Removes dataset"
+        "rm-dataset", parents=[parent_parser], description="Removes dataset"
     )
     rm_dataset_parser.add_argument("name", type=str, help="Dataset name")
     rm_dataset_parser.add_argument(
@@ -305,7 +310,9 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     dataset_stats_parser = subp.add_parser(
-        "dataset-stats", parents=[parent_parser], help="Shows basic dataset stats"
+        "dataset-stats",
+        parents=[parent_parser],
+        description="Shows basic dataset stats",
     )
     dataset_stats_parser.add_argument("name", type=str, help="Dataset name")
     dataset_stats_parser.add_argument(
@@ -330,7 +337,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     parse_merge_datasets = subp.add_parser(
-        "merge-datasets", parents=[parent_parser], help="Merges datasets"
+        "merge-datasets", parents=[parent_parser], description="Merges datasets"
     )
     parse_merge_datasets.add_argument(
         "--src",
@@ -360,7 +367,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     parse_ls = subp.add_parser(
-        "ls", parents=[parent_parser], help="List storage contents"
+        "ls", parents=[parent_parser], description="List storage contents"
     )
     add_sources_arg(parse_ls, nargs="*")
     parse_ls.add_argument(
@@ -378,7 +385,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     parse_du = subp.add_parser(
-        "du", parents=[parent_parser], help="Display space usage"
+        "du", parents=[parent_parser], description="Display space usage"
     )
     add_sources_arg(parse_du)
     parse_du.add_argument(
@@ -408,7 +415,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     parse_find = subp.add_parser(
-        "find", parents=[parent_parser], help="Search in a directory hierarchy"
+        "find", parents=[parent_parser], description="Search in a directory hierarchy"
     )
     add_sources_arg(parse_find)
     parse_find.add_argument(
@@ -461,20 +468,20 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     parse_index = subp.add_parser(
-        "index", parents=[parent_parser], help="Index storage location"
+        "index", parents=[parent_parser], description="Index storage location"
     )
     add_sources_arg(parse_index)
 
     subp.add_parser(
         "find-stale-storages",
         parents=[parent_parser],
-        help="Finds and marks stale storages",
+        description="Finds and marks stale storages",
     )
 
     show_parser = subp.add_parser(
         "show",
         parents=[parent_parser],
-        help="Create a new dataset with a query script",
+        description="Create a new dataset with a query script",
     )
     show_parser.add_argument("name", type=str, help="Dataset name")
     show_parser.add_argument(
@@ -489,7 +496,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     query_parser = subp.add_parser(
         "query",
         parents=[parent_parser],
-        help="Create a new dataset with a query script",
+        description="Create a new dataset with a query script",
     )
     query_parser.add_argument(
         "script", metavar="<script.py>", type=str, help="Filepath for script"
@@ -520,7 +527,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     )
 
     apply_udf_parser = subp.add_parser(
-        "apply-udf", parents=[parent_parser], help="Apply UDF"
+        "apply-udf", parents=[parent_parser], description="Apply UDF"
     )
     apply_udf_parser.add_argument("udf", type=str, help="UDF location")
     apply_udf_parser.add_argument("source", type=str, help="Source storage or dataset")
@@ -541,12 +548,14 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         "--udf-params", type=str, default=None, help="UDF class parameters"
     )
     subp.add_parser(
-        "clear-cache", parents=[parent_parser], help="Clear the local file cache"
+        "clear-cache", parents=[parent_parser], description="Clear the local file cache"
     )
     subp.add_parser(
-        "gc", parents=[parent_parser], help="Garbage collect temporary tables"
+        "gc", parents=[parent_parser], description="Garbage collect temporary tables"
     )
 
+    subp.add_parser("internal-run-udf", parents=[parent_parser])
+    subp.add_parser("internal-run-udf-worker", parents=[parent_parser])
     add_completion_parser(subp, [parent_parser])
     return parser
 
@@ -555,7 +564,7 @@ def add_completion_parser(subparsers, parents):
     parser = subparsers.add_parser(
         "completion",
         parents=parents,
-        help="Output shell completion script",
+        description="Output shell completion script",
     )
     parser.add_argument(
         "-s",
@@ -817,7 +826,7 @@ def show(
         .limit(limit)
         .offset(offset)
     )
-    records = query.to_records()
+    records = query.to_db_records()
     show_records(records, collapse_columns=not no_collapse)
 
 
@@ -901,27 +910,23 @@ def completion(shell: str) -> str:
     )
 
 
-def main(argv: Optional[list[str]] = None) -> int:  # noqa: C901, PLR0911, PLR0912, PLR0915
+def main(argv: Optional[list[str]] = None) -> int:  # noqa: C901, PLR0912, PLR0915
     # Required for Windows multiprocessing support
     freeze_support()
 
     parser = get_parser()
     args = parser.parse_args(argv)
 
-    if args.internal_run_udf:
+    if args.command == "internal-run-udf":
         from datachain.query.dispatch import udf_entrypoint
 
         return udf_entrypoint()
 
-    if args.internal_run_udf_worker:
+    if args.command == "internal-run-udf-worker":
         from datachain.query.dispatch import udf_worker_entrypoint
 
         return udf_worker_entrypoint()
 
-    if args.command is None:
-        parser.print_help()
-        return 1
-
     from .catalog import get_catalog
 
     logger.addHandler(logging.StreamHandler())

datachain/data_storage/metastore.py

@@ -40,6 +40,7 @@ from datachain.error import (
     StorageNotFoundError,
     TableMissingError,
 )
+from datachain.job import Job
 from datachain.storage import Storage, StorageStatus, StorageURI
 from datachain.utils import JSONSerialize, is_expired
 
@@ -67,6 +68,7 @@ class AbstractMetastore(ABC, Serializable):
     storage_class: type[Storage] = Storage
     dataset_class: type[DatasetRecord] = DatasetRecord
     dependency_class: type[DatasetDependency] = DatasetDependency
+    job_class: type[Job] = Job
 
     def __init__(
         self,
@@ -377,6 +379,9 @@ class AbstractMetastore(ABC, Serializable):
     # Jobs
     #
 
+    def list_jobs_by_ids(self, ids: list[str], conn=None) -> Iterator["Job"]:
+        raise NotImplementedError
+
     @abstractmethod
     def create_job(
         self,
@@ -1467,6 +1472,10 @@ class AbstractDBMetastore(AbstractMetastore):
             Column("metrics", JSON, nullable=False),
         ]
 
+    @cached_property
+    def _job_fields(self) -> list[str]:
+        return [c.name for c in self._jobs_columns() if c.name]  # type: ignore[attr-defined]
+
     @cached_property
     def _jobs(self) -> "Table":
         return Table(self.JOBS_TABLE, self.db.metadata, *self._jobs_columns())
@@ -1484,6 +1493,21 @@ class AbstractDBMetastore(AbstractMetastore):
             return self._jobs.update()
         return self._jobs.update().where(*where)
 
+    def _parse_job(self, rows) -> Job:
+        return Job.parse(*rows)
+
+    def _parse_jobs(self, rows) -> Iterator["Job"]:
+        for _, g in groupby(rows, lambda r: r[0]):
+            yield self._parse_job(*list(g))
+
+    def _jobs_query(self):
+        return self._jobs_select(*[getattr(self._jobs.c, f) for f in self._job_fields])
+
+    def list_jobs_by_ids(self, ids: list[str], conn=None) -> Iterator["Job"]:
+        """List jobs by ids."""
+        query = self._jobs_query().where(self._jobs.c.id.in_(ids))
+        yield from self._parse_jobs(self.db.execute(query, conn=conn))
+
     def create_job(
         self,
         name: str,

datachain/data_storage/warehouse.py

@@ -390,7 +390,9 @@ class AbstractWarehouse(ABC, Serializable):
         expressions: tuple[_ColumnsClauseArgument[Any], ...] = (
             sa.func.count(table.c.sys__id),
         )
-        if "size" in table.columns:
+        if "file__size" in table.columns:
+            expressions = (*expressions, sa.func.sum(table.c.file__size))
+        elif "size" in table.columns:
             expressions = (*expressions, sa.func.sum(table.c.size))
         query = select(*expressions)
         ((nrows, *rest),) = self.db.execute(query)

datachain/job.py

@@ -0,0 +1,56 @@
+import json
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Optional, TypeVar
+
+J = TypeVar("J", bound="Job")
+
+
+@dataclass
+class Job:
+    id: str
+    name: str
+    status: int
+    created_at: datetime
+    query: str
+    query_type: int
+    workers: int
+    params: dict[str, str]
+    metrics: dict[str, Any]
+    finished_at: Optional[datetime] = None
+    python_version: Optional[str] = None
+    error_message: str = ""
+    error_stack: str = ""
+
+    @classmethod
+    def parse(
+        cls: type[J],
+        id: str,
+        name: str,
+        status: int,
+        created_at: datetime,
+        finished_at: Optional[datetime],
+        query: str,
+        query_type: int,
+        workers: int,
+        python_version: Optional[str],
+        error_message: str,
+        error_stack: str,
+        params: str,
+        metrics: str,
+    ) -> "Job":
+        return cls(
+            id,
+            name,
+            status,
+            created_at,
+            query,
+            query_type,
+            workers,
+            json.loads(params),
+            json.loads(metrics),
+            finished_at,
+            python_version,
+            error_message,
+            error_stack,
+        )

datachain/lib/arrow.py

@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Optional
 
 import pyarrow as pa
 from pyarrow.dataset import dataset
+from tqdm import tqdm
 
 from datachain.lib.file import File, IndexedFile
 from datachain.lib.udf import Generator
@@ -13,33 +14,44 @@ if TYPE_CHECKING:
 
 
 class ArrowGenerator(Generator):
-    def __init__(self, schema: Optional["pa.Schema"] = None, **kwargs):
+    def __init__(
+        self,
+        schema: Optional["pa.Schema"] = None,
+        nrows: Optional[int] = None,
+        **kwargs,
+    ):
         """
         Generator for getting rows from tabular files.
 
         Parameters:
 
         schema : Optional pyarrow schema for validation.
+        nrows : Optional row limit.
         kwargs: Parameters to pass to pyarrow.dataset.dataset.
         """
         super().__init__()
         self.schema = schema
+        self.nrows = nrows
         self.kwargs = kwargs
 
     def process(self, file: File):
         path = file.get_path()
         ds = dataset(path, filesystem=file.get_fs(), schema=self.schema, **self.kwargs)
         index = 0
-        for record_batch in ds.to_batches():
-            for record in record_batch.to_pylist():
-                source = IndexedFile(file=file, index=index)
-                yield [source, *record.values()]
-                index += 1
+        with tqdm(desc="Parsed by pyarrow", unit=" rows") as pbar:
+            for record_batch in ds.to_batches():
+                for record in record_batch.to_pylist():
+                    source = IndexedFile(file=file, index=index)
+                    yield [source, *record.values()]
+                    index += 1
+                    if self.nrows and index >= self.nrows:
+                        return
+                pbar.update(len(record_batch))
 
 
 def infer_schema(chain: "DataChain", **kwargs) -> pa.Schema:
     schemas = []
-    for file in chain.iterate_one("file"):
+    for file in chain.collect("file"):
         ds = dataset(file.get_path(), filesystem=file.get_fs(), **kwargs)  # type: ignore[union-attr]
         schemas.append(ds.schema)
     return pa.unify_schemas(schemas)