datachain 0.8.3__py3-none-any.whl → 0.8.5__py3-none-any.whl

datachain/func/conditional.py

@@ -1,9 +1,15 @@
 from typing import Union
 
+from sqlalchemy import case as sql_case
+from sqlalchemy.sql.elements import BinaryExpression
+
+from datachain.lib.utils import DataChainParamsError
 from datachain.sql.functions import conditional
 
 from .func import ColT, Func
 
+CaseT = Union[int, float, complex, bool, str]
+
 
 def greatest(*args: Union[ColT, float]) -> Func:
     """
@@ -79,3 +85,71 @@ def least(*args: Union[ColT, float]) -> Func:
     return Func(
         "least", inner=conditional.least, cols=cols, args=func_args, result_type=int
     )
+
+
+def case(*args: tuple[BinaryExpression, CaseT], else_=None) -> Func:
+    """
+    Returns the case function that produces case expression which has a list of
+    conditions and corresponding results. Results can only be python primitives
+    like string, numbes or booleans. Result type is inferred from condition results.
+
+    Args:
+        args (tuple(BinaryExpression, value(str | int | float | complex | bool):
+            - Tuple of binary expression and values pair which corresponds to one
+            case condition - value
+        else_ (str | int | float | complex | bool): else value in case expression
+
+    Returns:
+        Func: A Func object that represents the case function.
+
+    Example:
+        ```py
+        dc.mutate(
+            res=func.case((C("num") > 0, "P"), (C("num") < 0, "N"), else_="Z"),
+        )
+        ```
+    """
+    supported_types = [int, float, complex, str, bool]
+
+    type_ = type(else_) if else_ else None
+
+    if not args:
+        raise DataChainParamsError("Missing statements")
+
+    for arg in args:
+        if type_ and not isinstance(arg[1], type_):
+            raise DataChainParamsError("Statement values must be of the same type")
+        type_ = type(arg[1])
+
+    if type_ not in supported_types:
+        raise DataChainParamsError(
+            f"Only python literals ({supported_types}) are supported for values"
+        )
+
+    kwargs = {"else_": else_}
+    return Func("case", inner=sql_case, args=args, kwargs=kwargs, result_type=type_)
+
+
+def ifelse(condition: BinaryExpression, if_val: CaseT, else_val: CaseT) -> Func:
+    """
+    Returns the ifelse function that produces if expression which has a condition
+    and values for true and false outcome. Results can only be python primitives
+    like string, numbes or booleans. Result type is inferred from the values.
+
+    Args:
+        condition: BinaryExpression - condition which is evaluated
+        if_val: (str | int | float | complex | bool): value for true condition outcome
+        else_val: (str | int | float | complex | bool): value for false condition
+         outcome
+
+    Returns:
+        Func: A Func object that represents the ifelse function.
+
+    Example:
+        ```py
+        dc.mutate(
+            res=func.ifelse(C("num") > 0, "P", "N"),
+        )
+        ```
+    """
+    return case((condition, if_val), else_=else_val)

datachain/func/func.py

@@ -35,6 +35,7 @@ class Func(Function):
         inner: Callable,
         cols: Optional[Sequence[ColT]] = None,
         args: Optional[Sequence[Any]] = None,
+        kwargs: Optional[dict[str, Any]] = None,
         result_type: Optional["DataType"] = None,
         is_array: bool = False,
         is_window: bool = False,
@@ -45,6 +46,7 @@ class Func(Function):
         self.inner = inner
         self.cols = cols or []
         self.args = args or []
+        self.kwargs = kwargs or {}
         self.result_type = result_type
         self.is_array = is_array
         self.is_window = is_window
@@ -63,6 +65,7 @@ class Func(Function):
             self.inner,
             self.cols,
             self.args,
+            self.kwargs,
             self.result_type,
             self.is_array,
             self.is_window,
@@ -333,6 +336,7 @@ class Func(Function):
             self.inner,
             self.cols,
             self.args,
+            self.kwargs,
             self.result_type,
             self.is_array,
             self.is_window,
@@ -387,7 +391,7 @@ class Func(Function):
             return col
 
         cols = [get_col(col) for col in self._db_cols]
-        func_col = self.inner(*cols, *self.args)
+        func_col = self.inner(*cols, *self.args, **self.kwargs)
 
         if self.is_window:
             if not self.window:

datachain/lib/arrow.py

@@ -91,7 +91,9 @@ class ArrowGenerator(Generator):
                 yield from record_batch.to_pylist()
 
         it = islice(iter_records(), self.nrows)
-        with tqdm(it, desc="Parsed by pyarrow", unit="rows", total=self.nrows) as pbar:
+        with tqdm(
+            it, desc="Parsed by pyarrow", unit="rows", total=self.nrows, leave=False
+        ) as pbar:
             for index, record in enumerate(pbar):
                 yield self._process_record(
                     record, file, index, hf_schema, use_datachain_schema
@@ -149,6 +151,10 @@ def infer_schema(chain: "DataChain", **kwargs) -> pa.Schema:
     for file in chain.collect("file"):
         ds = dataset(file.get_path(), filesystem=file.get_fs(), **kwargs)  # type: ignore[union-attr]
         schemas.append(ds.schema)
+    if not schemas:
+        raise ValueError(
+            "Cannot infer schema (no files to process or can't access them)"
+        )
     return pa.unify_schemas(schemas)
 
 

datachain/lib/dc.py

@@ -451,6 +451,7 @@ class DataChain:
             return dc
 
         if update or not list_ds_exists:
+            # disable prefetch for listing, as it pre-downloads all files
             (
                 cls.from_records(
                     DataChain.DEFAULT_FILE_RECORD,
@@ -458,6 +459,7 @@ class DataChain:
                     settings=settings,
                     in_memory=in_memory,
                 )
+                .settings(prefetch=0)
                 .gen(
                     list_bucket(list_uri, cache, client_config=client_config),
                     output={f"{object_name}": File},
@@ -1534,7 +1536,7 @@ class DataChain:
 
         Example:
             ```py
-            diff = persons.diff(
+            res = persons.compare(
                 new_persons,
                 on=["id"],
                 right_on=["other_id"],
@@ -1547,9 +1549,9 @@ class DataChain:
             )
             ```
         """
-        from datachain.lib.diff import compare as chain_compare
+        from datachain.diff import _compare
 
-        return chain_compare(
+        return _compare(
             self,
             other,
             on,
@@ -1882,6 +1884,9 @@ class DataChain:
                     "`nrows` only supported for csv and json formats.",
                 )
 
+        if "file" not in self.schema or not self.count():
+            raise DatasetPrepareError(self.name, "no files to parse.")
+
         schema = None
         col_names = output if isinstance(output, Sequence) else None
         if col_names or not output:

datachain/lib/file.py

@@ -269,10 +269,21 @@ class File(DataModel):
         client = self._catalog.get_client(self.source)
         client.download(self, callback=self._download_cb)
 
-    async def _prefetch(self) -> None:
-        if self._caching_enabled:
-            client = self._catalog.get_client(self.source)
-            await client._download(self, callback=self._download_cb)
+    async def _prefetch(self, download_cb: Optional["Callback"] = None) -> bool:
+        from datachain.client.hf import HfClient
+
+        if self._catalog is None:
+            raise RuntimeError("cannot prefetch file because catalog is not setup")
+
+        client = self._catalog.get_client(self.source)
+        if client.protocol == HfClient.protocol:
+            return False
+
+        await client._download(self, callback=download_cb or self._download_cb)
+        self._set_stream(
+            self._catalog, caching_enabled=True, download_cb=DEFAULT_CALLBACK
+        )
+        return True
 
     def get_local_path(self) -> Optional[str]:
         """Return path to a file in a local cache.
@@ -364,7 +375,7 @@ class File(DataModel):
 
         try:
             info = client.fs.info(client.get_full_path(self.path))
-            converted_info = client.info_to_file(info, self.source)
+            converted_info = client.info_to_file(info, self.path)
             return type(self)(
                 path=self.path,
                 source=self.source,

datachain/lib/hf.py

@@ -95,7 +95,7 @@ class HFGenerator(Generator):
         ds = self.ds_dict[split]
         if split:
             desc += f" split '{split}'"
-        with tqdm(desc=desc, unit=" rows") as pbar:
+        with tqdm(desc=desc, unit=" rows", leave=False) as pbar:
             for row in ds:
                 output_dict = {}
                 if split and "split" in self.output_schema.model_fields:

datachain/lib/listing.py

@@ -85,6 +85,24 @@ def ls(
     return dc.filter(pathfunc.parent(_file_c("path")) == path.lstrip("/").rstrip("/*"))
 
 
+def _isfile(client: "Client", path: str) -> bool:
+    """
+    Returns True if uri points to a file
+    """
+    try:
+        info = client.fs.info(path)
+        name = info.get("name")
+        # case for special simulated directories on some clouds
+        # e.g. Google creates a zero byte file with the same name as the
+        # directory with a trailing slash at the end
+        if not name or name.endswith("/"):
+            return False
+
+        return info["type"] == "file"
+    except:  # noqa: E722
+        return False
+
+
 def parse_listing_uri(uri: str, cache, client_config) -> tuple[Optional[str], str, str]:
     """
     Parsing uri and returns listing dataset name, listing uri and listing path
@@ -94,7 +112,7 @@ def parse_listing_uri(uri: str, cache, client_config) -> tuple[Optional[str], st
     storage_uri, path = Client.parse_url(uri)
     telemetry.log_param("client", client.PREFIX)
 
-    if not uri.endswith("/") and client.fs.isfile(uri):
+    if not uri.endswith("/") and _isfile(client, uri):
         return None, f'{storage_uri}/{path.lstrip("/")}', path
     if uses_glob(path):
         lst_uri_path = posixpath.dirname(path)

datachain/lib/pytorch.py

@@ -1,5 +1,8 @@
 import logging
-from collections.abc import Iterator
+import os
+import weakref
+from collections.abc import Generator, Iterable, Iterator
+from contextlib import closing
 from typing import TYPE_CHECKING, Any, Callable, Optional
 
 from PIL import Image
@@ -9,15 +12,19 @@ from torch.utils.data import IterableDataset, get_worker_info
 from torchvision.transforms import v2
 
 from datachain import Session
-from datachain.asyn import AsyncMapper
+from datachain.cache import get_temp_cache
 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
+from datachain.progress import CombinedDownloadCallback
+from datachain.query.dataset import get_download_callback
 
 if TYPE_CHECKING:
     from torchvision.transforms.v2 import Transform
 
+    from datachain.cache import DataChainCache as Cache
+
 
 logger = logging.getLogger("datachain")
 
@@ -75,6 +82,19 @@ class PytorchDataset(IterableDataset):
         if (prefetch := dc_settings.prefetch) is not None:
             self.prefetch = prefetch
 
+        self._cache = catalog.cache
+        self._prefetch_cache: Optional[Cache] = None
+        if prefetch and not self.cache:
+            tmp_dir = catalog.cache.tmp_dir
+            assert tmp_dir
+            self._prefetch_cache = get_temp_cache(tmp_dir, prefix="prefetch-")
+            self._cache = self._prefetch_cache
+            weakref.finalize(self, self._prefetch_cache.destroy)
+
+    def close(self) -> None:
+        if self._prefetch_cache:
+            self._prefetch_cache.destroy()
+
     def _init_catalog(self, catalog: "Catalog"):
         # For compatibility with multiprocessing,
         # we can only store params in __init__(), as Catalog isn't picklable
@@ -89,9 +109,15 @@ class PytorchDataset(IterableDataset):
         ms = ms_cls(*ms_args, **ms_kwargs)
         wh_cls, wh_args, wh_kwargs = self._wh_params
         wh = wh_cls(*wh_args, **wh_kwargs)
-        return Catalog(ms, wh, **self._catalog_params)
+        catalog = Catalog(ms, wh, **self._catalog_params)
+        catalog.cache = self._cache
+        return catalog
 
-    def _rows_iter(self, total_rank: int, total_workers: int):
+    def _row_iter(
+        self,
+        total_rank: int,
+        total_workers: int,
+    ) -> Generator[tuple[Any, ...], None, None]:
         catalog = self._get_catalog()
         session = Session("PyTorch", catalog=catalog)
         ds = DataChain.from_dataset(
@@ -104,16 +130,34 @@ class PytorchDataset(IterableDataset):
         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()
-        yield from map(self._process_row, rows)
+    def _iter_with_prefetch(self) -> Generator[tuple[Any], None, None]:
+        from datachain.lib.udf import _prefetch_inputs
 
-    def _process_row(self, row_features):
+        total_rank, total_workers = self.get_rank_and_workers()
+        download_cb = CombinedDownloadCallback()
+        if os.getenv("DATACHAIN_SHOW_PREFETCH_PROGRESS"):
+            download_cb = get_download_callback(
+                f"{total_rank}/{total_workers}",
+                position=total_rank,
+                leave=True,
+            )
+
+        rows = self._row_iter(total_rank, total_workers)
+        rows = _prefetch_inputs(
+            rows,
+            self.prefetch,
+            download_cb=download_cb,
+            after_prefetch=download_cb.increment_file_count,
+        )
+
+        with download_cb, closing(rows):
+            yield from rows
+
+    def __iter__(self) -> Iterator[list[Any]]:
+        with closing(self._iter_with_prefetch()) as rows:
+            yield from map(self._process_row, rows)
+
+    def _process_row(self, row_features: Iterable[Any]) -> list[Any]:
         row = []
         for fr in row_features:
             if hasattr(fr, "read"):

datachain/lib/signal_schema.py

@@ -13,13 +13,14 @@ from typing import (  # noqa: UP035
     Final,
     List,
     Literal,
+    Mapping,
     Optional,
     Union,
     get_args,
     get_origin,
 )
 
-from pydantic import BaseModel, create_model
+from pydantic import BaseModel, Field, create_model
 from sqlalchemy import ColumnElement
 from typing_extensions import Literal as LiteralEx
 
@@ -85,8 +86,31 @@ class SignalResolvingTypeError(SignalResolvingError):
         )
 
 
+class CustomType(BaseModel):
+    schema_version: int = Field(ge=1, le=2, strict=True)
+    name: str
+    fields: dict[str, str]
+    bases: list[tuple[str, str, Optional[str]]]
+
+    @classmethod
+    def deserialize(cls, data: dict[str, Any], type_name: str) -> "CustomType":
+        version = data.get("schema_version", 1)
+
+        if version == 1:
+            data = {
+                "schema_version": 1,
+                "name": type_name,
+                "fields": data,
+                "bases": [],
+            }
+
+        return cls(**data)
+
+
 def create_feature_model(
-    name: str, fields: dict[str, Union[type, tuple[type, Any]]]
+    name: str,
+    fields: Mapping[str, Union[type, None, tuple[type, Any]]],
+    base: Optional[type] = None,
 ) -> type[BaseModel]:
     """
     This gets or returns a dynamic feature model for use in restoring a model
@@ -98,7 +122,7 @@ def create_feature_model(
     name = name.replace("@", "_")
     return create_model(
         name,
-        __base__=DataModel,  # type: ignore[call-overload]
+        __base__=base or DataModel,  # type: ignore[call-overload]
         # These are tuples for each field of: annotation, default (if any)
         **{
             field_name: anno if isinstance(anno, tuple) else (anno, None)
@@ -156,7 +180,7 @@ class SignalSchema:
         return SignalSchema(signals)
 
     @staticmethod
-    def _serialize_custom_model_fields(
+    def _serialize_custom_model(
         version_name: str, fr: type[BaseModel], custom_types: dict[str, Any]
     ) -> str:
         """This serializes any custom type information to the provided custom_types
@@ -165,12 +189,23 @@ class SignalSchema:
             # This type is already stored in custom_types.
             return version_name
         fields = {}
+
         for field_name, info in fr.model_fields.items():
             field_type = info.annotation
             # All fields should be typed.
             assert field_type
             fields[field_name] = SignalSchema._serialize_type(field_type, custom_types)
-        custom_types[version_name] = fields
+
+        bases: list[tuple[str, str, Optional[str]]] = []
+        for type_ in fr.__mro__:
+            model_store_name = (
+                ModelStore.get_name(type_) if issubclass(type_, DataModel) else None
+            )
+            bases.append((type_.__name__, type_.__module__, model_store_name))
+
+        ct = CustomType(schema_version=2, name=version_name, fields=fields, bases=bases)
+        custom_types[version_name] = ct.model_dump()
+
         return version_name
 
     @staticmethod
@@ -184,15 +219,12 @@ class SignalSchema:
             if st is None or not ModelStore.is_pydantic(st):
                 continue
             # Register and save feature types.
-            ModelStore.register(st)
             st_version_name = ModelStore.get_name(st)
             if st is fr:
                 # If the main type is Pydantic, then use the ModelStore version name.
                 type_name = st_version_name
             # Save this type to custom_types.
-            SignalSchema._serialize_custom_model_fields(
-                st_version_name, st, custom_types
-            )
+            SignalSchema._serialize_custom_model(st_version_name, st, custom_types)
         return type_name
 
     def serialize(self) -> dict[str, Any]:
@@ -215,7 +247,7 @@ class SignalSchema:
                 depth += 1
             elif c == "]":
                 if depth == 0:
-                    raise TypeError(
+                    raise ValueError(
                         "Extra closing square bracket when parsing subtype list"
                     )
                 depth -= 1
@@ -223,16 +255,51 @@ class SignalSchema:
                 subtypes.append(type_name[start:i].strip())
                 start = i + 1
         if depth > 0:
-            raise TypeError("Unclosed square bracket when parsing subtype list")
+            raise ValueError("Unclosed square bracket when parsing subtype list")
         subtypes.append(type_name[start:].strip())
         return subtypes
 
     @staticmethod
-    def _resolve_type(type_name: str, custom_types: dict[str, Any]) -> Optional[type]:  # noqa: PLR0911
+    def _deserialize_custom_type(
+        type_name: str, custom_types: dict[str, Any]
+    ) -> Optional[type]:
+        """Given a type name like MyType@v1 gets a type from ModelStore or recreates
+        it based on the information from the custom types dict that includes fields and
+        bases."""
+        model_name, version = ModelStore.parse_name_version(type_name)
+        fr = ModelStore.get(model_name, version)
+        if fr:
+            return fr
+
+        if type_name in custom_types:
+            ct = CustomType.deserialize(custom_types[type_name], type_name)
+
+            fields = {
+                field_name: SignalSchema._resolve_type(field_type_str, custom_types)
+                for field_name, field_type_str in ct.fields.items()
+            }
+
+            base_model = None
+            for base in ct.bases:
+                _, _, model_store_name = base
+                if model_store_name:
+                    model_name, version = ModelStore.parse_name_version(
+                        model_store_name
+                    )
+                    base_model = ModelStore.get(model_name, version)
+                    if base_model:
+                        break
+
+            return create_feature_model(type_name, fields, base=base_model)
+
+        return None
+
+    @staticmethod
+    def _resolve_type(type_name: str, custom_types: dict[str, Any]) -> Optional[type]:
         """Convert a string-based type back into a python type."""
         type_name = type_name.strip()
         if not type_name:
-            raise TypeError("Type cannot be empty")
+            raise ValueError("Type cannot be empty")
         if type_name == "NoneType":
             return None
 
@@ -240,14 +307,14 @@ class SignalSchema:
         subtypes: Optional[tuple[Optional[type], ...]] = None
         if bracket_idx > -1:
             if bracket_idx == 0:
-                raise TypeError("Type cannot start with '['")
+                raise ValueError("Type cannot start with '['")
             close_bracket_idx = type_name.rfind("]")
             if close_bracket_idx == -1:
-                raise TypeError("Unclosed square bracket when parsing type")
+                raise ValueError("Unclosed square bracket when parsing type")
             if close_bracket_idx < bracket_idx:
-                raise TypeError("Square brackets are out of order when parsing type")
+                raise ValueError("Square brackets are out of order when parsing type")
             if close_bracket_idx == bracket_idx + 1:
-                raise TypeError("Empty square brackets when parsing type")
+                raise ValueError("Empty square brackets when parsing type")
             subtype_names = SignalSchema._split_subtypes(
                 type_name[bracket_idx + 1 : close_bracket_idx]
             )
@@ -267,18 +334,10 @@ class SignalSchema:
                 return fr[subtypes]  # type: ignore[index]
             return fr  # type: ignore[return-value]
 
-        model_name, version = ModelStore.parse_name_version(type_name)
-        fr = ModelStore.get(model_name, version)
+        fr = SignalSchema._deserialize_custom_type(type_name, custom_types)
         if fr:
             return fr
 
-        if type_name in custom_types:
-            fields = custom_types[type_name]
-            fields = {
-                field_name: SignalSchema._resolve_type(field_type_str, custom_types)
-                for field_name, field_type_str in fields.items()
-            }
-            return create_feature_model(type_name, fields)
         # This can occur if a third-party or custom type is used, which is not available
         # when deserializing.
         warnings.warn(
@@ -317,7 +376,7 @@ class SignalSchema:
                         stacklevel=2,
                     )
                     continue
-            except TypeError as err:
+            except ValueError as err:
                 raise SignalSchemaError(
                     f"cannot deserialize '{signal}': {err}"
                 ) from err
@@ -662,6 +721,9 @@ class SignalSchema:
                 stacklevel=2,
             )
             return "Any"
+        if ModelStore.is_pydantic(type_):
+            ModelStore.register(type_)
+            return ModelStore.get_name(type_)
         return type_.__name__
 
     @staticmethod