datachain 0.3.13__py3-none-any.whl → 0.3.15__py3-none-any.whl

datachain/lib/file.py

@@ -1,5 +1,6 @@
 import io
 import json
+import logging
 import os
 import posixpath
 from abc import ABC, abstractmethod
@@ -15,6 +16,9 @@ from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 from PIL import Image
 from pydantic import Field, field_validator
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
 from datachain.cache import UniqueId
 from datachain.client.fileslice import FileSlice
 from datachain.lib.data_model import DataModel
@@ -25,6 +29,8 @@ from datachain.utils import TIME_ZERO
 if TYPE_CHECKING:
     from datachain.catalog import Catalog
 
+logger = logging.getLogger("datachain")
+
 # how to create file path when exporting
 ExportPlacement = Literal["filename", "etag", "fullpath", "checksum"]
 
@@ -251,14 +257,18 @@ class File(DataModel):
         dump = self.model_dump()
         return UniqueId(*(dump[k] for k in self._unique_id_keys))
 
-    def get_local_path(self) -> Optional[str]:
+    def get_local_path(self, download: bool = False) -> Optional[str]:
         """Returns path to a file in a local cache.
         Return None if file is not cached. Throws an exception if cache is not setup."""
         if self._catalog is None:
             raise RuntimeError(
                 "cannot resolve local file path because catalog is not setup"
             )
-        return self._catalog.cache.get_path(self.get_uid())
+        uid = self.get_uid()
+        if download:
+            client = self._catalog.get_client(self.source)
+            client.download(uid, callback=self._download_cb)
+        return self._catalog.cache.get_path(uid)
 
     def get_file_suffix(self):
         """Returns last part of file name with `.`."""
@@ -313,6 +323,70 @@ class File(DataModel):
         """Returns `fsspec` filesystem for the file."""
         return self._catalog.get_client(self.source).fs
 
+    def resolve(self) -> "Self":
+        """
+        Resolve a File object by checking its existence and updating its metadata.
+
+        Returns:
+            File: The resolved File object with updated metadata.
+        """
+        if self._catalog is None:
+            raise RuntimeError("Cannot resolve file: catalog is not set")
+
+        try:
+            client = self._catalog.get_client(self.source)
+        except NotImplementedError as e:
+            raise RuntimeError(
+                f"Unsupported protocol for file source: {self.source}"
+            ) from e
+
+        try:
+            info = client.fs.info(client.get_full_path(self.path))
+            converted_info = client.info_to_file(info, self.source)
+            return type(self)(
+                path=self.path,
+                source=self.source,
+                size=converted_info.size,
+                etag=converted_info.etag,
+                version=converted_info.version,
+                is_latest=converted_info.is_latest,
+                last_modified=converted_info.last_modified,
+                location=self.location,
+            )
+        except (FileNotFoundError, PermissionError, OSError) as e:
+            logger.warning("File system error when resolving %s: %s", self.path, str(e))
+
+        return type(self)(
+            path=self.path,
+            source=self.source,
+            size=0,
+            etag="",
+            version="",
+            is_latest=True,
+            last_modified=TIME_ZERO,
+            location=self.location,
+        )
+
+
+def resolve(file: File) -> File:
+    """
+    Resolve a File object by checking its existence and updating its metadata.
+
+    This function is a wrapper around the File.resolve() method, designed to be
+    used as a mapper in DataChain operations.
+
+    Args:
+        file (File): The File object to resolve.
+
+    Returns:
+        File: The resolved File object with updated metadata.
+
+    Raises:
+        RuntimeError: If the file's catalog is not set or if
+        the file source protocol is unsupported.
+    """
+    return file.resolve()
+
 
 class TextFile(File):
     """`DataModel` for reading text files."""

datachain/lib/hf.py

@@ -15,7 +15,7 @@ try:
         Value,
         load_dataset,
     )
-    from datasets.features.features import string_to_arrow
+    from datasets.features.features import Features, string_to_arrow
     from datasets.features.image import image_to_bytes
 
 except ImportError as exc:
@@ -36,6 +36,7 @@ from datachain.lib.data_model import DataModel, DataType, dict_to_data_model
 from datachain.lib.udf import Generator
 
 if TYPE_CHECKING:
+    import pyarrow as pa
     from pydantic import BaseModel
 
 
@@ -71,6 +72,15 @@ class HFGenerator(Generator):
         *args,
         **kwargs,
     ):
