datachain 0.7.8__py3-none-any.whl → 0.7.10__py3-none-any.whl

datachain/cli.py

@@ -16,7 +16,7 @@ from tabulate import tabulate
 from datachain import Session, utils
 from datachain.cli_utils import BooleanOptionalAction, CommaSeparatedArgs, KeyValueArgs
 from datachain.config import Config
-from datachain.error import DataChainError
+from datachain.error import DataChainError, DatasetNotFoundError
 from datachain.lib.dc import DataChain
 from datachain.studio import (
     edit_studio_dataset,
@@ -1056,7 +1056,10 @@ def rm_dataset(
     all, local, studio = _determine_flavors(studio, local, all, token)
 
     if all or local:
-        catalog.remove_dataset(name, version=version, force=force)
+        try:
+            catalog.remove_dataset(name, version=version, force=force)
+        except DatasetNotFoundError:
+            print("Dataset not found in local", file=sys.stderr)
 
     if (all or studio) and token:
         remove_studio_dataset(team, name, version, force)
@@ -1077,7 +1080,10 @@ def edit_dataset(
     all, local, studio = _determine_flavors(studio, local, all, token)
 
     if all or local:
-        catalog.edit_dataset(name, new_name, description, labels)
+        try:
+            catalog.edit_dataset(name, new_name, description, labels)
+        except DatasetNotFoundError:
+            print("Dataset not found in local", file=sys.stderr)
 
     if (all or studio) and token:
         edit_studio_dataset(team, name, new_name, description, labels)

datachain/client/fsspec.py

@@ -172,7 +172,7 @@ class Client(ABC):
         return url == cls.PREFIX
 
     @classmethod
-    def get_uri(cls, name) -> "StorageURI":
+    def get_uri(cls, name: str) -> "StorageURI":
         from datachain.dataset import StorageURI
 
         return StorageURI(f"{cls.PREFIX}{name}")
@@ -278,7 +278,9 @@ class Client(ABC):
     ) -> None:
         await self._fetch_nested(start_prefix, result_queue)
 
-    async def _fetch_dir(self, prefix, pbar, result_queue: ResultQueue) -> set[str]:
+    async def _fetch_dir(
+        self, prefix: str, pbar, result_queue: ResultQueue
+    ) -> set[str]:
         path = f"{self.name}/{prefix}"
         infos = await self.ls_dir(path)
         files = []

datachain/client/local.py

@@ -12,6 +12,7 @@ from datachain.lib.file import File
 from .fsspec import Client
 
 if TYPE_CHECKING:
+    from datachain.cache import DataChainCache
     from datachain.dataset import StorageURI
 
 
@@ -21,7 +22,11 @@ class FileClient(Client):
     protocol = "file"
 
     def __init__(
-        self, name: str, fs_kwargs: dict[str, Any], cache, use_symlinks: bool = False
+        self,
+        name: str,
+        fs_kwargs: dict[str, Any],
+        cache: "DataChainCache",
+        use_symlinks: bool = False,
     ) -> None:
         super().__init__(name, fs_kwargs, cache)
         self.use_symlinks = use_symlinks
@@ -30,7 +35,7 @@ class FileClient(Client):
         raise TypeError("Signed urls are not implemented for local file system")
 
     @classmethod
-    def get_uri(cls, name) -> "StorageURI":
+    def get_uri(cls, name: str) -> "StorageURI":
         from datachain.dataset import StorageURI
 
         return StorageURI(f'{cls.PREFIX}/{name.removeprefix("/")}')
@@ -77,7 +82,7 @@ class FileClient(Client):
         return bucket, path
 
     @classmethod
-    def from_name(cls, name: str, cache, kwargs) -> "FileClient":
+    def from_name(cls, name: str, cache: "DataChainCache", kwargs) -> "FileClient":
         use_symlinks = kwargs.pop("use_symlinks", False)
         return cls(name, kwargs, cache, use_symlinks=use_symlinks)
 
@@ -85,7 +90,7 @@ class FileClient(Client):
     def from_source(
         cls,
         uri: str,
-        cache,
+        cache: "DataChainCache",
         use_symlinks: bool = False,
         **kwargs,
     ) -> "FileClient":

datachain/data_storage/metastore.py

@@ -725,9 +725,10 @@ class AbstractDBMetastore(AbstractMetastore):
 
     def list_datasets(self) -> Iterator["DatasetListRecord"]:
         """Lists all datasets."""
-        yield from self._parse_dataset_list(
-            self.db.execute(self._base_list_datasets_query())
+        query = self._base_list_datasets_query().order_by(
+            self._datasets.c.name, self._datasets_versions.c.version
         )
+        yield from self._parse_dataset_list(self.db.execute(query))
 
     def list_datasets_by_prefix(
         self, prefix: str, conn=None

datachain/func/__init__.py

@@ -17,8 +17,9 @@ from .aggregate import (
 )
 from .array import cosine_distance, euclidean_distance, length, sip_hash_64
 from .conditional import greatest, least
-from .numeric import bit_and, bit_or, bit_xor, int_hash_64
+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
 from .window import window
 
 __all__ = [
@@ -26,8 +27,10 @@ __all__ = [
     "array",
     "avg",
     "bit_and",
+    "bit_hamming_distance",
     "bit_or",
     "bit_xor",
+    "byte_hamming_distance",
     "case",
     "collect",
     "concat",

datachain/func/numeric.py

@@ -160,3 +160,49 @@ def int_hash_64(col: Union[ColT, int]) -> Func:
     return Func(
         "int_hash_64", inner=numeric.int_hash_64, cols=cols, args=args, result_type=int
     )
+
+
+def bit_hamming_distance(*args: Union[ColT, int]) -> Func:
+    """
+    Computes the Hamming distance between the bit representations of two integer values.
+
+    The Hamming distance is the number of positions at which the corresponding bits
+    are different. This function returns the dissimilarity between the integers,
+    where 0 indicates identical integers and values closer to the number of bits
+    in the integer indicate higher dissimilarity.
+
+    Args:
+        args (str | int): Two integers to compute the Hamming distance between.
+            If a str is provided, it is assumed to be the name of the column.
+            If an int is provided, it is assumed to be an integer literal.
+
+    Returns:
+        Func: A Func object that represents the Hamming distance function.
+
+    Example:
+        ```py
+        dc.mutate(
+            ham_dist=func.bit_hamming_distance("embed1", 123456),
+        )
+        ```
+
+    Notes:
+        - Result column will always be of type int.
+    """
+    cols, func_args = [], []
+    for arg in args:
+        if isinstance(arg, int):
+            func_args.append(arg)
+        else:
+            cols.append(arg)
+
+    if len(cols) + len(func_args) != 2:
+        raise ValueError("bit_hamming_distance() requires exactly two arguments")
+
+    return Func(
+        "bit_hamming_distance",
+        inner=numeric.bit_hamming_distance,
+        cols=cols,
+        args=func_args,
+        result_type=int,
+    )

datachain/func/string.py

@@ -152,3 +152,49 @@ def regexp_replace(col: Union[str, Func], regex: str, replacement: str) -> Func:
         args = None
 
     return Func("regexp_replace", inner=inner, cols=cols, args=args, result_type=str)
+
+
+def byte_hamming_distance(*args: Union[str, Func]) -> Func:
+    """
+    Computes the Hamming distance between two strings.
+
+    The Hamming distance is the number of positions at which the corresponding
+    characters are different. This function returns the dissimilarity between
+    the strings, where 0 indicates identical strings and values closer to the length
+    of the strings indicate higher dissimilarity.
+
+    Args:
+        args (str | literal): Two strings to compute the Hamming distance between.
+            If a str is provided, it is assumed to be the name of the column.
+            If a Literal is provided, it is assumed to be a string literal.
+
+    Returns:
+        Func: A Func object that represents the Hamming distance function.
+
+    Example:
+        ```py
+        dc.mutate(
+            ham_dist=func.byte_hamming_distance("file.phash", literal("hello")),
+        )
+        ```
+
+    Notes:
+        - Result column will always be of type int.
+    """
+    cols, func_args = [], []
+    for arg in args:
+        if get_origin(arg) is literal:
+            func_args.append(arg)
+        else:
+            cols.append(arg)
+
+    if len(cols) + len(func_args) != 2:
+        raise ValueError("byte_hamming_distance() requires exactly two arguments")
+
+    return Func(
+        "byte_hamming_distance",
+        inner=string.byte_hamming_distance,
+        cols=cols,
+        args=func_args,
+        result_type=int,
+    )

datachain/lib/convert/flatten.py

@@ -1,19 +1,21 @@
+from collections.abc import Generator
+
 from pydantic import BaseModel
 
 from datachain.lib.model_store import ModelStore
 
 
-def flatten(obj: BaseModel):
+def flatten(obj: BaseModel) -> tuple:
     return tuple(_flatten_fields_values(obj.model_fields, obj))
 
 
-def flatten_list(obj_list):
+def flatten_list(obj_list: list[BaseModel]) -> tuple:
     return tuple(
         val for obj in obj_list for val in _flatten_fields_values(obj.model_fields, obj)
     )
 
 
-def _flatten_list_field(value: list):
+def _flatten_list_field(value: list) -> list:
     assert isinstance(value, list)
     if value and ModelStore.is_pydantic(type(value[0])):
         return [val.model_dump() for val in value]
@@ -22,7 +24,7 @@ def _flatten_list_field(value: list):
     return value
 
 
-def _flatten_fields_values(fields, obj: BaseModel):
+def _flatten_fields_values(fields: dict, obj: BaseModel) -> Generator:
     for name, f_info in fields.items():
         anno = f_info.annotation
         # Optimization: Access attributes directly to skip the model_dump() call.
@@ -40,5 +42,5 @@ def _flatten_fields_values(fields, obj: BaseModel):
             yield value
 
 
-def _flatten(obj):
+def _flatten(obj: BaseModel) -> tuple:
     return tuple(_flatten_fields_values(obj.model_fields, obj))

datachain/lib/convert/unflatten.py

@@ -9,12 +9,12 @@ from pydantic import BaseModel
 from datachain.query.schema import DEFAULT_DELIMITER
 
 
-def unflatten_to_json(model: type[BaseModel], row: Sequence[Any], pos=0) -> dict:
+def unflatten_to_json(model: type[BaseModel], row: Sequence[Any], pos: int = 0) -> dict:
     return unflatten_to_json_pos(model, row, pos)[0]
 
 
 def unflatten_to_json_pos(
-    model: type[BaseModel], row: Sequence[Any], pos=0
+    model: type[BaseModel], row: Sequence[Any], pos: int = 0
 ) -> tuple[dict, int]:
     res = {}
     for name, f_info in model.model_fields.items():

datachain/lib/convert/values_to_tuples.py

@@ -11,7 +11,7 @@ from datachain.lib.utils import DataChainParamsError
 
 
 class ValuesToTupleError(DataChainParamsError):
-    def __init__(self, ds_name, msg):
+    def __init__(self, ds_name: str, msg: str):
         if ds_name:
             ds_name = f"' {ds_name}'"
         super().__init__(f"Cannot convert signals for dataset{ds_name}: {msg}")

datachain/lib/dc.py

@@ -1446,6 +1446,7 @@ class DataChain:
             tokenizer=tokenizer,
             tokenizer_kwargs=tokenizer_kwargs,
             num_samples=num_samples,
+            dc_settings=chain._settings,
         )
 
     def remove_file_signals(self) -> "Self":  # noqa: D102

datachain/lib/pytorch.py

@@ -10,8 +10,10 @@ from torchvision.transforms import v2
 from tqdm import tqdm
 
 from datachain import Session
+from datachain.asyn import AsyncMapper
 from datachain.catalog import Catalog, get_catalog
 from datachain.lib.dc import DataChain
+from datachain.lib.settings import Settings
 from datachain.lib.text import convert_text
 
 if TYPE_CHECKING:
@@ -30,6 +32,8 @@ def label_to_int(value: str, classes: list) -> int:
 
 
 class PytorchDataset(IterableDataset):
+    prefetch: int = 2
+
     def __init__(
         self,
         name: str,
@@ -39,6 +43,7 @@ class PytorchDataset(IterableDataset):
         tokenizer: Optional[Callable] = None,
         tokenizer_kwargs: Optional[dict[str, Any]] = None,
         num_samples: int = 0,
+        dc_settings: Optional[Settings] = None,
     ):
         """
         Pytorch IterableDataset that streams DataChain datasets.
@@ -66,6 +71,11 @@ class PytorchDataset(IterableDataset):
             catalog = get_catalog()
         self._init_catalog(catalog)
 
+        dc_settings = dc_settings or Settings()
+        self.cache = dc_settings.cache
+        if (prefetch := dc_settings.prefetch) is not None:
+            self.prefetch = prefetch
+
     def _init_catalog(self, catalog: "Catalog"):
         # For compatibility with multiprocessing,
         # we can only store params in __init__(), as Catalog isn't picklable
@@ -82,51 +92,58 @@ class PytorchDataset(IterableDataset):
         wh = wh_cls(*wh_args, **wh_kwargs)
         return Catalog(ms, wh, **self._catalog_params)
 
-    def __iter__(self) -> Iterator[Any]:
-        if self.catalog is None:
-            self.catalog = self._get_catalog()
-        session = Session.get(catalog=self.catalog)
-        total_rank, total_workers = self.get_rank_and_workers()
+    def _rows_iter(self, total_rank: int, total_workers: int):
+        catalog = self._get_catalog()
+        session = Session("PyTorch", catalog=catalog)
         ds = DataChain.from_dataset(
             name=self.name, version=self.version, session=session
-        )
+        ).settings(cache=self.cache, prefetch=self.prefetch)
         ds = ds.remove_file_signals()
 
         if self.num_samples > 0:
             ds = ds.sample(self.num_samples)
         ds = ds.chunk(total_rank, total_workers)
+        yield from ds.collect()
+
+    def __iter__(self) -> Iterator[Any]:
+        total_rank, total_workers = self.get_rank_and_workers()
+        rows = self._rows_iter(total_rank, total_workers)
+        if self.prefetch > 0:
+            from datachain.lib.udf import _prefetch_input
+
+            rows = AsyncMapper(_prefetch_input, rows, workers=self.prefetch).iterate()
+
         desc = f"Parsed PyTorch dataset for rank={total_rank} worker"
-        with tqdm(desc=desc, unit=" rows") as pbar:
-            for row_features in ds.collect():
-                row = []
-                for fr in row_features:
-                    if hasattr(fr, "read"):
-                        row.append(fr.read())  # type: ignore[unreachable]
-                    else:
-                        row.append(fr)
-                # Apply transforms
-                if self.transform:
-                    try:
-                        if isinstance(self.transform, v2.Transform):
-                            row = self.transform(row)
-                        for i, val in enumerate(row):
-                            if isinstance(val, Image.Image):
-                                row[i] = self.transform(val)
-                    except ValueError:
-                        logger.warning(
-                            "Skipping transform due to unsupported data types."
-                        )
-                        self.transform = None
-                if self.tokenizer:
-                    for i, val in enumerate(row):
-                        if isinstance(val, str) or (
-                            isinstance(val, list) and isinstance(val[0], str)
-                        ):
-                            row[i] = convert_text(
-                                val, self.tokenizer, self.tokenizer_kwargs
-                            ).squeeze(0)  # type: ignore[union-attr]
-                yield row
-                pbar.update(1)
+        with tqdm(rows, desc=desc, unit=" rows", position=total_rank) as rows_it:
+            yield from map(self._process_row, rows_it)
+
+    def _process_row(self, row_features):
+        row = []
+        for fr in row_features:
+            if hasattr(fr, "read"):
+                row.append(fr.read())  # type: ignore[unreachable]
+            else:
+                row.append(fr)
+        # Apply transforms
+        if self.transform:
+            try:
+                if isinstance(self.transform, v2.Transform):
+                    row = self.transform(row)
+                for i, val in enumerate(row):
+                    if isinstance(val, Image.Image):
+                        row[i] = self.transform(val)
+            except ValueError:
+                logger.warning("Skipping transform due to unsupported data types.")
+                self.transform = None
+        if self.tokenizer:
+            for i, val in enumerate(row):
+                if isinstance(val, str) or (
+                    isinstance(val, list) and isinstance(val[0], str)
+                ):
+                    row[i] = convert_text(
+                        val, self.tokenizer, self.tokenizer_kwargs
+                    ).squeeze(0)  # type: ignore[union-attr]
+        return row
 
     @staticmethod
     def get_rank_and_workers() -> tuple[int, int]:

datachain/lib/utils.py

@@ -28,7 +28,7 @@ class DataChainParamsError(DataChainError):
 
 
 class DataChainColumnError(DataChainParamsError):
-    def __init__(self, col_name, msg):
+    def __init__(self, col_name: str, msg: str):
         super().__init__(f"Error for column {col_name}: {msg}")
 
 

datachain/query/dataset.py

@@ -215,7 +215,7 @@ class DatasetDiffOperation(Step):
         Should return select query that calculates desired diff between dataset queries
         """
 
-    def apply(self, query_generator, temp_tables: list[str]):
+    def apply(self, query_generator, temp_tables: list[str]) -> "StepResult":
         source_query = query_generator.exclude(("sys__id",))
         target_query = self.dq.apply_steps().select()
         temp_tables.extend(self.dq.temp_table_names)

datachain/remote/studio.py

@@ -119,18 +119,27 @@ class StudioClient:
                 "\tpip install 'datachain[remote]'"
             ) from None
 
-    def _send_request_msgpack(self, route: str, data: dict[str, Any]) -> Response[Any]:
+    def _send_request_msgpack(
+        self, route: str, data: dict[str, Any], method: Optional[str] = "POST"
+    ) -> Response[Any]:
         import msgpack
         import requests
 
-        response = requests.post(
-            f"{self.url}/{route}",
-            json={**data, "team_name": self.team},
+        kwargs = (
+            {"params": {**data, "team_name": self.team}}
+            if method == "GET"
+            else {"json": {**data, "team_name": self.team}}
+        )
+
+        response = requests.request(
+            method=method,  # type: ignore[arg-type]
+            url=f"{self.url}/{route}",
             headers={
                 "Content-Type": "application/json",
                 "Authorization": f"token {self.token}",
             },
             timeout=self.timeout,
+            **kwargs,  # type: ignore[arg-type]
         )
         ok = response.ok
         if not ok:
@@ -148,7 +157,9 @@ class StudioClient:
         return Response(response_data, ok, message)
 
     @retry_with_backoff(retries=5)
-    def _send_request(self, route: str, data: dict[str, Any]) -> Response[Any]:
+    def _send_request(
+        self, route: str, data: dict[str, Any], method: Optional[str] = "POST"
+    ) -> Response[Any]:
         """
         Function that communicate Studio API.
         It will raise an exception, and try to retry, if 5xx status code is
@@ -157,14 +168,21 @@ class StudioClient:
         """
         import requests
 
-        response = requests.post(
-            f"{self.url}/{route}",
-            json={**data, "team_name": self.team},
+        kwargs = (
+            {"params": {**data, "team_name": self.team}}
+            if method == "GET"
+            else {"json": {**data, "team_name": self.team}}
+        )
+
+        response = requests.request(
+            method=method,  # type: ignore[arg-type]
+            url=f"{self.url}/{route}",
             headers={
                 "Content-Type": "application/json",
                 "Authorization": f"token {self.token}",
             },
             timeout=self.timeout,
+            **kwargs,  # type: ignore[arg-type]
         )
         try:
             response.raise_for_status()
@@ -222,7 +240,7 @@ class StudioClient:
             yield path, response
 
     def ls_datasets(self) -> Response[LsData]:
-        return self._send_request("datachain/ls-datasets", {})
+        return self._send_request("datachain/datasets", {}, method="GET")
 
     def edit_dataset(
         self,
@@ -232,20 +250,14 @@ class StudioClient:
         labels: Optional[list[str]] = None,
     ) -> Response[DatasetInfoData]:
         body = {
+            "new_name": new_name,
             "dataset_name": name,
+            "description": description,
+            "labels": labels,
         }
 
-        if new_name is not None:
-            body["new_name"] = new_name
-
-        if description is not None:
-            body["description"] = description
-
-        if labels is not None:
-            body["labels"] = labels  # type: ignore[assignment]
-
         return self._send_request(
-            "datachain/edit-dataset",
+            "datachain/datasets",
             body,
         )
 
@@ -256,12 +268,13 @@ class StudioClient:
         force: Optional[bool] = False,
     ) -> Response[DatasetInfoData]:
         return self._send_request(
-            "datachain/rm-dataset",
+            "datachain/datasets",
             {
                 "dataset_name": name,
                 "version": version,
                 "force": force,
             },
+            method="DELETE",
         )
 
     def dataset_info(self, name: str) -> Response[DatasetInfoData]:
@@ -272,7 +285,9 @@ class StudioClient:
 
             return dataset_info
 
-        response = self._send_request("datachain/dataset-info", {"dataset_name": name})
+        response = self._send_request(
+            "datachain/datasets/info", {"dataset_name": name}, method="GET"
+        )
         if response.ok:
             response.data = _parse_dataset_info(response.data)
         return response
@@ -282,14 +297,16 @@ class StudioClient:
     ) -> Response[DatasetRowsData]:
         req_data = {"dataset_name": name, "dataset_version": version}
         return self._send_request_msgpack(
-            "datachain/dataset-rows",
+            "datachain/datasets/rows",
             {**req_data, "offset": offset, "limit": DATASET_ROWS_CHUNK_SIZE},
+            method="GET",
         )
 
     def dataset_stats(self, name: str, version: int) -> Response[DatasetStatsData]:
         response = self._send_request(
-            "datachain/dataset-stats",
+            "datachain/datasets/stats",
             {"dataset_name": name, "dataset_version": version},
+            method="GET",
         )
         if response.ok:
             response.data = DatasetStats(**response.data)
@@ -299,16 +316,18 @@ class StudioClient:
         self, name: str, version: int
     ) -> Response[DatasetExportSignedUrls]:
         return self._send_request(
-            "datachain/dataset-export",
+            "datachain/datasets/export",
             {"dataset_name": name, "dataset_version": version},
+            method="GET",
         )
 
     def dataset_export_status(
         self, name: str, version: int
     ) -> Response[DatasetExportStatus]:
         return self._send_request(
-            "datachain/dataset-export-status",
+            "datachain/datasets/export-status",
             {"dataset_name": name, "dataset_version": version},
+            method="GET",
         )
 
     def upload_file(self, file_name: str, content: bytes) -> Response[FileUploadData]:

datachain/sql/functions/numeric.py

@@ -35,9 +35,21 @@ class int_hash_64(GenericFunction):  # noqa: N801
     inherit_cache = True
 
 
+class bit_hamming_distance(GenericFunction):  # noqa: N801
+    """
+    Returns the Hamming distance between two integers.
+    """
+
+    type = Int64()
+    package = "numeric"
+    name = "hamming_distance"
+    inherit_cache = True
+
+
 compiler_not_implemented(bit_and)
 compiler_not_implemented(bit_or)
 compiler_not_implemented(bit_xor)
 compiler_not_implemented(bit_rshift)
 compiler_not_implemented(bit_lshift)
 compiler_not_implemented(int_hash_64)
+compiler_not_implemented(bit_hamming_distance)