datachain 0.3.14__tar.gz → 0.3.15__tar.gz

{datachain-0.3.14/src/datachain.egg-info → datachain-0.3.15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.3.14
+Version: 0.3.15
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0

{datachain-0.3.14 → datachain-0.3.15}/src/datachain/catalog/catalog.py

@@ -621,10 +621,6 @@ class Catalog:
                 code_ast.body[-1:] = new_expressions
         return code_ast
 
-    def parse_url(self, uri: str, **config: Any) -> tuple[Client, str]:
-        config = config or self.client_config
-        return Client.parse_url(uri, self.cache, **config)
-
     def get_client(self, uri: StorageURI, **config: Any) -> Client:
         """
         Return the client corresponding to the given source `uri`.
@@ -651,17 +647,16 @@ class Catalog:
         partial_path: Optional[str]
 
         client_config = client_config or self.client_config
-        client, path = self.parse_url(source, **client_config)
+        uri, path = Client.parse_url(source)
+        client = Client.get_client(source, self.cache, **client_config)
         stem = os.path.basename(os.path.normpath(path))
         prefix = (
             posixpath.dirname(path)
             if glob.has_magic(stem) or client.fs.isfile(source)
             else path
         )
-        storage_dataset_name = Storage.dataset_name(
-            client.uri, posixpath.join(prefix, "")
-        )
-        source_metastore = self.metastore.clone(client.uri)
+        storage_dataset_name = Storage.dataset_name(uri, posixpath.join(prefix, ""))
+        source_metastore = self.metastore.clone(uri)
 
         columns = [
             Column("path", String),
@@ -675,15 +670,13 @@ class Catalog:
         ]
 
         if skip_indexing:
-            source_metastore.create_storage_if_not_registered(client.uri)
-            storage = source_metastore.get_storage(client.uri)
-            source_metastore.init_partial_id(client.uri)
-            partial_id = source_metastore.get_next_partial_id(client.uri)
+            source_metastore.create_storage_if_not_registered(uri)
+            storage = source_metastore.get_storage(uri)
+            source_metastore.init_partial_id(uri)
+            partial_id = source_metastore.get_next_partial_id(uri)
 
-            source_metastore = self.metastore.clone(
-                uri=client.uri, partial_id=partial_id
-            )
-            source_metastore.init(client.uri)
+            source_metastore = self.metastore.clone(uri=uri, partial_id=partial_id)
+            source_metastore.init(uri)
 
             source_warehouse = self.warehouse.clone()
             dataset = self.create_dataset(
@@ -701,20 +694,16 @@ class Catalog:
             in_progress,
             partial_id,
             partial_path,
-        ) = source_metastore.register_storage_for_indexing(
-            client.uri, force_update, prefix
-        )
+        ) = source_metastore.register_storage_for_indexing(uri, force_update, prefix)
         if in_progress:
             raise PendingIndexingError(f"Pending indexing operation: uri={storage.uri}")
 
         if not need_index:
             assert partial_id is not None
             assert partial_path is not None
-            source_metastore = self.metastore.clone(
-                uri=client.uri, partial_id=partial_id
-            )
+            source_metastore = self.metastore.clone(uri=uri, partial_id=partial_id)
             source_warehouse = self.warehouse.clone()
-            dataset = self.get_dataset(Storage.dataset_name(client.uri, partial_path))
+            dataset = self.get_dataset(Storage.dataset_name(uri, partial_path))
             lst = Listing(storage, source_metastore, source_warehouse, client, dataset)
             logger.debug(
                 "Using cached listing %s. Valid till: %s",
@@ -731,11 +720,11 @@ class Catalog:
 
             return lst, path
 
-        source_metastore.init_partial_id(client.uri)
-        partial_id = source_metastore.get_next_partial_id(client.uri)
+        source_metastore.init_partial_id(uri)
+        partial_id = source_metastore.get_next_partial_id(uri)
 
-        source_metastore.init(client.uri)
-        source_metastore = self.metastore.clone(uri=client.uri, partial_id=partial_id)
+        source_metastore.init(uri)
+        source_metastore = self.metastore.clone(uri=uri, partial_id=partial_id)
 
         source_warehouse = self.warehouse.clone()
 
@@ -1370,7 +1359,7 @@ class Catalog:
 
     def signed_url(self, source: str, path: str, client_config=None) -> str:
         client_config = client_config or self.client_config
-        client, _ = self.parse_url(source, **client_config)
+        client = Client.get_client(source, self.cache, **client_config)
         return client.url(path)
 
     def export_dataset_table(

{datachain-0.3.14 → datachain-0.3.15}/src/datachain/client/fsspec.py

@@ -116,15 +116,16 @@ class Client(ABC):
         return DATA_SOURCE_URI_PATTERN.match(name) is not None
 
     @staticmethod
-    def parse_url(
-        source: str,
-        cache: DataChainCache,
-        **kwargs,
-    ) -> tuple["Client", str]:
+    def parse_url(source: str) -> tuple[StorageURI, str]:
+        cls = Client.get_implementation(source)
+        storage_name, rel_path = cls.split_url(source)
+        return cls.get_uri(storage_name), rel_path
+
+    @staticmethod
+    def get_client(source: str, cache: DataChainCache, **kwargs) -> "Client":
         cls = Client.get_implementation(source)
-        storage_url, rel_path = cls.split_url(source)
-        client = cls.from_name(storage_url, cache, kwargs)
-        return client, rel_path
+        storage_url, _ = cls.split_url(source)
+        return cls.from_name(storage_url, cache, kwargs)
 
     @classmethod
     def create_fs(cls, **kwargs) -> "AbstractFileSystem":

{datachain-0.3.14 → datachain-0.3.15}/src/datachain/dataset.py

@@ -112,7 +112,7 @@ class DatasetDependency:
 
         if is_listing_dataset(dataset_name):
             dependency_type = DatasetDependencyType.STORAGE  # type: ignore[arg-type]
-            dependency_name = listing_uri_from_name(dataset_name)
+            dependency_name, _ = Client.parse_url(listing_uri_from_name(dataset_name))
 
         return cls(
             id,

{datachain-0.3.14 → datachain-0.3.15}/src/datachain/lib/arrow.py

@@ -13,8 +13,10 @@ from datachain.lib.model_store import ModelStore
 from datachain.lib.udf import Generator
 
 if TYPE_CHECKING:
+    from datasets.features.features import Features
     from pydantic import BaseModel
 
+    from datachain.lib.data_model import DataType
     from datachain.lib.dc import DataChain
 
 
@@ -46,7 +48,10 @@ class ArrowGenerator(Generator):
         self.kwargs = kwargs
 
     def process(self, file: File):
-        if self.nrows:
+        if file._caching_enabled:
+            path = file.get_local_path(download=True)
+            ds = dataset(path, schema=self.input_schema, **self.kwargs)
+        elif self.nrows:
             path = _nrows_file(file, self.nrows)
             ds = dataset(path, schema=self.input_schema, **self.kwargs)
         else:
@@ -54,6 +59,7 @@ class ArrowGenerator(Generator):
             ds = dataset(
                 path, filesystem=file.get_fs(), schema=self.input_schema, **self.kwargs
             )
+        hf_schema = _get_hf_schema(ds.schema)
         index = 0
         with tqdm(desc="Parsed by pyarrow", unit=" rows") as pbar:
             for record_batch in ds.to_batches():
@@ -62,9 +68,17 @@ class ArrowGenerator(Generator):
                     if self.output_schema:
                         fields = self.output_schema.model_fields
                         vals_dict = {}
-                        for (field, field_info), val in zip(fields.items(), vals):
-                            if ModelStore.is_pydantic(field_info.annotation):
-                                vals_dict[field] = field_info.annotation(**val)  # type: ignore[misc]
+                        for i, ((field, field_info), val) in enumerate(
+                            zip(fields.items(), vals)
+                        ):
+                            anno = field_info.annotation
+                            if hf_schema:
+                                from datachain.lib.hf import convert_feature
+
+                                feat = list(hf_schema[0].values())[i]
+                                vals_dict[field] = convert_feature(val, feat, anno)
+                            elif ModelStore.is_pydantic(anno):
+                                vals_dict[field] = anno(**val)  # type: ignore[misc]
                             else:
                                 vals_dict[field] = val
                         vals = [self.output_schema(**vals_dict)]
@@ -91,26 +105,36 @@ def schema_to_output(schema: pa.Schema, col_names: Optional[Sequence[str]] = Non
             "Error generating output from Arrow schema - "
             f"Schema has {len(schema)} columns but got {len(col_names)} column names."
         )
-    default_column = 0
+    if not col_names:
+        col_names = schema.names
+    columns = _convert_col_names(col_names)  # type: ignore[arg-type]
+    hf_schema = _get_hf_schema(schema)
+    if hf_schema:
+        return {
+            column: hf_type for hf_type, column in zip(hf_schema[1].values(), columns)
+        }
     output = {}
-    for i, field in enumerate(schema):
-        if col_names:
-            column = col_names[i]
-        else:
-            column = field.name
-        column = column.lower()
-        column = re.sub("[^0-9a-z_]+", "", column)
-        if not column:
-            column = f"c{default_column}"
-            default_column += 1
+    for field, column in zip(schema, columns):
         dtype = arrow_type_mapper(field.type, column)  # type: ignore[assignment]
         if field.nullable and not ModelStore.is_pydantic(dtype):
             dtype = Optional[dtype]  # type: ignore[assignment]
         output[column] = dtype
-
     return output
 
 
+def _convert_col_names(col_names: Sequence[str]) -> list[str]:
+    default_column = 0
+    converted_col_names = []
+    for column in col_names:
+        column = column.lower()
+        column = re.sub("[^0-9a-z_]+", "", column)
+        if not column:
+            column = f"c{default_column}"
+            default_column += 1
+        converted_col_names.append(column)
+    return converted_col_names
+
+
 def arrow_type_mapper(col_type: pa.DataType, column: str = "") -> type:  # noqa: PLR0911
     """Convert pyarrow types to basic types."""
     from datetime import datetime
@@ -156,3 +180,14 @@ def _nrows_file(file: File, nrows: int) -> str:
                 writer.write(line)
                 writer.write("\n")
     return tf.name
+
+
+def _get_hf_schema(
+    schema: "pa.Schema",
+) -> Optional[tuple["Features", dict[str, "DataType"]]]:
+    if schema.metadata and b"huggingface" in schema.metadata:
+        from datachain.lib.hf import get_output_schema, schema_from_arrow
+
+        features = schema_from_arrow(schema)
+        return features, get_output_schema(features)
+    return None

{datachain-0.3.14 → datachain-0.3.15}/src/datachain/lib/dc.py

@@ -408,7 +408,11 @@ class DataChain(DatasetQuery):
                     in_memory=in_memory,
                 )
                 .gen(
-                    list_bucket(list_uri, client_config=session.catalog.client_config),
+                    list_bucket(
+                        list_uri,
+                        session.catalog.cache,
+                        client_config=session.catalog.client_config,
+                    ),
                     output={f"{object_name}": File},
                 )
                 .save(list_dataset_name, listing=True)
@@ -1523,7 +1527,8 @@ class DataChain(DatasetQuery):
             output = {"split": str}
 
         model_name = model_name or object_name or ""
-        output = output | get_output_schema(next(iter(ds_dict.values())), model_name)
+        hf_features = next(iter(ds_dict.values())).features
+        output = output | get_output_schema(hf_features, model_name)
         model = dict_to_data_model(model_name, output)
         if object_name:
             output = {object_name: model}

{datachain-0.3.14 → datachain-0.3.15}/src/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-0.3.14 → datachain-0.3.15}/src/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-0.3.14 → datachain-0.3.15}/src/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,7 +28,8 @@ 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()):
             yield from entries
 
@@ -76,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-0.3.14 → datachain-0.3.15}/src/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-0.3.14 → datachain-0.3.15}/src/datachain/lib/signal_schema.py

@@ -386,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]]:

datachain-0.3.15/src/datachain/lib/tar.py

@@ -0,0 +1,33 @@
+import hashlib
+import tarfile
+from collections.abc import Iterator
+
+from datachain.lib.file import File, TarVFile
+
+
+def build_tar_member(parent: File, info: tarfile.TarInfo) -> File:
+    new_parent = parent.get_full_name()
+    etag_string = "-".join([parent.etag, info.name, str(info.mtime)])
+    etag = hashlib.md5(etag_string.encode(), usedforsecurity=False).hexdigest()
+    return File(
+        source=parent.source,
+        path=f"{new_parent}/{info.name}",
+        version=parent.version,
+        size=info.size,
+        etag=etag,
+        location=[
+            {
+                "vtype": TarVFile.get_vtype(),
+                "parent": parent.model_dump_custom(),
+                "size": info.size,
+                "offset": info.offset_data,
+            }
+        ],
+    )
+
+
+def process_tar(file: File) -> Iterator[File]:
+    with file.open() as fd:
+        with tarfile.open(fileobj=fd) as tar:
+            for entry in tar.getmembers():
+                yield build_tar_member(file, entry)