datachain 0.6.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

datachain/remote/studio.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import os
 from collections.abc import Iterable, Iterator
 from datetime import datetime, timedelta, timezone
 from struct import unpack
@@ -10,8 +11,10 @@ from typing import (
     TypeVar,
 )
 
+from datachain.config import Config
 from datachain.dataset import DatasetStats
-from datachain.utils import retry_with_backoff
+from datachain.error import DataChainError
+from datachain.utils import STUDIO_URL, retry_with_backoff
 
 T = TypeVar("T")
 LsData = Optional[list[dict[str, Any]]]
@@ -54,14 +57,54 @@ class Response(Generic[T]):
 
 
 class StudioClient:
-    def __init__(
-        self, url: str, username: str, token: str, timeout: float = 3600.0
-    ) -> None:
+    def __init__(self, timeout: float = 3600.0, team: Optional[str] = None) -> None:
         self._check_dependencies()
-        self.url = url.rstrip("/")
-        self.username = username
-        self.token = token
         self.timeout = timeout
+        self._config = None
+        self._team = team
+
+    @property
+    def token(self) -> str:
+        token = os.environ.get("DVC_STUDIO_TOKEN") or self.config.get("token")
+
+        if not token:
+            raise DataChainError(
+                "Studio token is not set. Use `datachain studio login` "
+                "or environment variable `DVC_STUDIO_TOKEN` to set it."
+            )
+
+        return token
+
+    @property
+    def url(self) -> str:
+        return (
+            os.environ.get("DVC_STUDIO_URL") or self.config.get("url") or STUDIO_URL
+        ) + "/api"
+
+    @property
+    def config(self) -> dict:
+        if self._config is None:
+            self._config = Config().read().get("studio", {})
+        return self._config  # type: ignore [return-value]
+
+    @property
+    def team(self) -> str:
+        if self._team is None:
+            self._team = self._get_team()
+        return self._team
+
+    def _get_team(self) -> str:
+        team = os.environ.get("DVC_STUDIO_TEAM") or self.config.get("team")
+
+        if not team:
+            raise DataChainError(
+                "Studio team is not set. "
+                "Use `datachain studio team <team_name>` "
+                "or environment variable `DVC_STUDIO_TEAM` to set it."
+                "You can also set it in the config file as team under studio."
+            )
+
+        return team
 
     def _check_dependencies(self) -> None:
         try:
