datachain 0.20.4__py3-none-any.whl → 0.21.0__py3-none-any.whl

datachain/lib/settings.py

@@ -14,16 +14,12 @@ class Settings:
         workers=None,
         min_task_size=None,
         prefetch=None,
-        namespace=None,
-        project=None,
     ):
         self._cache = cache
         self.parallel = parallel
         self._workers = workers
         self.min_task_size = min_task_size
         self.prefetch = prefetch
-        self.namespace = namespace
-        self.project = project
 
         if not isinstance(cache, bool) and cache is not None:
             raise SettingsError(
@@ -71,10 +67,6 @@ class Settings:
             res["workers"] = self.workers
         if self.min_task_size is not None:
             res["min_task_size"] = self.min_task_size
-        if self.namespace is not None:
-            res["namespace"] = self.namespace
-        if self.project is not None:
-            res["project"] = self.project
         return res
 
     def add(self, settings: "Settings"):
@@ -82,7 +74,5 @@ class Settings:
         self.parallel = settings.parallel or self.parallel
         self._workers = settings._workers or self._workers
         self.min_task_size = settings.min_task_size or self.min_task_size
-        self.namespace = settings.namespace or self.namespace
-        self.project = settings.project or self.project
         if settings.prefetch is not None:
             self.prefetch = settings.prefetch

datachain/lib/tar.py

@@ -6,11 +6,12 @@ from datachain.lib.file import File, TarVFile
 
 
 def build_tar_member(parent: File, info: tarfile.TarInfo) -> File:
+    new_parent = parent.get_full_name()
     etag_string = "-".join([parent.etag, info.name, str(info.mtime)])
     etag = hashlib.md5(etag_string.encode(), usedforsecurity=False).hexdigest()
     return File(
         source=parent.source,
-        path=f"{parent.path}/{info.name}",
+        path=f"{new_parent}/{info.name}",
         version=parent.version,
         size=info.size,
         etag=etag,

datachain/lib/udf_signature.py

@@ -16,7 +16,7 @@ class UdfSignatureError(DataChainParamsError):
 
 
 @dataclass
-class UdfSignature:  # noqa: PLW1641
+class UdfSignature:
     func: Union[Callable, UDFBase]
     params: dict[str, Union[DataType, Any]]
     output_schema: SignalSchema

datachain/lib/webdataset.py

@@ -34,29 +34,29 @@ warnings.filterwarnings(
 
 
 class WDSError(DataChainError):
-    def __init__(self, tar_name: str, message: str):
-        super().__init__(f"WebDataset error '{tar_name}': {message}")
+    def __init__(self, tar_stream, message: str):
+        super().__init__(f"WebDataset error '{tar_stream.get_full_name()}': {message}")
 
 
 class CoreFileDuplicationError(WDSError):
-    def __init__(self, tar_name: str, file1: str, file2: str):
+    def __init__(self, tar_stream, file1: str, file2: str):
         super().__init__(
-            tar_name, f"duplication of files with core extensions: {file1}, {file2}"
+            tar_stream, f"duplication of files with core extensions: {file1}, {file2}"
         )
 
 
 class CoreFileNotFoundError(WDSError):
-    def __init__(self, tar_name: str, extensions: Sequence[str], stem: str):
+    def __init__(self, tar_stream, extensions, stem):
         super().__init__(
-            tar_name,
+            tar_stream,
             f"no files with the extensions '{','.join(extensions)}'"
             f" were found for file stem {stem}",
         )
 
 
 class UnknownFileExtensionError(WDSError):
-    def __init__(self, tar_name, name: str, ext: str):
-        super().__init__(tar_name, f"unknown extension '{ext}' for file '{name}'")
+    def __init__(self, tar_stream, name, ext):
+        super().__init__(tar_stream, f"unknown extension '{ext}' for file '{name}'")
 
 
 class WDSBasic(DataModel):
@@ -113,10 +113,10 @@ class Builder:
     def __init__(
         self,
         tar_stream: File,
-        core_extensions: Sequence[str],
+        core_extensions: list[str],
         wds_class: type[WDSBasic],
-        tar: tarfile.TarFile,
-        encoding: str = "utf-8",
+        tar,
+        encoding="utf-8",
     ):
         self._core_extensions = core_extensions
         self._tar_stream = tar_stream
@@ -145,20 +145,18 @@ class Builder:
         if ext in self._core_extensions:
             if self.state.core_file is not None:
                 raise CoreFileDuplicationError(
-                    self._tar_stream.name, file.name, self.state.core_file.name
+                    self._tar_stream, file.name, self.state.core_file.name
                 )
             self.state.core_file = file
         elif ext in self.state.data:
             raise WDSError(
-                self._tar_stream.name,
+                self._tar_stream,
                 f"file with extension '.{ext}' already exists in the archive",
             )
         else:
             type_ = self._get_type(ext)
             if type_ is None:
-                raise UnknownFileExtensionError(
-                    self._tar_stream.name, fstream.name, ext
-                )
+                raise UnknownFileExtensionError(self._tar_stream, fstream.name, ext)
 
             if issubclass(type_, WDSReadableSubclass):
                 reader = type_._reader
@@ -167,7 +165,7 @@ class Builder:
 
             if reader is None:
                 raise WDSError(
-                    self._tar_stream.name,
+                    self._tar_stream,
                     f"unable to find a reader for type {type_}, extension .{ext}",
                 )
             self.state.data[ext] = reader(self, file)
@@ -175,7 +173,7 @@ class Builder:
     def produce(self):
         if self.state.core_file is None:
             raise CoreFileNotFoundError(
-                self._tar_stream.name, self._core_extensions, self.state.stem
+                self._tar_stream, self._core_extensions, self.state.stem
             )
 
         file = build_tar_member(self._tar_stream, self.state.core_file)
@@ -196,13 +194,7 @@ class Builder:
         return anno
 
 
-def get_tar_groups(
-    stream: File,
-    tar: tarfile.TarFile,
-    core_extensions: Sequence[str],
-    spec: type[WDSBasic],
-    encoding: str = "utf-8",
-) -> Iterator[WDSBasic]:
+def get_tar_groups(stream, tar, core_extensions, spec, encoding="utf-8"):
     builder = Builder(stream, core_extensions, spec, tar, encoding)
 
     for item in sorted(tar.getmembers(), key=lambda m: Path(m.name).stem):
@@ -218,11 +210,9 @@ def get_tar_groups(
 
 
 def process_webdataset(
-    core_extensions: Sequence[str] = ("jpg", "png"),
-    spec: type[WDSBasic] = WDSAllFile,
-    encoding: str = "utf-8",
-) -> Callable[[File], Iterator]:
-    def wds_func(file: File) -> Iterator[spec]:  # type: ignore[valid-type]
+    core_extensions: Sequence[str] = ("jpg", "png"), spec=WDSAllFile, encoding="utf-8"
+) -> Callable:
+    def wds_func(file: File) -> Iterator[spec]:
         with file.open() as fd:
             with tarfile.open(fileobj=fd) as tar:
                 yield from get_tar_groups(file, tar, core_extensions, spec, encoding)

datachain/listing.py

@@ -66,9 +66,7 @@ class Listing:
     @cached_property
     def dataset(self) -> "DatasetRecord":
         assert self.dataset_name
-        return self.metastore.get_dataset(
-            self.dataset_name, self.metastore.listing_project.id
-        )
+        return self.metastore.get_dataset(self.dataset_name)
 
     @cached_property
     def dataset_rows(self):

datachain/query/dataset.py

@@ -41,13 +41,12 @@ from datachain.data_storage.schema import (
     partition_col_names,
     partition_columns,
 )
-from datachain.dataset import DatasetDependency, DatasetStatus, RowDict
+from datachain.dataset import DATASET_PREFIX, 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
@@ -84,7 +83,7 @@ PartitionByType = Union[
     Function, ColumnElement, Sequence[Union[Function, ColumnElement]]
 ]
 JoinPredicateType = Union[str, ColumnClause, ColumnElement]
-DatasetDependencyType = tuple["DatasetRecord", str]
+DatasetDependencyType = tuple[str, str]
 
 logger = logging.getLogger("datachain")
 
@@ -170,17 +169,18 @@ class QueryStep:
     """A query that returns all rows from specific dataset version"""
 
     catalog: "Catalog"
-    dataset: "DatasetRecord"
+    dataset_name: str
     dataset_version: str
 
     def apply(self) -> "StepResult":
         def q(*columns):
             return sqlalchemy.select(*columns)
 
-        dr = self.catalog.warehouse.dataset_rows(self.dataset, self.dataset_version)
+        dataset = self.catalog.get_dataset(self.dataset_name)
+        dr = self.catalog.warehouse.dataset_rows(dataset, self.dataset_version)
 
         return step_result(
-            q, dr.columns, dependencies=[(self.dataset, self.dataset_version)]
+            q, dr.columns, dependencies=[(self.dataset_name, self.dataset_version)]
         )
 
 
@@ -1095,8 +1095,6 @@ 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,
@@ -1130,38 +1128,33 @@ class DatasetQuery:
         if version:
             self.version = version
 
-        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
+        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
         elif fallback_to_studio and is_token_set():
             self._set_starting_step(
-                self.catalog.get_dataset_with_remote_fallback(
-                    name,
-                    namespace_name=namespace_name,
-                    project_name=project_name,
-                    version=version,
-                )
+                self.catalog.get_dataset_with_remote_fallback(name, version)
             )
         else:
-            project = self.catalog.metastore.get_project(project_name, namespace_name)
-            self._set_starting_step(self.catalog.get_dataset(name, project=project))
+            self._set_starting_step(self.catalog.get_dataset(name))
 
     def _set_starting_step(self, ds: "DatasetRecord") -> None:
         if not self.version:
             self.version = ds.latest_version
 
-        self.starting_step = QueryStep(self.catalog, ds, self.version)
+        self.starting_step = QueryStep(self.catalog, ds.name, 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())
@@ -1169,6 +1162,21 @@ 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(
@@ -1649,8 +1657,6 @@ 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()
@@ -1670,36 +1676,26 @@ class DatasetQuery:
 
     def _add_dependencies(self, dataset: "DatasetRecord", version: str):
         dependencies: set[DatasetDependencyType] = set()
-        for dep_dataset, dep_dataset_version in self.dependencies:
-            if Session.is_temp_dataset(dep_dataset.name):
+        for dep_name, dep_version in self.dependencies:
+            if Session.is_temp_dataset(dep_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_dataset.name,
-                    dep_dataset_version,
-                    dep_dataset.project,
-                    indirect=False,
+                    dep_name, dep_version, indirect=False
                 ):
                     if dep:
-                        dep_project = self.catalog.metastore.get_project(
-                            dep.project, dep.namespace
-                        )
-                        dependencies.add(
-                            (
-                                self.catalog.get_dataset(dep.name, dep_project),
-                                dep.version,
-                            )
-                        )
+                        dependencies.add((dep.name, dep.version))
             else:
-                dependencies.add((dep_dataset, dep_dataset_version))
+                dependencies.add((dep_name, dep_version))
 
-        for dep_dataset, dep_dataset_version in dependencies:
+        for dep_name, dep_version in dependencies:
+            # ds_dependency_name, ds_dependency_version = dependency
             self.catalog.metastore.add_dataset_dependency(
-                dataset,
+                dataset.name,
                 version,
-                dep_dataset,
-                dep_dataset_version,
+                dep_name,
+                dep_version,
             )
 
     def exec(self) -> "Self":
@@ -1715,7 +1711,6 @@ 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,
@@ -1724,13 +1719,8 @@ 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, project).has_version(version)
-            ):
+            if name and version and self.catalog.get_dataset(name).has_version(version):
                 raise RuntimeError(f"Dataset {name} already has version {version}")
         except DatasetNotFoundError:
             pass
@@ -1755,7 +1745,6 @@ class DatasetQuery:
 
             dataset = self.catalog.create_dataset(
                 name,
-                project,
                 version=version,
                 feature_schema=feature_schema,
                 columns=columns,
@@ -1781,25 +1770,11 @@ class DatasetQuery:
 
             if dependencies:
                 # overriding 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.dependencies = {(dep.name, dep.version) for dep in dependencies}
             self._add_dependencies(dataset, version)  # type: ignore [arg-type]
         finally:
             self.cleanup()
-        return self.__class__(
-            name=name,
-            namespace_name=project.namespace.name,
-            project_name=project.name,
-            version=version,
-            catalog=self.catalog,
-        )
+        return self.__class__(name=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, dataset.project, force=True)
+                self.catalog.remove_dataset(dataset.name, force=True)
         # suppress error when metastore has been reset during testing
         except TableMissingError:
             pass

datachain/remote/studio.py

@@ -17,7 +17,6 @@ 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
 
@@ -37,33 +36,13 @@ 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(get_studio_env_variable("TOKEN"))
+        bool(os.environ.get("DVC_STUDIO_TOKEN"))
         or Config().read().get("studio", {}).get("token") is not None
     )
 
@@ -99,12 +78,12 @@ class StudioClient:
 
     @property
     def token(self) -> str:
-        token = get_studio_env_variable("TOKEN") or self.config.get("token")
+        token = os.environ.get("DVC_STUDIO_TOKEN") or self.config.get("token")
 
         if not token:
             raise DataChainError(
                 "Studio token is not set. Use `datachain auth login` "
-                "or environment variable `DATACHAIN_STUDIO_TOKEN` to set it."
+                "or environment variable `DVC_STUDIO_TOKEN` to set it."
             )
 
         return token
@@ -112,8 +91,8 @@ class StudioClient:
     @property
     def url(self) -> str:
         return (
-            get_studio_env_variable("URL") or self.config.get("url") or STUDIO_URL
-        ).rstrip("/") + "/api"
+            os.environ.get("DVC_STUDIO_URL") or self.config.get("url") or STUDIO_URL
+        ) + "/api"
 
     @property
     def config(self) -> dict:
@@ -128,13 +107,13 @@ class StudioClient:
         return self._team
 
     def _get_team(self) -> str:
-        team = get_studio_env_variable("TEAM") or self.config.get("team")
+        team = os.environ.get("DVC_STUDIO_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 `DATACHAIN_STUDIO_TEAM` to set it. "
+                "or environment variable `DVC_STUDIO_TEAM` to set it. "
                 "You can also set `studio.team` in the config file."
             )
 
@@ -312,17 +291,13 @@ 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,
-            "name": name,
-            "namespace": namespace,
-            "project": project,
+            "dataset_name": name,
             "description": description,
             "attrs": attrs,
         }
