datachain 0.7.5__py3-none-any.whl → 0.7.7__py3-none-any.whl

datachain/cli.py

@@ -233,6 +233,67 @@ def add_studio_parser(subparsers, parent_parser) -> None:
         help="The team to list datasets for. By default, it will use team from config.",
     )
 
+    studio_run_help = "Run a job in Studio"
+    studio_run_description = "This command runs a job in Studio."
+
+    studio_run_parser = studio_subparser.add_parser(
+        "run",
+        parents=[parent_parser],
+        description=studio_run_description,
+        help=studio_run_help,
+    )
+
+    studio_run_parser.add_argument(
+        "query_file",
+        action="store",
+        help="The query file to run.",
+    )
+
+    studio_run_parser.add_argument(
+        "--team",
+        action="store",
+        default=None,
+        help="The team to run a job for. By default, it will use team from config.",
+    )
+    studio_run_parser.add_argument(
+        "--env-file",
+        action="store",
+        help="File containing environment variables to set for the job.",
+    )
+
+    studio_run_parser.add_argument(
+        "--env",
+        nargs="+",
+        help="Environment variable. Can be specified multiple times. Format: KEY=VALUE",
+    )
+
+    studio_run_parser.add_argument(
+        "--workers",
+        type=int,
+        help="Number of workers to use for the job.",
+    )
+    studio_run_parser.add_argument(
+        "--files",
+        nargs="+",
+        help="Files to include in the job.",
+    )
+    studio_run_parser.add_argument(
+        "--python-version",
+        action="store",
+        help="Python version to use for the job (e.g. '3.9', '3.10', '3.11').",
+    )
+    studio_run_parser.add_argument(
+        "--req-file",
+        action="store",
+        help="File containing Python package requirements.",
+    )
+
+    studio_run_parser.add_argument(
+        "--req",
+        nargs="+",
+        help="Python package requirement. Can be specified multiple times.",
+    )
+
 
 def get_parser() -> ArgumentParser:  # noqa: PLR0915
     try:
@@ -358,7 +419,18 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
 
     add_studio_parser(subp, parent_parser)
 
