datachain 0.7.4__py3-none-any.whl → 0.7.6__py3-none-any.whl

datachain/catalog/catalog.py

@@ -42,7 +42,6 @@ from datachain.dataset import (
     DatasetRecord,
     DatasetStats,
     DatasetStatus,
-    RowDict,
     StorageURI,
     create_dataset_uri,
     parse_dataset_uri,
@@ -69,13 +68,11 @@ from .datasource import DataSource
 
 if TYPE_CHECKING:
     from datachain.data_storage import (
-        AbstractIDGenerator,
         AbstractMetastore,
         AbstractWarehouse,
     )
     from datachain.dataset import DatasetListVersion
     from datachain.job import Job
-    from datachain.lib.file import File
     from datachain.listing import Listing
 
 logger = logging.getLogger("datachain")
@@ -127,8 +124,10 @@ class DatasetRowsFetcher(NodesThreadPool):
         self,
         metastore: "AbstractMetastore",
         warehouse: "AbstractWarehouse",
-        dataset_name: str,
-        dataset_version: int,
+        remote_ds_name: str,
+        remote_ds_version: int,
+        local_ds_name: str,
+        local_ds_version: int,
         schema: dict[str, Union[SQLType, type[SQLType]]],
         max_threads: int = PULL_DATASET_MAX_THREADS,
     ):
@@ -136,8 +135,10 @@ class DatasetRowsFetcher(NodesThreadPool):
         self._check_dependencies()
         self.metastore = metastore
         self.warehouse = warehouse
-        self.dataset_name = dataset_name
-        self.dataset_version = dataset_version
+        self.remote_ds_name = remote_ds_name
+        self.remote_ds_version = remote_ds_version
+        self.local_ds_name = local_ds_name
+        self.local_ds_version = local_ds_version
         self.schema = schema
         self.last_status_check: Optional[float] = None
         self.studio_client = StudioClient()
@@ -171,7 +172,7 @@ class DatasetRowsFetcher(NodesThreadPool):
         Checks are done every PULL_DATASET_CHECK_STATUS_INTERVAL seconds
         """
         export_status_response = self.studio_client.dataset_export_status(
-            self.dataset_name, self.dataset_version
+            self.remote_ds_name, self.remote_ds_version
         )
         if not export_status_response.ok:
             raise_remote_error(export_status_response.message)
@@ -203,7 +204,7 @@ class DatasetRowsFetcher(NodesThreadPool):
 
         # metastore and warehouse are not thread safe
         with self.metastore.clone() as metastore, self.warehouse.clone() as warehouse:
-            dataset = metastore.get_dataset(self.dataset_name)
+            local_ds = metastore.get_dataset(self.local_ds_name)
 
             urls = list(urls)
             while urls:
@@ -227,7 +228,7 @@ class DatasetRowsFetcher(NodesThreadPool):
                     df = df.drop("sys__id", axis=1)
 
                     inserted = warehouse.insert_dataset_rows(
-                        df, dataset, self.dataset_version
+                        df, local_ds, self.local_ds_version
                     )
                     self.increase_counter(inserted)  # type: ignore [arg-type]
                     urls.remove(url)
@@ -520,7 +521,6 @@ def find_column_to_str(  # noqa: PLR0911
 class Catalog:
     def __init__(
         self,
-        id_generator: "AbstractIDGenerator",
         metastore: "AbstractMetastore",
         warehouse: "AbstractWarehouse",
         cache_dir=None,
@@ -533,7 +533,6 @@ class Catalog:
     ):
         datachain_dir = DataChainDir(cache=cache_dir, tmp=tmp_dir)
         datachain_dir.init()
-        self.id_generator = id_generator
         self.metastore = metastore
         self._warehouse = warehouse
         self.cache = DataChainCache(datachain_dir.cache, datachain_dir.tmp)
@@ -567,7 +566,6 @@ class Catalog:
     def copy(self, cache=True, db=True):
         result = copy(self)
         if not db:
-            result.id_generator = None
             result.metastore = None
             result._warehouse = None
             result.warehouse = None
@@ -967,7 +965,6 @@ class Catalog:
         are cleaned up as soon as they are no longer needed.
         """
         self.warehouse.cleanup_tables(names)