+        """
+        Generator for chain from huggingface datasets.
+
+        Parameters:
+
+        ds : Path or name of the dataset to read from Hugging Face Hub,
+            or an instance of `datasets.Dataset`-like object.
+        output_schema : Pydantic model for validation.
+        """
         super().__init__()
         self.ds = ds
         self.output_schema = output_schema
@@ -92,7 +102,7 @@ class HFGenerator(Generator):
                     output_dict["split"] = split
                 for name, feat in ds.features.items():
                     anno = self.output_schema.model_fields[name].annotation
-                    output_dict[name] = _convert_feature(row[name], feat, anno)
+                    output_dict[name] = convert_feature(row[name], feat, anno)
                 yield self.output_schema(**output_dict)
                 pbar.update(1)
 
@@ -106,7 +116,7 @@ def stream_splits(ds: Union[str, HFDatasetType], *args, **kwargs):
     return {"": ds}
 
 
-def _convert_feature(val: Any, feat: Any, anno: Any) -> Any:
+def convert_feature(val: Any, feat: Any, anno: Any) -> Any:  # noqa: PLR0911
     if isinstance(feat, (Value, Array2D, Array3D, Array4D, Array5D)):
         return val
     if isinstance(feat, ClassLabel):
@@ -117,20 +127,23 @@ def _convert_feature(val: Any, feat: Any, anno: Any) -> Any:
             for sname in val:
                 sfeat = feat.feature[sname]
                 sanno = anno.model_fields[sname].annotation
-                sdict[sname] = [_convert_feature(v, sfeat, sanno) for v in val[sname]]
+                sdict[sname] = [convert_feature(v, sfeat, sanno) for v in val[sname]]
             return anno(**sdict)
         return val
     if isinstance(feat, Image):
+        if isinstance(val, dict):
+            return HFImage(img=val["bytes"])
         return HFImage(img=image_to_bytes(val))
     if isinstance(feat, Audio):
         return HFAudio(**val)
 
 
 def get_output_schema(
-    ds: Union[Dataset, IterableDataset], model_name: str = ""
+    features: Features, model_name: str = "", stream: bool = True
 ) -> dict[str, DataType]:
+    """Generate UDF output schema from huggingface datasets features."""
     fields_dict = {}
-    for name, val in ds.features.items():
+    for name, val in features.items():
         fields_dict[name] = _feature_to_chain_type(name, val)  # type: ignore[assignment]
     return fields_dict  # type: ignore[return-value]
 
@@ -165,3 +178,7 @@ def _feature_to_chain_type(name: str, val: Any) -> type:  # noqa: PLR0911
     if isinstance(val, Audio):
         return HFAudio
     raise TypeError(f"Unknown huggingface datasets type {type(val)}")
+
+
+def schema_from_arrow(schema: "pa.Schema"):
+    return Features.from_arrow_schema(schema)

datachain/lib/listing.py

@@ -20,7 +20,7 @@ LISTING_TTL = 4 * 60 * 60  # cached listing lasts 4 hours
 LISTING_PREFIX = "lst__"  # listing datasets start with this name
 
 
