datachain 0.8.4__tar.gz → 0.8.6__tar.gz

{datachain-0.8.4 → datachain-0.8.6}/.gitignore

@@ -140,3 +140,6 @@ cython_debug/
 .vscode/
 .datachain/
 .dvcx/
+
+# pt files produced by ultralytics examples
+*.pt

{datachain-0.8.4 → datachain-0.8.6}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: datachain
-Version: 0.8.4
+Version: 0.8.6
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -96,10 +96,10 @@ Requires-Dist: defusedxml; extra == "examples"
 Requires-Dist: accelerate; extra == "examples"
 Requires-Dist: unstructured_ingest[embed-huggingface]; extra == "examples"
 Requires-Dist: unstructured[pdf]<0.16.12; extra == "examples"
-Requires-Dist: pdfplumber==0.11.4; extra == "examples"
+Requires-Dist: pdfplumber==0.11.5; extra == "examples"
 Requires-Dist: huggingface_hub[hf_transfer]; extra == "examples"
 Requires-Dist: onnx==1.16.1; extra == "examples"
-Requires-Dist: ultralytics==8.3.55; extra == "examples"
+Requires-Dist: ultralytics==8.3.58; extra == "examples"
 
 ================
 |logo| DataChain
@@ -134,7 +134,7 @@ Use Cases
 1. **ETL.** Pythonic framework for describing and running unstructured data transformations
    and enrichments, applying models to data, including LLMs.
 2. **Analytics.** DataChain dataset is a table that combines all the information about data
-   objects in one place + it provides dataframe-like API and vecrorized engine to do analytics
+   objects in one place + it provides dataframe-like API and vectorized engine to do analytics
    on these tables at scale.
 3. **Versioning.** DataChain doesn't store, require moving or copying data (unlike DVC).
    Perfect use case is a bucket with thousands or millions of images, videos, audio, PDFs.
@@ -270,7 +270,7 @@ DataChain Studio Platform
 
 `DataChain Studio`_ is a proprietary solution for teams that offers:
 
-- **Centralized dataset registry** to manage data, code and dependency
+- **Centralized dataset registry** to manage data, code and
   dependencies in one place.
 - **Data Lineage** for data sources as well as derivative dataset.
 - **UI for Multimodal Data** like images, videos, and PDFs.

{datachain-0.8.4 → datachain-0.8.6}/README.rst

@@ -31,7 +31,7 @@ Use Cases
 1. **ETL.** Pythonic framework for describing and running unstructured data transformations
    and enrichments, applying models to data, including LLMs.
 2. **Analytics.** DataChain dataset is a table that combines all the information about data
-   objects in one place + it provides dataframe-like API and vecrorized engine to do analytics
+   objects in one place + it provides dataframe-like API and vectorized engine to do analytics
    on these tables at scale.
 3. **Versioning.** DataChain doesn't store, require moving or copying data (unlike DVC).
    Perfect use case is a bucket with thousands or millions of images, videos, audio, PDFs.
@@ -167,7 +167,7 @@ DataChain Studio Platform
 
 `DataChain Studio`_ is a proprietary solution for teams that offers:
 
-- **Centralized dataset registry** to manage data, code and dependency
+- **Centralized dataset registry** to manage data, code and
   dependencies in one place.
 - **Data Lineage** for data sources as well as derivative dataset.
 - **UI for Multimodal Data** like images, videos, and PDFs.

{datachain-0.8.4 → datachain-0.8.6}/docs/index.md

@@ -42,7 +42,7 @@ database for easy and efficient querying.
     including LLMs.
 2.  **Analytics.** DataChain dataset is a table that combines all the
     information about data objects in one place + it provides
-    dataframe-like API and vecrorized engine to do analytics on these
+    dataframe-like API and vectorized engine to do analytics on these
     tables at scale.
 3.  **Versioning.** DataChain doesn't store, require moving or copying
     data (unlike DVC). Perfect use case is a bucket with thousands or

{datachain-0.8.4 → datachain-0.8.6}/examples/computer_vision/ultralytics-bbox.py

@@ -1,3 +1,8 @@
+import os
+
+os.environ["YOLO_VERBOSE"] = "false"
+
+
 from io import BytesIO
 
 from PIL import Image

{datachain-0.8.4 → datachain-0.8.6}/examples/computer_vision/ultralytics-pose.py

@@ -1,3 +1,8 @@
+import os
+
+os.environ["YOLO_VERBOSE"] = "false"
+
+
 from io import BytesIO
 
 from PIL import Image

