datachain 0.8.3__py3-none-any.whl → 0.8.5__py3-none-any.whl

datachain/cli/parser/studio.py

@@ -0,0 +1,126 @@
+def add_studio_parser(subparsers, parent_parser) -> None:
+    studio_help = "Commands to authenticate DataChain with Iterative Studio"
+    studio_description = (
+        "Authenticate DataChain with Studio and set the token. "
+        "Once this token has been properly configured,\n"
+        "DataChain will utilize it for seamlessly sharing datasets\n"
+        "and using Studio features from CLI"
+    )
+
+    studio_parser = subparsers.add_parser(
+        "studio",
+        parents=[parent_parser],
+        description=studio_description,
+        help=studio_help,
+    )
+    studio_subparser = studio_parser.add_subparsers(
+        dest="cmd",
+        help="Use `DataChain studio CMD --help` to display command-specific help.",
+        required=True,
+    )
+
+    studio_login_help = "Authenticate DataChain with Studio host"
+    studio_login_description = (
+        "By default, this command authenticates the DataChain with Studio\n"
+        "using default scopes and assigns a random name as the token name."
+    )
+    login_parser = studio_subparser.add_parser(
+        "login",
+        parents=[parent_parser],
+        description=studio_login_description,
+        help=studio_login_help,
+    )
+
+    login_parser.add_argument(
+        "-H",
+        "--hostname",
+        action="store",
+        default=None,
+        help="The hostname of the Studio instance to authenticate with.",
+    )
+    login_parser.add_argument(
+        "-s",
+        "--scopes",
+        action="store",
+        default=None,
+        help="The scopes for the authentication token. ",
+    )
+
+    login_parser.add_argument(
+        "-n",
+        "--name",
+        action="store",
+        default=None,
+        help="The name of the authentication token. It will be used to\n"
+        "identify token shown in Studio profile.",
+    )
+
+    login_parser.add_argument(
+        "--no-open",
+        action="store_true",
+        default=False,
+        help="Use authentication flow based on user code.\n"
+        "You will be presented with user code to enter in browser.\n"
+        "DataChain will also use this if it cannot launch browser on your behalf.",
+    )
+
+    studio_logout_help = "Logout user from Studio"
+    studio_logout_description = "This removes the studio token from your global config."
+
+    studio_subparser.add_parser(
+        "logout",
+        parents=[parent_parser],
+        description=studio_logout_description,
+        help=studio_logout_help,
+    )
+
+    studio_team_help = "Set the default team for DataChain"
+    studio_team_description = (
+        "Set the default team for DataChain to use when interacting with Studio."
+    )
+
+    team_parser = studio_subparser.add_parser(
+        "team",
+        parents=[parent_parser],
+        description=studio_team_description,
+        help=studio_team_help,
+    )
+    team_parser.add_argument(
+        "team_name",
+        action="store",
+        help="The name of the team to set as the default.",
+    )
+    team_parser.add_argument(
+        "--global",
+        action="store_true",
+        default=False,
+        help="Set the team globally for all DataChain projects.",
+    )
+
+    studio_token_help = "View the token datachain uses to contact Studio"  # noqa: S105 # nosec B105
+
+    studio_subparser.add_parser(
+        "token",
+        parents=[parent_parser],
+        description=studio_token_help,
+        help=studio_token_help,
+    )
+
+    studio_ls_dataset_help = "List the available datasets from Studio"
+    studio_ls_dataset_description = (
+        "This command lists all the datasets available in Studio.\n"
+        "It will show the dataset name and the number of versions available."
+    )
+
+    ls_dataset_parser = studio_subparser.add_parser(
+        "dataset",
+        parents=[parent_parser],
+        description=studio_ls_dataset_description,
+        help=studio_ls_dataset_help,
+    )
+    ls_dataset_parser.add_argument(
+        "--team",
+        action="store",
+        default=None,
+        help="The team to list datasets for. By default, it will use team from config.",
+    )

datachain/cli/parser/utils.py

