datachain 0.21.1__py3-none-any.whl → 0.22.0__py3-none-any.whl

datachain/query/dataset.py

@@ -41,12 +41,13 @@ from datachain.data_storage.schema import (
     partition_col_names,
     partition_columns,
 )
-from datachain.dataset import DATASET_PREFIX, DatasetDependency, DatasetStatus, RowDict
+from datachain.dataset import DatasetDependency, DatasetStatus, RowDict
 from datachain.error import DatasetNotFoundError, QueryScriptCancelError
 from datachain.func.base import Function
 from datachain.lib.listing import is_listing_dataset, listing_dataset_expired
 from datachain.lib.udf import UDFAdapter, _get_cache
 from datachain.progress import CombinedDownloadCallback, TqdmCombinedDownloadCallback
+from datachain.project import Project
 from datachain.query.schema import C, UDFParamSpec, normalize_param
 from datachain.query.session import Session
 from datachain.query.udf import UdfInfo
@@ -83,7 +84,7 @@ PartitionByType = Union[
     Function, ColumnElement, Sequence[Union[Function, ColumnElement]]
 ]
 JoinPredicateType = Union[str, ColumnClause, ColumnElement]
-DatasetDependencyType = tuple[str, str]
+DatasetDependencyType = tuple["DatasetRecord", str]
 
 logger = logging.getLogger("datachain")
 
@@ -169,18 +170,17 @@ class QueryStep:
     """A query that returns all rows from specific dataset version"""
 
     catalog: "Catalog"
-    dataset_name: str
+    dataset: "DatasetRecord"
     dataset_version: str
 
     def apply(self) -> "StepResult":
         def q(*columns):
             return sqlalchemy.select(*columns)
 
-        dataset = self.catalog.get_dataset(self.dataset_name)
-        dr = self.catalog.warehouse.dataset_rows(dataset, self.dataset_version)
+        dr = self.catalog.warehouse.dataset_rows(self.dataset, self.dataset_version)
 
         return step_result(
-            q, dr.columns, dependencies=[(self.dataset_name, self.dataset_version)]
+            q, dr.columns, dependencies=[(self.dataset, self.dataset_version)]
         )
 
 
@@ -1095,6 +1095,8 @@ class DatasetQuery:
         self,
         name: str,
         version: Optional[str] = None,
+        project_name: Optional[str] = None,
+        namespace_name: Optional[str] = None,
         catalog: Optional["Catalog"] = None,
         session: Optional[Session] = None,
         indexing_column_types: Optional[dict[str, Any]] = None,
@@ -1128,33 +1130,38 @@ class DatasetQuery:
         if version:
             self.version = version
 
-        if is_listing_dataset(name):
-            if version:
-                # this listing dataset should already be listed as we specify
-                # exact version
-                self._set_starting_step(self.catalog.get_dataset(name))
-            else:
-                # not setting query step yet as listing dataset might not exist at
-                # this point
-                self.list_ds_name = name
+        namespace_name = namespace_name or self.catalog.metastore.default_namespace_name
+        project_name = project_name or self.catalog.metastore.default_project_name
+
+        if is_listing_dataset(name) and not version:
+            # not setting query step yet as listing dataset might not exist at
+            # this point
+            self.list_ds_name = name
         elif fallback_to_studio and is_token_set():
             self._set_starting_step(
-                self.catalog.get_dataset_with_remote_fallback(name, version)
+                self.catalog.get_dataset_with_remote_fallback(
+                    name,
+                    namespace_name=namespace_name,
+                    project_name=project_name,
+                    version=version,
+                )
             )
         else:
-            self._set_starting_step(self.catalog.get_dataset(name))
+            project = self.catalog.metastore.get_project(project_name, namespace_name)
+            self._set_starting_step(self.catalog.get_dataset(name, project=project))
 
     def _set_starting_step(self, ds: "DatasetRecord") -> None:
         if not self.version:
             self.version = ds.latest_version
 
