datachain 0.9.0__tar.gz → 0.10.0__tar.gz

{datachain-0.9.0 → datachain-0.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: datachain
-Version: 0.9.0
+Version: 0.10.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -175,7 +175,7 @@ high confidence scores.
 
     likely_cats = annotated.filter((Column("meta.inference.confidence") > 0.93) \
                                    & (Column("meta.inference.class_") == "cat"))
-    likely_cats.export_files("high-confidence-cats/", signal="file")
+    likely_cats.to_storage("high-confidence-cats/", signal="file")
 
 
 Example: LLM based text-file evaluation
@@ -216,7 +216,7 @@ Python code:
     )
 
     successful_chain = chain.filter(Column("is_success") == True)
-    successful_chain.export_files("./output_mistral")
+    successful_chain.to_storage("./output_mistral")
 
     print(f"{successful_chain.count()} files were exported")
 

{datachain-0.9.0 → datachain-0.10.0}/README.rst

@@ -68,7 +68,7 @@ high confidence scores.
 
     likely_cats = annotated.filter((Column("meta.inference.confidence") > 0.93) \
                                    & (Column("meta.inference.class_") == "cat"))
-    likely_cats.export_files("high-confidence-cats/", signal="file")
+    likely_cats.to_storage("high-confidence-cats/", signal="file")
 
 
 Example: LLM based text-file evaluation
@@ -109,7 +109,7 @@ Python code:
     )
 
     successful_chain = chain.filter(Column("is_success") == True)
-    successful_chain.export_files("./output_mistral")
+    successful_chain.to_storage("./output_mistral")
 
     print(f"{successful_chain.count()} files were exported")
 

{datachain-0.9.0 → datachain-0.10.0}/docs/quick-start.md

@@ -47,7 +47,7 @@ annotated = images_id.merge(meta, on="id", right_on="meta.id")
 
 likely_cats = annotated.filter((Column("meta.inference.confidence") > 0.93) \
                                & (Column("meta.inference.class_") == "cat"))
-likely_cats.export_files("high-confidence-cats/", signal="file")
+likely_cats.to_storage("high-confidence-cats/", signal="file")
 ```
 
 ## Data curation with a local AI model
@@ -85,7 +85,7 @@ chain = (
 )
 
 positive_chain = chain.filter(Column("is_positive") == True)
-positive_chain.export_files("./output")
+positive_chain.to_storage("./output")
 
 print(f"{positive_chain.count()} files were exported")
 ```

{datachain-0.9.0 → datachain-0.10.0}/src/datachain/client/fsspec.py

@@ -390,7 +390,7 @@ class Client(ABC):
         )  # type: ignore[return-value]
 
     def upload(self, data: bytes, path: str) -> "File":
-        full_path = self.get_full_path(path)
+        full_path = path if path.startswith(self.PREFIX) else self.get_full_path(path)
 
         parent = posixpath.dirname(full_path)
         self.fs.makedirs(parent, exist_ok=True)

{datachain-0.9.0 → datachain-0.10.0}/src/datachain/lib/dc.py

@@ -411,6 +411,7 @@ class DataChain:
         object_name: str = "file",
         update: bool = False,
         anon: bool = False,
+        client_config: Optional[dict] = None,
     ) -> "Self":
         """Get data from a storage as a list of file with all file attributes.
         It returns the chain itself as usual.
@@ -423,15 +424,32 @@ class DataChain:
             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
             chain = DataChain.from_storage("s3://my-bucket/my-dir")
             ```
+
+            With AWS S3-compatible storage
+            ```py
+            chain = DataChain.from_storage(
+                "s3://my-bucket/my-dir",
+                client_config = {"aws_endpoint_url": "<minio-endpoint-url>"}
+            )
+            ```
+
+            Pass existing session
+            ```py
+            session = Session.get()
+            chain = DataChain.from_storage("s3://my-bucket/my-dir", session=session)
+            ```
         """
         file_type = get_file_type(type)
 
-        client_config = {"anon": True} if anon else None
+        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
@@ -481,25 +499,56 @@ class DataChain:
         version: Optional[int] = None,
         session: Optional[Session] = None,
         settings: Optional[dict] = None,