@@ -0,0 +1,63 @@
+from argparse import Action, ArgumentParser, ArgumentTypeError
+from typing import Union
+
+from datachain.cli.utils import CommaSeparatedArgs
+
+FIND_COLUMNS = ["du", "name", "path", "size", "type"]
+
+
+def find_columns_type(
+    columns_str: str,
+    default_colums_str: str = "path",
+) -> list[str]:
+    if not columns_str:
+        columns_str = default_colums_str
+
+    return [parse_find_column(c) for c in columns_str.split(",")]
+
+
+def parse_find_column(column: str) -> str:
+    column_lower = column.strip().lower()
+    if column_lower in FIND_COLUMNS:
+        return column_lower
+    raise ArgumentTypeError(
+        f"Invalid column for find: '{column}' Options are: {','.join(FIND_COLUMNS)}"
+    )
+
+
+def add_sources_arg(parser: ArgumentParser, nargs: Union[str, int] = "+") -> Action:
+    return parser.add_argument(
+        "sources",
+        type=str,
+        nargs=nargs,
+        help="Data sources - paths to cloud storage dirs",
+    )
+
+
+def add_show_args(parser: ArgumentParser) -> None:
+    parser.add_argument(
+        "--limit",
+        action="store",
+        default=10,
+        type=int,
+        help="Number of rows to show",
+    )
+    parser.add_argument(
+        "--offset",
+        action="store",
+        default=0,
+        type=int,
+        help="Number of rows to offset",
+    )
+    parser.add_argument(
+        "--columns",
+        default=[],
+        action=CommaSeparatedArgs,
+        help="Columns to show",
+    )
+    parser.add_argument(
+        "--no-collapse",
+        action="store_true",
+        default=False,
+        help="Do not collapse the columns",
+    )

datachain/{cli_utils.py → cli/utils.py}

@@ -1,4 +1,8 @@
-from argparse import SUPPRESS, Action, ArgumentError, _AppendAction
+import logging
+from argparse import SUPPRESS, Action, ArgumentError, Namespace, _AppendAction
+from typing import Optional
+
+from datachain.error import DataChainError
 
 
 class BooleanOptionalAction(Action):
@@ -70,3 +74,25 @@ class KeyValueArgs(_AppendAction):  # pylint: disable=protected-access
             items[key.strip()] = value
 
         setattr(namespace, self.dest, items)
+
+
+def get_logging_level(args: Namespace) -> int:
+    if args.quiet:
+        return logging.CRITICAL
+    if args.verbose:
+        return logging.DEBUG
+    return logging.INFO
+
+
+def determine_flavors(studio: bool, local: bool, all: bool, token: Optional[str]):
+    if studio and not token:
+        raise DataChainError(
+            "Not logged in to Studio. Log in with 'datachain studio login'."
+        )
+
+    if local or studio:
+        all = False
+
+    all = all and not (local or studio)
+
+    return all, local, studio

datachain/client/azure.py