-        self.starting_step = QueryStep(self.catalog, ds.name, self.version)
+        self.starting_step = QueryStep(self.catalog, ds, self.version)
 
         # at this point we know our starting dataset so setting up schemas
         self.feature_schema = ds.get_version(self.version).feature_schema
         self.column_types = copy(ds.schema)
         if "sys__id" in self.column_types:
             self.column_types.pop("sys__id")
+        self.project = ds.project
 
     def __iter__(self):
         return iter(self.db_results())
@@ -1162,21 +1169,6 @@ class DatasetQuery:
     def __or__(self, other):
         return self.union(other)
 
-    def pull_dataset(self, name: str, version: Optional[str] = None) -> "DatasetRecord":
-        print("Dataset not found in local catalog, trying to get from studio")
-
-        remote_ds_uri = f"{DATASET_PREFIX}{name}"
-        if version:
-            remote_ds_uri += f"@v{version}"
-
-        self.catalog.pull_dataset(
-            remote_ds_uri=remote_ds_uri,
-            local_ds_name=name,
-            local_ds_version=version,
-        )
-
-        return self.catalog.get_dataset(name)
-
     @staticmethod
     def get_table() -> "TableClause":
         table_name = "".join(
@@ -1657,6 +1649,8 @@ class DatasetQuery:
         workers: Union[bool, int] = False,
         min_task_size: Optional[int] = None,
         partition_by: Optional[PartitionByType] = None,
+        namespace: Optional[str] = None,
+        project: Optional[str] = None,
         cache: bool = False,
     ) -> "Self":
         query = self.clone()
@@ -1676,26 +1670,36 @@ class DatasetQuery:
 
     def _add_dependencies(self, dataset: "DatasetRecord", version: str):
         dependencies: set[DatasetDependencyType] = set()
-        for dep_name, dep_version in self.dependencies:
-            if Session.is_temp_dataset(dep_name):
+        for dep_dataset, dep_dataset_version in self.dependencies:
+            if Session.is_temp_dataset(dep_dataset.name):
                 # temp dataset are created for optimization and they will be removed
                 # afterwards. Therefore, we should not put them as dependencies, but
                 # their own direct dependencies
                 for dep in self.catalog.get_dataset_dependencies(
-                    dep_name, dep_version, indirect=False
+                    dep_dataset.name,
+                    dep_dataset_version,
+                    dep_dataset.project,
+                    indirect=False,
                 ):
                     if dep:
-                        dependencies.add((dep.name, dep.version))
+                        dep_project = self.catalog.metastore.get_project(
+                            dep.project, dep.namespace
+                        )
+                        dependencies.add(
+                            (
+                                self.catalog.get_dataset(dep.name, dep_project),
+                                dep.version,
+                            )
+                        )
             else:
-                dependencies.add((dep_name, dep_version))
+                dependencies.add((dep_dataset, dep_dataset_version))
 
-        for dep_name, dep_version in dependencies:
-            # ds_dependency_name, ds_dependency_version = dependency
+        for dep_dataset, dep_dataset_version in dependencies:
             self.catalog.metastore.add_dataset_dependency(
-                dataset.name,
+                dataset,
                 version,
-                dep_name,
-                dep_version,
+                dep_dataset,
+                dep_dataset_version,
             )
 
     def exec(self) -> "Self":
@@ -1711,6 +1715,7 @@ class DatasetQuery:
         self,
         name: Optional[str] = None,
         version: Optional[str] = None,
+        project: Optional[Project] = None,
         feature_schema: Optional[dict] = None,
         dependencies: Optional[list[DatasetDependency]] = None,
         description: Optional[str] = None,
@@ -1719,8 +1724,13 @@ class DatasetQuery:
         **kwargs,
     ) -> "Self":
         """Save the query as a dataset."""
+        project = project or self.catalog.metastore.default_project
         try:
-            if name and version and self.catalog.get_dataset(name).has_version(version):
+            if (
+                name
+                and version
+                and self.catalog.get_dataset(name, project).has_version(version)
+            ):
                 raise RuntimeError(f"Dataset {name} already has version {version}")
         except DatasetNotFoundError:
             pass
