datachain 0.8.9__py3-none-any.whl → 0.8.11__py3-none-any.whl

datachain/cli/parser/studio.py

@@ -1,31 +1,32 @@
-def add_studio_parser(subparsers, parent_parser) -> None:
-    studio_help = "Manage Studio authentication"
-    studio_description = (
-        "Manage authentication and settings for Studio. "
-        "Configure tokens for sharing datasets and using Studio features."
-    )
+def add_auth_parser(subparsers, parent_parser) -> None:
+    from dvc_studio_client.auth import AVAILABLE_SCOPES
+
+    auth_help = "Manage Studio authentication"
+    auth_description = "Manage authentication and settings for Studio. "
 
-    studio_parser = subparsers.add_parser(
-        "studio",
+    auth_parser = subparsers.add_parser(
+        "auth",
         parents=[parent_parser],
-        description=studio_description,
-        help=studio_help,
+        description=auth_description,
+        help=auth_help,
     )
-    studio_subparser = studio_parser.add_subparsers(
+    auth_subparser = auth_parser.add_subparsers(
         dest="cmd",
-        help="Use `datachain studio CMD --help` to display command-specific help",
+        help="Use `datachain auth CMD --help` to display command-specific help",
     )
 
-    studio_login_help = "Authenticate with Studio"
-    studio_login_description = (
+    auth_login_help = "Authenticate with Studio"
+    auth_login_description = (
         "Authenticate with Studio using default scopes. "
-        "A random name will be assigned as the token name if not specified."
+        "A random name will be assigned if the token name is not specified."
     )
-    login_parser = studio_subparser.add_parser(
+
+    allowed_scopes = ", ".join(AVAILABLE_SCOPES)
+    login_parser = auth_subparser.add_parser(
         "login",
         parents=[parent_parser],
-        description=studio_login_description,
-        help=studio_login_help,
+        description=auth_login_description,
+        help=auth_login_help,
     )
 
     login_parser.add_argument(
@@ -40,7 +41,7 @@ def add_studio_parser(subparsers, parent_parser) -> None:
         "--scopes",
         action="store",
         default=None,
-        help="Authentication token scopes",
+        help=f"Authentication token scopes. Allowed scopes: {allowed_scopes}",
     )
 
     login_parser.add_argument(
@@ -58,26 +59,26 @@ def add_studio_parser(subparsers, parent_parser) -> None:
         help="Use code-based authentication without browser",
     )
 
-    studio_logout_help = "Log out from Studio"
-    studio_logout_description = (
+    auth_logout_help = "Log out from Studio"
+    auth_logout_description = (
         "Remove the Studio authentication token from global config."
     )
 
-    studio_subparser.add_parser(
+    auth_subparser.add_parser(
         "logout",
         parents=[parent_parser],
-        description=studio_logout_description,
-        help=studio_logout_help,
+        description=auth_logout_description,
+        help=auth_logout_help,
     )
 
-    studio_team_help = "Set default team for Studio operations"
-    studio_team_description = "Set the default team for Studio operations."
+    auth_team_help = "Set default team for Studio operations"
+    auth_team_description = "Set the default team for Studio operations."
 
-    team_parser = studio_subparser.add_parser(
+    team_parser = auth_subparser.add_parser(
         "team",
         parents=[parent_parser],
-        description=studio_team_description,
-        help=studio_team_help,
+        description=auth_team_description,
+        help=auth_team_help,
     )
     team_parser.add_argument(
         "team_name",
@@ -91,12 +92,12 @@ def add_studio_parser(subparsers, parent_parser) -> None:
         help="Set team globally for all projects",
     )
 
-    studio_token_help = "View Studio authentication token"  # noqa: S105
-    studio_token_description = "Display the current authentication token for Studio."  # noqa: S105
+    auth_token_help = "View Studio authentication token"  # noqa: S105
+    auth_token_description = "Display the current authentication token for Studio."  # noqa: S105
 
-    studio_subparser.add_parser(
+    auth_subparser.add_parser(
         "token",
         parents=[parent_parser],
-        description=studio_token_description,
-        help=studio_token_help,
+        description=auth_token_description,
+        help=auth_token_help,
     )

datachain/cli/parser/utils.py

@@ -30,7 +30,25 @@ def add_sources_arg(parser: ArgumentParser, nargs: Union[str, int] = "+") -> Act
         "sources",
         type=str,
         nargs=nargs,
-        help="Data sources - paths to cloud storage directories",
+        help="Data sources - paths to source storage directories or files",
+    )
+
+
+def add_anon_arg(parser: ArgumentParser) -> None:
+    parser.add_argument(
+        "--anon",
+        action="store_true",
+        help="Use anonymous access to storage",
+    )
+
+
+def add_update_arg(parser: ArgumentParser) -> None:
+    parser.add_argument(
+        "-u",
+        "--update",
+        action="count",
+        default=0,
+        help="Update cached list of files for the sources",
     )
 
 

datachain/cli/utils.py

@@ -87,7 +87,7 @@ def get_logging_level(args: Namespace) -> int:
 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'."
+            "Not logged in to Studio. Log in with 'datachain auth login'."
         )
 
     if local or studio:

datachain/client/fsspec.py

@@ -3,6 +3,7 @@ import functools
 import logging
 import multiprocessing
 import os
+import posixpath
 import re
 import sys
 from abc import ABC, abstractmethod
@@ -25,7 +26,7 @@ from fsspec.asyn import get_loop, sync
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 from tqdm.auto import tqdm
 
-from datachain.cache import DataChainCache
+from datachain.cache import Cache
 from datachain.client.fileslice import FileWrapper
 from datachain.error import ClientError as DataChainClientError
 from datachain.nodes_fetcher import NodesFetcher
@@ -74,9 +75,7 @@ class Client(ABC):
     PREFIX: ClassVar[str]
     protocol: ClassVar[str]
 
-    def __init__(
-        self, name: str, fs_kwargs: dict[str, Any], cache: DataChainCache
-    ) -> None:
+    def __init__(self, name: str, fs_kwargs: dict[str, Any], cache: Cache) -> None:
         self.name = name
         self.fs_kwargs = fs_kwargs
         self._fs: Optional[AbstractFileSystem] = None
@@ -122,7 +121,7 @@ class Client(ABC):
         return cls.get_uri(storage_name), rel_path
 
     @staticmethod
-    def get_client(source: str, cache: DataChainCache, **kwargs) -> "Client":
+    def get_client(source: str, cache: Cache, **kwargs) -> "Client":
         cls = Client.get_implementation(source)
         storage_url, _ = cls.split_url(source)
         if os.name == "nt":
@@ -145,7 +144,7 @@ class Client(ABC):
     def from_name(
         cls,
         name: str,
-        cache: DataChainCache,
+        cache: Cache,
         kwargs: dict[str, Any],
     ) -> "Client":
         return cls(name, kwargs, cache)
@@ -154,7 +153,7 @@ class Client(ABC):
     def from_source(
         cls,
         uri: "StorageURI",
-        cache: DataChainCache,
+        cache: Cache,
         **kwargs,
     ) -> "Client":
         return cls(cls.FS_CLASS._strip_protocol(uri), kwargs, cache)
@@ -390,8 +389,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":
+    def upload(self, data: bytes, path: str) -> "File":
         full_path = self.get_full_path(path)
+
+        parent = posixpath.dirname(full_path)
+        self.fs.makedirs(parent, exist_ok=True)
+
         self.fs.pipe_file(full_path, data)
         file_info = self.fs.info(full_path)
         return self.info_to_file(file_info, path)

datachain/client/local.py

@@ -12,7 +12,7 @@ from datachain.lib.file import File
 from .fsspec import Client
 
 if TYPE_CHECKING:
-    from datachain.cache import DataChainCache
+    from datachain.cache import Cache
     from datachain.dataset import StorageURI
 
 
@@ -25,7 +25,7 @@ class FileClient(Client):
         self,
         name: str,
         fs_kwargs: dict[str, Any],
-        cache: "DataChainCache",
+        cache: "Cache",
         use_symlinks: bool = False,
     ) -> None:
         super().__init__(name, fs_kwargs, cache)
@@ -82,7 +82,7 @@ class FileClient(Client):
         return bucket, path
 
     @classmethod
-    def from_name(cls, name: str, cache: "DataChainCache", kwargs) -> "FileClient":
+    def from_name(cls, name: str, cache: "Cache", kwargs) -> "FileClient":
         use_symlinks = kwargs.pop("use_symlinks", False)
         return cls(name, kwargs, cache, use_symlinks=use_symlinks)
 
@@ -90,7 +90,7 @@ class FileClient(Client):
     def from_source(
         cls,
         uri: str,
-        cache: "DataChainCache",
+        cache: "Cache",
         use_symlinks: bool = False,
         **kwargs,
     ) -> "FileClient":

datachain/data_storage/schema.py

@@ -200,7 +200,7 @@ class DataTable:
         columns: Sequence["sa.Column"] = (),
         metadata: Optional["sa.MetaData"] = None,
     ):
-        # copy columns, since re-using the same objects from another table
+        # copy columns, since reusing the same objects from another table
         # may raise an error
         columns = cls.sys_columns() + [cls.copy_column(c) for c in columns]
         columns = dedup_columns(columns)

datachain/data_storage/sqlite.py

@@ -19,6 +19,7 @@ from sqlalchemy import MetaData, Table, UniqueConstraint, exists, select
 from sqlalchemy.dialects import sqlite
 from sqlalchemy.schema import CreateIndex, CreateTable, DropTable
 from sqlalchemy.sql import func
+from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
 from sqlalchemy.sql.expression import bindparam, cast
 from sqlalchemy.sql.selectable import Select
 from tqdm.auto import tqdm
@@ -40,7 +41,6 @@ if TYPE_CHECKING:
     from sqlalchemy.schema import SchemaItem
     from sqlalchemy.sql._typing import _FromClauseArgument, _OnClauseArgument
     from sqlalchemy.sql.elements import ColumnElement
-    from sqlalchemy.sql.selectable import Join
     from sqlalchemy.types import TypeEngine
 
     from datachain.lib.file import File
@@ -654,16 +654,47 @@ class SQLiteWarehouse(AbstractWarehouse):
         right: "_FromClauseArgument",
         onclause: "_OnClauseArgument",
         inner: bool = True,
-    ) -> "Join":
+        full: bool = False,
+        columns=None,
+    ) -> "Select":
         """
         Join two tables together.
         """
-        return sqlalchemy.join(
-            left,
-            right,
-            onclause,
-            isouter=not inner,
+        if not full:
+            join_query = sqlalchemy.join(
+                left,
+                right,
+                onclause,
+                isouter=not inner,
+            )
+            return sqlalchemy.select(*columns).select_from(join_query)
+
+        left_right_join = sqlalchemy.select(*columns).select_from(
+            sqlalchemy.join(left, right, onclause, isouter=True)
         )
+        right_left_join = sqlalchemy.select(*columns).select_from(
+            sqlalchemy.join(right, left, onclause, isouter=True)
+        )
+
+        def add_left_rows_filter(exp: BinaryExpression):
+            """
+            Adds filter to right_left_join to remove unmatched left table rows by
+            getting column names that need to be NULL from BinaryExpressions in onclause
+            """
+            return right_left_join.where(
+                getattr(left.c, exp.left.name) == None  # type: ignore[union-attr] # noqa: E711
+            )
+
+        if isinstance(onclause, BinaryExpression):
+            right_left_join = add_left_rows_filter(onclause)
+
+        if isinstance(onclause, BooleanClauseList):
+            for c in onclause.get_children():
+                if isinstance(c, BinaryExpression):
+                    right_left_join = add_left_rows_filter(c)
+
+        union = sqlalchemy.union(left_right_join, right_left_join).subquery()
+        return sqlalchemy.select(*union.c).select_from(union)
 
     def create_pre_udf_table(self, query: "Select") -> "Table":
         """

datachain/data_storage/warehouse.py

@@ -31,7 +31,7 @@ if TYPE_CHECKING:
         _FromClauseArgument,
         _OnClauseArgument,
     )
-    from sqlalchemy.sql.selectable import Join, Select
+    from sqlalchemy.sql.selectable import Select
     from sqlalchemy.types import TypeEngine
 
     from datachain.data_storage import schema
@@ -873,7 +873,7 @@ class AbstractWarehouse(ABC, Serializable):
         right: "_FromClauseArgument",
         onclause: "_OnClauseArgument",
         inner: bool = True,
-    ) -> "Join":
+    ) -> "Select":
         """
         Join two tables together.
         """

datachain/dataset.py

@@ -91,7 +91,7 @@ class DatasetDependency:
         if self.type == DatasetDependencyType.DATASET:
             return self.name
 
-        list_dataset_name, _, _ = parse_listing_uri(self.name.strip("/"), None, {})
+        list_dataset_name, _, _ = parse_listing_uri(self.name.strip("/"), {})
         assert list_dataset_name
         return list_dataset_name
 

datachain/error.py

@@ -1,3 +1,15 @@
+import botocore.errorfactory
+import botocore.exceptions
+import gcsfs.retry
+
+REMOTE_ERRORS = (
+    gcsfs.retry.HttpError,  # GCS
+    OSError,  # GCS
+    botocore.exceptions.BotoCoreError,  # S3
+    ValueError,  # Azure
+)
+
+
 class DataChainError(RuntimeError):
     pass
 

datachain/func/__init__.py

@@ -16,7 +16,7 @@ from .aggregate import (
     sum,
 )
 from .array import cosine_distance, euclidean_distance, length, sip_hash_64
-from .conditional import case, greatest, ifelse, least
+from .conditional import case, greatest, ifelse, isnone, 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
@@ -42,6 +42,7 @@ __all__ = [
     "greatest",
     "ifelse",
     "int_hash_64",
+    "isnone",
     "least",
     "length",
     "literal",

datachain/func/conditional.py

@@ -1,14 +1,15 @@
-from typing import Union
+from typing import Optional, Union
 
+from sqlalchemy import ColumnElement
 from sqlalchemy import case as sql_case
-from sqlalchemy.sql.elements import BinaryExpression
 
 from datachain.lib.utils import DataChainParamsError
+from datachain.query.schema import Column
 from datachain.sql.functions import conditional
 
 from .func import ColT, Func
 
-CaseT = Union[int, float, complex, bool, str]
+CaseT = Union[int, float, complex, bool, str, Func]
 
 
 def greatest(*args: Union[ColT, float]) -> Func:
@@ -87,17 +88,21 @@ def least(*args: Union[ColT, float]) -> Func:
     )
 
 
-def case(*args: tuple[BinaryExpression, CaseT], else_=None) -> Func:
+def case(
+    *args: tuple[Union[ColumnElement, Func], CaseT], else_: Optional[CaseT] = None
+) -> Func:
     """
     Returns the case function that produces case expression which has a list of
-    conditions and corresponding results. Results can only be python primitives
-    like string, numbes or booleans. Result type is inferred from condition results.
+    conditions and corresponding results. Results can be python primitives like string,
+    numbers or booleans but can also be other nested function (including case function).
+    Result type is inferred from condition results.
 
     Args:
-        args (tuple(BinaryExpression, value(str | int | float | complex | bool):
-            - Tuple of binary expression and values pair which corresponds to one
-            case condition - value
-        else_ (str | int | float | complex | bool): else value in case expression
+        args (tuple((ColumnElement, Func), (str | int | float | complex | bool, Func))):
+            Tuple of condition and values pair.
+        else_ (str | int | float | complex | bool, Func): optional else value in case
+            expression. If omitted, and no case conditions are satisfied, the result
+            will be None (NULL in DB).
 
     Returns:
         Func: A Func object that represents the case function.
@@ -111,15 +116,24 @@ def case(*args: tuple[BinaryExpression, CaseT], else_=None) -> Func:
     """
     supported_types = [int, float, complex, str, bool]
 
-    type_ = type(else_) if else_ else None
+    def _get_type(val):
+        if isinstance(val, Func):
+            # nested functions
+            return val.result_type
+        return type(val)
 
     if not args:
         raise DataChainParamsError("Missing statements")
 
+    type_ = _get_type(else_) if else_ is not None else None
+
     for arg in args:
-        if type_ and not isinstance(arg[1], type_):
-            raise DataChainParamsError("Statement values must be of the same type")
-        type_ = type(arg[1])
+        arg_type = _get_type(arg[1])
+        if type_ and arg_type != type_:
+            raise DataChainParamsError(
+                f"Statement values must be of the same type, got {type_} and {arg_type}"
+            )
+        type_ = arg_type
 
     if type_ not in supported_types:
         raise DataChainParamsError(
@@ -127,20 +141,25 @@ def case(*args: tuple[BinaryExpression, CaseT], else_=None) -> Func:
         )
 
     kwargs = {"else_": else_}
-    return Func("case", inner=sql_case, args=args, kwargs=kwargs, result_type=type_)
+
+    return Func("case", inner=sql_case, cols=args, kwargs=kwargs, result_type=type_)
 
 
-def ifelse(condition: BinaryExpression, if_val: CaseT, else_val: CaseT) -> Func:
+def ifelse(
+    condition: Union[ColumnElement, Func], if_val: CaseT, else_val: CaseT
+) -> Func:
     """
     Returns the ifelse function that produces if expression which has a condition
-    and values for true and false outcome. Results can only be python primitives
-    like string, numbes or booleans. Result type is inferred from the values.
+    and values for true and false outcome. Results can be one of python primitives
+    like string, numbers or booleans, but can also be nested functions.
+    Result type is inferred from the values.
 
     Args:
-        condition: BinaryExpression - condition which is evaluated
-        if_val: (str | int | float | complex | bool): value for true condition outcome
-        else_val: (str | int | float | complex | bool): value for false condition
-         outcome
+        condition (ColumnElement, Func):  Condition which is evaluated.
+        if_val (str | int | float | complex | bool, Func): Value for true
+            condition outcome.
+        else_val (str | int | float | complex | bool, Func): Value for false condition
+            outcome.
 
     Returns:
         Func: A Func object that represents the ifelse function.
@@ -148,8 +167,33 @@ def ifelse(condition: BinaryExpression, if_val: CaseT, else_val: CaseT) -> Func:
     Example:
         ```py
         dc.mutate(
-            res=func.ifelse(C("num") > 0, "P", "N"),
+            res=func.ifelse(isnone("col"), "EMPTY", "NOT_EMPTY")
         )
         ```
     """
     return case((condition, if_val), else_=else_val)
+
+
+def isnone(col: Union[str, Column]) -> Func:
+    """
+    Returns True if column value is None, otherwise False.
+
+    Args:
+        col (str | Column): Column to check if it's None or not.
+            If a string is provided, it is assumed to be the name of the column.
+
+    Returns:
+        Func: A Func object that represents the conditional to check if column is None.
+
+    Example:
+        ```py
+        dc.mutate(test=ifelse(isnone("col"), "EMPTY", "NOT_EMPTY"))
+        ```
+    """
+    from datachain import C
+
+    if isinstance(col, str):
+        # if string, it is assumed to be the name of the column
+        col = C(col)
+
+    return case((col.is_(None) if col is not None else True, True), else_=False)

datachain/func/func.py

@@ -23,7 +23,7 @@ if TYPE_CHECKING:
     from .window import Window
 
 
-ColT = Union[str, ColumnElement, "Func"]
+ColT = Union[str, ColumnElement, "Func", tuple]
 
 
 class Func(Function):
@@ -78,7 +78,7 @@ class Func(Function):
         return (
             [
                 col
-                if isinstance(col, (Func, BindParameter, Case, Comparator))
+                if isinstance(col, (Func, BindParameter, Case, Comparator, tuple))
                 else ColumnMeta.to_db_name(
                     col.name if isinstance(col, ColumnElement) else col
                 )
@@ -381,17 +381,24 @@ class Func(Function):
         col_type = self.get_result_type(signals_schema)
         sql_type = python_to_sql(col_type)
 
-        def get_col(col: ColT) -> ColT:
+        def get_col(col: ColT, string_as_literal=False) -> ColT:
+            # string_as_literal is used only for conditionals like `case()` where
+            # literals are nested inside ColT as we have tuples of condition - values
+            # and if user wants to set some case value as column, explicit `C("col")`
+            # syntax must be used to distinguish from literals
+            if isinstance(col, tuple):
+                return tuple(get_col(x, string_as_literal=True) for x in col)
             if isinstance(col, Func):
                 return col.get_column(signals_schema, table=table)
-            if isinstance(col, str):
+            if isinstance(col, str) and not string_as_literal:
                 column = Column(col, sql_type)
                 column.table = table
                 return column
             return col
 
         cols = [get_col(col) for col in self._db_cols]
-        func_col = self.inner(*cols, *self.args, **self.kwargs)
+        kwargs = {k: get_col(v, string_as_literal=True) for k, v in self.kwargs.items()}
+        func_col = self.inner(*cols, *self.args, **kwargs)
 
         if self.is_window:
             if not self.window:
@@ -416,6 +423,11 @@ class Func(Function):
 
 
 def get_db_col_type(signals_schema: "SignalSchema", col: ColT) -> "DataType":
+    if isinstance(col, tuple):
+        raise DataChainParamsError(
+            "Cannot get type from tuple, please provide type hint to the function"
+        )
+
     if isinstance(col, Func):
         return col.get_result_type(signals_schema)
 

datachain/lib/convert/python_to_sql.py

@@ -52,15 +52,15 @@ def python_to_sql(typ):  # noqa: PLR0911
 
     args = get_args(typ)
     if inspect.isclass(orig) and (issubclass(list, orig) or issubclass(tuple, orig)):
-        if args is None or len(args) != 1:
+        if args is None:
             raise TypeError(f"Cannot resolve type '{typ}' for flattening features")
 
         args0 = args[0]
         if ModelStore.is_pydantic(args0):
             return Array(JSON())
 
-        next_type = python_to_sql(args0)
-        return Array(next_type)
+        list_type = list_of_args_to_type(args)
+        return Array(list_type)
 
     if orig is Annotated:
         # Ignoring annotations
@@ -82,6 +82,18 @@ def python_to_sql(typ):  # noqa: PLR0911
     raise TypeError(f"Cannot recognize type {typ}")
 
 
+def list_of_args_to_type(args) -> SQLType:
+    first_type = python_to_sql(args[0])
+    for next_arg in args[1:]:
+        try:
+            next_type = python_to_sql(next_arg)
+            if next_type != first_type:
+                return JSON()
+        except TypeError:
+            return JSON()
+    return first_type
+
+
 def _is_json_inside_union(orig, args) -> bool:
     if orig == Union and len(args) >= 2:
         # List in JSON: Union[dict, list[dict]]