@@ -31,8 +31,12 @@ class AzureClient(Client):
         Generate a signed URL for the given path.
         """
         version_id = kwargs.pop("version_id", None)
+        content_disposition = kwargs.pop("content_disposition", None)
         result = self.fs.sign(
-            self.get_full_path(path, version_id), expiration=expires, **kwargs
+            self.get_full_path(path, version_id),
+            expiration=expires,
+            content_disposition=content_disposition,
+            **kwargs,
         )
         return result + (f"&versionid={version_id}" if version_id else "")
 
@@ -42,7 +46,7 @@ class AzureClient(Client):
             prefix = prefix.lstrip(DELIMITER) + DELIMITER
         found = False
         try:
-            with tqdm(desc=f"Listing {self.uri}", unit=" objects") as pbar:
+            with tqdm(desc=f"Listing {self.uri}", unit=" objects", leave=False) as pbar:
                 async with self.fs.service_client.get_container_client(
                     container=self.name
                 ) as container_client:

datachain/client/fsspec.py

@@ -215,7 +215,7 @@ class Client(ABC):
         info = await self.fs._info(
             self.get_full_path(file.path, file.version), **kwargs
         )
-        return self.info_to_file(info, "").etag
+        return self.info_to_file(info, file.path).etag
 
     def get_file_info(self, path: str, version_id: Optional[str] = None) -> "File":
         info = self.fs.info(self.get_full_path(path, version_id), version_id=version_id)
@@ -249,7 +249,7 @@ class Client(ABC):
         await main_task
 
     async def _fetch_nested(self, start_prefix: str, result_queue: ResultQueue) -> None:
-        progress_bar = tqdm(desc=f"Listing {self.uri}", unit=" objects")
+        progress_bar = tqdm(desc=f"Listing {self.uri}", unit=" objects", leave=False)
         loop = get_loop()
 
         queue: asyncio.Queue[str] = asyncio.Queue()
@@ -343,7 +343,7 @@ class Client(ABC):
         return self.version_path(f"{self.PREFIX}{self.name}/{rel_path}", version_id)
 
     @abstractmethod
-    def info_to_file(self, v: dict[str, Any], parent: str) -> "File": ...
+    def info_to_file(self, v: dict[str, Any], path: str) -> "File": ...
 
     def fetch_nodes(
         self,
@@ -390,6 +390,12 @@ class Client(ABC):
             self.fs.open(self.get_full_path(file.path, file.version)), cb
         )  # type: ignore[return-value]
 
+    def upload(self, path: str, data: bytes) -> "File":
+        full_path = self.get_full_path(path)
+        self.fs.pipe_file(full_path, data)
+        file_info = self.fs.info(full_path)
+        return self.info_to_file(file_info, path)
+
     def download(self, file: "File", *, callback: Callback = DEFAULT_CALLBACK) -> None:
         sync(get_loop(), functools.partial(self._download, file, callback=callback))
 

datachain/client/gcs.py

@@ -39,11 +39,15 @@ class GCSClient(Client):
         (see https://cloud.google.com/storage/docs/access-public-data#api-link).
         """
         version_id = kwargs.pop("version_id", None)
+        content_disposition = kwargs.pop("content_disposition", None)
         if self.fs.storage_options.get("token") == "anon":
             query = f"?generation={version_id}" if version_id else ""
             return f"https://storage.googleapis.com/{self.name}/{path}{query}"
         return self.fs.sign(
-            self.get_full_path(path, version_id), expiration=expires, **kwargs
+            self.get_full_path(path, version_id),
+            expiration=expires,
+            response_disposition=content_disposition,
+            **kwargs,
         )
 
     @staticmethod
@@ -83,7 +87,7 @@ class GCSClient(Client):
         self, page_queue: PageQueue, result_queue: ResultQueue
     ) -> bool:
         found = False
-        with tqdm(desc=f"Listing {self.uri}", unit=" objects") as pbar:
+        with tqdm(desc=f"Listing {self.uri}", unit=" objects", leave=False) as pbar:
             while (page := await page_queue.get()) is not None:
                 if page:
                     found = True

datachain/client/s3.py

@@ -51,6 +51,21 @@ class ClientS3(Client):
 
         return cast(S3FileSystem, super().create_fs(**kwargs))
 
+    def url(self, path: str, expires: int = 3600, **kwargs) -> str:
+        """
+        Generate a signed URL for the given path.
+        """
+        version_id = kwargs.pop("version_id", None)
+        content_disposition = kwargs.pop("content_disposition", None)
+        if content_disposition:
+            kwargs["ResponseContentDisposition"] = content_disposition
+
+        return self.fs.sign(
+            self.get_full_path(path, version_id),
+            expiration=expires,
+            **kwargs,
+        )
+
     async def _fetch_flat(self, start_prefix: str, result_queue: ResultQueue) -> None:
         async def get_pages(it, page_queue):
             try:
@@ -61,7 +76,7 @@ class ClientS3(Client):
 
         async def process_pages(page_queue, result_queue):
             found = False
