datachain 0.14.0__tar.gz → 0.14.1__tar.gz

{datachain-0.14.0/src/datachain.egg-info → datachain-0.14.1}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.14.0
+Version: 0.14.1
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
-License: Apache-2.0
+License-Expression: Apache-2.0
 Project-URL: Documentation, https://datachain.dvc.ai
 Project-URL: Issues, https://github.com/iterative/datachain/issues
 Project-URL: Source, https://github.com/iterative/datachain

{datachain-0.14.0 → datachain-0.14.1}/pyproject.toml

@@ -1,12 +1,13 @@
 [build-system]
-requires = ["setuptools>=48", "setuptools_scm[toml]>=6.3.1"]
+requires = ["setuptools>=77", "setuptools_scm[toml]>=6.3.1"]
 build-backend = "setuptools.build_meta"
 
 [project]
 name = "datachain"
 description = "Wrangle unstructured AI data at scale"
 readme = "README.rst"
-license = {text = "Apache-2.0"}
+license = "Apache-2.0"
+license-files = ["LICENSE"]
 authors = [{name = "Dmitry Petrov", email = "support@dvc.org"}]
 classifiers = [
   "Programming Language :: Python :: 3",

{datachain-0.14.0 → datachain-0.14.1}/src/datachain/catalog/catalog.py

@@ -588,7 +588,7 @@ class Catalog:
 
         from_storage(
             source, session=self.session, update=update, object_name=object_name
-        )
+        ).exec()
 
         list_ds_name, list_uri, list_path, _ = get_listing(
             source, self.session, update=update

{datachain-0.14.0 → datachain-0.14.1}/src/datachain/client/fsspec.py

@@ -89,9 +89,9 @@ class Client(ABC):
         from .local import FileClient
         from .s3 import ClientS3
 
-        protocol = urlparse(str(url)).scheme
+        protocol = urlparse(os.fspath(url)).scheme
 
-        if not protocol or _is_win_local_path(str(url)):
+        if not protocol or _is_win_local_path(os.fspath(url)):
             return FileClient
         if protocol == ClientS3.protocol:
             return ClientS3
@@ -122,7 +122,7 @@ class Client(ABC):
         source: Union[str, os.PathLike[str]], cache: Cache, **kwargs
     ) -> "Client":
         cls = Client.get_implementation(source)
-        storage_url, _ = cls.split_url(str(source))
+        storage_url, _ = cls.split_url(os.fspath(source))
         if os.name == "nt":
             storage_url = storage_url.removeprefix("/")
 

{datachain-0.14.0 → datachain-0.14.1}/src/datachain/lib/dc/json.py