{datachain-0.8.4 → datachain-0.8.6}/examples/computer_vision/ultralytics-segment.py

@@ -1,3 +1,8 @@
+import os
+
+os.environ["YOLO_VERBOSE"] = "false"
+
+
 from io import BytesIO
 
 from PIL import Image

{datachain-0.8.4 → datachain-0.8.6}/examples/get_started/torch-loader.py

@@ -56,7 +56,7 @@ class CNN(nn.Module):
 if __name__ == "__main__":
     ds = (
         DataChain.from_storage(STORAGE, type="image")
-        .settings(cache=True, prefetch=25)
+        .settings(prefetch=25)
         .filter(C("file.path").glob("*.jpg"))
         .map(
             label=lambda path: label_to_int(basename(path)[:3], CLASSES),
@@ -68,7 +68,7 @@ if __name__ == "__main__":
     train_loader = DataLoader(
         ds.to_pytorch(transform=transform),
         batch_size=25,
-        num_workers=max(4, os.cpu_count() or 2),
+        num_workers=min(4, os.cpu_count() or 2),
         persistent_workers=True,
         multiprocessing_context=multiprocessing.get_context("spawn"),
     )

{datachain-0.8.4 → datachain-0.8.6}/pyproject.toml

@@ -109,10 +109,10 @@ examples = [
   "accelerate",
   "unstructured_ingest[embed-huggingface]",
   "unstructured[pdf]<0.16.12",
-  "pdfplumber==0.11.4",
+  "pdfplumber==0.11.5",
   "huggingface_hub[hf_transfer]",
   "onnx==1.16.1",
-  "ultralytics==8.3.55"
+  "ultralytics==8.3.58"
 ]
 
 [project.urls]

{datachain-0.8.4 → datachain-0.8.6}/src/datachain/asyn.py

@@ -8,12 +8,14 @@ from collections.abc import (
     Iterable,
     Iterator,
 )
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, wait
 from heapq import heappop, heappush
 from typing import Any, Callable, Generic, Optional, TypeVar
 
 from fsspec.asyn import get_loop
 
+from datachain.utils import safe_closing
+
 ASYNC_WORKERS = 20
 
 InputT = TypeVar("InputT", contravariant=True)  # noqa: PLC0105
@@ -56,6 +58,7 @@ class AsyncMapper(Generic[InputT, ResultT]):
         self.pool = ThreadPoolExecutor(workers)
         self._tasks: set[asyncio.Task] = set()
         self._shutdown_producer = threading.Event()
+        self._producer_is_shutdown = threading.Event()
 
     def start_task(self, coro: Coroutine) -> asyncio.Task:
         task = self.loop.create_task(coro)
@@ -64,11 +67,16 @@ class AsyncMapper(Generic[InputT, ResultT]):
         return task
 
     def _produce(self) -> None:
-        for item in self.iterable:
-            if self._shutdown_producer.is_set():
-                return
-            fut = asyncio.run_coroutine_threadsafe(self.work_queue.put(item), self.loop)
-            fut.result()  # wait until the item is in the queue
+        try:
+            with safe_closing(self.iterable):
+                for item in self.iterable:
+                    if self._shutdown_producer.is_set():
+                        return
+                    coro = self.work_queue.put(item)
+                    fut = asyncio.run_coroutine_threadsafe(coro, self.loop)
+                    fut.result()  # wait until the item is in the queue
+        finally:
+            self._producer_is_shutdown.set()
 
     async def produce(self) -> None:
         await self.to_thread(self._produce)
@@ -179,6 +187,8 @@ class AsyncMapper(Generic[InputT, ResultT]):
             self.shutdown_producer()
             if not async_run.done():
                 async_run.cancel()
+                wait([async_run])
+            self._producer_is_shutdown.wait()
 
     def __iter__(self):
         return self.iterate()

{datachain-0.8.4 → datachain-0.8.6}/src/datachain/cache.py

@@ -1,8 +1,12 @@
 import os
+from collections.abc import Iterator
+from contextlib import contextmanager
+from tempfile import mkdtemp
 from typing import TYPE_CHECKING, Optional
 
 from dvc_data.hashfile.db.local import LocalHashFileDB
 from dvc_objects.fs.local import LocalFileSystem
+from dvc_objects.fs.utils import remove
 from fsspec.callbacks import Callback, TqdmCallback
 
 from .progress import Tqdm
@@ -20,6 +24,23 @@ def try_scandir(path):
         pass
 
 
+def get_temp_cache(tmp_dir: str, prefix: Optional[str] = None) -> "DataChainCache":
+    cache_dir = mkdtemp(prefix=prefix, dir=tmp_dir)
+    return DataChainCache(cache_dir, tmp_dir=tmp_dir)
+
+
+@contextmanager
+def temporary_cache(
+    tmp_dir: str, prefix: Optional[str] = None, delete: bool = True
+) -> Iterator["DataChainCache"]:
+    cache = get_temp_cache(tmp_dir, prefix=prefix)
+    try:
+        yield cache
+    finally:
+        if delete:
+            cache.destroy()
+
+
 class DataChainCache:
     def __init__(self, cache_dir: str, tmp_dir: str):
         self.odb = LocalHashFileDB(
@@ -28,6 +49,9 @@ class DataChainCache:
             tmp_dir=tmp_dir,
         )
 
+    def __eq__(self, other) -> bool:
+        return self.odb == other.odb
+
     @property
     def cache_dir(self):
         return self.odb.path
@@ -63,7 +87,7 @@ class DataChainCache:
         if size < 0:
             size = await client.get_size(from_path, version_id=file.version)
         cb = callback or TqdmCallback(
-            tqdm_kwargs={"desc": odb_fs.name(from_path), "bytes": True},
+            tqdm_kwargs={"desc": odb_fs.name(from_path), "bytes": True, "leave": False},
             tqdm_cls=Tqdm,
             size=size,
         )
@@ -82,20 +106,18 @@ class DataChainCache:
             os.unlink(tmp_info)
 
     def store_data(self, file: "File", contents: bytes) -> None:
-        checksum = file.get_hash()
-        dst = self.path_from_checksum(checksum)
-        if not os.path.exists(dst):
-            # Create the file only if it's not already in cache
-            os.makedirs(os.path.dirname(dst), exist_ok=True)
-            with open(dst, mode="wb") as f:
-                f.write(contents)
-
-    def clear(self):
+        self.odb.add_bytes(file.get_hash(), contents)
+
+    def clear(self) -> None:
         """
         Completely clear the cache.
         """
         self.odb.clear()
 
+    def destroy(self) -> None:
+        # `clear` leaves the prefix directory structure intact.
+        remove(self.cache_dir)
+
     def get_total_size(self) -> int:
         total = 0
         for subdir in try_scandir(self.odb.path):

{datachain-0.8.4 → datachain-0.8.6}/src/datachain/catalog/catalog.py

@@ -405,6 +405,7 @@ def get_download_bar(bar_format: str, total_size: int):
         unit_scale=True,
         unit_divisor=1000,
         total=total_size,
+        leave=False,
     )
 
 
@@ -429,6 +430,7 @@ def instantiate_node_groups(
             unit_scale=True,
             unit_divisor=1000,
             total=total_files,
+            leave=False,
         )
     )
 