@@ -1745,6 +1755,7 @@ class DatasetQuery:
 
             dataset = self.catalog.create_dataset(
                 name,
+                project,
                 version=version,
                 feature_schema=feature_schema,
                 columns=columns,
@@ -1770,11 +1781,25 @@ class DatasetQuery:
 
             if dependencies:
                 # overriding dependencies
-                self.dependencies = {(dep.name, dep.version) for dep in dependencies}
+                self.dependencies = set()
+                for dep in dependencies:
+                    dep_project = self.catalog.metastore.get_project(
+                        dep.project, dep.namespace
+                    )
+                    self.dependencies.add(
+                        (self.catalog.get_dataset(dep.name, dep_project), dep.version)
+                    )
+
             self._add_dependencies(dataset, version)  # type: ignore [arg-type]
         finally:
             self.cleanup()
-        return self.__class__(name=name, version=version, catalog=self.catalog)
+        return self.__class__(
+            name=name,
+            namespace_name=project.namespace.name,
+            project_name=project.name,
+            version=version,
+            catalog=self.catalog,
+        )
 
     @property
     def is_ordered(self) -> bool:

datachain/query/session.py

@@ -108,7 +108,7 @@ class Session:
         prefix = self.get_temp_prefix()
         try:
             for dataset in list(self.catalog.metastore.list_datasets_by_prefix(prefix)):
-                self.catalog.remove_dataset(dataset.name, force=True)
+                self.catalog.remove_dataset(dataset.name, dataset.project, force=True)
         # suppress error when metastore has been reset during testing
         except TableMissingError:
             pass

datachain/remote/studio.py

@@ -17,6 +17,7 @@ import websockets
 from requests.exceptions import HTTPError, Timeout
 
 from datachain.config import Config
+from datachain.dataset import DatasetRecord
 from datachain.error import DataChainError
 from datachain.utils import STUDIO_URL, retry_with_backoff
 
@@ -36,13 +37,33 @@ logger = logging.getLogger("datachain")
 DATASET_ROWS_CHUNK_SIZE = 8192
 
 
+def get_studio_env_variable(name: str) -> Any:
+    """
+    Get the value of a DataChain Studio environment variable.
+    It first checks for the variable prefixed with 'DATACHAIN_STUDIO_',
+    then checks for the deprecated 'DVC_STUDIO_' prefix.
+    If neither is set, it returns the provided default value.
+    """
+    if (value := os.environ.get(f"DATACHAIN_STUDIO_{name}")) is not None:
+        return value
+    if (value := os.environ.get(f"DVC_STUDIO_{name}")) is not None:  # deprecated
+        logger.warning(
+            "Environment variable 'DVC_STUDIO_%s' is deprecated, "
+            "use 'DATACHAIN_STUDIO_%s' instead.",
+            name,
+            name,
+        )
+        return value
+    return None
+
+
 def _is_server_error(status_code: int) -> bool:
     return str(status_code).startswith("5")
 
 
 def is_token_set() -> bool:
     return (
-        bool(os.environ.get("DVC_STUDIO_TOKEN"))
+        bool(get_studio_env_variable("TOKEN"))
         or Config().read().get("studio", {}).get("token") is not None
     )
 
@@ -78,12 +99,12 @@ class StudioClient:
 
     @property
     def token(self) -> str:
-        token = os.environ.get("DVC_STUDIO_TOKEN") or self.config.get("token")
+        token = get_studio_env_variable("TOKEN") or self.config.get("token")
 
         if not token:
             raise DataChainError(
                 "Studio token is not set. Use `datachain auth login` "
-                "or environment variable `DVC_STUDIO_TOKEN` to set it."
+                "or environment variable `DATACHAIN_STUDIO_TOKEN` to set it."
             )
 
         return token
@@ -91,8 +112,8 @@ class StudioClient:
     @property
     def url(self) -> str:
         return (
-            os.environ.get("DVC_STUDIO_URL") or self.config.get("url") or STUDIO_URL
-        ) + "/api"
+            get_studio_env_variable("URL") or self.config.get("url") or STUDIO_URL
+        ).rstrip("/") + "/api"
 
     @property
     def config(self) -> dict:
