datachain 0.11.0__py3-none-any.whl → 0.12.0__py3-none-any.whl

datachain/lib/dc.py

@@ -22,7 +22,7 @@ import orjson
 import sqlalchemy
 from pydantic import BaseModel
 from sqlalchemy.sql.functions import GenericFunction
-from sqlalchemy.sql.sqltypes import NullType
+from tqdm import tqdm
 
 from datachain.dataset import DatasetRecord
 from datachain.func import literal
@@ -32,7 +32,14 @@ from datachain.lib.convert.python_to_sql import python_to_sql
 from datachain.lib.convert.values_to_tuples import values_to_tuples
 from datachain.lib.data_model import DataModel, DataType, DataValue, dict_to_data_model
 from datachain.lib.dataset_info import DatasetInfo
-from datachain.lib.file import ArrowRow, File, FileType, get_file_type
+from datachain.lib.file import (
+    EXPORT_FILES_MAX_THREADS,
+    ArrowRow,
+    File,
+    FileExporter,
+    FileType,
+    get_file_type,
+)
 from datachain.lib.file import ExportPlacement as FileExportPlacement
 from datachain.lib.listing import get_file_info, get_listing, list_bucket, ls
 from datachain.lib.listing_info import ListingInfo
@@ -47,7 +54,6 @@ from datachain.query import Session
 from datachain.query.dataset import DatasetQuery, PartitionByType
 from datachain.query.schema import DEFAULT_DELIMITER, Column, ColumnMeta
 from datachain.sql.functions import path as pathfunc
-from datachain.telemetry import telemetry
 from datachain.utils import batched_it, inside_notebook, row_to_nested_dict
 
 if TYPE_CHECKING:
@@ -65,7 +71,6 @@ _T = TypeVar("_T")
 D = TypeVar("D", bound="DataChain")
 UDFObjT = TypeVar("UDFObjT", bound=UDFBase)
 
-
 DEFAULT_PARQUET_CHUNK_SIZE = 100_000
 
 
@@ -208,7 +213,7 @@ class DataChain:
         from mistralai.client import MistralClient
         from mistralai.models.chat_completion import ChatMessage
 