@@ -64,7 +64,7 @@ def from_json(
     from .storage import from_storage
 
     if schema_from == "auto":
-        schema_from = str(path)
+        schema_from = os.fspath(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]

datachain-0.14.1/src/datachain/lib/dc/storage.py

@@ -0,0 +1,170 @@
+import os.path
+from typing import (
+    TYPE_CHECKING,
+    Optional,
+    Union,
+)
+
+from datachain.lib.file import (
+    FileType,
+    get_file_type,
+)
+from datachain.lib.listing import (
+    get_file_info,
+    get_listing,
+    list_bucket,
+    ls,
+)
+from datachain.query import Session
+
+if TYPE_CHECKING:
+    from .datachain import DataChain
+
+
+def from_storage(
+    uri: Union[str, os.PathLike[str], list[str], list[os.PathLike[str]]],
+    *,
+    type: FileType = "binary",
+    session: Optional[Session] = None,
+    settings: Optional[dict] = None,
+    in_memory: bool = False,
+    recursive: Optional[bool] = True,
+    object_name: str = "file",
+    update: bool = False,
+    anon: bool = False,
+    client_config: Optional[dict] = None,
+) -> "DataChain":
+    """Get data from storage(s) as a list of file with all file attributes.
+    It returns the chain itself as usual.
+
+    Parameters:
+        uri : storage URI with directory or list of URIs.
+            URIs must start with storage prefix such
+            as `s3://`, `gs://`, `az://` or "file:///"
+        type : read file as "binary", "text", or "image" data. Default is "binary".
+        recursive : search recursively for the given path.
+        object_name : Created object column name.
+        update : force storage reindexing. Default is False.
+        anon : If True, we will treat cloud bucket as public one
+        client_config : Optional client configuration for the storage client.
+
+    Returns:
+        DataChain: A DataChain object containing the file information.
+
+    Examples:
+        Simple call from s3:
+        ```python
+        import datachain as dc
+        chain = dc.from_storage("s3://my-bucket/my-dir")
+        ```
+
+        Multiple URIs:
+        ```python
+        chain = dc.from_storage([
+            "s3://bucket1/dir1",
+            "s3://bucket2/dir2"
+        ])
+        ```
+
+        With AWS S3-compatible storage:
+        ```python
+        chain = dc.from_storage(
+            "s3://my-bucket/my-dir",
+            client_config = {"aws_endpoint_url": "<minio-endpoint-url>"}
+        )
+        ```
+
+        Pass existing session
+        ```py
+        session = Session.get()
+        chain = dc.from_storage([
+            "path/to/dir1",
+            "path/to/dir2"
+        ], session=session, recursive=True)
+        ```
+
+    Note:
+        When using multiple URIs with `update=True`, the function optimizes by
+        avoiding redundant updates for URIs pointing to the same storage location.
+    """
+    from .datachain import DataChain
+    from .datasets import from_dataset
+    from .records import from_records
+    from .values import from_values
+
+    file_type = get_file_type(type)
+
+    if anon:
+        client_config = (client_config or {}) | {"anon": True}
+    session = Session.get(session, client_config=client_config, in_memory=in_memory)
+    cache = session.catalog.cache
+    client_config = session.catalog.client_config
+
+    uris = uri if isinstance(uri, (list, tuple)) else [uri]
+
+    if not uris:
+        raise ValueError("No URIs provided")
+
+    storage_chain = None
+    listed_ds_name = set()
+    file_values = []
+
+    for single_uri in uris:
+        list_ds_name, list_uri, list_path, list_ds_exists = get_listing(
+            single_uri, session, update=update
+        )
+
+        # list_ds_name is None if object is a file, we don't want to use cache
+        # or do listing in that case - just read that single object
+        if not list_ds_name:
+            file_values.append(
+                get_file_info(list_uri, cache, client_config=client_config)
+            )
+            continue
+
+        dc = from_dataset(list_ds_name, session=session, settings=settings)
+        dc.signals_schema = dc.signals_schema.mutate({f"{object_name}": file_type})
+
+        if update or not list_ds_exists:
+
+            def lst_fn(ds_name, lst_uri):
+                # disable prefetch for listing, as it pre-downloads all files
+                (
+                    from_records(
+                        DataChain.DEFAULT_FILE_RECORD,
+                        session=session,
+                        settings=settings,
+                        in_memory=in_memory,
+                    )
+                    .settings(prefetch=0)
+                    .gen(
+                        list_bucket(lst_uri, cache, client_config=client_config),
+                        output={f"{object_name}": file_type},
+                    )
+                    .save(ds_name, listing=True)
+                )
+
+            dc._query.add_before_steps(
+                lambda ds_name=list_ds_name, lst_uri=list_uri: lst_fn(ds_name, lst_uri)
+            )
+
+        chain = ls(dc, list_path, recursive=recursive, object_name=object_name)
+
+        storage_chain = storage_chain.union(chain) if storage_chain else chain
+        listed_ds_name.add(list_ds_name)
+
+    if file_values:
+        file_chain = from_values(
+            session=session,
+            settings=settings,
+            in_memory=in_memory,
+            file=file_values,
+        )
+        file_chain.signals_schema = file_chain.signals_schema.mutate(
+            {f"{object_name}": file_type}
+        )
+        storage_chain = storage_chain.union(file_chain) if storage_chain else file_chain
+
+    assert storage_chain is not None
+
+    return storage_chain

{datachain-0.14.0 → datachain-0.14.1}/src/datachain/query/dataset.py

@@ -47,6 +47,7 @@ from datachain.error import (
     QueryScriptCancelError,
 )
 from datachain.func.base import Function
+from datachain.lib.listing import is_listing_dataset
 from datachain.lib.udf import UDFAdapter, _get_cache
 from datachain.progress import CombinedDownloadCallback, TqdmCombinedDownloadCallback
 from datachain.query.schema import C, UDFParamSpec, normalize_param
@@ -151,13 +152,6 @@ def step_result(
     )
 
 
-class StartingStep(ABC):
-    """An initial query processing step, referencing a data source."""
-
-    @abstractmethod
-    def apply(self) -> "StepResult": ...
-
-
 @frozen
 class Step(ABC):
     """A query processing step (filtering, mutation, etc.)"""
@@ -170,7 +164,7 @@ class Step(ABC):
 
 
 @frozen
-class QueryStep(StartingStep):
+class QueryStep:
     catalog: "Catalog"
     dataset_name: str
     dataset_version: int
@@ -1097,26 +1091,42 @@ class DatasetQuery:
         self.temp_table_names: list[str] = []
         self.dependencies: set[DatasetDependencyType] = set()
         self.table = self.get_table()
-        self.starting_step: StartingStep
+        self.starting_step: Optional[QueryStep] = None
         self.name: Optional[str] = None
         self.version: Optional[int] = None
         self.feature_schema: Optional[dict] = None
         self.column_types: Optional[dict[str, Any]] = None
+        self.before_steps: list[Callable] = []
 
-        self.name = name
+        self.list_ds_name: Optional[str] = None
 
-        if fallback_to_studio and is_token_set():
-            ds = self.catalog.get_dataset_with_remote_fallback(name, version)
+        self.name = name
+        self.dialect = self.catalog.warehouse.db.dialect
+        if version:
+            self.version = version
+
+        if is_listing_dataset(name):
+            # not setting query step yet as listing dataset might not exist at
+            # this point
+            self.list_ds_name = name
+        elif fallback_to_studio and is_token_set():
+            self._set_starting_step(
+                self.catalog.get_dataset_with_remote_fallback(name, version)
+            )
         else:
-            ds = self.catalog.get_dataset(name)
+            self._set_starting_step(self.catalog.get_dataset(name))
+
+    def _set_starting_step(self, ds: "DatasetRecord") -> None:
+        if not self.version:
+            self.version = ds.latest_version
 
-        self.version = version or ds.latest_version
+        self.starting_step = QueryStep(self.catalog, ds.name, self.version)
+
+        # at this point we know our starting dataset so setting up schemas
         self.feature_schema = ds.get_version(self.version).feature_schema
         self.column_types = copy(ds.schema)
         if "sys__id" in self.column_types:
             self.column_types.pop("sys__id")
-        self.starting_step = QueryStep(self.catalog, name, self.version)
-        self.dialect = self.catalog.warehouse.db.dialect
 
     def __iter__(self):
         return iter(self.db_results())
@@ -1180,11 +1190,23 @@ class DatasetQuery:
         col.table = self.table
         return col
 
+    def add_before_steps(self, fn: Callable) -> None:
+        """
+        Setting custom function to be run before applying steps
+        """
+        self.before_steps.append(fn)
+
     def apply_steps(self) -> QueryGenerator:
         """
         Apply the steps in the query and return the resulting
         sqlalchemy.SelectBase.
         """
+        for fn in self.before_steps:
+            fn()
+
+        if self.list_ds_name:
+            # at this point we know what is our starting listing dataset name
+            self._set_starting_step(self.catalog.get_dataset(self.list_ds_name))  # type: ignore [arg-type]
         query = self.clone()
 
         index = os.getenv("DATACHAIN_QUERY_CHUNK_INDEX", self._chunk_index)
@@ -1203,6 +1225,7 @@ class DatasetQuery:
             query = query.filter(C.sys__rand % total == index)
             query.steps = query.steps[-1:] + query.steps[:-1]
 
+        assert query.starting_step
         result = query.starting_step.apply()
         self.dependencies.update(result.dependencies)
 

{datachain-0.14.0 → datachain-0.14.1/src/datachain.egg-info}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.14.0
+Version: 0.14.1
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
-License: Apache-2.0
+License-Expression: Apache-2.0
 Project-URL: Documentation, https://datachain.dvc.ai
 Project-URL: Issues, https://github.com/iterative/datachain/issues
 Project-URL: Source, https://github.com/iterative/datachain

{datachain-0.14.0 → datachain-0.14.1}/tests/benchmarks/test_ls.py

@@ -2,5 +2,5 @@ from datachain.cli import ls
 
 
 def test_ls(benchmark, tmp_dir):
-    bucket = "s3://noaa-bathymetry-pds/"
+    bucket = "s3://noaa-dcdb-bathymetry-pds/"
     benchmark.pedantic(ls, args=([bucket],), kwargs={"client_config": {"anon": True}})

{datachain-0.14.0 → datachain-0.14.1}/tests/func/test_datachain.py

@@ -9,6 +9,7 @@ import uuid
 from collections.abc import Iterator
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
+from unittest.mock import patch
 
 import numpy as np
 import pandas as pd
@@ -152,7 +153,7 @@ def test_from_storage_partials(cloud_test_catalog):
         return name
 
     dogs_uri = f"{src_uri}/dogs"
-    dc.from_storage(dogs_uri, session=session)
+    dc.from_storage(dogs_uri, session=session).exec()
     assert _get_listing_datasets(session) == [
         f"{_list_dataset_name(dogs_uri)}@v1",
     ]
@@ -162,7 +163,7 @@ def test_from_storage_partials(cloud_test_catalog):
         f"{_list_dataset_name(dogs_uri)}@v1",
     ]
 