@@ -80,7 +123,7 @@ class StudioClient:
 
         response = requests.post(
             f"{self.url}/{route}",
-            json={**data, "team_name": self.username},
+            json={**data, "team_name": self.team},
             headers={
                 "Content-Type": "application/json",
                 "Authorization": f"token {self.token}",
@@ -108,7 +151,7 @@ class StudioClient:
 
         response = requests.post(
             f"{self.url}/{route}",
-            json={**data, "team_name": self.username},
+            json={**data, "team_name": self.team},
             headers={
                 "Content-Type": "application/json",
                 "Authorization": f"token {self.token}",
@@ -174,6 +217,9 @@ class StudioClient:
             response = self._send_request_msgpack("ls", {"source": path})
             yield path, response
 
+    def ls_datasets(self) -> Response[LsData]:
+        return self._send_request("datachain/ls-datasets", {})
+
     def dataset_info(self, name: str) -> Response[DatasetInfoData]:
         def _parse_dataset_info(dataset_info):
             _parse_dates(dataset_info, ["created_at", "finished_at"])
@@ -182,7 +228,7 @@ class StudioClient:
 
             return dataset_info
 
-        response = self._send_request("dataset-info", {"dataset_name": name})
+        response = self._send_request("datachain/dataset-info", {"dataset_name": name})
         if response.ok:
             response.data = _parse_dataset_info(response.data)
         return response
@@ -192,13 +238,14 @@ class StudioClient:
     ) -> Response[DatasetRowsData]:
         req_data = {"dataset_name": name, "dataset_version": version}
         return self._send_request_msgpack(
-            "dataset-rows",
+            "datachain/dataset-rows",
             {**req_data, "offset": offset, "limit": DATASET_ROWS_CHUNK_SIZE},
         )
 
     def dataset_stats(self, name: str, version: int) -> Response[DatasetStatsData]:
         response = self._send_request(
-            "dataset-stats", {"dataset_name": name, "dataset_version": version}
+            "datachain/dataset-stats",
+            {"dataset_name": name, "dataset_version": version},
         )
         if response.ok:
             response.data = DatasetStats(**response.data)
@@ -208,12 +255,14 @@ class StudioClient:
         self, name: str, version: int
     ) -> Response[DatasetExportSignedUrls]:
         return self._send_request(
-            "dataset-export", {"dataset_name": name, "dataset_version": version}
+            "datachain/dataset-export",
+            {"dataset_name": name, "dataset_version": version},
         )
 
     def dataset_export_status(
         self, name: str, version: int
     ) -> Response[DatasetExportStatus]:
         return self._send_request(
-            "dataset-export-status", {"dataset_name": name, "dataset_version": version}
+            "datachain/dataset-export-status",
+            {"dataset_name": name, "dataset_version": version},
         )

datachain/sql/functions/__init__.py

@@ -1,7 +1,7 @@
 from sqlalchemy.sql.expression import func
 
 from . import array, path, string
-from .array import avg
+from .aggregate import avg
 from .conditional import greatest, least
 from .random import rand
 

datachain/sql/functions/aggregate.py

@@ -0,0 +1,47 @@
+from sqlalchemy.sql.functions import GenericFunction, ReturnTypeFromArgs
+
+from datachain.sql.types import Float, String
+from datachain.sql.utils import compiler_not_implemented
+
+
+class avg(GenericFunction):  # noqa: N801
+    """
+    Returns the average of the column.
+    """
+
+    type = Float()
+    package = "array"
+    name = "avg"
+    inherit_cache = True
+
+
+class group_concat(GenericFunction):  # noqa: N801
+    """
+    Returns the concatenated string of the column.
+    """
+
+    type = String()
+    package = "array"
+    name = "group_concat"
+    inherit_cache = True
+
+
+class any_value(ReturnTypeFromArgs):  # noqa: N801
+    """
+    Returns first value of the column.
+    """
+
+    inherit_cache = True
+
+
+class collect(ReturnTypeFromArgs):  # noqa: N801
+    """
+    Returns an array of the column.
+    """
+
+    inherit_cache = True
+
+
+compiler_not_implemented(avg)
+compiler_not_implemented(group_concat)
+compiler_not_implemented(any_value)

datachain/sql/functions/array.py

@@ -44,15 +44,7 @@ class sip_hash_64(GenericFunction):  # noqa: N801
     inherit_cache = True
 
 
-class avg(GenericFunction):  # noqa: N801
-    type = Float()
-    package = "array"
-    name = "avg"
-    inherit_cache = True
-
-
 compiler_not_implemented(cosine_distance)
 compiler_not_implemented(euclidean_distance)
 compiler_not_implemented(length)
 compiler_not_implemented(sip_hash_64)
-compiler_not_implemented(avg)

datachain/sql/sqlite/base.py

@@ -14,7 +14,7 @@ from sqlalchemy.sql.elements import literal
 from sqlalchemy.sql.expression import case
 from sqlalchemy.sql.functions import func
 
-from datachain.sql.functions import array, conditional, random, string
+from datachain.sql.functions import aggregate, array, conditional, random, string
 from datachain.sql.functions import path as sql_path
 from datachain.sql.selectable import Values, base_values_compiler
 from datachain.sql.sqlite.types import (
@@ -84,7 +84,10 @@ def setup():
     compiles(conditional.least, "sqlite")(compile_least)
     compiles(Values, "sqlite")(compile_values)
     compiles(random.rand, "sqlite")(compile_rand)
-    compiles(array.avg, "sqlite")(compile_avg)
+    compiles(aggregate.avg, "sqlite")(compile_avg)
+    compiles(aggregate.group_concat, "sqlite")(compile_group_concat)
+    compiles(aggregate.any_value, "sqlite")(compile_any_value)
+    compiles(aggregate.collect, "sqlite")(compile_collect)
 
     if load_usearch_extension(sqlite3.connect(":memory:")):
         compiles(array.cosine_distance, "sqlite")(compile_cosine_distance_ext)
@@ -400,6 +403,21 @@ def compile_avg(element, compiler, **kwargs):
     return compiler.process(func.avg(*element.clauses.clauses), **kwargs)
 
 
+def compile_group_concat(element, compiler, **kwargs):
+    return compiler.process(func.aggregate_strings(*element.clauses.clauses), **kwargs)
+
+
+def compile_any_value(element, compiler, **kwargs):
+    # use bare column to return any value from the group,
+    # this is documented behavior for sqlite,
+    # see https://www.sqlite.org/lang_select.html#bare_columns_in_an_aggregate_query
+    return compiler.process(*element.clauses.clauses, **kwargs)
+
+
+def compile_collect(element, compiler, **kwargs):
+    return compiler.process(func.json_group_array(*element.clauses.clauses), **kwargs)
+
+
 def load_usearch_extension(conn) -> bool:
     try:
         # usearch is part of the vector optional dependencies

datachain/studio.py

@@ -0,0 +1,129 @@
+import os
+from typing import TYPE_CHECKING
+
+from datachain.catalog.catalog import raise_remote_error
+from datachain.config import Config, ConfigLevel
+from datachain.error import DataChainError
+from datachain.remote.studio import StudioClient
+from datachain.utils import STUDIO_URL
+
+if TYPE_CHECKING:
+    from argparse import Namespace
+
+POST_LOGIN_MESSAGE = (
+    "Once you've logged in, return here "
+    "and you'll be ready to start using DataChain with Studio."
+)
+
+
+def process_studio_cli_args(args: "Namespace"):
+    if args.cmd == "login":
+        return login(args)
+    if args.cmd == "logout":
+        return logout()
+    if args.cmd == "token":
+        return token()
+    if args.cmd == "datasets":
+        return list_datasets(args)
+    if args.cmd == "team":
+        return set_team(args)
+    raise DataChainError(f"Unknown command '{args.cmd}'.")
+
+
+def set_team(args: "Namespace"):
+    level = ConfigLevel.GLOBAL if args.__dict__.get("global") else ConfigLevel.LOCAL
+    config = Config(level)
+    with config.edit() as conf:
+        studio_conf = conf.get("studio", {})
+        studio_conf["team"] = args.team_name
+        conf["studio"] = studio_conf
+
+    print(f"Set default team to '{args.team_name}' in {config.config_file()}")
+
+
+def login(args: "Namespace"):
+    from dvc_studio_client.auth import StudioAuthError, get_access_token
+
+    config = Config().read().get("studio", {})
+    name = args.name
+    hostname = (
+        args.hostname
+        or os.environ.get("DVC_STUDIO_URL")
+        or config.get("url")
+        or STUDIO_URL
+    )
+    scopes = args.scopes
+
+    if config.get("url", hostname) == hostname and "token" in config:
+        raise DataChainError(
+            "Token already exists. "
+            "To login with a different token, "
+            "logout using `datachain studio logout`."
+        )
+
+    open_browser = not args.no_open
+    try:
+        _, access_token = get_access_token(
+            token_name=name,
+            hostname=hostname,
+            scopes=scopes,
+            open_browser=open_browser,
+            client_name="DataChain",
+            post_login_message=POST_LOGIN_MESSAGE,
+        )
+    except StudioAuthError as exc:
+        raise DataChainError(f"Failed to authenticate with Studio: {exc}") from exc
+
+    config_path = save_config(hostname, access_token)
+    print(f"Authentication complete. Saved token to {config_path}.")
+    return 0
+
+
+def logout():
+    with Config(ConfigLevel.GLOBAL).edit() as conf:
+        token = conf.get("studio", {}).get("token")
+        if not token:
+            raise DataChainError(
+                "Not logged in to Studio. Log in with 'datachain studio login'."
+            )
+
+        del conf["studio"]["token"]
+
+    print("Logged out from Studio. (you can log back in with 'datachain studio login')")
+
+
+def token():
+    config = Config().read().get("studio", {})
+    token = config.get("token")
+    if not token:
+        raise DataChainError(
+            "Not logged in to Studio. Log in with 'datachain studio login'."
+        )
+
+    print(token)
+
+
+def list_datasets(args: "Namespace"):
+    client = StudioClient(team=args.team)
+    response = client.ls_datasets()
+    if not response.ok:
+        raise_remote_error(response.message)
+    if not response.data:
+        print("No datasets found.")
+        return
+    for d in response.data:
+        name = d.get("name")
+        for v in d.get("versions", []):
+            version = v.get("version")
+            print(f"{name} (v{version})")
+
+
+def save_config(hostname, token):
+    config = Config(ConfigLevel.GLOBAL)
+    with config.edit() as conf:
+        studio_conf = conf.get("studio", {})
+        studio_conf["url"] = hostname
+        studio_conf["token"] = token
+        conf["studio"] = studio_conf
+
+    return config.config_file()

datachain/utils.py

@@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
 from uuid import UUID
 
 import cloudpickle
+import platformdirs
 from dateutil import tz
 from dateutil.parser import isoparse
 from pydantic import BaseModel
@@ -25,6 +26,13 @@ if TYPE_CHECKING:
 NUL = b"\0"
 TIME_ZERO = datetime.fromtimestamp(0, tz=timezone.utc)
 
+APPNAME = "datachain"
+APPAUTHOR = "iterative"
+ENV_DATACHAIN_SYSTEM_CONFIG_DIR = "DATACHAIN_SYSTEM_CONFIG_DIR"
+ENV_DATACHAIN_GLOBAL_CONFIG_DIR = "DATACHAIN_GLOBAL_CONFIG_DIR"
+STUDIO_URL = "https://studio.dvc.ai"
+
+
 T = TypeVar("T", bound="DataChainDir")
 
 
@@ -33,6 +41,7 @@ class DataChainDir:
     CACHE = "cache"
     TMP = "tmp"
     DB = "db"
+    CONFIG = "config"
     ENV_VAR = "DATACHAIN_DIR"
     ENV_VAR_DATACHAIN_ROOT = "DATACHAIN_ROOT_DIR"
 
@@ -42,6 +51,7 @@ class DataChainDir:
         cache: Optional[str] = None,
         tmp: Optional[str] = None,
         db: Optional[str] = None,
+        config: Optional[str] = None,
     ) -> None:
         self.root = osp.abspath(root) if root is not None else self.default_root()
         self.cache = (
@@ -51,12 +61,24 @@ class DataChainDir:
             osp.abspath(tmp) if tmp is not None else osp.join(self.root, self.TMP)
         )
         self.db = osp.abspath(db) if db is not None else osp.join(self.root, self.DB)
+        self.config = (
+            osp.abspath(config)
+            if config is not None
+            else osp.join(self.root, self.CONFIG)
+        )
+        self.config = (
+            osp.abspath(config)
+            if config is not None
+            else osp.join(self.root, self.CONFIG)
+        )
 
     def init(self):
         os.makedirs(self.root, exist_ok=True)
         os.makedirs(self.cache, exist_ok=True)
         os.makedirs(self.tmp, exist_ok=True)
         os.makedirs(osp.split(self.db)[0], exist_ok=True)
+        os.makedirs(osp.split(self.config)[0], exist_ok=True)
+        os.makedirs(osp.split(self.config)[0], exist_ok=True)
 
     @classmethod
     def default_root(cls) -> str:
@@ -82,6 +104,18 @@ class DataChainDir:
         return instance
 
 
+def system_config_dir():
+    return os.getenv(ENV_DATACHAIN_SYSTEM_CONFIG_DIR) or platformdirs.site_config_dir(
+        APPNAME, APPAUTHOR
+    )
+
+
+def global_config_dir():
+    return os.getenv(ENV_DATACHAIN_GLOBAL_CONFIG_DIR) or platformdirs.user_config_dir(
+        APPNAME, APPAUTHOR
+    )
+
+
 def human_time_to_int(time: str) -> Optional[int]:
     if not time:
         return None
@@ -421,3 +455,27 @@ def env2bool(var, undefined=False):
     if var is None:
         return undefined
     return bool(re.search("1|y|yes|true", var, flags=re.IGNORECASE))
+
+
+def nested_dict_path_set(
+    data: dict[str, Any], path: Sequence[str], value: Any
+) -> dict[str, Any]:
+    """Sets a value inside a nested dict based on the list of dict keys as a path,
+    and will create sub-dicts as needed to set the value."""
+    sub_data = data
+    for element in path[:-1]:
+        if element not in sub_data:
+            sub_data[element] = {}
+        sub_data = sub_data[element]
+    sub_data[path[len(path) - 1]] = value
+    return data
+
+
+def row_to_nested_dict(
+    headers: Iterable[Sequence[str]], row: Iterable[Any]
+) -> dict[str, Any]:
+    """Converts a row to a nested dict based on the provided headers."""
+    result: dict[str, Any] = {}
+    for h, v in zip(headers, row):
+        nested_dict_path_set(result, h, v)
+    return result

{datachain-0.6.0.dist-info → datachain-0.6.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.6.0
+Version: 0.6.2
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -19,7 +19,7 @@ License-File: LICENSE
 Requires-Dist: pyyaml
 Requires-Dist: tomlkit
 Requires-Dist: tqdm
-Requires-Dist: numpy
+Requires-Dist: numpy <3,>=1
 Requires-Dist: pandas >=2.0.0
 Requires-Dist: pyarrow
 Requires-Dist: typing-extensions
@@ -38,15 +38,16 @@ Requires-Dist: orjson >=3.10.5
 Requires-Dist: pydantic <3,>=2
 Requires-Dist: jmespath >=1.0
 Requires-Dist: datamodel-code-generator >=0.25
-Requires-Dist: Pillow <11,>=10.0.0
+Requires-Dist: Pillow <12,>=10.0.0
 Requires-Dist: msgpack <2,>=1.0.4
 Requires-Dist: psutil
 Requires-Dist: huggingface-hub
 Requires-Dist: iterative-telemetry >=0.0.9
-Requires-Dist: numpy <2,>=1 ; sys_platform == "win32"
+Requires-Dist: platformdirs
+Requires-Dist: dvc-studio-client <1,>=0.21
 Provides-Extra: dev
 Requires-Dist: datachain[docs,tests] ; extra == 'dev'
-Requires-Dist: mypy ==1.11.2 ; extra == 'dev'
+Requires-Dist: mypy ==1.12.1 ; extra == 'dev'
 Requires-Dist: types-python-dateutil ; extra == 'dev'
 Requires-Dist: types-pytz ; extra == 'dev'
 Requires-Dist: types-PyYAML ; extra == 'dev'
@@ -63,7 +64,7 @@ Requires-Dist: datachain[tests] ; extra == 'examples'
 Requires-Dist: numpy <2,>=1 ; extra == 'examples'
 Requires-Dist: defusedxml ; extra == 'examples'
 Requires-Dist: accelerate ; extra == 'examples'
-Requires-Dist: unstructured[embed-huggingface,pdf] ; extra == 'examples'
+Requires-Dist: unstructured[embed-huggingface,pdf] <0.16.0 ; extra == 'examples'
 Requires-Dist: pdfplumber ==0.11.4 ; extra == 'examples'
 Requires-Dist: huggingface-hub[hf_transfer] ; extra == 'examples'
 Requires-Dist: onnx ==1.16.1 ; extra == 'examples'