@@ -107,13 +128,13 @@ class StudioClient:
         return self._team
 
     def _get_team(self) -> str:
-        team = os.environ.get("DVC_STUDIO_TEAM") or self.config.get("team")
+        team = get_studio_env_variable("TEAM") or self.config.get("team")
 
         if not team:
             raise DataChainError(
                 "Studio team is not set. "
                 "Use `datachain auth team <team_name>` "
-                "or environment variable `DVC_STUDIO_TEAM` to set it. "
+                "or environment variable `DATACHAIN_STUDIO_TEAM` to set it. "
                 "You can also set `studio.team` in the config file."
             )
 
@@ -291,13 +312,17 @@ class StudioClient:
     def edit_dataset(
         self,
         name: str,
+        namespace: str,
+        project: str,
         new_name: Optional[str] = None,
         description: Optional[str] = None,
         attrs: Optional[list[str]] = None,
     ) -> Response[DatasetInfoData]:
         body = {
             "new_name": new_name,
-            "dataset_name": name,
+            "name": name,
+            "namespace": namespace,
+            "project": project,
             "description": description,
             "attrs": attrs,
         }
@@ -310,44 +335,44 @@ class StudioClient:
     def rm_dataset(
         self,
         name: str,
+        namespace: str,
+        project: str,
         version: Optional[str] = None,
         force: Optional[bool] = False,
     ) -> Response[DatasetInfoData]:
         return self._send_request(
             "datachain/datasets",
             {
-                "dataset_name": name,
-                "dataset_version": version,
+                "name": name,
+                "namespace": namespace,
+                "project": project,
+                "version": version,
                 "force": force,
             },
             method="DELETE",
         )
 
-    def dataset_info(self, name: str) -> Response[DatasetInfoData]:
+    def dataset_info(
+        self, namespace: str, project: str, name: str
+    ) -> Response[DatasetInfoData]:
         def _parse_dataset_info(dataset_info):
             _parse_dates(dataset_info, ["created_at", "finished_at"])
             for version in dataset_info.get("versions"):
                 _parse_dates(version, ["created_at"])
+            _parse_dates(dataset_info.get("project"), ["created_at"])
+            _parse_dates(dataset_info.get("project").get("namespace"), ["created_at"])
 
             return dataset_info
 
         response = self._send_request(
-            "datachain/datasets/info", {"dataset_name": name}, method="GET"
+            "datachain/datasets/info",
+            {"namespace": namespace, "project": project, "name": name},
+            method="GET",
         )
         if response.ok:
             response.data = _parse_dataset_info(response.data)
         return response
 
-    def dataset_rows_chunk(
-        self, name: str, version: str, offset: int
-    ) -> Response[DatasetRowsData]:
-        req_data = {"dataset_name": name, "dataset_version": version}
-        return self._send_request_msgpack(
-            "datachain/datasets/rows",
-            {**req_data, "offset": offset, "limit": DATASET_ROWS_CHUNK_SIZE},
-            method="GET",
-        )
-
     def dataset_job_versions(self, job_id: str) -> Response[DatasetJobVersionsData]:
         return self._send_request(
             "datachain/datasets/dataset_job_versions",
@@ -356,20 +381,30 @@ class StudioClient:
         )
 
     def export_dataset_table(
-        self, name: str, version: str
+        self, dataset: DatasetRecord, version: str
     ) -> Response[DatasetExportSignedUrls]:
         return self._send_request(
             "datachain/datasets/export",
-            {"dataset_name": name, "dataset_version": version},
+            {
+                "namespace": dataset.project.namespace.name,
+                "project": dataset.project.name,
+                "name": dataset.name,
+                "version": version,
+            },
             method="GET",
         )
 
     def dataset_export_status(
-        self, name: str, version: str
+        self, dataset: DatasetRecord, version: str
     ) -> Response[DatasetExportStatus]:
         return self._send_request(
             "datachain/datasets/export-status",
-            {"dataset_name": name, "dataset_version": version},
+            {
+                "namespace": dataset.project.namespace.name,
+                "project": dataset.project.name,
+                "name": dataset.name,
+                "version": version,
+            },
             method="GET",
         )
 