-def list_bucket(uri: str, client_config=None) -> Callable:
+def list_bucket(uri: str, cache, client_config=None) -> Callable:
     """
     Function that returns another generator function that yields File objects
     from bucket where each File represents one bucket entry.
@@ -28,10 +28,10 @@ def list_bucket(uri: str, client_config=None) -> Callable:
 
     def list_func() -> Iterator[File]:
         config = client_config or {}
-        client, path = Client.parse_url(uri, None, **config)  # type: ignore[arg-type]
+        client = Client.get_client(uri, cache, **config)  # type: ignore[arg-type]
+        _, path = Client.parse_url(uri)
         for entries in iter_over_async(client.scandir(path.rstrip("/")), get_loop()):
-            for entry in entries:
-                yield entry.to_file(client.uri)
+            yield from entries
 
     return list_func
 
@@ -77,16 +77,17 @@ def parse_listing_uri(uri: str, cache, client_config) -> tuple[str, str, str]:
     """
     Parsing uri and returns listing dataset name, listing uri and listing path
     """
-    client, path = Client.parse_url(uri, cache, **client_config)
+    client = Client.get_client(uri, cache, **client_config)
+    storage_uri, path = Client.parse_url(uri)
 
     # clean path without globs
     lst_uri_path = (
         posixpath.dirname(path) if uses_glob(path) or client.fs.isfile(uri) else path
     )
 
-    lst_uri = f"{client.uri}/{lst_uri_path.lstrip('/')}"
+    lst_uri = f"{storage_uri}/{lst_uri_path.lstrip('/')}"
     ds_name = (
-        f"{LISTING_PREFIX}{client.uri}/{posixpath.join(lst_uri_path, '').lstrip('/')}"
+        f"{LISTING_PREFIX}{storage_uri}/{posixpath.join(lst_uri_path, '').lstrip('/')}"
     )
 
     return ds_name, lst_uri, path

datachain/lib/listing_info.py

@@ -13,8 +13,8 @@ class ListingInfo(DatasetInfo):
 
     @property
     def storage_uri(self) -> str:
-        client, _ = Client.parse_url(self.uri, None)  # type: ignore[arg-type]
-        return client.uri
+        uri, _ = Client.parse_url(self.uri)
+        return uri
 
     @property
     def expires(self) -> Optional[datetime]:

datachain/lib/model_store.py

@@ -1,6 +1,6 @@
 import inspect
 import logging
-from typing import ClassVar, Optional
+from typing import Any, ClassVar, Optional
 
 from pydantic import BaseModel
 
@@ -69,7 +69,7 @@ class ModelStore:
             del cls.store[fr.__name__][version]
 
     @staticmethod
-    def is_pydantic(val):
+    def is_pydantic(val: Any) -> bool:
         return (
             not hasattr(val, "__origin__")
             and inspect.isclass(val)

datachain/lib/pytorch.py

@@ -7,6 +7,7 @@ from torch import float32
 from torch.distributed import get_rank, get_world_size
 from torch.utils.data import IterableDataset, get_worker_info
 from torchvision.transforms import v2
+from tqdm import tqdm
 
 from datachain.catalog import Catalog, get_catalog
 from datachain.lib.dc import DataChain
@@ -93,33 +94,38 @@ class PytorchDataset(IterableDataset):
         if self.num_samples > 0:
             ds = ds.sample(self.num_samples)
         ds = ds.chunk(total_rank, total_workers)
-        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)
+        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, 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
+                        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)
 
     @staticmethod
     def get_rank_and_workers() -> tuple[int, int]:

datachain/lib/signal_schema.py

@@ -4,11 +4,14 @@ from collections.abc import Iterator, Sequence
 from dataclasses import dataclass
 from datetime import datetime
 from inspect import isclass
-from typing import (
+from typing import (  # noqa: UP035
     TYPE_CHECKING,
     Annotated,
     Any,
     Callable,
+    Dict,
+    Final,
+    List,
     Literal,
     Optional,
     Union,
@@ -42,8 +45,13 @@ NAMES_TO_TYPES = {
     "dict": dict,
     "bytes": bytes,
     "datetime": datetime,
-    "Literal": Literal,
+    "Final": Final,
     "Union": Union,
+    "Optional": Optional,
+    "List": list,
+    "Dict": dict,
+    "Literal": Any,
+    "Any": Any,
 }
 
 
@@ -146,35 +154,11 @@ class SignalSchema:
         return SignalSchema(signals)
 
     @staticmethod
-    def _get_name_original_type(fr_type: type) -> tuple[str, type]:
-        """Returns the name of and the original type for the given type,
-        based on whether the type is Optional or not."""
-        orig = get_origin(fr_type)
-        args = get_args(fr_type)
-        # Check if fr_type is Optional
-        if orig == Union and len(args) == 2 and (type(None) in args):
-            fr_type = args[0]
-            orig = get_origin(fr_type)
-        if orig in (Literal, LiteralEx):
-            # Literal has no __name__ in Python 3.9
-            type_name = "Literal"
-        elif orig == Union:
-            # Union also has no __name__ in Python 3.9
-            type_name = "Union"
-        else:
-            type_name = str(fr_type.__name__)  # type: ignore[union-attr]
-        return type_name, fr_type
-
-    @staticmethod
-    def serialize_custom_model_fields(
-        name: str, fr: type, custom_types: dict[str, Any]
+    def _serialize_custom_model_fields(
+        version_name: str, fr: type[BaseModel], custom_types: dict[str, Any]
     ) -> str:
         """This serializes any custom type information to the provided custom_types
-        dict, and returns the name of the type provided."""
-        if hasattr(fr, "__origin__") or not issubclass(fr, BaseModel):
-            # Don't store non-feature types.
-            return name
-        version_name = ModelStore.get_name(fr)
+        dict, and returns the name of the type serialized."""
         if version_name in custom_types:
             # This type is already stored in custom_types.
             return version_name
@@ -183,37 +167,102 @@ class SignalSchema:
             field_type = info.annotation
             # All fields should be typed.
             assert field_type