-    parse_pull = subp.add_parser(
+    datasets_parser = subp.add_parser(
+        "datasets",
+        aliases=["ds"],
+        parents=[parent_parser],
+        description="Commands for managing datasers",
+    )
+    datasets_subparser = datasets_parser.add_subparsers(
+        dest="datasets_cmd",
+        help="Use `datachain datasets CMD --help` to display command specific help",
+    )
+
+    parse_pull = datasets_subparser.add_parser(
         "pull",
         parents=[parent_parser],
         description="Pull specific dataset version from SaaS",
@@ -400,9 +472,21 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         "--edatachain-file",
         help="Use a different filename for the resulting .edatachain file",
     )
+    parse_pull.add_argument(
+        "--local-name",
+        action="store",
+        default=None,
+        help="Name of the local dataset",
+    )
+    parse_pull.add_argument(
+        "--local-version",
+        action="store",
+        default=None,
+        help="Version of the local dataset",
+    )
 
-    parse_edit_dataset = subp.add_parser(
-        "edit-dataset", parents=[parent_parser], description="Edit dataset metadata"
+    parse_edit_dataset = datasets_subparser.add_parser(
+        "edit", parents=[parent_parser], description="Edit dataset metadata"
     )
     parse_edit_dataset.add_argument("name", type=str, help="Dataset name")
     parse_edit_dataset.add_argument(
@@ -447,8 +531,8 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help="The team to edit a dataset. By default, it will use team from config.",
     )
 
-    datasets_parser = subp.add_parser(
-        "datasets", parents=[parent_parser], description="List datasets"
+    datasets_parser = datasets_subparser.add_parser(
+        "ls", parents=[parent_parser], description="List datasets"
     )
     datasets_parser.add_argument(
         "--studio",
@@ -477,8 +561,8 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help="The team to list datasets for. By default, it will use team from config.",
     )
 
-    rm_dataset_parser = subp.add_parser(
-        "rm-dataset", parents=[parent_parser], description="Removes dataset"
+    rm_dataset_parser = datasets_subparser.add_parser(
+        "rm", parents=[parent_parser], description="Removes dataset", aliases=["remove"]
     )
     rm_dataset_parser.add_argument("name", type=str, help="Dataset name")
     rm_dataset_parser.add_argument(
@@ -521,8 +605,8 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help="The team to delete a dataset. By default, it will use team from config.",
     )
 
-    dataset_stats_parser = subp.add_parser(
-        "dataset-stats",
+    dataset_stats_parser = datasets_subparser.add_parser(
+        "stats",
         parents=[parent_parser],
         description="Shows basic dataset stats",
     )
@@ -1203,27 +1287,59 @@ def main(argv: Optional[list[str]] = None) -> int:  # noqa: C901, PLR0912, PLR09
                 edatachain=args.edatachain,
                 edatachain_file=args.edatachain_file,
             )
-        elif args.command == "pull":
-            catalog.pull_dataset(
-                args.dataset,
-                args.output,
-                no_cp=args.no_cp,
-                force=bool(args.force),
-                edatachain=args.edatachain,
-                edatachain_file=args.edatachain_file,
-            )
-        elif args.command == "edit-dataset":
-            edit_dataset(
-                catalog,
-                args.name,
-                new_name=args.new_name,
-                description=args.description,
-                labels=args.labels,
-                studio=args.studio,
-                local=args.local,
-                all=args.all,
-                team=args.team,
-            )
+        elif args.command in ("datasets", "ds"):
+            if args.datasets_cmd == "pull":
+                catalog.pull_dataset(
+                    args.dataset,
+                    args.output,
+                    local_ds_name=args.local_name,
+                    local_ds_version=args.local_version,
+                    no_cp=args.no_cp,
+                    force=bool(args.force),
+                    edatachain=args.edatachain,
+                    edatachain_file=args.edatachain_file,
+                )
+            elif args.datasets_cmd == "edit":
+                edit_dataset(
+                    catalog,
+                    args.name,
+                    new_name=args.new_name,
+                    description=args.description,
+                    labels=args.labels,
+                    studio=args.studio,
+                    local=args.local,
+                    all=args.all,
+                    team=args.team,
+                )
+            elif args.datasets_cmd == "ls":
+                datasets(
+                    catalog=catalog,
+                    studio=args.studio,
+                    local=args.local,
+                    all=args.all,
+                    team=args.team,
+                )
+            elif args.datasets_cmd in ("rm", "remove"):
+                rm_dataset(
+                    catalog,
+                    args.name,
+                    version=args.version,
+                    force=args.force,
+                    studio=args.studio,
+                    local=args.local,
+                    all=args.all,
+                    team=args.team,
+                )
+            elif args.datasets_cmd == "stats":
+                dataset_stats(
+                    catalog,
+                    args.name,
+                    args.version,
+                    show_bytes=args.bytes,
+                    si=args.si,
+                )
+            else:
+                raise Exception(f"Unexpected command {args.datasets_cmd}")
         elif args.command == "ls":
             ls(
                 args.sources,
@@ -1235,14 +1351,7 @@ def main(argv: Optional[list[str]] = None) -> int:  # noqa: C901, PLR0912, PLR09
                 update=bool(args.update),
                 client_config=client_config,
             )
-        elif args.command == "datasets":
-            datasets(
-                catalog=catalog,
-                studio=args.studio,
-                local=args.local,
-                all=args.all,
-                team=args.team,
-            )
+
         elif args.command == "show":
             show(
                 catalog,
@@ -1254,25 +1363,7 @@ def main(argv: Optional[list[str]] = None) -> int:  # noqa: C901, PLR0912, PLR09
                 no_collapse=args.no_collapse,
                 schema=args.schema,
             )
-        elif args.command == "rm-dataset":
-            rm_dataset(
-                catalog,
-                args.name,
-                version=args.version,
-                force=args.force,
-                studio=args.studio,
-                local=args.local,
-                all=args.all,
-                team=args.team,
-            )
-        elif args.command == "dataset-stats":
-            dataset_stats(
-                catalog,
-                args.name,
-                args.version,
-                show_bytes=args.bytes,
-                si=args.si,
-            )
+
         elif args.command == "du":
             du(
                 catalog,

datachain/data_storage/__init__.py

@@ -1,12 +1,9 @@
-from .id_generator import AbstractDBIDGenerator, AbstractIDGenerator
 from .job import JobQueryType, JobStatus
 from .metastore import AbstractDBMetastore, AbstractMetastore
 from .warehouse import AbstractWarehouse
 
 __all__ = [
-    "AbstractDBIDGenerator",
     "AbstractDBMetastore",
-    "AbstractIDGenerator",
     "AbstractMetastore",
     "AbstractWarehouse",
     "JobQueryType",

datachain/data_storage/metastore.py

@@ -45,7 +45,7 @@ if TYPE_CHECKING:
     from sqlalchemy import Delete, Insert, Select, Update
     from sqlalchemy.schema import SchemaItem
 
-    from datachain.data_storage import AbstractIDGenerator, schema
+    from datachain.data_storage import schema
     from datachain.data_storage.db_engine import DatabaseEngine
 
 logger = logging.getLogger("datachain")
@@ -304,16 +304,10 @@ class AbstractDBMetastore(AbstractMetastore):
     DATASET_DEPENDENCY_TABLE = "datasets_dependencies"
     JOBS_TABLE = "jobs"
 
-    id_generator: "AbstractIDGenerator"
     db: "DatabaseEngine"
 
-    def __init__(
-        self,
-        id_generator: "AbstractIDGenerator",
-        uri: Optional[StorageURI] = None,
-    ):
+    def __init__(self, uri: Optional[StorageURI] = None):
         uri = uri or StorageURI("")
-        self.id_generator = id_generator
         super().__init__(uri)
 
     def close(self) -> None:
@@ -322,7 +316,6 @@ class AbstractDBMetastore(AbstractMetastore):
 
     def cleanup_tables(self, temp_table_names: list[str]) -> None:
         """Cleanup temp tables."""
-        self.id_generator.delete_uris(temp_table_names)
 
     @classmethod
     def _datasets_columns(cls) -> list["SchemaItem"]:

datachain/data_storage/sqlite.py

@@ -15,7 +15,6 @@ from typing import (
 )
 
 import sqlalchemy
-from packaging import version
 from sqlalchemy import MetaData, Table, UniqueConstraint, exists, select
 from sqlalchemy.dialects import sqlite
 from sqlalchemy.schema import CreateIndex, CreateTable, DropTable
@@ -27,7 +26,6 @@ from tqdm import tqdm
 import datachain.sql.sqlite
 from datachain.data_storage import AbstractDBMetastore, AbstractWarehouse
 from datachain.data_storage.db_engine import DatabaseEngine
-from datachain.data_storage.id_generator import AbstractDBIDGenerator
 from datachain.data_storage.schema import DefaultSchema
 from datachain.dataset import DatasetRecord, StorageURI
 from datachain.error import DataChainError
@@ -275,123 +273,16 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         self.execute_str(f"ALTER TABLE {comp_old_name} RENAME TO {comp_new_name}")
 
 
-class SQLiteIDGenerator(AbstractDBIDGenerator):
-    _db: "SQLiteDatabaseEngine"
-
-    def __init__(
-        self,
-        db: Optional["SQLiteDatabaseEngine"] = None,
-        table_prefix: Optional[str] = None,
-        skip_db_init: bool = False,
-        db_file: Optional[str] = None,
-        in_memory: bool = False,
-    ):
-        db_file = get_db_file_in_memory(db_file, in_memory)
-
-        db = db or SQLiteDatabaseEngine.from_db_file(db_file)
-
-        super().__init__(db, table_prefix, skip_db_init)
-
-    def clone(self) -> "SQLiteIDGenerator":
-        """Clones SQLiteIDGenerator implementation."""
-        return SQLiteIDGenerator(
-            self._db.clone(), self._table_prefix, skip_db_init=True
-        )
-
-    def clone_params(self) -> tuple[Callable[..., Any], list[Any], dict[str, Any]]:
-        """
-        Returns the function, args, and kwargs needed to instantiate a cloned copy
-        of this SQLiteIDGenerator implementation, for use in separate processes
-        or machines.
-        """
-        return (
-            SQLiteIDGenerator.init_after_clone,
-            [],
-            {
-                "db_clone_params": self._db.clone_params(),
-                "table_prefix": self._table_prefix,
-            },
-        )
-
-    @classmethod
-    def init_after_clone(
-        cls,
-        *,
-        db_clone_params: tuple[Callable, list, dict[str, Any]],
-        table_prefix: Optional[str] = None,
-    ) -> "SQLiteIDGenerator":
-        """
-        Initializes a new instance of this SQLiteIDGenerator implementation
-        using the given parameters, which were obtained from a call to clone_params.
-        """
-        (db_class, db_args, db_kwargs) = db_clone_params
-        return cls(
-            db=db_class(*db_args, **db_kwargs),
-            table_prefix=table_prefix,
-            skip_db_init=True,
-        )
-
-    @property
-    def db(self) -> "SQLiteDatabaseEngine":
-        return self._db
-
-    def init_id(self, uri: str) -> None:
-        """Initializes the ID generator for the given URI with zero last_id."""
-        self._db.execute(
-            sqlite.insert(self._table)
-            .values(uri=uri, last_id=0)
-            .on_conflict_do_nothing()
-        )
-
-    def get_next_ids(self, uri: str, count: int) -> range:
-        """Returns a range of IDs for the given URI."""
-
-        sqlite_version = version.parse(sqlite3.sqlite_version)
-        is_returning_supported = sqlite_version >= version.parse("3.35.0")
-        if is_returning_supported:
-            stmt = (
-                sqlite.insert(self._table)
-                .values(uri=uri, last_id=count)
-                .on_conflict_do_update(
-                    index_elements=["uri"],
-                    set_={"last_id": self._table.c.last_id + count},
-                )
-                .returning(self._table.c.last_id)
-            )
-            last_id = self._db.execute(stmt).fetchone()[0]
-        else:
-            # Older versions of SQLite are still the default under Ubuntu LTS,
-            # e.g. Ubuntu 20.04 LTS (Focal Fossa) uses 3.31.1
-            # Transactions ensure no concurrency conflicts
-            with self._db.transaction() as conn:
-                stmt_ins = (
-                    sqlite.insert(self._table)
-                    .values(uri=uri, last_id=count)
-                    .on_conflict_do_update(
-                        index_elements=["uri"],
-                        set_={"last_id": self._table.c.last_id + count},
-                    )
-                )
-                self._db.execute(stmt_ins, conn=conn)
-
-                stmt_sel = select(self._table.c.last_id).where(self._table.c.uri == uri)
-                last_id = self._db.execute(stmt_sel, conn=conn).fetchone()[0]
-
-        return range(last_id - count + 1, last_id + 1)
-
-
 class SQLiteMetastore(AbstractDBMetastore):
     """
     SQLite Metastore uses SQLite3 for storing indexed data locally.
     This is currently used for the local cli.
     """
 
-    id_generator: "SQLiteIDGenerator"
     db: "SQLiteDatabaseEngine"
 
     def __init__(
         self,
-        id_generator: "SQLiteIDGenerator",
         uri: Optional[StorageURI] = None,
         db: Optional["SQLiteDatabaseEngine"] = None,
         db_file: Optional[str] = None,
@@ -399,7 +290,7 @@ class SQLiteMetastore(AbstractDBMetastore):
     ):
         uri = uri or StorageURI("")
         self.schema: DefaultSchema = DefaultSchema()
-        super().__init__(id_generator, uri)
+        super().__init__(uri)
 
         # needed for dropping tables in correct order for tests because of
         # foreign keys
@@ -424,11 +315,7 @@ class SQLiteMetastore(AbstractDBMetastore):
         if not uri and self.uri:
             uri = self.uri
 
-        return SQLiteMetastore(
-            self.id_generator.clone(),
-            uri=uri,
-            db=self.db.clone(),
-        )
+        return SQLiteMetastore(uri=uri, db=self.db.clone())
 
     def clone_params(self) -> tuple[Callable[..., Any], list[Any], dict[str, Any]]:
         """
@@ -439,7 +326,6 @@ class SQLiteMetastore(AbstractDBMetastore):
             SQLiteMetastore.init_after_clone,
             [],
             {
-                "id_generator_clone_params": self.id_generator.clone_params(),
                 "uri": self.uri,
                 "db_clone_params": self.db.clone_params(),
             },
@@ -449,21 +335,11 @@ class SQLiteMetastore(AbstractDBMetastore):
     def init_after_clone(
         cls,
         *,
-        id_generator_clone_params: tuple[Callable, list, dict[str, Any]],
         uri: StorageURI,
         db_clone_params: tuple[Callable, list, dict[str, Any]],
     ) -> "SQLiteMetastore":
-        (
-            id_generator_class,
-            id_generator_args,
-            id_generator_kwargs,
-        ) = id_generator_clone_params
         (db_class, db_args, db_kwargs) = db_clone_params
-        return cls(
-            id_generator=id_generator_class(*id_generator_args, **id_generator_kwargs),
-            uri=uri,
-            db=db_class(*db_args, **db_kwargs),
-        )
+        return cls(uri=uri, db=db_class(*db_args, **db_kwargs))
 
     def _init_tables(self) -> None:
         """Initialize tables."""
@@ -518,7 +394,6 @@ class SQLiteWarehouse(AbstractWarehouse):
     This is currently used for the local cli.
     """
 
-    id_generator: "SQLiteIDGenerator"
     db: "SQLiteDatabaseEngine"
 
     # Cache for our defined column types to dialect specific TypeEngine relations
@@ -526,13 +401,12 @@ class SQLiteWarehouse(AbstractWarehouse):
 
     def __init__(
         self,
-        id_generator: "SQLiteIDGenerator",
         db: Optional["SQLiteDatabaseEngine"] = None,
         db_file: Optional[str] = None,
         in_memory: bool = False,
     ):
         self.schema: DefaultSchema = DefaultSchema()
-        super().__init__(id_generator)
+        super().__init__()
 
         db_file = get_db_file_in_memory(db_file, in_memory)
 
@@ -543,7 +417,7 @@ class SQLiteWarehouse(AbstractWarehouse):
         self.close()
 
     def clone(self, use_new_connection: bool = False) -> "SQLiteWarehouse":
-        return SQLiteWarehouse(self.id_generator.clone(), db=self.db.clone())
+        return SQLiteWarehouse(db=self.db.clone())
 
     def clone_params(self) -> tuple[Callable[..., Any], list[Any], dict[str, Any]]:
         """
@@ -553,29 +427,17 @@ class SQLiteWarehouse(AbstractWarehouse):
         return (
             SQLiteWarehouse.init_after_clone,
             [],
-            {
-                "id_generator_clone_params": self.id_generator.clone_params(),
-                "db_clone_params": self.db.clone_params(),
-            },
+            {"db_clone_params": self.db.clone_params()},
         )
 
     @classmethod
     def init_after_clone(
         cls,
         *,
-        id_generator_clone_params: tuple[Callable, list, dict[str, Any]],
         db_clone_params: tuple[Callable, list, dict[str, Any]],
     ) -> "SQLiteWarehouse":
-        (
-            id_generator_class,
-            id_generator_args,
-            id_generator_kwargs,
-        ) = id_generator_clone_params
         (db_class, db_args, db_kwargs) = db_clone_params
-        return cls(
-            id_generator=id_generator_class(*id_generator_args, **id_generator_kwargs),
-            db=db_class(*db_args, **db_kwargs),
-        )
+        return cls(db=db_class(*db_args, **db_kwargs))
 
     def _reflect_tables(self, filter_tables=None):
         """

datachain/data_storage/warehouse.py

@@ -34,7 +34,7 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.selectable import Join, Select
     from sqlalchemy.types import TypeEngine
 
-    from datachain.data_storage import AbstractIDGenerator, schema
+    from datachain.data_storage import schema
     from datachain.data_storage.db_engine import DatabaseEngine
     from datachain.data_storage.schema import DataTable
     from datachain.lib.file import File
@@ -69,13 +69,9 @@ class AbstractWarehouse(ABC, Serializable):
     UDF_TABLE_NAME_PREFIX = "udf_"
     TMP_TABLE_NAME_PREFIX = "tmp_"
 
-    id_generator: "AbstractIDGenerator"
     schema: "schema.Schema"
     db: "DatabaseEngine"
 
-    def __init__(self, id_generator: "AbstractIDGenerator"):
-        self.id_generator = id_generator
-
     def __enter__(self) -> "AbstractWarehouse":
         return self
 

datachain/dataset.py

@@ -488,6 +488,18 @@ class DatasetRecord:
             if v.version == version
         )
 
+    def get_version_by_uuid(self, uuid: str) -> DatasetVersion:
+        try:
+            return next(
+                v
+                for v in self.versions  # type: ignore [union-attr]
+                if v.uuid == uuid
+            )
+        except StopIteration:
+            raise DatasetVersionNotFoundError(
+                f"Dataset {self.name} does not have version with uuid {uuid}"
+            ) from None
+
     def remove_version(self, version: int) -> None:
         if not self.versions or not self.has_version(version):
             return
@@ -635,6 +647,9 @@ class DatasetListRecord:
             LISTING_PREFIX
         )
 
+    def has_version_with_uuid(self, uuid: str) -> bool:
+        return any(v.uuid == uuid for v in self.versions)
+
 
 class RowDict(dict):
     pass

datachain/func/__init__.py

@@ -1,4 +1,4 @@
-from sqlalchemy import literal
+from sqlalchemy import case, literal
 
 from . import array, path, random, string
 from .aggregate import (
@@ -24,6 +24,7 @@ __all__ = [
     "any_value",
     "array",
     "avg",
+    "case",
     "collect",
     "concat",
     "cosine_distance",

datachain/func/func.py

@@ -2,9 +2,11 @@ import inspect
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Callable, Optional, Union
 
-from sqlalchemy import BindParameter, ColumnElement, desc
+from sqlalchemy import BindParameter, Case, ColumnElement, desc
+from sqlalchemy.ext.hybrid import Comparator
 
 from datachain.lib.convert.python_to_sql import python_to_sql
+from datachain.lib.convert.sql_to_python import sql_to_python
 from datachain.lib.utils import DataChainColumnError, DataChainParamsError
 from datachain.query.schema import Column, ColumnMeta
 
@@ -71,7 +73,7 @@ class Func(Function):
         return (
             [
                 col
-                if isinstance(col, (Func, BindParameter))
+                if isinstance(col, (Func, BindParameter, Case, Comparator))
                 else ColumnMeta.to_db_name(
                     col.name if isinstance(col, ColumnElement) else col
                 )
@@ -273,6 +275,9 @@ def get_db_col_type(signals_schema: "SignalSchema", col: ColT) -> "DataType":
     if isinstance(col, Func):
         return col.get_result_type(signals_schema)
 
+    if isinstance(col, ColumnElement) and not hasattr(col, "name"):
+        return sql_to_python(col)
+
     return signals_schema.get_column_type(
         col.name if isinstance(col, ColumnElement) else col
     )

datachain/lib/dc.py

@@ -1150,7 +1150,7 @@ class DataChain:
     def group_by(
         self,
         *,
-        partition_by: Union[str, Func, Sequence[Union[str, Func]]],
+        partition_by: Optional[Union[str, Func, Sequence[Union[str, Func]]]] = None,
         **kwargs: Func,
     ) -> "Self":
         """Group rows by specified set of signals and return new signals
@@ -1167,10 +1167,10 @@ class DataChain:
             )
             ```
         """
-        if isinstance(partition_by, (str, Func)):
+        if partition_by is None:
+            partition_by = []
+        elif isinstance(partition_by, (str, Func)):
             partition_by = [partition_by]
-        if not partition_by:
-            raise ValueError("At least one column should be provided for partition_by")
 
         partition_by_columns: list[Column] = []
         signal_columns: list[Column] = []