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

datachain/cli/parser/job.py

@@ -0,0 +1,120 @@
+def add_jobs_parser(subparsers, parent_parser) -> None:
+    jobs_help = "Commands to handle the Job running with Iterative Studio"
+    jobs_description = (
+        "This will help us to run, cancel and view the status of the job in Studio. "
+    )
+    jobs_parser = subparsers.add_parser(
+        "job", parents=[parent_parser], description=jobs_description, help=jobs_help
+    )
+    jobs_subparser = jobs_parser.add_subparsers(
+        dest="cmd",
+        help="Use `DataChain studio CMD --help` to display command-specific help.",
+        required=True,
+    )
+
+    studio_run_help = "Run a job in Studio"
+    studio_run_description = "This command runs a job in Studio."
+
+    studio_run_parser = jobs_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.",
+    )
+
+    studio_cancel_help = "Cancel a job in Studio"
+    studio_cancel_description = "This command cancels a job in Studio."
+
+    studio_cancel_parser = jobs_subparser.add_parser(
+        "cancel",
+        parents=[parent_parser],
+        description=studio_cancel_description,
+        help=studio_cancel_help,
+    )
+
+    studio_cancel_parser.add_argument(
+        "job_id",
+        action="store",
+        help="The job ID to cancel.",
+    )
+    studio_cancel_parser.add_argument(
+        "--team",
+        action="store",
+        default=None,
+        help="The team to cancel a job for. By default, it will use team from config.",
+    )
+
+    studio_log_help = "Show the logs and latest status of Jobs in Studio"
+    studio_log_description = (
+        "This will display the logs and latest status of jobs in Studio"
+    )
+
+    studio_log_parser = jobs_subparser.add_parser(
+        "logs",
+        parents=[parent_parser],
+        description=studio_log_description,
+        help=studio_log_help,
+    )
+
+    studio_log_parser.add_argument(
+        "job_id",
+        action="store",
+        help="The job ID to show the logs.",
+    )
+    studio_log_parser.add_argument(
+        "--team",
+        action="store",
+        default=None,
+        help="The team to check the logs. By default, it will use team from config.",
+    )

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/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)
@@ -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/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, 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

datachain/func/conditional.py

@@ -1,5 +1,9 @@
 from typing import Union
 
+from sqlalchemy import case as sql_case
+from sqlalchemy.sql.elements import BinaryExpression
+
+from datachain.lib.utils import DataChainParamsError
 from datachain.sql.functions import conditional
 
 from .func import ColT, Func
@@ -79,3 +83,51 @@ def least(*args: Union[ColT, float]) -> Func:
     return Func(
         "least", inner=conditional.least, cols=cols, args=func_args, result_type=int
     )