@@ -534,6 +536,12 @@ def find_column_to_str(  # noqa: PLR0911
     return ""
 
 
+def clone_catalog_with_cache(catalog: "Catalog", cache: "DataChainCache") -> "Catalog":
+    clone = catalog.copy()
+    clone.cache = cache
+    return clone
+
+
 class Catalog:
     def __init__(
         self,
@@ -1242,10 +1250,17 @@ class Catalog:
         path: str,
         version_id: Optional[str] = None,
         client_config=None,
+        content_disposition: Optional[str] = None,
+        **kwargs,
     ) -> str:
         client_config = client_config or self.client_config
         client = Client.get_client(source, self.cache, **client_config)
-        return client.url(path, version_id=version_id)
+        return client.url(
+            path,
+            version_id=version_id,
+            content_disposition=content_disposition,
+            **kwargs,
+        )
 
     def export_dataset_table(
         self,
@@ -1437,6 +1452,7 @@ class Catalog:
             unit_scale=True,
             unit_divisor=1000,
             total=ds_stats.num_objects,  # type: ignore [union-attr]
+            leave=False,
         )
 
         schema = DatasetRecord.parse_schema(remote_ds_version.schema)

{datachain-0.8.4 → datachain-0.8.6}/src/datachain/client/azure.py

@@ -31,8 +31,12 @@ class AzureClient(Client):
         Generate a signed URL for the given path.
         """
         version_id = kwargs.pop("version_id", None)
+        content_disposition = kwargs.pop("content_disposition", None)
         result = self.fs.sign(
-            self.get_full_path(path, version_id), expiration=expires, **kwargs
+            self.get_full_path(path, version_id),
+            expiration=expires,
+            content_disposition=content_disposition,
+            **kwargs,
         )
         return result + (f"&versionid={version_id}" if version_id else "")
 
@@ -42,7 +46,7 @@ class AzureClient(Client):
             prefix = prefix.lstrip(DELIMITER) + DELIMITER
         found = False
         try:
-            with tqdm(desc=f"Listing {self.uri}", unit=" objects") as pbar:
+            with tqdm(desc=f"Listing {self.uri}", unit=" objects", leave=False) as pbar:
                 async with self.fs.service_client.get_container_client(
                     container=self.name
                 ) as container_client:

{datachain-0.8.4 → datachain-0.8.6}/src/datachain/client/fsspec.py

@@ -249,7 +249,7 @@ class Client(ABC):
         await main_task
 
     async def _fetch_nested(self, start_prefix: str, result_queue: ResultQueue) -> None:
-        progress_bar = tqdm(desc=f"Listing {self.uri}", unit=" objects")
+        progress_bar = tqdm(desc=f"Listing {self.uri}", unit=" objects", leave=False)
         loop = get_loop()
 
         queue: asyncio.Queue[str] = asyncio.Queue()

{datachain-0.8.4 → datachain-0.8.6}/src/datachain/client/gcs.py

@@ -39,11 +39,15 @@ class GCSClient(Client):
         (see https://cloud.google.com/storage/docs/access-public-data#api-link).
         """
         version_id = kwargs.pop("version_id", None)
+        content_disposition = kwargs.pop("content_disposition", None)
         if self.fs.storage_options.get("token") == "anon":
             query = f"?generation={version_id}" if version_id else ""
             return f"https://storage.googleapis.com/{self.name}/{path}{query}"
         return self.fs.sign(
-            self.get_full_path(path, version_id), expiration=expires, **kwargs
+            self.get_full_path(path, version_id),
+            expiration=expires,
+            response_disposition=content_disposition,
+            **kwargs,
         )
 
     @staticmethod
@@ -83,7 +87,7 @@ class GCSClient(Client):
         self, page_queue: PageQueue, result_queue: ResultQueue
     ) -> bool:
         found = False