-            field_type_name, field_type = SignalSchema._get_name_original_type(
-                field_type
-            )
-            # Serialize this type to custom_types if it is a custom type as well.
-            fields[field_name] = SignalSchema.serialize_custom_model_fields(
-                field_type_name, field_type, custom_types
-            )
+            fields[field_name] = SignalSchema._serialize_type(field_type, custom_types)
         custom_types[version_name] = fields
         return version_name
 
+    @staticmethod
+    def _serialize_type(fr: type, custom_types: dict[str, Any]) -> str:
+        """Serialize a given type to a string, including automatic ModelStore
+        registration, and save this type and subtypes to custom_types as well."""
+        subtypes: list[Any] = []
+        type_name = SignalSchema._type_to_str(fr, subtypes)
+        # Iterate over all subtypes (includes the input type).
+        for st in subtypes:
+            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
+            )
+        return type_name
+
     def serialize(self) -> dict[str, Any]:
         signals: dict[str, Any] = {}
         custom_types: dict[str, Any] = {}
         for name, fr_type in self.values.items():
-            if (fr := ModelStore.to_pydantic(fr_type)) is not None:
-                ModelStore.register(fr)
-                signals[name] = ModelStore.get_name(fr)
-                type_name, fr_type = SignalSchema._get_name_original_type(fr)
-            else:
-                type_name, fr_type = SignalSchema._get_name_original_type(fr_type)
-                signals[name] = type_name
-            self.serialize_custom_model_fields(type_name, fr_type, custom_types)
+            signals[name] = self._serialize_type(fr_type, custom_types)
         if custom_types:
             signals["_custom_types"] = custom_types
         return signals
 
     @staticmethod