+
+
+def case(
+    *args: tuple[BinaryExpression, Union[int, float, complex, bool, str]], else_=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.
+
+    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
+
+    Returns:
+        Func: A Func object that represents the case function.
+
+    Example:
+        ```py
+        dc.mutate(
+            res=func.case((C("num") > 0, "P"), (C("num") < 0, "N"), else_="Z"),
+        )
+        ```
+
+    Note:
+        - Result column will always be of the same type as the input columns.
+    """
+    supported_types = [int, float, complex, str, bool]
+
+    type_ = type(else_) if else_ else None
+
+    if not args:
+        raise DataChainParamsError("Missing case statements")
+
+    for arg in args:
+        if type_ and not isinstance(arg[1], type_):
+            raise DataChainParamsError("Case statement values must be of the same type")
+        type_ = type(arg[1])
+
+    if type_ not in supported_types:
+        raise DataChainParamsError(
+            f"Case supports only python literals ({supported_types}) for values"
+        )
+
+    kwargs = {"else_": else_}
+    return Func("case", inner=sql_case, args=args, kwargs=kwargs, result_type=type_)

datachain/func/func.py

@@ -35,6 +35,7 @@ class Func(Function):
         inner: Callable,
         cols: Optional[Sequence[ColT]] = None,
         args: Optional[Sequence[Any]] = None,
+        kwargs: Optional[dict[str, Any]] = None,
         result_type: Optional["DataType"] = None,
         is_array: bool = False,
         is_window: bool = False,
@@ -45,6 +46,7 @@ class Func(Function):
         self.inner = inner
         self.cols = cols or []
         self.args = args or []
+        self.kwargs = kwargs or {}
         self.result_type = result_type
         self.is_array = is_array
         self.is_window = is_window
@@ -63,6 +65,7 @@ class Func(Function):
             self.inner,
             self.cols,
             self.args,
+            self.kwargs,
             self.result_type,
             self.is_array,
             self.is_window,
@@ -333,6 +336,7 @@ class Func(Function):
             self.inner,
             self.cols,
             self.args,
+            self.kwargs,
             self.result_type,
             self.is_array,
             self.is_window,
@@ -387,7 +391,7 @@ class Func(Function):
             return col
 
         cols = [get_col(col) for col in self._db_cols]
-        func_col = self.inner(*cols, *self.args)
+        func_col = self.inner(*cols, *self.args, **self.kwargs)
 
         if self.is_window:
             if not self.window:

datachain/lib/arrow.py

@@ -149,6 +149,10 @@ def infer_schema(chain: "DataChain", **kwargs) -> pa.Schema:
     for file in chain.collect("file"):
         ds = dataset(file.get_path(), filesystem=file.get_fs(), **kwargs)  # type: ignore[union-attr]
         schemas.append(ds.schema)
+    if not schemas:
+        raise ValueError(
+            "Cannot infer schema (no files to process or can't access them)"
+        )
     return pa.unify_schemas(schemas)
 
 

datachain/lib/dc.py

@@ -1882,6 +1882,9 @@ class DataChain:
                     "`nrows` only supported for csv and json formats.",
                 )
 
+        if "file" not in self.schema or not self.count():
+            raise DatasetPrepareError(self.name, "no files to parse.")
+
         schema = None
         col_names = output if isinstance(output, Sequence) else None
         if col_names or not output:

datachain/lib/file.py

@@ -364,7 +364,7 @@ class File(DataModel):
 
         try:
             info = client.fs.info(client.get_full_path(self.path))
-            converted_info = client.info_to_file(info, self.source)
+            converted_info = client.info_to_file(info, self.path)
             return type(self)(
                 path=self.path,
                 source=self.source,

datachain/lib/listing.py

@@ -85,6 +85,24 @@ def ls(
     return dc.filter(pathfunc.parent(_file_c("path")) == path.lstrip("/").rstrip("/*"))
 
 
+def _isfile(client: "Client", path: str) -> bool:
+    """
+    Returns True if uri points to a file
+    """
+    try:
+        info = client.fs.info(path)
+        name = info.get("name")
+        # case for special simulated directories on some clouds
+        # e.g. Google creates a zero byte file with the same name as the
+        # directory with a trailing slash at the end
+        if not name or name.endswith("/"):
+            return False
+
+        return info["type"] == "file"
+    except:  # noqa: E722
+        return False
+
+
 def parse_listing_uri(uri: str, cache, client_config) -> tuple[Optional[str], str, str]:
     """
     Parsing uri and returns listing dataset name, listing uri and listing path
@@ -94,7 +112,7 @@ def parse_listing_uri(uri: str, cache, client_config) -> tuple[Optional[str], st
     storage_uri, path = Client.parse_url(uri)
     telemetry.log_param("client", client.PREFIX)
 
-    if not uri.endswith("/") and client.fs.isfile(uri):
+    if not uri.endswith("/") and _isfile(client, uri):
         return None, f'{storage_uri}/{path.lstrip("/")}', path
     if uses_glob(path):
         lst_uri_path = posixpath.dirname(path)