-        from datachain.dc import DataChain, Column
+        from datachain import DataChain, Column
 
         PROMPT = (
             "Was this bot dialog successful? "
@@ -401,7 +406,7 @@ class DataChain:
     @classmethod
     def from_storage(
         cls,
-        uri,
+        uri: Union[str, os.PathLike[str]],
         *,
         type: FileType = "binary",
         session: Optional[Session] = None,
@@ -543,6 +548,8 @@ class DataChain:
             )
             ```
         """
+        from datachain.telemetry import telemetry
+
         query = DatasetQuery(
             name=name,
             version=version,
@@ -566,7 +573,7 @@ class DataChain:
     @classmethod
     def from_json(
         cls,
-        path,
+        path: Union[str, os.PathLike[str]],
         type: FileType = "text",
         spec: Optional[DataType] = None,
         schema_from: Optional[str] = "auto",
@@ -603,7 +610,7 @@ class DataChain:
             ```
         """
         if schema_from == "auto":
-            schema_from = path
+            schema_from = str(path)
 
         def jmespath_to_name(s: str):
             name_end = re.search(r"\W", s).start() if re.search(r"\W", s) else len(s)  # type: ignore[union-attr]
@@ -694,9 +701,22 @@ class DataChain:
         in_memory: bool = False,
         object_name: str = "dataset",
         include_listing: bool = False,
+        studio: bool = False,
     ) -> "DataChain":
         """Generate chain with list of registered datasets.
 
+        Args:
+            session: Optional session instance. If not provided, uses default session.
+            settings: Optional dictionary of settings to configure the chain.
+            in_memory: If True, creates an in-memory session. Defaults to False.
+            object_name: Name of the output object in the chain. Defaults to "dataset".
+            include_listing: If True, includes listing datasets. Defaults to False.
+            studio: If True, returns datasets from Studio only,
+                otherwise returns all local datasets. Defaults to False.
+
+        Returns:
+            DataChain: A new DataChain instance containing dataset information.
+
         Example:
             ```py
             from datachain import DataChain
@@ -712,7 +732,7 @@ class DataChain:
         datasets = [
             DatasetInfo.from_models(d, v, j)
             for d, v, j in catalog.list_datasets_versions(
-                include_listing=include_listing
+                include_listing=include_listing, studio=studio
             )
         ]
 
@@ -1050,7 +1070,7 @@ class DataChain:
     def select(self, *args: str, _sys: bool = True) -> "Self":
         """Select only a specified set of signals."""
         new_schema = self.signals_schema.resolve(*args)
-        if _sys:
+        if self._sys and _sys:
             new_schema = SignalSchema({"sys": Sys}) | new_schema
         columns = new_schema.db_signals()
         return self._evolve(
@@ -1093,6 +1113,7 @@ class DataChain:
         partition_by_columns: list[Column] = []
         signal_columns: list[Column] = []
         schema_fields: dict[str, DataType] = {}
+        keep_columns: list[str] = []
 
         # validate partition_by columns and add them to the schema
         for col in partition_by:
@@ -1100,10 +1121,13 @@ class DataChain:
                 col_db_name = ColumnMeta.to_db_name(col)
                 col_type = self.signals_schema.get_column_type(col_db_name)
                 column = Column(col_db_name, python_to_sql(col_type))
+                if col not in keep_columns:
+                    keep_columns.append(col)
             elif isinstance(col, Function):
                 column = col.get_column(self.signals_schema)
                 col_db_name = column.name
                 col_type = column.type.python_type
+                schema_fields[col_db_name] = col_type
             else:
                 raise DataChainColumnError(
                     col,
@@ -1113,7 +1137,6 @@ class DataChain:
                     ),
                 )
             partition_by_columns.append(column)
-            schema_fields[col_db_name] = col_type
 
         # validate signal columns and add them to the schema
         if not kwargs:
@@ -1128,9 +1151,13 @@ class DataChain:
             signal_columns.append(column)
             schema_fields[col_name] = func.get_result_type(self.signals_schema)
 
+        signal_schema = SignalSchema(schema_fields)
+        if keep_columns:
+            signal_schema |= self.signals_schema.to_partial(*keep_columns)
+
         return self._evolve(
             query=self._query.group_by(signal_columns, partition_by_columns),
-            signal_schema=SignalSchema(schema_fields),
+            signal_schema=signal_schema,
         )
 
     def mutate(self, **kwargs) -> "Self":
@@ -1181,6 +1208,8 @@ class DataChain:
         )
         ```
         """
+        from sqlalchemy.sql.sqltypes import NullType
+
         primitives = (bool, str, int, float)
 
         for col_name, expr in kwargs.items():
@@ -1225,23 +1254,37 @@ class DataChain:
     @overload
     def collect_flatten(self) -> Iterator[tuple[Any, ...]]: ...
 
+    @overload
+    def collect_flatten(self, *, include_hidden: bool) -> Iterator[tuple[Any, ...]]: ...
+
     @overload
     def collect_flatten(
         self, *, row_factory: Callable[[list[str], tuple[Any, ...]], _T]
     ) -> Iterator[_T]: ...
 
-    def collect_flatten(self, *, row_factory=None):
+    @overload
+    def collect_flatten(
+        self,
+        *,
+        row_factory: Callable[[list[str], tuple[Any, ...]], _T],
+        include_hidden: bool,
+    ) -> Iterator[_T]: ...
+
+    def collect_flatten(self, *, row_factory=None, include_hidden: bool = True):
         """Yields flattened rows of values as a tuple.
 
         Args:
             row_factory : A callable to convert row to a custom format.
                           It should accept two arguments: a list of column names and
                           a tuple of row values.
+            include_hidden: Whether to include hidden signals from the schema.
         """
-        db_signals = self._effective_signals_schema.db_signals()
+        db_signals = self._effective_signals_schema.db_signals(
+            include_hidden=include_hidden
+        )
         with self._query.ordered_select(*db_signals).as_iterable() as rows:
             if row_factory:
-                rows = (row_factory(db_signals, r) for r in rows)
+                rows = (row_factory(db_signals, r) for r in rows)  # type: ignore[assignment]
             yield from rows
 
     def to_columnar_data_with_names(
@@ -1275,10 +1318,23 @@ class DataChain:
         self, *, row_factory: Callable[[list[str], tuple[Any, ...]], _T]
     ) -> list[_T]: ...
 
-    def results(self, *, row_factory=None):  # noqa: D102
+    @overload
+    def results(
+        self,
+        *,
+        row_factory: Callable[[list[str], tuple[Any, ...]], _T],
+        include_hidden: bool,
+    ) -> list[_T]: ...
+
+    @overload
+    def results(self, *, include_hidden: bool) -> list[tuple[Any, ...]]: ...
+
+    def results(self, *, row_factory=None, include_hidden=True):  # noqa: D102
         if row_factory is None:
-            return list(self.collect_flatten())
-        return list(self.collect_flatten(row_factory=row_factory))
+            return list(self.collect_flatten(include_hidden=include_hidden))
+        return list(
+            self.collect_flatten(row_factory=row_factory, include_hidden=include_hidden)
+        )
 
     def to_records(self) -> list[dict[str, Any]]:
         """Convert every row to a dictionary."""
@@ -1788,21 +1844,25 @@ class DataChain:
             **fr_map,
         )
 
-    def to_pandas(self, flatten=False) -> "pd.DataFrame":
+    def to_pandas(self, flatten=False, include_hidden=True) -> "pd.DataFrame":
         """Return a pandas DataFrame from the chain.
 
         Parameters:
             flatten : Whether to use a multiindex or flatten column names.
+            include_hidden : Whether to include hidden columns.
         """
         import pandas as pd
 
-        headers, max_length = self._effective_signals_schema.get_headers_with_length()
+        headers, max_length = self._effective_signals_schema.get_headers_with_length(
+            include_hidden=include_hidden
+        )
         if flatten or max_length < 2:
             columns = [".".join(filter(None, header)) for header in headers]
         else:
             columns = pd.MultiIndex.from_tuples(map(tuple, headers))
 
-        return pd.DataFrame.from_records(self.results(), columns=columns)
+        results = self.results(include_hidden=include_hidden)
+        return pd.DataFrame.from_records(results, columns=columns)
 
     def show(
         self,
@@ -1810,6 +1870,7 @@ class DataChain:
         flatten=False,
         transpose=False,
         truncate=True,
+        include_hidden=False,
     ) -> None:
         """Show a preview of the chain results.
 
@@ -1818,11 +1879,12 @@ class DataChain:
             flatten : Whether to use a multiindex or flatten column names.
             transpose : Whether to transpose rows and columns.
             truncate : Whether or not to truncate the contents of columns.
+            include_hidden : Whether to include hidden columns.
         """
         import pandas as pd
 
         dc = self.limit(limit) if limit > 0 else self  # type: ignore[misc]
-        df = dc.to_pandas(flatten)
+        df = dc.to_pandas(flatten, include_hidden=include_hidden)
 
         if df.empty:
             print("Empty result")
@@ -2495,22 +2557,28 @@ class DataChain:
 
     def to_storage(
         self,
-        output: str,
+        output: Union[str, os.PathLike[str]],
         signal: str = "file",
         placement: FileExportPlacement = "fullpath",
-        use_cache: bool = True,
         link_type: Literal["copy", "symlink"] = "copy",
+        num_threads: Optional[int] = EXPORT_FILES_MAX_THREADS,
+        anon: bool = False,
+        client_config: Optional[dict] = None,
     ) -> None:
-        """Export files from a specified signal to a directory.
+        """Export files from a specified signal to a directory. Files can be
+        exported to a local or cloud directory.
 
         Args:
             output: Path to the target directory for exporting files.
             signal: Name of the signal to export files from.
             placement: The method to use for naming exported files.
                 The possible values are: "filename", "etag", "fullpath", and "checksum".
-            use_cache: If `True`, cache the files before exporting.
             link_type: Method to use for exporting files.
                 Falls back to `'copy'` if symlinking fails.
+            num_threads : number of threads to use for exporting files.
+                By default it uses 5 threads.
+            anon: If true, we will treat cloud bucket as public one
+            client_config: Optional configuration for the destination storage client
 
         Example:
             Cross cloud transfer
@@ -2525,8 +2593,26 @@ class DataChain:
         ):
             raise ValueError("Files with the same name found")
 
-        for file in self.collect(signal):
-            file.export(output, placement, use_cache, link_type=link_type)  # type: ignore[union-attr]
+        if anon:
+            client_config = (client_config or {}) | {"anon": True}
+
+        progress_bar = tqdm(
+            desc=f"Exporting files to {output}: ",
+            unit=" files",
+            unit_scale=True,
+            unit_divisor=10,
+            total=self.count(),
+            leave=False,
+        )
+        file_exporter = FileExporter(
+            output,
+            placement,
+            self._settings.cache if self._settings else False,
+            link_type,
+            max_threads=num_threads or 1,
+            client_config=client_config,
+        )
+        file_exporter.run(self.collect(signal), progress_bar)
 
     def shuffle(self) -> "Self":
         """Shuffle the rows of the chain deterministically."""

datachain/lib/file.py

@@ -18,12 +18,12 @@ from urllib.request import url2pathname
 
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 from fsspec.utils import stringify_path
-from PIL import Image as PilImage
 from pydantic import Field, field_validator
 
 from datachain.client.fileslice import FileSlice
 from datachain.lib.data_model import DataModel
 from datachain.lib.utils import DataChainError
+from datachain.nodes_thread_pool import NodesThreadPool
 from datachain.sql.types import JSON, Boolean, DateTime, Int, String
 from datachain.utils import TIME_ZERO
 
@@ -43,6 +43,41 @@ logger = logging.getLogger("datachain")
 ExportPlacement = Literal["filename", "etag", "fullpath", "checksum"]
 
 FileType = Literal["binary", "text", "image", "video"]
+EXPORT_FILES_MAX_THREADS = 5
+
+
+class FileExporter(NodesThreadPool):
+    """Class that does file exporting concurrently with thread pool"""
+
+    def __init__(
+        self,
+        output: Union[str, os.PathLike[str]],
+        placement: ExportPlacement,
+        use_cache: bool,
+        link_type: Literal["copy", "symlink"],
+        max_threads: int = EXPORT_FILES_MAX_THREADS,
+        client_config: Optional[dict] = None,
+    ):
+        super().__init__(max_threads)
+        self.output = output
+        self.placement = placement
+        self.use_cache = use_cache
+        self.link_type = link_type
+        self.client_config = client_config
+
+    def done_task(self, done):
+        for task in done:
+            task.result()
+
+    def do_task(self, file):
+        file.export(
+            self.output,
+            self.placement,
+            self.use_cache,
+            link_type=self.link_type,
+            client_config=self.client_config,
+        )
+        self.increase_counter(1)
 
 
 class VFileError(DataChainError):
@@ -158,6 +193,7 @@ class File(DataModel):
         "last_modified": DateTime,
         "location": JSON,
     }
+    _hidden_fields: ClassVar[list[str]] = ["version", "source"]
 
     _unique_id_keys: ClassVar[list[str]] = [
         "source",
@@ -206,6 +242,30 @@ class File(DataModel):
         self._catalog = None
         self._caching_enabled: bool = False
 
+    def as_text_file(self) -> "TextFile":
+        """Convert the file to a `TextFile` object."""
+        if isinstance(self, TextFile):
+            return self
+        file = TextFile(**self.model_dump())
+        file._set_stream(self._catalog, caching_enabled=self._caching_enabled)
+        return file
+
+    def as_image_file(self) -> "ImageFile":
+        """Convert the file to a `ImageFile` object."""
+        if isinstance(self, ImageFile):
+            return self
+        file = ImageFile(**self.model_dump())
+        file._set_stream(self._catalog, caching_enabled=self._caching_enabled)
+        return file
+
+    def as_video_file(self) -> "VideoFile":
+        """Convert the file to a `VideoFile` object."""
+        if isinstance(self, VideoFile):
+            return self
+        file = VideoFile(**self.model_dump())
+        file._set_stream(self._catalog, caching_enabled=self._caching_enabled)
+        return file
+
     @classmethod
     def upload(
         cls, data: bytes, path: str, catalog: Optional["Catalog"] = None
@@ -255,24 +315,24 @@ class File(DataModel):
             ) as f:
                 yield io.TextIOWrapper(f) if mode == "r" else f
 
-    def read(self, length: int = -1):
-        """Returns file contents."""
+    def read_bytes(self, length: int = -1):
+        """Returns file contents as bytes."""
         with self.open() as stream:
             return stream.read(length)
 
-    def read_bytes(self):
-        """Returns file contents as bytes."""
-        return self.read()
-
     def read_text(self):
         """Returns file contents as text."""
         with self.open(mode="r") as stream:
             return stream.read()
 
-    def save(self, destination: str):
+    def read(self, length: int = -1):
+        """Returns file contents."""
+        return self.read_bytes(length)
+
+    def save(self, destination: str, client_config: Optional[dict] = None):
         """Writes it's content to destination"""
         destination = stringify_path(destination)
-        client: Client = self._catalog.get_client(destination)
+        client: Client = self._catalog.get_client(destination, **(client_config or {}))
 
         if client.PREFIX == "file://" and not destination.startswith(client.PREFIX):
             destination = Path(destination).absolute().as_uri()
@@ -296,17 +356,17 @@ class File(DataModel):
 
     def export(
         self,
-        output: str,
+        output: Union[str, os.PathLike[str]],
         placement: ExportPlacement = "fullpath",
         use_cache: bool = True,
         link_type: Literal["copy", "symlink"] = "copy",
+        client_config: Optional[dict] = None,
     ) -> None:
         """Export file to new location."""
-        if use_cache:
-            self._caching_enabled = use_cache
+        self._caching_enabled = use_cache
         dst = self.get_destination_path(output, placement)
         dst_dir = os.path.dirname(dst)
-        client: Client = self._catalog.get_client(dst_dir)
+        client: Client = self._catalog.get_client(dst_dir, **(client_config or {}))
         client.fs.makedirs(dst_dir, exist_ok=True)
 
         if link_type == "symlink":
@@ -316,7 +376,7 @@ class File(DataModel):
                 if exc.errno not in (errno.ENOTSUP, errno.EXDEV, errno.ENOSYS):
                     raise
 
-        self.save(dst)
+        self.save(dst, client_config=client_config)
 
     def _set_stream(
         self,
@@ -337,15 +397,10 @@ class File(DataModel):
         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
@@ -393,7 +448,9 @@ class File(DataModel):
             path = url2pathname(path)
         return path
 
-    def get_destination_path(self, output: str, placement: ExportPlacement) -> str:
+    def get_destination_path(
+        self, output: Union[str, os.PathLike[str]], placement: ExportPlacement
+    ) -> str:
         """
         Returns full destination path of a file for exporting to some output
         based on export placement
@@ -502,11 +559,11 @@ class TextFile(File):
         with self.open() as stream:
             return stream.read()
 
-    def save(self, destination: str):
+    def save(self, destination: str, client_config: Optional[dict] = None):
         """Writes it's content to destination"""
         destination = stringify_path(destination)
 
-        client: Client = self._catalog.get_client(destination)
+        client: Client = self._catalog.get_client(destination, **(client_config or {}))
         with client.fs.open(destination, mode="w") as f:
             f.write(self.read_text())
 
@@ -514,18 +571,36 @@ class TextFile(File):
 class ImageFile(File):
     """`DataModel` for reading image files."""
 
+    def get_info(self) -> "Image":
+        """
+        Retrieves metadata and information about the image file.
+
+        Returns:
+            Image: A Model containing image metadata such as width, height and format.
+        """
+        from .image import image_info
+
+        return image_info(self)
+
     def read(self):
         """Returns `PIL.Image.Image` object."""
+        from PIL import Image as PilImage
+
         fobj = super().read()
         return PilImage.open(BytesIO(fobj))
 
-    def save(self, destination: str):
+    def save(  # type: ignore[override]
+        self,
+        destination: str,
+        format: Optional[str] = None,
+        client_config: Optional[dict] = None,
+    ):
         """Writes it's content to destination"""
         destination = stringify_path(destination)
 
-        client: Client = self._catalog.get_client(destination)
+        client: Client = self._catalog.get_client(destination, **(client_config or {}))
         with client.fs.open(destination, mode="wb") as f:
-            self.read().save(f)
+            self.read().save(f, format=format)
 
 
 class Image(DataModel):

datachain/lib/image.py

@@ -1,17 +1,41 @@
 from typing import Callable, Optional, Union
 
 import torch
-from PIL import Image
+from PIL import Image as PILImage
+
+from datachain.lib.file import File, FileError, Image, ImageFile
+
+
+def image_info(file: Union[File, ImageFile]) -> Image:
+    """
+    Returns image file information.
+
+    Args:
+        file (ImageFile): Image file object.
+
+    Returns:
+        Image: Image file information.
+    """
+    try:
+        img = file.as_image_file().read()
+    except Exception as exc:
+        raise FileError(file, "unable to open image file") from exc
+
+    return Image(
+        width=img.width,
+        height=img.height,
+        format=img.format or "",
+    )
 
 
 def convert_image(
-    img: Image.Image,
+    img: PILImage.Image,
     mode: str = "RGB",
     size: Optional[tuple[int, int]] = None,
     transform: Optional[Callable] = None,
     encoder: Optional[Callable] = None,
     device: Optional[Union[str, torch.device]] = None,
-) -> Union[Image.Image, torch.Tensor]:
+) -> Union[PILImage.Image, torch.Tensor]:
     """
     Resize, transform, and otherwise convert an image.
 
@@ -47,13 +71,13 @@ def convert_image(
 
 
 def convert_images(
-    images: Union[Image.Image, list[Image.Image]],
+    images: Union[PILImage.Image, list[PILImage.Image]],
     mode: str = "RGB",
     size: Optional[tuple[int, int]] = None,
     transform: Optional[Callable] = None,
     encoder: Optional[Callable] = None,
     device: Optional[Union[str, torch.device]] = None,
-) -> Union[list[Image.Image], torch.Tensor]:
+) -> Union[list[PILImage.Image], torch.Tensor]:
     """
     Resize, transform, and otherwise convert one or more images.
 
@@ -65,7 +89,7 @@ def convert_images(
         encoder (Callable): Encode image using model.
         device (str or torch.device): Device to use.
     """
-    if isinstance(images, Image.Image):
+    if isinstance(images, PILImage.Image):
         images = [images]
 
     converted = [