-        fallback_to_remote: bool = True,
+        fallback_to_studio: bool = True,
     ) -> "Self":
         """Get data from a saved Dataset. It returns the chain itself.
+        If dataset or version is not found locally, it will try to pull it from Studio.
 
         Parameters:
             name : dataset name
             version : dataset version
+            session : Session to use for the chain.
+            settings : Settings to use for the chain.
+            fallback_to_studio : Try to pull dataset from Studio if not found locally.
+                Default is True.
 
         Example:
             ```py
             chain = DataChain.from_dataset("my_cats")
             ```
+
+            ```py
+            chain = DataChain.from_dataset("my_cats", fallback_to_studio=False)
+            ```
+
+            ```py
+            chain = DataChain.from_dataset("my_cats", version=1)
+            ```
+
+            ```py
+            session = Session.get(client_config={"aws_endpoint_url": "<minio-url>"})
+            settings = {
+                "cache": True,
+                "parallel": 4,
+                "workers": 4,
+                "min_task_size": 1000,
+                "prefetch": 10,
+            }
+            chain = DataChain.from_dataset(
+                name="my_cats",
+                version=1,
+                session=session,
+                settings=settings,
+                fallback_to_studio=True,
+            )
+            ```
         """
         query = DatasetQuery(
             name=name,
             version=version,
             session=session,
             indexing_column_types=File._datachain_column_types,
-            fallback_to_remote=fallback_to_remote,
+            fallback_to_studio=fallback_to_studio,
         )
         telemetry.send_event_once("class", "datachain_init", name=name, version=version)
         if settings:
@@ -2444,7 +2493,7 @@ class DataChain:
         self._setup = self._setup | kwargs
         return self
 