datachain/studio.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Optional
 import tabulate
 
 from datachain.config import Config, ConfigLevel
-from datachain.dataset import QUERY_DATASET_PREFIX
+from datachain.dataset import QUERY_DATASET_PREFIX, parse_dataset_name
 from datachain.error import DataChainError
 from datachain.remote.studio import StudioClient
 from datachain.utils import STUDIO_URL
@@ -102,11 +102,13 @@ def set_team(args: "Namespace"):
 def login(args: "Namespace"):
     from dvc_studio_client.auth import StudioAuthError, get_access_token
 
+    from datachain.remote.studio import get_studio_env_variable
+
     config = Config().read().get("studio", {})
     name = args.name
     hostname = (
         args.hostname
-        or os.environ.get("DVC_STUDIO_URL")
+        or get_studio_env_variable("URL")
         or config.get("url")
         or STUDIO_URL
     )
@@ -165,6 +167,11 @@ def token():
 
 
 def list_datasets(team: Optional[str] = None, name: Optional[str] = None):
+    def ds_full_name(ds: dict) -> str:
+        return (
+            f"{ds['project']['namespace']['name']}.{ds['project']['name']}.{ds['name']}"
+        )
+
     if name:
         yield from list_dataset_versions(team, name)
         return
@@ -181,18 +188,22 @@ def list_datasets(team: Optional[str] = None, name: Optional[str] = None):
 
     for d in response.data:
         name = d.get("name")
+        full_name = ds_full_name(d)
         if name and name.startswith(QUERY_DATASET_PREFIX):
             continue
 
         for v in d.get("versions", []):
             version = v.get("version")
-            yield (name, version)
+            yield (full_name, version)
 
 
 def list_dataset_versions(team: Optional[str] = None, name: str = ""):
     client = StudioClient(team=team)
 
-    response = client.dataset_info(name)
+    namespace_name, project_name, name = parse_dataset_name(name)
+    if not namespace_name or not project_name:
+        raise DataChainError(f"Missing namespace or project form dataset name {name}")
+    response = client.dataset_info(namespace_name, project_name, name)
 
     if not response.ok:
         raise DataChainError(response.message)
@@ -208,12 +219,16 @@ def list_dataset_versions(team: Optional[str] = None, name: str = ""):
 def edit_studio_dataset(
     team_name: Optional[str],
     name: str,
+    namespace: str,
+    project: str,
     new_name: Optional[str] = None,
     description: Optional[str] = None,
     attrs: Optional[list[str]] = None,
 ):
     client = StudioClient(team=team_name)
-    response = client.edit_dataset(name, new_name, description, attrs)
+    response = client.edit_dataset(
+        name, namespace, project, new_name, description, attrs
+    )
     if not response.ok:
         raise DataChainError(response.message)
 
@@ -223,11 +238,13 @@ def edit_studio_dataset(
 def remove_studio_dataset(
     team_name: Optional[str],
     name: str,
+    namespace: str,
+    project: str,
     version: Optional[str] = None,
     force: Optional[bool] = False,
 ):
     client = StudioClient(team=team_name)
-    response = client.rm_dataset(name, version, force)
+    response = client.rm_dataset(name, namespace, project, version, force)
     if not response.ok:
         raise DataChainError(response.message)
 

{datachain-0.21.1.dist-info → datachain-0.22.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.21.1
+Version: 0.22.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -94,7 +94,7 @@ Requires-Dist: scipy; extra == "tests"
 Requires-Dist: ultralytics; extra == "tests"
 Provides-Extra: dev
 Requires-Dist: datachain[docs,tests]; extra == "dev"
-Requires-Dist: mypy==1.16.0; extra == "dev"
+Requires-Dist: mypy==1.16.1; extra == "dev"
 Requires-Dist: types-python-dateutil; extra == "dev"
 Requires-Dist: types-pytz; extra == "dev"
 Requires-Dist: types-PyYAML; extra == "dev"