-        self.id_generator.delete_uris(names)
 
     def create_dataset_from_sources(
         self,
@@ -1101,6 +1098,13 @@ class Catalog:
     def get_dataset(self, name: str) -> DatasetRecord:
         return self.metastore.get_dataset(name)
 
+    def get_dataset_with_version_uuid(self, uuid: str) -> DatasetRecord:
+        """Returns dataset that contains version with specific uuid"""
+        for dataset in self.ls_datasets():
+            if dataset.has_version_with_uuid(uuid):
+                return self.get_dataset(dataset.name)
+        raise DatasetNotFoundError(f"Dataset with version uuid {uuid} not found.")
+
     def get_remote_dataset(self, name: str) -> DatasetRecord:
         studio_client = StudioClient()
 
@@ -1268,35 +1272,6 @@ class Catalog:
         dataset = self.get_dataset(name)
         return self.update_dataset(dataset, **update_data)
 
-    def get_file_from_row(
-        self, dataset_name: str, dataset_version: int, row: RowDict, signal_name: str
-    ) -> "File":
-        """
-        Function that returns specific file signal from dataset row by name.
-        """
-        from datachain.lib.file import File
-        from datachain.lib.signal_schema import DEFAULT_DELIMITER, SignalSchema
-
-        version = self.get_dataset(dataset_name).get_version(dataset_version)
-        schema = SignalSchema.deserialize(version.feature_schema)
-
-        if signal_name not in schema.get_signals(File):
-            raise RuntimeError(
-                f"File signal with path {signal_name} not found in ",
-                f"dataset {dataset_name}@v{dataset_version} signals schema",
-            )
-
-        prefix = signal_name.replace(".", DEFAULT_DELIMITER) + DEFAULT_DELIMITER
-        file_signals = {
-            c_name.removeprefix(prefix): c_value
-            for c_name, c_value in row.items()
-            if c_name.startswith(prefix)
-            and DEFAULT_DELIMITER not in c_name.removeprefix(prefix)
-            and c_name.removeprefix(prefix) in File.model_fields
-        }
-
-        return File(**file_signals)
-
     def ls(
         self,
         sources: list[str],
@@ -1316,10 +1291,12 @@ class Catalog:
         for source in data_sources:  # type: ignore [union-attr]
             yield source, source.ls(fields)
 
-    def pull_dataset(
+    def pull_dataset(  # noqa: PLR0915
         self,
-        dataset_uri: str,
+        remote_ds_uri: str,
         output: Optional[str] = None,
+        local_ds_name: Optional[str] = None,
+        local_ds_version: Optional[int] = None,
         no_cp: bool = False,
         force: bool = False,
         edatachain: bool = False,
@@ -1327,105 +1304,112 @@ class Catalog:
         *,
         client_config=None,
     ) -> None:
-        # TODO add progress bar https://github.com/iterative/dvcx/issues/750
-        # TODO copy correct remote dates https://github.com/iterative/dvcx/issues/new
-        # TODO compare dataset stats on remote vs local pull to assert it's ok
-        def _instantiate_dataset():
+        def _instantiate(ds_uri: str) -> None:
             if no_cp:
                 return
+            assert output
             self.cp(
-                [dataset_uri],
+                [ds_uri],
                 output,
                 force=force,
                 no_edatachain_file=not edatachain,
                 edatachain_file=edatachain_file,
                 client_config=client_config,
             )
-            print(f"Dataset {dataset_uri} instantiated locally to {output}")
+            print(f"Dataset {ds_uri} instantiated locally to {output}")
 
         if not output and not no_cp:
             raise ValueError("Please provide output directory for instantiation")
 
-        client_config = client_config or self.client_config
-
         studio_client = StudioClient()
 
         try:
-            remote_dataset_name, version = parse_dataset_uri(dataset_uri)
+            remote_ds_name, version = parse_dataset_uri(remote_ds_uri)
         except Exception as e:
             raise DataChainError("Error when parsing dataset uri") from e
 
-        dataset = None
-        try:
-            dataset = self.get_dataset(remote_dataset_name)
-        except DatasetNotFoundError:
-            # we will create new one if it doesn't exist
-            pass
-
-        if dataset and version and dataset.has_version(version):
-            """No need to communicate with Studio at all"""
-            dataset_uri = create_dataset_uri(remote_dataset_name, version)
-            print(f"Local copy of dataset {dataset_uri} already present")
-            _instantiate_dataset()
-            return
+        remote_ds = self.get_remote_dataset(remote_ds_name)
 
-        remote_dataset = self.get_remote_dataset(remote_dataset_name)
-        # if version is not specified in uri, take the latest one
-        if not version:
-            version = remote_dataset.latest_version
-            print(f"Version not specified, pulling the latest one (v{version})")
-            # updating dataset uri with latest version
-            dataset_uri = create_dataset_uri(remote_dataset_name, version)
+        try:
+            # if version is not specified in uri, take the latest one
+            if not version:
+                version = remote_ds.latest_version
+                print(f"Version not specified, pulling the latest one (v{version})")
+                # updating dataset uri with latest version
+                remote_ds_uri = create_dataset_uri(remote_ds_name, version)
+            remote_ds_version = remote_ds.get_version(version)
+        except (DatasetVersionNotFoundError, StopIteration) as exc:
+            raise DataChainError(
+                f"Dataset {remote_ds_name} doesn't have version {version} on server"
+            ) from exc
 
-        assert version
+        local_ds_name = local_ds_name or remote_ds.name
+        local_ds_version = local_ds_version or remote_ds_version.version
+        local_ds_uri = create_dataset_uri(local_ds_name, local_ds_version)
 
-        if dataset and dataset.has_version(version):
-            print(f"Local copy of dataset {dataset_uri} already present")
-            _instantiate_dataset()
+        try:
+            # try to find existing dataset with the same uuid to avoid pulling again
+            existing_ds = self.get_dataset_with_version_uuid(remote_ds_version.uuid)
+            existing_ds_version = existing_ds.get_version_by_uuid(
+                remote_ds_version.uuid
+            )
+            existing_ds_uri = create_dataset_uri(
+                existing_ds.name, existing_ds_version.version
+            )
+            if existing_ds_uri == remote_ds_uri:
+                print(f"Local copy of dataset {remote_ds_uri} already present")
+            else:
+                print(
+                    f"Local copy of dataset {remote_ds_uri} already present as"
+                    f" dataset {existing_ds_uri}"
+                )
+            _instantiate(existing_ds_uri)
             return
+        except DatasetNotFoundError:
+            pass
 
         try:
-            remote_dataset_version = remote_dataset.get_version(version)
-        except (DatasetVersionNotFoundError, StopIteration) as exc:
-            raise DataChainError(
-                f"Dataset {remote_dataset_name} doesn't have version {version}"
-                " on server"
-            ) from exc
+            local_dataset = self.get_dataset(local_ds_name)
+            if local_dataset and local_dataset.has_version(local_ds_version):
+                raise DataChainError(
+                    f"Local dataset {local_ds_uri} already exists with different uuid,"
+                    " please choose different local dataset name or version"
+                )
+        except DatasetNotFoundError:
+            pass
 
-        stats_response = studio_client.dataset_stats(remote_dataset_name, version)
+        stats_response = studio_client.dataset_stats(
+            remote_ds_name, remote_ds_version.version
+        )
         if not stats_response.ok:
             raise_remote_error(stats_response.message)
-        dataset_stats = stats_response.data
+        ds_stats = stats_response.data
 
         dataset_save_progress_bar = tqdm(
-            desc=f"Saving dataset {dataset_uri} locally: ",
+            desc=f"Saving dataset {remote_ds_uri} locally: ",
             unit=" rows",
             unit_scale=True,
             unit_divisor=1000,
-            total=dataset_stats.num_objects,  # type: ignore [union-attr]
+            total=ds_stats.num_objects,  # type: ignore [union-attr]
         )
 
-        schema = DatasetRecord.parse_schema(remote_dataset_version.schema)
+        schema = DatasetRecord.parse_schema(remote_ds_version.schema)
 
-        columns = tuple(
-            sa.Column(name, typ) for name, typ in schema.items() if name != "sys__id"
-        )
-        # creating new dataset (version) locally
-        dataset = self.create_dataset(
-            remote_dataset_name,
-            version,
-            query_script=remote_dataset_version.query_script,
+        local_ds = self.create_dataset(
+            local_ds_name,
+            local_ds_version,
+            query_script=remote_ds_version.query_script,
             create_rows=True,
-            columns=columns,
-            feature_schema=remote_dataset_version.feature_schema,
+            columns=tuple(sa.Column(n, t) for n, t in schema.items() if n != "sys__id"),
+            feature_schema=remote_ds_version.feature_schema,
             validate_version=False,
-            uuid=remote_dataset_version.uuid,
+            uuid=remote_ds_version.uuid,
         )
 
         # asking remote to export dataset rows table to s3 and to return signed
         # urls of exported parts, which are in parquet format
         export_response = studio_client.export_dataset_table(
-            remote_dataset_name, version
+            remote_ds_name, remote_ds_version.version
         )
         if not export_response.ok:
             raise_remote_error(export_response.message)
@@ -1442,8 +1426,10 @@ class Catalog:
                 rows_fetcher = DatasetRowsFetcher(
                     metastore,
                     warehouse,
-                    dataset.name,
-                    version,
+                    remote_ds_name,
+                    remote_ds_version.version,
+                    local_ds_name,
+                    local_ds_version,
                     schema,
                 )
                 try:
@@ -1455,23 +1441,23 @@ class Catalog:
                         dataset_save_progress_bar,
                     )
                 except:
-                    self.remove_dataset(dataset.name, version)
+                    self.remove_dataset(local_ds_name, local_ds_version)
                     raise
 
-        dataset = self.metastore.update_dataset_status(
-            dataset,
+        local_ds = self.metastore.update_dataset_status(
+            local_ds,
             DatasetStatus.COMPLETE,
-            version=version,
-            error_message=remote_dataset.error_message,
-            error_stack=remote_dataset.error_stack,
-            script_output=remote_dataset.error_stack,
+            version=local_ds_version,
+            error_message=remote_ds.error_message,
+            error_stack=remote_ds.error_stack,
+            script_output=remote_ds.error_stack,
         )
-        self.update_dataset_version_with_warehouse_info(dataset, version)
+        self.update_dataset_version_with_warehouse_info(local_ds, local_ds_version)
 
         dataset_save_progress_bar.close()
-        print(f"Dataset {dataset_uri} saved locally")
+        print(f"Dataset {remote_ds_uri} saved locally")
 
-        _instantiate_dataset()
+        _instantiate(local_ds_uri)
 
     def clone(
         self,

datachain/catalog/loader.py

@@ -4,21 +4,16 @@ from typing import Any, Optional
 
 from datachain.catalog import Catalog
 from datachain.data_storage import (
-    AbstractIDGenerator,
     AbstractMetastore,
     AbstractWarehouse,
 )
 from datachain.data_storage.serializer import deserialize
 from datachain.data_storage.sqlite import (
-    SQLiteIDGenerator,
     SQLiteMetastore,
     SQLiteWarehouse,
 )
 from datachain.utils import get_envs_by_prefix
 
-ID_GENERATOR_SERIALIZED = "DATACHAIN__ID_GENERATOR"
-ID_GENERATOR_IMPORT_PATH = "DATACHAIN_ID_GENERATOR"
-ID_GENERATOR_ARG_PREFIX = "DATACHAIN_ID_GENERATOR_ARG_"
 METASTORE_SERIALIZED = "DATACHAIN__METASTORE"
 METASTORE_IMPORT_PATH = "DATACHAIN_METASTORE"
 METASTORE_ARG_PREFIX = "DATACHAIN_METASTORE_ARG_"
@@ -31,45 +26,7 @@ DISTRIBUTED_ARG_PREFIX = "DATACHAIN_DISTRIBUTED_ARG_"
 IN_MEMORY_ERROR_MESSAGE = "In-memory is only supported on SQLite"
 
 
-def get_id_generator(in_memory: bool = False) -> "AbstractIDGenerator":
-    id_generator_serialized = os.environ.get(ID_GENERATOR_SERIALIZED)
-    if id_generator_serialized:
-        id_generator_obj = deserialize(id_generator_serialized)
-        if not isinstance(id_generator_obj, AbstractIDGenerator):
-            raise RuntimeError(
-                "Deserialized ID generator is not an instance of AbstractIDGenerator: "
-                f"{id_generator_obj}"
-            )
-        return id_generator_obj
-
-    id_generator_import_path = os.environ.get(ID_GENERATOR_IMPORT_PATH)
-    id_generator_arg_envs = get_envs_by_prefix(ID_GENERATOR_ARG_PREFIX)
-    # Convert env variable names to keyword argument names by lowercasing them
-    id_generator_args: dict[str, Any] = {
-        k.lower(): v for k, v in id_generator_arg_envs.items()
-    }
-
-    if not id_generator_import_path:
-        id_generator_args["in_memory"] = in_memory
-        return SQLiteIDGenerator(**id_generator_args)
-    if in_memory:
-        raise RuntimeError(IN_MEMORY_ERROR_MESSAGE)
-    # ID generator paths are specified as (for example):
-    # datachain.data_storage.SQLiteIDGenerator
-    if "." not in id_generator_import_path:
-        raise RuntimeError(
-            f"Invalid {ID_GENERATOR_IMPORT_PATH} import path:"
-            f"{id_generator_import_path}"
-        )
-    module_name, _, class_name = id_generator_import_path.rpartition(".")
-    id_generator = import_module(module_name)
-    id_generator_class = getattr(id_generator, class_name)
-    return id_generator_class(**id_generator_args)
-
-
-def get_metastore(
-    id_generator: Optional["AbstractIDGenerator"], in_memory: bool = False
-) -> "AbstractMetastore":
+def get_metastore(in_memory: bool = False) -> "AbstractMetastore":
     metastore_serialized = os.environ.get(METASTORE_SERIALIZED)
     if metastore_serialized:
         metastore_obj = deserialize(metastore_serialized)
@@ -80,9 +37,6 @@ def get_metastore(
             )
         return metastore_obj
 
-    if id_generator is None:
-        id_generator = get_id_generator()
-
     metastore_import_path = os.environ.get(METASTORE_IMPORT_PATH)
     metastore_arg_envs = get_envs_by_prefix(METASTORE_ARG_PREFIX)
     # Convert env variable names to keyword argument names by lowercasing them
@@ -91,10 +45,8 @@ def get_metastore(
     }
 
     if not metastore_import_path:
-        if not isinstance(id_generator, SQLiteIDGenerator):
-            raise ValueError("SQLiteMetastore can only be used with SQLiteIDGenerator")
         metastore_args["in_memory"] = in_memory
-        return SQLiteMetastore(id_generator, **metastore_args)
+        return SQLiteMetastore(**metastore_args)
     if in_memory:
         raise RuntimeError(IN_MEMORY_ERROR_MESSAGE)
     # Metastore paths are specified as (for example):
@@ -106,12 +58,10 @@ def get_metastore(
     module_name, _, class_name = metastore_import_path.rpartition(".")
     metastore = import_module(module_name)
     metastore_class = getattr(metastore, class_name)
-    return metastore_class(id_generator, **metastore_args)
+    return metastore_class(**metastore_args)
 
 
-def get_warehouse(
-    id_generator: Optional["AbstractIDGenerator"], in_memory: bool = False
-) -> "AbstractWarehouse":
+def get_warehouse(in_memory: bool = False) -> "AbstractWarehouse":
     warehouse_serialized = os.environ.get(WAREHOUSE_SERIALIZED)
     if warehouse_serialized:
         warehouse_obj = deserialize(warehouse_serialized)
@@ -122,9 +72,6 @@ def get_warehouse(
             )
         return warehouse_obj
 
-    if id_generator is None:
-        id_generator = get_id_generator()
-
     warehouse_import_path = os.environ.get(WAREHOUSE_IMPORT_PATH)
     warehouse_arg_envs = get_envs_by_prefix(WAREHOUSE_ARG_PREFIX)
     # Convert env variable names to keyword argument names by lowercasing them
@@ -133,10 +80,8 @@ def get_warehouse(
     }
 
     if not warehouse_import_path:
-        if not isinstance(id_generator, SQLiteIDGenerator):
-            raise ValueError("SQLiteWarehouse can only be used with SQLiteIDGenerator")
         warehouse_args["in_memory"] = in_memory
-        return SQLiteWarehouse(id_generator, **warehouse_args)
+        return SQLiteWarehouse(**warehouse_args)
     if in_memory:
         raise RuntimeError(IN_MEMORY_ERROR_MESSAGE)
     # Warehouse paths are specified as (for example):
@@ -148,7 +93,7 @@ def get_warehouse(
     module_name, _, class_name = warehouse_import_path.rpartition(".")
     warehouse = import_module(module_name)
     warehouse_class = getattr(warehouse, class_name)
-    return warehouse_class(id_generator, **warehouse_args)
+    return warehouse_class(**warehouse_args)
 
 
 def get_distributed_class(**kwargs):
@@ -188,11 +133,9 @@ def get_catalog(
     and name of variable after, e.g. if it accepts team_id as kwargs
     we can provide DATACHAIN_METASTORE_ARG_TEAM_ID=12345 env variable.
     """
-    id_generator = get_id_generator(in_memory=in_memory)
     return Catalog(
-        id_generator=id_generator,
-        metastore=get_metastore(id_generator, in_memory=in_memory),
-        warehouse=get_warehouse(id_generator, in_memory=in_memory),
+        metastore=get_metastore(in_memory=in_memory),
+        warehouse=get_warehouse(in_memory=in_memory),
         client_config=client_config,
         in_memory=in_memory,
     )