-        with tqdm(desc=f"Listing {self.uri}", unit=" objects") as pbar:
+        with tqdm(desc=f"Listing {self.uri}", unit=" objects", leave=False) as pbar:
             while (page := await page_queue.get()) is not None:
                 if page:
                     found = True

{datachain-0.8.4 → datachain-0.8.6}/src/datachain/client/s3.py

@@ -1,4 +1,5 @@
 import asyncio
+import os
 from typing import Any, Optional, cast
 from urllib.parse import parse_qs, urlsplit, urlunsplit
 
@@ -31,9 +32,11 @@ class ClientS3(Client):
         if "aws_token" in kwargs:
             kwargs.setdefault("token", kwargs.pop("aws_token"))
 
-        # caching bucket regions to use the right one in signed urls, otherwise
-        # it tries to randomly guess and creates wrong signature
-        kwargs.setdefault("cache_regions", True)
+        # remove this `if` when https://github.com/fsspec/s3fs/pull/929 lands
+        if not os.environ.get("AWS_REGION") and not os.environ.get("AWS_ENDPOINT_URL"):
+            # caching bucket regions to use the right one in signed urls, otherwise
+            # it tries to randomly guess and creates wrong signature
+            kwargs.setdefault("cache_regions", True)
 
         # We want to use newer v4 signature version since regions added after
         # 2014 are not going to support v2 which is the older one.
@@ -51,6 +54,21 @@ class ClientS3(Client):
 
         return cast(S3FileSystem, super().create_fs(**kwargs))
 
+    def url(self, path: str, expires: int = 3600, **kwargs) -> str:
+        """
+        Generate a signed URL for the given path.
+        """
+        version_id = kwargs.pop("version_id", None)
+        content_disposition = kwargs.pop("content_disposition", None)
+        if content_disposition:
+            kwargs["ResponseContentDisposition"] = content_disposition
+
+        return self.fs.sign(
+            self.get_full_path(path, version_id),
+            expiration=expires,
+            **kwargs,
+        )
+
     async def _fetch_flat(self, start_prefix: str, result_queue: ResultQueue) -> None:
         async def get_pages(it, page_queue):
             try:
@@ -61,7 +79,7 @@ class ClientS3(Client):
 
         async def process_pages(page_queue, result_queue):
             found = False
-            with tqdm(desc=f"Listing {self.uri}", unit=" objects") as pbar:
+            with tqdm(desc=f"Listing {self.uri}", unit=" objects", leave=False) as pbar:
                 while (res := await page_queue.get()) is not None:
                     if res:
                         found = True