@@ -335,44 +310,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",
             {
-                "name": name,
-                "namespace": namespace,
-                "project": project,
-                "version": version,
+                "dataset_name": name,
+                "dataset_version": version,
                 "force": force,
             },
             method="DELETE",
         )
 
-    def dataset_info(
-        self, namespace: str, project: str, name: str
-    ) -> Response[DatasetInfoData]:
+    def dataset_info(self, 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",
-            {"namespace": namespace, "project": project, "name": name},
-            method="GET",
+            "datachain/datasets/info", {"dataset_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",
@@ -381,30 +356,20 @@ class StudioClient:
         )
 
     def export_dataset_table(
-        self, dataset: DatasetRecord, version: str
+        self, name: str, version: str
     ) -> Response[DatasetExportSignedUrls]:
         return self._send_request(
             "datachain/datasets/export",
-            {
-                "namespace": dataset.project.namespace.name,
-                "project": dataset.project.name,
-                "name": dataset.name,
-                "version": version,
-            },
+            {"dataset_name": name, "dataset_version": version},
             method="GET",
         )
 
     def dataset_export_status(
-        self, dataset: DatasetRecord, version: str
+        self, name: str, version: str
     ) -> Response[DatasetExportStatus]:
         return self._send_request(
             "datachain/datasets/export-status",
-            {
-                "namespace": dataset.project.namespace.name,
-                "project": dataset.project.name,
-                "name": dataset.name,
-                "version": version,
-            },
+            {"dataset_name": name, "dataset_version": version},
             method="GET",
         )