-    def _resolve_type(type_name: str, custom_types: dict[str, Any]) -> Optional[type]:
+    def _split_subtypes(type_name: str) -> list[str]:
+        """This splits a list of subtypes, including proper square bracket handling."""
+        start = 0
+        depth = 0
+        subtypes = []
+        for i, c in enumerate(type_name):
+            if c == "[":
+                depth += 1
+            elif c == "]":
+                if depth == 0:
+                    raise TypeError(
+                        "Extra closing square bracket when parsing subtype list"
+                    )
+                depth -= 1
+            elif c == "," and depth == 0:
+                subtypes.append(type_name[start:i].strip())
+                start = i + 1
+        if depth > 0:
+            raise TypeError("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
         """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")
+        if type_name == "NoneType":
+            return None
+
+        bracket_idx = type_name.find("[")
+        subtypes: Optional[tuple[Optional[type], ...]] = None
+        if bracket_idx > -1:
+            if bracket_idx == 0:
+                raise TypeError("Type cannot start with '['")
+            close_bracket_idx = type_name.rfind("]")
+            if close_bracket_idx == -1:
+                raise TypeError("Unclosed square bracket when parsing type")
+            if close_bracket_idx < bracket_idx:
+                raise TypeError("Square brackets are out of order when parsing type")
+            if close_bracket_idx == bracket_idx + 1:
+                raise TypeError("Empty square brackets when parsing type")
+            subtype_names = SignalSchema._split_subtypes(
+                type_name[bracket_idx + 1 : close_bracket_idx]
+            )
+            # Types like Union require the parameters to be a tuple of types.
+            subtypes = tuple(
+                SignalSchema._resolve_type(st, custom_types) for st in subtype_names
+            )
+            type_name = type_name[:bracket_idx].strip()
+
         fr = NAMES_TO_TYPES.get(type_name)
         if fr:
+            if subtypes:
+                if len(subtypes) == 1:
+                    # Types like Optional require there to be only one argument.
+                    return fr[subtypes[0]]  # type: ignore[index]
+                # Other types like Union require the parameters to be a tuple of types.
+                return fr[subtypes]  # type: ignore[index]
             return fr  # type: ignore[return-value]
 
         model_name, version = ModelStore.parse_name_version(type_name)
@@ -228,7 +277,14 @@ class SignalSchema:
                 for field_name, field_type_str in fields.items()
             }
             return create_feature_model(type_name, fields)
-        return None
+        # This can occur if a third-party or custom type is used, which is not available
+        # when deserializing.
+        warnings.warn(
+            f"Could not resolve type: '{type_name}'.",
+            SignalSchemaWarning,
+            stacklevel=2,
+        )
+        return Any  # type: ignore[return-value]
 
     @staticmethod
     def deserialize(schema: dict[str, Any]) -> "SignalSchema":
@@ -242,9 +298,14 @@ class SignalSchema:
                 # This entry is used as a lookup for custom types,
                 # and is not an actual field.
                 continue
+            if not isinstance(type_name, str):
+                raise SignalSchemaError(
+                    f"cannot deserialize '{type_name}': "
+                    "serialized types must be a string"
+                )
             try:
                 fr = SignalSchema._resolve_type(type_name, custom_types)
-                if fr is None:
+                if fr is Any:
                     # Skip if the type is not found, so all data can be displayed.
                     warnings.warn(
                         f"In signal '{signal}': "
@@ -258,7 +319,7 @@ class SignalSchema:
                 raise SignalSchemaError(
                     f"cannot deserialize '{signal}': {err}"
                 ) from err
-            signals[signal] = fr
+            signals[signal] = fr  # type: ignore[assignment]
 
         return SignalSchema(signals)
 
@@ -325,11 +386,20 @@ class SignalSchema:
             else:
                 json, pos = unflatten_to_json_pos(fr, row, pos)  # type: ignore[union-attr]
                 obj = fr(**json)
-                if isinstance(obj, File):
-                    obj._set_stream(catalog, caching_enabled=cache)
+                SignalSchema._set_file_stream(obj, catalog, cache)
                 res.append(obj)
         return res
 
+    @staticmethod
+    def _set_file_stream(
+        obj: BaseModel, catalog: "Catalog", cache: bool = False
+    ) -> None:
+        if isinstance(obj, File):
+            obj._set_stream(catalog, caching_enabled=cache)
+        for field, finfo in obj.model_fields.items():
+            if ModelStore.is_pydantic(finfo.annotation):
+                SignalSchema._set_file_stream(getattr(obj, field), catalog, cache)
+
     def db_signals(
         self, name: Optional[str] = None, as_columns=False
     ) -> Union[list[str], list[Column]]:
@@ -509,31 +579,58 @@ class SignalSchema:
         return self.values.pop(name)
 
     @staticmethod
-    def _type_to_str(type_):  # noqa: PLR0911
+    def _type_to_str(type_: Optional[type], subtypes: Optional[list] = None) -> str:  # noqa: PLR0911
+        """Convert a type to a string-based representation."""
+        if type_ is None:
+            return "NoneType"
+
         origin = get_origin(type_)
 
         if origin == Union:
             args = get_args(type_)
-            formatted_types = ", ".join(SignalSchema._type_to_str(arg) for arg in args)
+            formatted_types = ", ".join(
+                SignalSchema._type_to_str(arg, subtypes) for arg in args
+            )
             return f"Union[{formatted_types}]"
         if origin == Optional:
             args = get_args(type_)
-            type_str = SignalSchema._type_to_str(args[0])
+            type_str = SignalSchema._type_to_str(args[0], subtypes)
             return f"Optional[{type_str}]"
-        if origin is list:
+        if origin in (list, List):  # noqa: UP006
             args = get_args(type_)
-            type_str = SignalSchema._type_to_str(args[0])
+            type_str = SignalSchema._type_to_str(args[0], subtypes)
             return f"list[{type_str}]"
-        if origin is dict:
+        if origin in (dict, Dict):  # noqa: UP006
             args = get_args(type_)
-            type_str = SignalSchema._type_to_str(args[0]) if len(args) > 0 else ""
-            vals = f", {SignalSchema._type_to_str(args[1])}" if len(args) > 1 else ""
+            type_str = (
+                SignalSchema._type_to_str(args[0], subtypes) if len(args) > 0 else ""
+            )
+            vals = (
+                f", {SignalSchema._type_to_str(args[1], subtypes)}"
+                if len(args) > 1
+                else ""
+            )
             return f"dict[{type_str}{vals}]"
         if origin == Annotated:
             args = get_args(type_)
-            return SignalSchema._type_to_str(args[0])
-        if origin in (Literal, LiteralEx):
+            return SignalSchema._type_to_str(args[0], subtypes)
+        if origin in (Literal, LiteralEx) or type_ in (Literal, LiteralEx):
             return "Literal"
+        if Any in (origin, type_):
+            return "Any"
+        if Final in (origin, type_):
+            return "Final"
+        if subtypes is not None:
+            # Include this type in the list of all subtypes, if requested.
+            subtypes.append(type_)
+        if not hasattr(type_, "__name__"):
+            # This can happen for some third-party or custom types, mostly on Python 3.9
+            warnings.warn(
+                f"Unable to determine name of type '{type_}'.",
+                SignalSchemaWarning,
+                stacklevel=2,
+            )
+            return "Any"
         return type_.__name__
 
     @staticmethod