{datachain-0.8.4 → datachain-0.8.6}/src/datachain/data_storage/db_engine.py

@@ -79,6 +79,15 @@ class DatabaseEngine(ABC, Serializable):
         conn: Optional[Any] = None,
     ) -> Iterator[tuple[Any, ...]]: ...
 
+    def get_table(self, name: str) -> "Table":
+        table = self.metadata.tables.get(name)
+        if table is None:
+            sa.Table(name, self.metadata, autoload_with=self.engine)
+            # ^^^ This table may not be correctly initialised on some dialects
+            # Grab it from metadata instead.
+            table = self.metadata.tables[name]
+        return table
+
     @abstractmethod
     def executemany(
         self, query, params, cursor: Optional[Any] = None

{datachain-0.8.4 → datachain-0.8.6}/src/datachain/data_storage/schema.py

@@ -16,7 +16,6 @@ from datachain.sql.functions import path as pathfunc
 from datachain.sql.types import Int, SQLType, UInt64
 
 if TYPE_CHECKING:
-    from sqlalchemy import Engine
     from sqlalchemy.engine.interfaces import Dialect
     from sqlalchemy.sql.base import (
         ColumnCollection,
@@ -25,6 +24,8 @@ if TYPE_CHECKING:
     )
     from sqlalchemy.sql.elements import ColumnElement
 
+    from datachain.data_storage.db_engine import DatabaseEngine
+
 
 DEFAULT_DELIMITER = "__"
 
@@ -150,14 +151,12 @@ class DataTable:
     def __init__(
         self,
         name: str,
-        engine: "Engine",
-        metadata: Optional["sa.MetaData"] = None,
+        engine: "DatabaseEngine",
         column_types: Optional[dict[str, SQLType]] = None,
         object_name: str = "file",
     ):
         self.name: str = name
         self.engine = engine
-        self.metadata: sa.MetaData = metadata if metadata is not None else sa.MetaData()
         self.column_types: dict[str, SQLType] = column_types or {}
         self.object_name = object_name
 
@@ -211,12 +210,7 @@ class DataTable:
         return sa.Table(name, metadata, *columns)
 
     def get_table(self) -> "sa.Table":
-        table = self.metadata.tables.get(self.name)
-        if table is None:
-            sa.Table(self.name, self.metadata, autoload_with=self.engine)
-            # ^^^ This table may not be correctly initialised on some dialects
-            # Grab it from metadata instead.
-            table = self.metadata.tables[self.name]
+        table = self.engine.get_table(self.name)
 
         column_types = self.column_types | {c.name: c.type for c in self.sys_columns()}
         # adjusting types for custom columns to be instances of SQLType if possible

{datachain-0.8.4 → datachain-0.8.6}/src/datachain/data_storage/sqlite.py

@@ -186,6 +186,12 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         self.db_file = db_file
         self.is_closed = False
 
+    def get_table(self, name: str) -> Table:
+        if self.is_closed:
+            # Reconnect in case of being closed previously.
+            self._reconnect()
+        return super().get_table(name)
+
     @retry_sqlite_locks
     def execute(
         self,
@@ -670,7 +676,7 @@ class SQLiteWarehouse(AbstractWarehouse):
         ]
         table = self.create_udf_table(columns)
 
-        with tqdm(desc="Preparing", unit=" rows") as pbar:
+        with tqdm(desc="Preparing", unit=" rows", leave=False) as pbar:
             self.copy_table(table, query, progress_cb=pbar.update)
 
         return table

{datachain-0.8.4 → datachain-0.8.6}/src/datachain/data_storage/warehouse.py

@@ -191,8 +191,7 @@ class AbstractWarehouse(ABC, Serializable):
         table_name = self.dataset_table_name(dataset.name, version)
         return self.schema.dataset_row_cls(
             table_name,
-            self.db.engine,
-            self.db.metadata,
+            self.db,
             dataset.get_schema(version),
             object_name=object_name,
         )
@@ -904,8 +903,11 @@ class AbstractWarehouse(ABC, Serializable):
         This should be implemented to ensure that the provided tables
         are cleaned up as soon as they are no longer needed.
         """
-        with tqdm(desc="Cleanup", unit=" tables") as pbar:
-            for name in set(names):
+        to_drop = set(names)
+        with tqdm(
+            desc="Cleanup", unit=" tables", total=len(to_drop), leave=False
+        ) as pbar:
+            for name in to_drop:
                 self.db.drop_table(Table(name, self.db.metadata), if_exists=True)
                 pbar.update(1)