-    dc.from_storage(src_uri, session=session)
+    dc.from_storage(src_uri, session=session).exec()
     assert _get_listing_datasets(session) == sorted(
         [
             f"{_list_dataset_name(dogs_uri)}@v1",
@@ -170,7 +171,7 @@ def test_from_storage_partials(cloud_test_catalog):
         ]
     )
 
-    dc.from_storage(f"{src_uri}/cats", session=session)
+    dc.from_storage(f"{src_uri}/cats", session=session).exec()
     assert _get_listing_datasets(session) == sorted(
         [
             f"{_list_dataset_name(dogs_uri)}@v1",
@@ -196,14 +197,14 @@ def test_from_storage_partials_with_update(cloud_test_catalog):
         return name
 
     uri = f"{src_uri}/cats"
-    dc.from_storage(uri, session=session)
+    dc.from_storage(uri, session=session).exec()
     assert _get_listing_datasets(session) == sorted(
         [
             f"{_list_dataset_name(uri)}@v1",
         ]
     )
 
-    dc.from_storage(uri, session=session, update=True)
+    dc.from_storage(uri, session=session, update=True).exec()
     assert _get_listing_datasets(session) == sorted(
         [
             f"{_list_dataset_name(uri)}@v1",
@@ -369,6 +370,85 @@ def test_export_images_files(test_session, tmp_dir, tmp_path, use_cache):
         assert images_equal(img["data"], exported_img)
 
 
+@pytest.mark.parametrize("use_cache", [True, False])
+def test_from_storage_multiple_uris_files(test_session, tmp_dir, tmp_path, use_cache):
+    images = [
+        {"name": "img1.jpg", "data": Image.new(mode="RGB", size=(64, 64))},
+        {"name": "img2.jpg", "data": Image.new(mode="RGB", size=(128, 128))},
+    ]
+
+    for img in images:
+        img["data"].save(tmp_path / img["name"])
+
+    dc.from_storage(
+        [
+            f"file://{tmp_path}/img1.jpg",
+            f"file://{tmp_path}/img2.jpg",
+        ],
+        session=test_session,
+        anon=True,
+        update=True,
+    ).to_storage(tmp_dir / "output", placement="filename")
+
+    for img in images:
+        exported_img = Image.open(tmp_dir / "output" / img["name"])
+        assert images_equal(img["data"], exported_img)
+
+    chain = dc.from_storage(
+        [
+            f"file://{tmp_path}/img1.jpg",
+            f"file://{tmp_path}/img2.jpg",
+            f"file://{tmp_dir}/output/*",
+        ]
+    )
+    assert chain.count() == 4
+
+    chain = dc.from_storage([f"file://{tmp_dir}/output/*"])
+    assert chain.count() == 2
+
+
+@pytest.mark.parametrize(
+    "cloud_type",
+    ["s3", "azure", "gs"],
+    indirect=True,
+)
+def test_from_storage_multiple_uris_cache(cloud_test_catalog):
+    ctc = cloud_test_catalog
+    src_uri = ctc.src_uri
+    session = ctc.session
+
+    with pytest.raises(ValueError):
+        dc.from_storage([])  # No URIs provided
+
+    with patch(
+        "datachain.lib.dc.storage.get_listing", wraps=dc.lib.listing.get_listing
+    ) as mock_get_listing:
+        chain = dc.from_storage(
+            [
+                f"{src_uri}/cats",
+                f"{src_uri}/dogs",
+                f"{src_uri}/cats/cat*",
+                f"{src_uri}/dogs/dog*",
+            ],
+            session=session,
+            update=True,
+        ).exec()
+        assert chain.count() == 11
+
+        files = chain.collect("file")
+        assert {f.name for f in files} == {
+            "cat1",
+            "cat2",
+            "dog1",
+            "dog2",
+            "dog3",
+            "dog4",
+        }
+
+        # Verify from_records was called exactly twice
+        assert mock_get_listing.call_count == 4  # TODO FIX THIS
+
+
 def test_from_storage_path_object(test_session, tmp_dir, tmp_path):
     images = [
         {"name": "img1.jpg", "data": Image.new(mode="RGB", size=(64, 64))},

{datachain-0.14.0 → datachain-0.14.1}/tests/func/test_ls.py

@@ -32,7 +32,7 @@ def test_ls_no_args(cloud_test_catalog, cloud_type, capsys):
     catalog = session.catalog
     src = cloud_test_catalog.src_uri
 
-    dc.from_storage(src, session=session).collect()
+    dc.from_storage(src, session=session).exec()
     ls([], catalog=catalog)
     captured = capsys.readouterr()
     assert captured.out == f"{src}/@v1\n"

{datachain-0.14.0 → datachain-0.14.1}/tests/unit/lib/test_datachain.py

@@ -339,7 +339,7 @@ def test_listings(test_session, tmp_dir):
     df.to_parquet(tmp_dir / "df.parquet")
 
     uri = tmp_dir.as_uri()
-    dc.from_storage(uri, session=test_session)
+    dc.from_storage(uri, session=test_session).exec()
 
     # check that listing is not returned as normal dataset
     assert not any(
@@ -370,13 +370,13 @@ def test_listings_reindex(test_session, tmp_dir):
 
     uri = tmp_dir.as_uri()
 
-    dc.from_storage(uri, session=test_session)
+    dc.from_storage(uri, session=test_session).exec()
     assert len(list(dc.listings(session=test_session).collect("listing"))) == 1
 
-    dc.from_storage(uri, session=test_session)
+    dc.from_storage(uri, session=test_session).exec()
     assert len(list(dc.listings(session=test_session).collect("listing"))) == 1
 
-    dc.from_storage(uri, session=test_session, update=True)
+    dc.from_storage(uri, session=test_session, update=True).exec()
     listings = list(dc.listings(session=test_session).collect("listing"))
     assert len(listings) == 2
     listings.sort(key=lambda lst: lst.version)

datachain-0.14.0/src/datachain/lib/dc/storage.py

@@ -1,118 +0,0 @@
-import os.path
-from typing import (
-    TYPE_CHECKING,
-    Optional,
-    Union,
-)
-
-from datachain.lib.file import (
-    File,
-    FileType,
-    get_file_type,
-)
-from datachain.lib.listing import get_file_info, get_listing, list_bucket, ls
-from datachain.query import Session
-
-if TYPE_CHECKING:
-    from .datachain import DataChain
-
-
-def from_storage(
-    uri: Union[str, os.PathLike[str]],
-    *,
-    type: FileType = "binary",
-    session: Optional[Session] = None,
-    settings: Optional[dict] = None,
-    in_memory: bool = False,
-    recursive: Optional[bool] = True,
-    object_name: str = "file",
-    update: bool = False,
-    anon: bool = False,
-    client_config: Optional[dict] = None,
-) -> "DataChain":
-    """Get data from a storage as a list of file with all file attributes.
-    It returns the chain itself as usual.
-
-    Parameters:
-        uri : storage URI with directory. URI must start with storage prefix such
-            as `s3://`, `gs://`, `az://` or "file:///"
-        type : read file as "binary", "text", or "image" data. Default is "binary".
-        recursive : search recursively for the given path.
-        object_name : Created object column name.
-        update : force storage reindexing. Default is False.
-        anon : If True, we will treat cloud bucket as public one
-        client_config : Optional client configuration for the storage client.
-
-    Example:
-        Simple call from s3
-        ```py
-        import datachain as dc
-        chain = dc.from_storage("s3://my-bucket/my-dir")
-        ```
-
-        With AWS S3-compatible storage
-        ```py
-        import datachain as dc
-        chain = dc.from_storage(
-            "s3://my-bucket/my-dir",
-            client_config = {"aws_endpoint_url": "<minio-endpoint-url>"}
-        )
-        ```
-
-        Pass existing session
-        ```py
-        session = Session.get()
-        import datachain as dc
-        chain = dc.from_storage("s3://my-bucket/my-dir", session=session)
-        ```
-    """
-    from .datachain import DataChain
-    from .datasets import from_dataset
-    from .records import from_records
-    from .values import from_values
-
-    file_type = get_file_type(type)
-
-    if anon:
-        client_config = (client_config or {}) | {"anon": True}
-    session = Session.get(session, client_config=client_config, in_memory=in_memory)
-    cache = session.catalog.cache
-    client_config = session.catalog.client_config
-
-    list_ds_name, list_uri, list_path, list_ds_exists = get_listing(
-        uri, session, update=update
-    )
-
-    # ds_name is None if object is a file, we don't want to use cache
-    # or do listing in that case - just read that single object
-    if not list_ds_name:
-        dc = from_values(
-            session=session,
-            settings=settings,
-            in_memory=in_memory,
-            file=[get_file_info(list_uri, cache, client_config=client_config)],
-        )
-        dc.signals_schema = dc.signals_schema.mutate({f"{object_name}": file_type})
-        return dc
-
-    if update or not list_ds_exists:
-        # disable prefetch for listing, as it pre-downloads all files
-        (
-            from_records(
-                DataChain.DEFAULT_FILE_RECORD,
-                session=session,
-                settings=settings,
-                in_memory=in_memory,
-            )
-            .settings(prefetch=0)
-            .gen(
-                list_bucket(list_uri, cache, client_config=client_config),
-                output={f"{object_name}": File},
-            )
-            .save(list_ds_name, listing=True)
-        )
-
-    dc = from_dataset(list_ds_name, session=session, settings=settings)
-    dc.signals_schema = dc.signals_schema.mutate({f"{object_name}": file_type})
-
-    return ls(dc, list_path, recursive=recursive, object_name=object_name)