-            with tqdm(desc=f"Listing {self.uri}", unit=" objects") as pbar:
+            with tqdm(desc=f"Listing {self.uri}", unit=" objects", leave=False) as pbar:
                 while (res := await page_queue.get()) is not None:
                     if res:
                         found = True

datachain/data_storage/db_engine.py

@@ -79,6 +79,15 @@ class DatabaseEngine(ABC, Serializable):
         conn: Optional[Any] = None,
     ) -> Iterator[tuple[Any, ...]]: ...
 
+    def get_table(self, name: str) -> "Table":
+        table = self.metadata.tables.get(name)
+        if table is None:
+            sa.Table(name, self.metadata, autoload_with=self.engine)
+            # ^^^ This table may not be correctly initialised on some dialects
+            # Grab it from metadata instead.
+            table = self.metadata.tables[name]
+        return table
+
     @abstractmethod
     def executemany(
         self, query, params, cursor: Optional[Any] = None

datachain/data_storage/schema.py

@@ -16,7 +16,6 @@ from datachain.sql.functions import path as pathfunc
 from datachain.sql.types import Int, SQLType, UInt64
 
 if TYPE_CHECKING:
-    from sqlalchemy import Engine
     from sqlalchemy.engine.interfaces import Dialect
     from sqlalchemy.sql.base import (
         ColumnCollection,
@@ -25,6 +24,8 @@ if TYPE_CHECKING:
     )
     from sqlalchemy.sql.elements import ColumnElement
 
+    from datachain.data_storage.db_engine import DatabaseEngine
+
 
 DEFAULT_DELIMITER = "__"
 
@@ -150,14 +151,12 @@ class DataTable:
     def __init__(
         self,
         name: str,
-        engine: "Engine",
-        metadata: Optional["sa.MetaData"] = None,
+        engine: "DatabaseEngine",
         column_types: Optional[dict[str, SQLType]] = None,
         object_name: str = "file",
     ):
         self.name: str = name
         self.engine = engine
-        self.metadata: sa.MetaData = metadata if metadata is not None else sa.MetaData()
         self.column_types: dict[str, SQLType] = column_types or {}
         self.object_name = object_name
 
@@ -211,12 +210,7 @@ class DataTable:
         return sa.Table(name, metadata, *columns)
 
     def get_table(self) -> "sa.Table":
-        table = self.metadata.tables.get(self.name)
-        if table is None:
-            sa.Table(self.name, self.metadata, autoload_with=self.engine)
-            # ^^^ This table may not be correctly initialised on some dialects
-            # Grab it from metadata instead.
-            table = self.metadata.tables[self.name]
+        table = self.engine.get_table(self.name)
 
         column_types = self.column_types | {c.name: c.type for c in self.sys_columns()}
         # adjusting types for custom columns to be instances of SQLType if possible

datachain/data_storage/sqlite.py

@@ -186,6 +186,12 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         self.db_file = db_file
         self.is_closed = False
 
+    def get_table(self, name: str) -> Table:
+        if self.is_closed:
+            # Reconnect in case of being closed previously.
+            self._reconnect()
+        return super().get_table(name)
+
     @retry_sqlite_locks
     def execute(
         self,
@@ -670,7 +676,7 @@ class SQLiteWarehouse(AbstractWarehouse):
         ]
         table = self.create_udf_table(columns)
 
-        with tqdm(desc="Preparing", unit=" rows") as pbar:
+        with tqdm(desc="Preparing", unit=" rows", leave=False) as pbar:
             self.copy_table(table, query, progress_cb=pbar.update)
 
         return table

datachain/data_storage/warehouse.py

@@ -191,8 +191,7 @@ class AbstractWarehouse(ABC, Serializable):
         table_name = self.dataset_table_name(dataset.name, version)
         return self.schema.dataset_row_cls(
             table_name,
-            self.db.engine,
-            self.db.metadata,
+            self.db,
             dataset.get_schema(version),
             object_name=object_name,
         )
@@ -904,8 +903,11 @@ class AbstractWarehouse(ABC, Serializable):
         This should be implemented to ensure that the provided tables
         are cleaned up as soon as they are no longer needed.
         """
-        with tqdm(desc="Cleanup", unit=" tables") as pbar:
-            for name in set(names):
+        to_drop = set(names)
+        with tqdm(
+            desc="Cleanup", unit=" tables", total=len(to_drop), leave=False
+        ) as pbar:
+            for name in to_drop:
                 self.db.drop_table(Table(name, self.db.metadata), if_exists=True)
                 pbar.update(1)
 

datachain/{lib/diff.py → diff/__init__.py}

@@ -1,6 +1,7 @@
 import random
 import string
 from collections.abc import Sequence
+from enum import Enum
 from typing import TYPE_CHECKING, Optional, Union
 
 import sqlalchemy as sa
@@ -16,7 +17,22 @@ if TYPE_CHECKING:
 C = Column
 
 
-def compare(  # noqa: PLR0912, PLR0915, C901
+def get_status_col_name() -> str:
+    """Returns new unique status col name"""
+    return "diff_" + "".join(
+        random.choice(string.ascii_letters)  # noqa: S311
+        for _ in range(10)
+    )
+
+
+class CompareStatus(str, Enum):
+    ADDED = "A"
+    DELETED = "D"
+    MODIFIED = "M"
+    SAME = "S"
+
+
+def _compare(  # noqa: PLR0912, PLR0915, C901
     left: "DataChain",
     right: "DataChain",
     on: Union[str, Sequence[str]],
@@ -72,13 +88,10 @@ def compare(  # noqa: PLR0912, PLR0915, C901
             "At least one of added, deleted, modified, same flags must be set"
         )
 
-    # we still need status column for internal implementation even if not
-    # needed in output
     need_status_col = bool(status_col)
-    status_col = status_col or "diff_" + "".join(
-        random.choice(string.ascii_letters)  # noqa: S311
-        for _ in range(10)
-    )
+    # we still need status column for internal implementation even if not
+    # needed in the output
+    status_col = status_col or get_status_col_name()
 
     # calculate on and compare column names
     right_on = right_on or on
@@ -112,7 +125,7 @@ def compare(  # noqa: PLR0912, PLR0915, C901
                 for c in [f"{_rprefix(c, rc)}{rc}" for c, rc in zip(on, right_on)]
             ]
         )
-        diff_cond.append((added_cond, "A"))
+        diff_cond.append((added_cond, CompareStatus.ADDED))
     if modified and compare:
         modified_cond = sa.or_(
             *[
@@ -120,7 +133,7 @@ def compare(  # noqa: PLR0912, PLR0915, C901
                 for c, rc in zip(compare, right_compare)  # type: ignore[arg-type]
             ]
         )
-        diff_cond.append((modified_cond, "M"))
+        diff_cond.append((modified_cond, CompareStatus.MODIFIED))
     if same and compare:
         same_cond = sa.and_(
             *[
@@ -128,9 +141,11 @@ def compare(  # noqa: PLR0912, PLR0915, C901
                 for c, rc in zip(compare, right_compare)  # type: ignore[arg-type]
             ]
         )
-        diff_cond.append((same_cond, "S"))
+        diff_cond.append((same_cond, CompareStatus.SAME))
 
-    diff = sa.case(*diff_cond, else_=None if compare else "M").label(status_col)
+    diff = sa.case(*diff_cond, else_=None if compare else CompareStatus.MODIFIED).label(
+        status_col
+    )
     diff.type = String()
 
     left_right_merge = left.merge(
@@ -145,7 +160,7 @@ def compare(  # noqa: PLR0912, PLR0915, C901
         )
     )
 
-    diff_col = sa.literal("D").label(status_col)
+    diff_col = sa.literal(CompareStatus.DELETED).label(status_col)
     diff_col.type = String()
 
     right_left_merge = right.merge(
@@ -195,3 +210,92 @@ def compare(  # noqa: PLR0912, PLR0915, C901
         res = res.select_except(C(status_col))
 
     return left._evolve(query=res, signal_schema=schema)
+
+
+def compare_and_split(
+    left: "DataChain",
+    right: "DataChain",
+    on: Union[str, Sequence[str]],
+    right_on: Optional[Union[str, Sequence[str]]] = None,
+    compare: Optional[Union[str, Sequence[str]]] = None,
+    right_compare: Optional[Union[str, Sequence[str]]] = None,
+    added: bool = True,
+    deleted: bool = True,
+    modified: bool = True,
+    same: bool = False,
+) -> dict[str, "DataChain"]:
+    """Comparing two chains and returning multiple chains, one for each of `added`,
+    `deleted`, `modified` and `same` status. Result is returned in form of
+    dictionary where each item represents one of the statuses and key values
+    are `A`, `D`, `M`, `S` corresponding. Note that status column is not in the
+    resulting chains.
+
+    Parameters:
+        left: Chain to calculate diff on.
+        right: Chain to calculate diff from.
+        on: Column or list of columns to match on. If both chains have the
+            same columns then this column is enough for the match. Otherwise,
+            `right_on` parameter has to specify the columns for the other chain.
+            This value is used to find corresponding row in other dataset. If not
+            found there, row is considered as added (or removed if vice versa), and
+            if found then row can be either modified or same.
+        right_on: Optional column or list of columns
+            for the `other` to match.
+        compare: Column or list of columns to compare on. If both chains have
+            the same columns then this column is enough for the compare. Otherwise,
+            `right_compare` parameter has to specify the columns for the other
+            chain. This value is used to see if row is modified or same. If
+            not set, all columns will be used for comparison
+        right_compare: Optional column or list of columns
+                for the `other` to compare to.
+        added (bool): Whether to return chain containing only added rows.
+        deleted (bool): Whether to return chain containing only deleted rows.
+        modified (bool): Whether to return chain containing only modified rows.
+        same (bool): Whether to return chain containing only same rows.
+
+    Example:
+        ```py
+        chains = compare(
+            persons,
+            new_persons,
+            on=["id"],
+            right_on=["other_id"],
+            compare=["name"],
+            added=True,
+            deleted=True,
+            modified=True,
+            same=True,
+        )
+        ```
+    """
+    status_col = get_status_col_name()
+
+    res = _compare(
+        left,
+        right,
+        on,
+        right_on=right_on,
+        compare=compare,
+        right_compare=right_compare,
+        added=added,
+        deleted=deleted,
+        modified=modified,
+        same=same,
+        status_col=status_col,
+    )
+
+    chains = {}
+
+    def filter_by_status(compare_status) -> "DataChain":
+        return res.filter(C(status_col) == compare_status).select_except(status_col)
+
+    if added:
+        chains[CompareStatus.ADDED.value] = filter_by_status(CompareStatus.ADDED)
+    if deleted:
+        chains[CompareStatus.DELETED.value] = filter_by_status(CompareStatus.DELETED)
+    if modified:
+        chains[CompareStatus.MODIFIED.value] = filter_by_status(CompareStatus.MODIFIED)
+    if same:
+        chains[CompareStatus.SAME.value] = filter_by_status(CompareStatus.SAME)
+
+    return chains

datachain/func/__init__.py

@@ -1,4 +1,4 @@
-from sqlalchemy import case, literal
+from sqlalchemy import literal
 
 from . import array, path, random, string
 from .aggregate import (
@@ -16,7 +16,7 @@ from .aggregate import (
     sum,
 )
 from .array import cosine_distance, euclidean_distance, length, sip_hash_64
-from .conditional import greatest, least
+from .conditional import case, greatest, ifelse, least
 from .numeric import bit_and, bit_hamming_distance, bit_or, bit_xor, int_hash_64
 from .random import rand
 from .string import byte_hamming_distance
@@ -40,6 +40,7 @@ __all__ = [
     "euclidean_distance",
     "first",
     "greatest",
+    "ifelse",
     "int_hash_64",
     "least",
     "length",