-    def export_files(
+    def to_storage(
         self,
         output: str,
         signal: str = "file",
@@ -2462,6 +2511,13 @@ class DataChain:
             use_cache: If `True`, cache the files before exporting.
             link_type: Method to use for exporting files.
                 Falls back to `'copy'` if symlinking fails.
+
+        Example:
+            Cross cloud transfer
+            ```py
+            ds = DataChain.from_storage("s3://mybucket")
+            ds.to_storage("gs://mybucket", placement="filename")
+            ```
         """
         if placement == "filename" and (
             self._query.distinct(pathfunc.name(C(f"{signal}__path"))).count()

{datachain-0.9.0 → datachain-0.10.0}/src/datachain/lib/file.py

@@ -17,6 +17,7 @@ from urllib.parse import unquote, urlparse
 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
 
@@ -214,10 +215,13 @@ class File(DataModel):
 
             catalog = get_catalog()
 
-        parent, name = posixpath.split(path)
+        from datachain.client.fsspec import Client
 
-        client = catalog.get_client(parent)
-        file = client.upload(data, name)
+        client_cls = Client.get_implementation(path)
+        source, rel_path = client_cls.split_url(path)
+
+        client = catalog.get_client(client_cls.get_uri(source))
+        file = client.upload(data, rel_path)
         if not isinstance(file, cls):
             file = cls(**file.model_dump())
         file._set_stream(catalog)
@@ -267,8 +271,9 @@ class File(DataModel):
 
     def save(self, destination: str):
         """Writes it's content to destination"""
-        with open(destination, mode="wb") as f:
-            f.write(self.read())
+        destination = stringify_path(destination)
+        client: Client = self._catalog.get_client(str(destination))
+        client.upload(self.read(), str(destination))
 
     def _symlink_to(self, destination: str):
         if self.location:
@@ -282,6 +287,7 @@ class File(DataModel):
             source = self.get_path()
         else:
             raise OSError(errno.EXDEV, "can't link across filesystems")
+
         return os.symlink(source, destination)
 
     def export(
@@ -296,7 +302,8 @@ class File(DataModel):
             self._caching_enabled = use_cache
         dst = self.get_destination_path(output, placement)
         dst_dir = os.path.dirname(dst)
-        os.makedirs(dst_dir, exist_ok=True)
+        client: Client = self._catalog.get_client(dst_dir)
+        client.fs.makedirs(dst_dir, exist_ok=True)
 
         if link_type == "symlink":
             try:
@@ -493,7 +500,10 @@ class TextFile(File):
 
     def save(self, destination: str):
         """Writes it's content to destination"""
-        with open(destination, mode="w") as f:
+        destination = stringify_path(destination)
+
+        client: Client = self._catalog.get_client(destination)
+        with client.fs.open(destination, mode="w") as f:
             f.write(self.read_text())
 
 
@@ -507,7 +517,11 @@ class ImageFile(File):
 
     def save(self, destination: str):
         """Writes it's content to destination"""
-        self.read().save(destination)
+        destination = stringify_path(destination)
+
+        client: Client = self._catalog.get_client(destination)
+        with client.fs.open(destination, mode="wb") as f:
+            self.read().save(f)
 
 
 class Image(DataModel):

{datachain-0.9.0 → datachain-0.10.0}/src/datachain/query/dataset.py

@@ -1085,7 +1085,7 @@ class DatasetQuery:
         session: Optional[Session] = None,
         indexing_column_types: Optional[dict[str, Any]] = None,
         in_memory: bool = False,
-        fallback_to_remote: bool = True,
+        fallback_to_studio: bool = True,
     ) -> None:
         self.session = Session.get(session, catalog=catalog, in_memory=in_memory)
         self.catalog = catalog or self.session.catalog
@@ -1103,7 +1103,7 @@ class DatasetQuery:
 
         self.name = name
 
-        if fallback_to_remote and is_token_set():
+        if fallback_to_studio and is_token_set():
             ds = self.catalog.get_dataset_with_remote_fallback(name, version)
         else:
             ds = self.catalog.get_dataset(name)

{datachain-0.9.0 → datachain-0.10.0}/src/datachain/query/session.py

@@ -139,21 +139,33 @@ class Session:
 
         # Access the active (most recent) context from the stack
         if cls.SESSION_CONTEXTS:
-            return cls.SESSION_CONTEXTS[-1]
+            session = cls.SESSION_CONTEXTS[-1]
 
-        if cls.GLOBAL_SESSION_CTX is None:
+        elif cls.GLOBAL_SESSION_CTX is None:
             cls.GLOBAL_SESSION_CTX = Session(
                 cls.GLOBAL_SESSION_NAME,
                 catalog,
                 client_config=client_config,
                 in_memory=in_memory,
             )
+            session = cls.GLOBAL_SESSION_CTX
 
             atexit.register(cls._global_cleanup)
             cls.ORIGINAL_EXCEPT_HOOK = sys.excepthook
             sys.excepthook = cls.except_hook
+        else:
+            session = cls.GLOBAL_SESSION_CTX
 
-        return cls.GLOBAL_SESSION_CTX
+        if client_config and session.catalog.client_config != client_config:
+            session = Session(
+                "session" + uuid4().hex[:4],
+                catalog,
+                client_config=client_config,
+                in_memory=in_memory,
+            )
+            session.__enter__()
+
+        return session
 
     @staticmethod
     def except_hook(exc_type, exc_value, exc_traceback):

{datachain-0.9.0 → datachain-0.10.0}/src/datachain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: datachain
-Version: 0.9.0
+Version: 0.10.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -175,7 +175,7 @@ high confidence scores.
 
     likely_cats = annotated.filter((Column("meta.inference.confidence") > 0.93) \
                                    & (Column("meta.inference.class_") == "cat"))
-    likely_cats.export_files("high-confidence-cats/", signal="file")
+    likely_cats.to_storage("high-confidence-cats/", signal="file")
 
 
 Example: LLM based text-file evaluation
@@ -216,7 +216,7 @@ Python code:
     )
 
     successful_chain = chain.filter(Column("is_success") == True)
-    successful_chain.export_files("./output_mistral")
+    successful_chain.to_storage("./output_mistral")
 
     print(f"{successful_chain.count()} files were exported")
 

{datachain-0.9.0 → datachain-0.10.0}/src/datachain.egg-info/SOURCES.txt

@@ -232,6 +232,7 @@ tests/func/__init__.py
 tests/func/fake-service-account-credentials.json
 tests/func/test_catalog.py
 tests/func/test_client.py
+tests/func/test_cloud_transfer.py
 tests/func/test_data_storage.py
 tests/func/test_datachain.py
 tests/func/test_datachain_merge.py

{datachain-0.9.0 → datachain-0.10.0}/tests/conftest.py

@@ -472,9 +472,9 @@ def cloud_server_credentials(cloud_server, monkeypatch):
 
 def get_cloud_test_catalog(cloud_server, tmp_path, metastore, warehouse):
     cache_dir = tmp_path / ".datachain" / "cache"
-    cache_dir.mkdir(parents=True)
+    cache_dir.mkdir(parents=True, exist_ok=True)
     tmpfile_dir = tmp_path / ".datachain" / "tmp"
-    tmpfile_dir.mkdir()
+    tmpfile_dir.mkdir(exist_ok=True)
 
     catalog = Catalog(
         metastore=metastore,

datachain-0.10.0/tests/func/test_cloud_transfer.py

@@ -0,0 +1,68 @@
+import pytest
+
+from datachain import Session
+from datachain.lib.dc import DataChain
+from tests.conftest import get_cloud_test_catalog, make_cloud_server
+
+
+def test_cross_cloud_transfer(
+    request,
+    tmp_upath_factory,
+    tree,
+    tmp_path,
+    metastore,
+    warehouse,
+):
+    disabled_remotes = request.config.getoption("--disable-remotes") or []
+
+    if any(remote in disabled_remotes for remote in ["azure", "gs", "all"]):
+        pytest.skip("Skipping all tests for azure, gs or all remotes")
+
+    azure_path = tmp_upath_factory.mktemp("azure", version_aware=False)
+    azure_server = make_cloud_server(azure_path, "azure", tree)
+
+    gcloud_path = tmp_upath_factory.mktemp("gs", version_aware=False)
+    gcloud_server = make_cloud_server(gcloud_path, "gs", tree)
+
+    # Initialize cloud catalogs
+    azure_catalog = get_cloud_test_catalog(azure_server, tmp_path, metastore, warehouse)
+    gcloud_catalog = get_cloud_test_catalog(
+        gcloud_server, tmp_path, metastore, warehouse
+    )
+
+    # Define test file paths
+    test_filename = "image_1.jpg"
+    test_content = b"bytes"
+
+    source_dir = f"{azure_catalog.src_uri}/source-test-images"
+    source_file = f"{source_dir}/{test_filename}"
+
+    dest_dir = f"{gcloud_catalog.src_uri}/destination-test-images"
+    dest_file = f"{dest_dir}/{test_filename}"
+
+    # Get cloud clients
+    azure_client = azure_catalog.catalog.get_client(source_file)
+    gcloud_client = gcloud_catalog.catalog.get_client(dest_file)
+
+    try:
+        # Create test file in Azure
+        with azure_client.fs.open(source_file, "wb") as f:
+            f.write(test_content)
+
+        # Perform cross-cloud transfer
+        combined_config = azure_server.client_config | gcloud_server.client_config
+        with Session("testSession", client_config=combined_config):
+            datachain = DataChain.from_storage(source_dir)
+            datachain.to_storage(dest_dir, placement="filename")
+
+        # Verify transfer
+        with gcloud_client.fs.open(dest_file, "rb") as f:
+            assert f.read() == test_content
+
+    finally:
+        # Cleanup
+        try:
+            azure_client.fs.rm(source_dir, recursive=True)
+            gcloud_client.fs.rm(dest_dir, recursive=True)
+        except FileNotFoundError:
+            pass

{datachain-0.9.0 → datachain-0.10.0}/tests/func/test_datachain.py

@@ -64,6 +64,16 @@ def test_catalog_anon(tmp_dir, catalog, anon):
     assert chain.session.catalog.client_config.get("anon", False) is anon
 
 
+def test_from_storage_client_config(tmp_dir, catalog):
+    dc = DataChain.from_storage(tmp_dir.as_uri())
+    assert dc.session.catalog.client_config == {}  # Default client config is set.
+
+    dc = DataChain.from_storage(tmp_dir.as_uri(), client_config={"anon": True})
+    assert dc.session.catalog.client_config == {
+        "anon": True
+    }  # New client config is set.
+
+
 def test_from_storage(cloud_test_catalog):
     ctc = cloud_test_catalog
     dc = DataChain.from_storage(ctc.src_uri, session=ctc.session)
@@ -292,20 +302,20 @@ def test_read_file(cloud_test_catalog, use_cache):
 @pytest.mark.parametrize("use_cache", [True, False])
 @pytest.mark.parametrize("file_type", ["", "binary", "text"])
 @pytest.mark.parametrize("cloud_type", ["file"], indirect=True)
-def test_export_files(
+def test_to_storage(
     tmp_dir, cloud_test_catalog, test_session, placement, use_map, use_cache, file_type
 ):
     ctc = cloud_test_catalog
     df = DataChain.from_storage(ctc.src_uri, type=file_type, session=test_session)
     if use_map:
-        df.export_files(tmp_dir / "output", placement=placement, use_cache=use_cache)
+        df.to_storage(tmp_dir / "output", placement=placement, use_cache=use_cache)
         df.map(
             res=lambda file: file.export(
                 tmp_dir / "output", placement=placement, use_cache=use_cache
             )
         ).exec()
     else:
-        df.export_files(tmp_dir / "output", placement=placement)
+        df.to_storage(tmp_dir / "output", placement=placement)
 
     expected = {
         "description": "Cats and Dogs",
@@ -341,14 +351,14 @@ def test_export_images_files(test_session, tmp_dir, tmp_path, use_cache):
             ImageFile(path=img["name"], source=f"file://{tmp_path}") for img in images
         ],
         session=test_session,
-    ).export_files(tmp_dir / "output", placement="filename", use_cache=use_cache)
+    ).to_storage(tmp_dir / "output", placement="filename", use_cache=use_cache)
 
     for img in images:
         exported_img = Image.open(tmp_dir / "output" / img["name"])
         assert images_equal(img["data"], exported_img)
 
 
-def test_export_files_filename_placement_not_unique_files(tmp_dir, test_session):
+def test_to_storage_files_filename_placement_not_unique_files(tmp_dir, test_session):
     data = b"some\x00data\x00is\x48\x65\x6c\x57\x6f\x72\x6c\x64\xff\xffheRe"
     bucket_name = "mybucket"
     files = ["dir1/a.json", "dir1/dir2/a.json"]
@@ -364,7 +374,7 @@ def test_export_files_filename_placement_not_unique_files(tmp_dir, test_session)
 
     df = DataChain.from_storage((tmp_dir / bucket_name).as_uri(), session=test_session)
     with pytest.raises(ValueError):
-        df.export_files(tmp_dir / "output", placement="filename")
+        df.to_storage(tmp_dir / "output", placement="filename")
 
 
 def test_show(capsys, test_session):

{datachain-0.9.0 → datachain-0.10.0}/tests/func/test_file.py

@@ -50,16 +50,18 @@ def test_upload(cloud_test_catalog):
 
     src_uri = ctc.src_uri
     filename = "image_1.jpg"
-    source = f"{src_uri}/upload-test-images"
+    dest = f"{src_uri}/upload-test-images"
     catalog = ctc.catalog
 
     img_bytes = b"bytes"
 
-    f = File.upload(img_bytes, f"{source}/{filename}", catalog)
+    f = File.upload(img_bytes, f"{dest}/{filename}", catalog)
 
-    assert f.path == filename
-    assert f.source == source
+    client = catalog.get_client(src_uri)
+    source, rel_path = client.split_url(f"{dest}/{filename}")
+
+    assert f.path == rel_path
+    assert f.source == client.get_uri(source)
     assert f.read() == img_bytes
 
-    client = catalog.get_client(src_uri)
-    client.fs.rm(source, recursive=True)
+    client.fs.rm(dest, recursive=True)

{datachain-0.9.0 → datachain-0.10.0}/tests/func/test_pull.py

@@ -295,7 +295,7 @@ def test_datachain_from_dataset_pull(
         ds = DataChain.from_dataset(
             name="dogs",
             version=1,
-            fallback_to_remote=True,
+            fallback_to_studio=True,
         )
 
     assert ds.dataset.name == "dogs"