datachain 0.14.3__tar.gz → 0.14.4__tar.gz

{datachain-0.14.3/src/datachain.egg-info → datachain-0.14.4}/PKG-INFO

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

{datachain-0.14.3 → datachain-0.14.4}/src/datachain/lib/dataset_info.py

@@ -12,6 +12,7 @@ from datachain.dataset import (
 )
 from datachain.job import Job
 from datachain.lib.data_model import DataModel
+from datachain.query.session import Session
 from datachain.utils import TIME_ZERO
 
 if TYPE_CHECKING:
@@ -32,6 +33,10 @@ class DatasetInfo(DataModel):
     error_message: str = Field(default="")
     error_stack: str = Field(default="")
 
+    @property
+    def is_temp(self) -> bool:
+        return Session.is_temp_dataset(self.name)
+
     @staticmethod
     def _validate_dict(
         v: Optional[Union[str, dict]],

{datachain-0.14.3 → datachain-0.14.4}/src/datachain/lib/dc/datasets.py

@@ -140,6 +140,8 @@ def datasets(
         )
     ]
 
+    datasets_values = [d for d in datasets_values if not d.is_temp]
+
     return read_values(
         session=session,
         settings=settings,

{datachain-0.14.3 → datachain-0.14.4}/src/datachain/lib/udf.py

@@ -16,7 +16,6 @@ from datachain.lib.convert.flatten import flatten
 from datachain.lib.data_model import DataValue
 from datachain.lib.file import File
 from datachain.lib.utils import AbstractUDF, DataChainError, DataChainParamsError
-from datachain.progress import CombinedDownloadCallback
 from datachain.query.batch import (
     Batch,
     BatchingStrategy,
@@ -327,8 +326,9 @@ def _prefetch_inputs(
 
     if after_prefetch is None:
         after_prefetch = noop
-        if isinstance(download_cb, CombinedDownloadCallback):
-            after_prefetch = download_cb.increment_file_count
+        if download_cb and hasattr(download_cb, "increment_file_count"):
+            increment_file_count: Callable[[], None] = download_cb.increment_file_count
+            after_prefetch = increment_file_count
 
     f = partial(_prefetch_input, download_cb=download_cb, after_prefetch=after_prefetch)
     mapper = AsyncMapper(f, prepared_inputs, workers=prefetch)

{datachain-0.14.3 → datachain-0.14.4}/src/datachain/query/dispatch.py

@@ -4,9 +4,8 @@ from itertools import chain
 from multiprocessing import cpu_count
 from sys import stdin
 from threading import Timer
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Literal, Optional
 
-import attrs
 import multiprocess
 from cloudpickle import load, loads
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
@@ -60,6 +59,7 @@ def udf_entrypoint() -> int:
     query = udf_info["query"]
     rows_total = udf_info["rows_total"]
     batching = udf_info["batching"]
+    is_generator = udf_info["is_generator"]
     n_workers = udf_info["processes"]
     if n_workers is True:
         n_workers = None  # Use default number of CPUs (cores)
@@ -72,17 +72,20 @@ def udf_entrypoint() -> int:
     ) as udf_inputs:
         download_cb = get_download_callback()
         processed_cb = get_processed_callback()
+        generated_cb = get_generated_callback(is_generator)
         try:
             dispatch.run_udf_parallel(
                 udf_inputs,
                 rows_total=rows_total,
                 n_workers=n_workers,
-                processed_cb=processed_cb,
                 download_cb=download_cb,
+                processed_cb=processed_cb,
+                generated_cb=generated_cb,
             )
         finally:
             download_cb.close()
             processed_cb.close()
+            generated_cb.close()
 
     return 0
 
@@ -128,7 +131,6 @@ class UDFDispatcher:
             self.done_queue,
             self.query,
             self.table,
-            self.is_generator,
             self.is_batching,
             self.cache,
             self.udf_fields,
@@ -152,16 +154,14 @@ class UDFDispatcher:
         for _ in range(n_workers):
             put_into_queue(task_queue, STOP_SIGNAL)
 
-    def create_input_queue(self):
-        return self.ctx.Queue()
-
     def run_udf_parallel(  # noqa: C901, PLR0912
         self,
         input_rows: Iterable[RowsOutput],
         rows_total: int,
         n_workers: Optional[int] = None,
-        processed_cb: Callback = DEFAULT_CALLBACK,
         download_cb: Callback = DEFAULT_CALLBACK,
+        processed_cb: Callback = DEFAULT_CALLBACK,
+        generated_cb: Callback = DEFAULT_CALLBACK,
     ) -> None:
         n_workers = get_n_workers_from_arg(n_workers)
 
@@ -214,6 +214,8 @@ class UDFDispatcher:
                     download_cb.relative_update(downloaded)
                 if processed := result.get("processed"):
                     processed_cb.relative_update(processed)
+                if generated := result.get("generated"):
+                    generated_cb.relative_update(generated)
 
                 status = result["status"]
                 if status in (OK_STATUS, NOTIFY_STATUS):
@@ -260,46 +262,61 @@ class UDFDispatcher:
                 p.join()
 
 
-class WorkerCallback(Callback):
-    def __init__(self, queue: "multiprocess.Queue"):
+class DownloadCallback(Callback):
+    def __init__(self, queue: "multiprocess.Queue") -> None:
         self.queue = queue
         super().__init__()
 
     def relative_update(self, inc: int = 1) -> None:
+        # This callback is used to notify the size of the downloaded files
+        pass
+
+    def increment_file_count(self, inc: int = 1) -> None:
         put_into_queue(self.queue, {"status": NOTIFY_STATUS, "downloaded": inc})
 
 
 class ProcessedCallback(Callback):
-    def __init__(self):
-        self.processed_rows: Optional[int] = None
+    def __init__(
+        self,
+        name: Literal["processed", "generated"],
+        queue: "multiprocess.Queue",
+    ) -> None:
+        self.name = name
+        self.queue = queue
         super().__init__()
 
     def relative_update(self, inc: int = 1) -> None:
-        self.processed_rows = inc
+        put_into_queue(self.queue, {"status": NOTIFY_STATUS, self.name: inc})
 
 
-@attrs.define
 class UDFWorker:
-    catalog: "Catalog"
-    udf: "UDFAdapter"
-    task_queue: "multiprocess.Queue"
-    done_queue: "multiprocess.Queue"
-    query: "Select"
-    table: "Table"
-    is_generator: bool
-    is_batching: bool
-    cache: bool
-    udf_fields: Sequence[str]
-    cb: Callback = attrs.field()
-
-    @cb.default
-    def _default_callback(self) -> WorkerCallback:
-        return WorkerCallback(self.done_queue)
+    def __init__(
+        self,
+        catalog: "Catalog",
+        udf: "UDFAdapter",
+        task_queue: "multiprocess.Queue",
+        done_queue: "multiprocess.Queue",
+        query: "Select",
+        table: "Table",
+        is_batching: bool,
+        cache: bool,
+        udf_fields: Sequence[str],
+    ) -> None:
+        self.catalog = catalog
+        self.udf = udf
+        self.task_queue = task_queue
+        self.done_queue = done_queue
+        self.query = query
+        self.table = table
+        self.is_batching = is_batching
+        self.cache = cache
+        self.udf_fields = udf_fields
+
+        self.download_cb = DownloadCallback(self.done_queue)
+        self.processed_cb = ProcessedCallback("processed", self.done_queue)
+        self.generated_cb = ProcessedCallback("generated", self.done_queue)
 
     def run(self) -> None:
-        processed_cb = ProcessedCallback()
-        generated_cb = get_generated_callback(self.is_generator)
-
         prefetch = self.udf.prefetch
         with _get_cache(self.catalog.cache, prefetch, use_cache=self.cache) as _cache:
             catalog = clone_catalog_with_cache(self.catalog, _cache)
@@ -308,29 +325,22 @@ class UDFWorker:
                 self.get_inputs(),
                 catalog,
                 self.cache,
-                download_cb=self.cb,
-                processed_cb=processed_cb,
+                download_cb=self.download_cb,
+                processed_cb=self.processed_cb,
             )
             with safe_closing(udf_results):
                 process_udf_outputs(
                     catalog.warehouse,
                     self.table,
-                    self.notify_and_process(udf_results, processed_cb),
+                    self.notify_and_process(udf_results),
                     self.udf,
-                    cb=generated_cb,
+                    cb=self.generated_cb,
                 )
+        put_into_queue(self.done_queue, {"status": FINISHED_STATUS})
 
-        put_into_queue(
-            self.done_queue,
-            {"status": FINISHED_STATUS, "processed": processed_cb.processed_rows},
-        )
-
-    def notify_and_process(self, udf_results, processed_cb):
+    def notify_and_process(self, udf_results):
         for row in udf_results:
-            put_into_queue(
-                self.done_queue,
-                {"status": OK_STATUS, "processed": processed_cb.processed_rows},
-            )
+            put_into_queue(self.done_queue, {"status": OK_STATUS})
             yield row
 
     def get_inputs(self):

{datachain-0.14.3 → datachain-0.14.4}/src/datachain/query/session.py

@@ -100,6 +100,10 @@ class Session:
     def get_temp_prefix(self) -> str:
         return f"{self.DATASET_PREFIX}{self.name}_"
 
+    @classmethod
+    def is_temp_dataset(cls, name) -> bool:
+        return name.startswith(cls.DATASET_PREFIX)
+
     def _cleanup_temp_datasets(self) -> None:
         prefix = self.get_temp_prefix()
         try:

{datachain-0.14.3 → datachain-0.14.4/src/datachain.egg-info}/PKG-INFO

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

{datachain-0.14.3 → datachain-0.14.4}/tests/func/test_datachain.py

@@ -538,6 +538,17 @@ def test_show(capsys, test_session):
         assert f"{i} {first_name[i]}" in normalized_output
 
 
+def test_show_without_temp_datasets(capsys, test_session):
+    dc.read_values(
+        key=[1, 2, 3, 4], session=test_session
+    ).save()  # creates temp dataset
+    dc.datasets().show()
+    captured = capsys.readouterr()
+    normalized_output = re.sub(r"\s+", " ", captured.out)
+    print(normalized_output)
+    assert "Empty result" in normalized_output
+
+
 def test_class_method_deprecated(capsys, test_session):
     with pytest.warns(DeprecationWarning):
         dc.DataChain.from_values(key=["a", "b", "c"], session=test_session)

{datachain-0.14.3 → datachain-0.14.4}/tests/unit/test_session.py

@@ -47,6 +47,18 @@ def test_session_empty_name():
     assert name.startswith(Session.GLOBAL_SESSION_NAME + "_")
 
 
+@pytest.mark.parametrize(
+    "name,is_temp",
+    (
+        ("session_global_456b5d_0cda3b", True),
+        ("session_TestSession_456b5d_0cda3b", True),
+        ("cats", False),
+    ),
+)
+def test_is_temp_dataset(name, is_temp):
+    assert Session.is_temp_dataset(name) is is_temp
+
+
 def test_ephemeral_dataset_lifecycle(catalog):
     session_name = "asd3d4"
     with Session(session_name, catalog=catalog) as session: