datachain 0.14.2__tar.gz → 0.14.3__tar.gz

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

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

{datachain-0.14.2 → datachain-0.14.3}/src/datachain/catalog/loader.py

@@ -7,6 +7,7 @@ from datachain.utils import get_envs_by_prefix
 if TYPE_CHECKING:
     from datachain.catalog import Catalog
     from datachain.data_storage import AbstractMetastore, AbstractWarehouse
+    from datachain.query.udf import AbstractUDFDistributor
 
 METASTORE_SERIALIZED = "DATACHAIN__METASTORE"
 METASTORE_IMPORT_PATH = "DATACHAIN_METASTORE"
@@ -15,7 +16,6 @@ WAREHOUSE_SERIALIZED = "DATACHAIN__WAREHOUSE"
 WAREHOUSE_IMPORT_PATH = "DATACHAIN_WAREHOUSE"
 WAREHOUSE_ARG_PREFIX = "DATACHAIN_WAREHOUSE_ARG_"
 DISTRIBUTED_IMPORT_PATH = "DATACHAIN_DISTRIBUTED"
-DISTRIBUTED_ARG_PREFIX = "DATACHAIN_DISTRIBUTED_ARG_"
 
 IN_MEMORY_ERROR_MESSAGE = "In-memory is only supported on SQLite"
 
@@ -100,27 +100,22 @@ def get_warehouse(in_memory: bool = False) -> "AbstractWarehouse":
     return warehouse_class(**warehouse_args)
 
 
-def get_distributed_class(**kwargs):
+def get_udf_distributor_class() -> type["AbstractUDFDistributor"]:
     distributed_import_path = os.environ.get(DISTRIBUTED_IMPORT_PATH)
-    distributed_arg_envs = get_envs_by_prefix(DISTRIBUTED_ARG_PREFIX)
-    # Convert env variable names to keyword argument names by lowercasing them
-    distributed_args = {k.lower(): v for k, v in distributed_arg_envs.items()}
 
     if not distributed_import_path:
         raise RuntimeError(
             f"{DISTRIBUTED_IMPORT_PATH} import path is required "
             "for distributed UDF processing."
         )
-    # Distributed class paths are specified as (for example):
-    # module.classname
+    # Distributed class paths are specified as (for example): module.classname
     if "." not in distributed_import_path:
         raise RuntimeError(
             f"Invalid {DISTRIBUTED_IMPORT_PATH} import path: {distributed_import_path}"
         )
     module_name, _, class_name = distributed_import_path.rpartition(".")
     distributed = import_module(module_name)
-    distributed_class = getattr(distributed, class_name)
-    return distributed_class(**distributed_args | kwargs)
+    return getattr(distributed, class_name)
 
 
 def get_catalog(

{datachain-0.14.2 → datachain-0.14.3}/src/datachain/data_storage/warehouse.py

@@ -199,6 +199,15 @@ class AbstractWarehouse(ABC, Serializable):
     # Query Execution
     #
 
+    def query_count(self, query: sa.sql.selectable.Select) -> int:
+        """Count the number of rows in a query."""
+        count_query = sa.select(func.count(1)).select_from(query.subquery())
+        return next(self.db.execute(count_query))[0]
+
+    def table_rows_count(self, table) -> int:
+        count_query = sa.select(func.count(1)).select_from(table)
+        return next(self.db.execute(count_query))[0]
+
     def dataset_select_paginated(
         self,
         query,

{datachain-0.14.2 → datachain-0.14.3}/src/datachain/query/dataset.py

@@ -55,10 +55,12 @@ from datachain.lib.udf import UDFAdapter, _get_cache
 from datachain.progress import CombinedDownloadCallback, TqdmCombinedDownloadCallback
 from datachain.query.schema import C, UDFParamSpec, normalize_param
 from datachain.query.session import Session
+from datachain.query.udf import UdfInfo
 from datachain.sql.functions.random import rand
 from datachain.utils import (
     batched,
     determine_processes,
+    determine_workers,
     filtered_cloudpickle_dumps,
     get_datachain_executable,
     safe_closing,
@@ -74,7 +76,6 @@ if TYPE_CHECKING:
     from datachain.data_storage import AbstractWarehouse
     from datachain.dataset import DatasetRecord
     from datachain.lib.udf import UDFAdapter, UDFResult
-    from datachain.query.udf import UdfInfo
 
     P = ParamSpec("P")
 
@@ -414,20 +415,15 @@ class UDFStep(Step, ABC):
     def populate_udf_table(self, udf_table: "Table", query: Select) -> None:
         from datachain.catalog import QUERY_SCRIPT_CANCELED_EXIT_CODE
 
-        use_partitioning = self.partition_by is not None
-        batching = self.udf.get_batching(use_partitioning)
-        workers = self.workers
-        if (
-            not workers
-            and os.environ.get("DATACHAIN_DISTRIBUTED")
-            and os.environ.get("DATACHAIN_SETTINGS_WORKERS")
-        ):
-            # Enable distributed processing by default if the module is available,
-            # and a default number of workers is provided.
-            workers = True
+        rows_total = self.catalog.warehouse.query_count(query)
+        if rows_total == 0:
+            return
 
-        processes = determine_processes(self.parallel)
+        workers = determine_workers(self.workers, rows_total=rows_total)
+        processes = determine_processes(self.parallel, rows_total=rows_total)
 
+        use_partitioning = self.partition_by is not None
+        batching = self.udf.get_batching(use_partitioning)
         udf_fields = [str(c.name) for c in query.selected_columns]
 
         prefetch = self.udf.prefetch
@@ -441,23 +437,24 @@ class UDFStep(Step, ABC):
                             "distributed processing."
                         )
 
-                    from datachain.catalog.loader import get_distributed_class
-
-                    distributor = get_distributed_class(
-                        min_task_size=self.min_task_size
-                    )
-                    distributor(
-                        self.udf,
-                        catalog,
-                        udf_table,
-                        query,
-                        workers,
-                        processes,
+                    from datachain.catalog.loader import get_udf_distributor_class
+
+                    udf_distributor_class = get_udf_distributor_class()
+                    udf_distributor = udf_distributor_class(
+                        catalog=catalog,
+                        table=udf_table,
+                        query=query,
+                        udf_data=filtered_cloudpickle_dumps(self.udf),
+                        batching=batching,
+                        workers=workers,
+                        processes=processes,
                         udf_fields=udf_fields,
+                        rows_total=rows_total,
+                        use_cache=self.cache,
                         is_generator=self.is_generator,
-                        use_partitioning=use_partitioning,
-                        cache=self.cache,
+                        min_task_size=self.min_task_size,
                     )
+                    udf_distributor()
                 elif processes:
                     # Parallel processing (faster for more CPU-heavy UDFs)
                     if catalog.in_memory:
@@ -465,19 +462,21 @@ class UDFStep(Step, ABC):
                             "In-memory databases cannot be used "
                             "with parallel processing."
                         )
-                    udf_info: UdfInfo = {
-                        "udf_data": filtered_cloudpickle_dumps(self.udf),
-                        "catalog_init": catalog.get_init_params(),
-                        "metastore_clone_params": catalog.metastore.clone_params(),
-                        "warehouse_clone_params": catalog.warehouse.clone_params(),
-                        "table": udf_table,
-                        "query": query,
-                        "udf_fields": udf_fields,
-                        "batching": batching,
-                        "processes": processes,
-                        "is_generator": self.is_generator,
-                        "cache": self.cache,
-                    }
+
+                    udf_info = UdfInfo(
+                        udf_data=filtered_cloudpickle_dumps(self.udf),
+                        catalog_init=catalog.get_init_params(),
+                        metastore_clone_params=catalog.metastore.clone_params(),
+                        warehouse_clone_params=catalog.warehouse.clone_params(),
+                        table=udf_table,
+                        query=query,
+                        udf_fields=udf_fields,
+                        batching=batching,
+                        processes=processes,
+                        is_generator=self.is_generator,
+                        cache=self.cache,
+                        rows_total=rows_total,
+                    )
 
                     # Run the UDFDispatcher in another process to avoid needing
                     # if __name__ == '__main__': in user scripts

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

@@ -11,11 +11,10 @@ import multiprocess
 from cloudpickle import load, loads
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 from multiprocess import get_context
-from sqlalchemy.sql import func
 
 from datachain.catalog import Catalog
 from datachain.catalog.catalog import clone_catalog_with_cache
-from datachain.catalog.loader import get_distributed_class
+from datachain.catalog.loader import get_udf_distributor_class
 from datachain.lib.udf import _get_cache
 from datachain.query.batch import RowsOutput, RowsOutputBatch
 from datachain.query.dataset import (
@@ -59,6 +58,7 @@ def udf_entrypoint() -> int:
     dispatch = UDFDispatcher(udf_info)
 
     query = udf_info["query"]
+    rows_total = udf_info["rows_total"]
     batching = udf_info["batching"]
     n_workers = udf_info["processes"]
     if n_workers is True:
@@ -67,12 +67,6 @@ def udf_entrypoint() -> int:
     wh_cls, wh_args, wh_kwargs = udf_info["warehouse_clone_params"]
     warehouse: AbstractWarehouse = wh_cls(*wh_args, **wh_kwargs)
 
-    total_rows = next(
-        warehouse.db.execute(
-            query.with_only_columns(func.count(query.c.sys__id)).order_by(None)
-        )
-    )[0]
-
     with contextlib.closing(
         batching(warehouse.dataset_select_paginated, query, ids_only=True)
     ) as udf_inputs:
@@ -81,7 +75,7 @@ def udf_entrypoint() -> int:
         try:
             dispatch.run_udf_parallel(
                 udf_inputs,
-                total_rows=total_rows,
+                rows_total=rows_total,
                 n_workers=n_workers,
                 processed_cb=processed_cb,
                 download_cb=download_cb,
@@ -94,7 +88,7 @@ def udf_entrypoint() -> int:
 
 
 def udf_worker_entrypoint() -> int:
-    return get_distributed_class().run_worker()
+    return get_udf_distributor_class().run_worker()
 
 
 class UDFDispatcher:
@@ -164,14 +158,14 @@ class UDFDispatcher:
     def run_udf_parallel(  # noqa: C901, PLR0912
         self,
         input_rows: Iterable[RowsOutput],
-        total_rows: int,
+        rows_total: int,
         n_workers: Optional[int] = None,
         processed_cb: Callback = DEFAULT_CALLBACK,
         download_cb: Callback = DEFAULT_CALLBACK,
     ) -> None:
         n_workers = get_n_workers_from_arg(n_workers)
 
-        input_batch_size = total_rows // n_workers
+        input_batch_size = rows_total // n_workers
         if input_batch_size == 0:
             input_batch_size = 1
         elif input_batch_size > DEFAULT_BATCH_SIZE:

datachain-0.14.3/src/datachain/query/udf.py

@@ -0,0 +1,49 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Callable, Optional, TypedDict, Union
+
+if TYPE_CHECKING:
+    from sqlalchemy import Select, Table
+
+    from datachain.catalog import Catalog
+    from datachain.query.batch import BatchingStrategy
+
+
+class UdfInfo(TypedDict):
+    udf_data: bytes
+    catalog_init: dict[str, Any]
+    metastore_clone_params: tuple[Callable[..., Any], list[Any], dict[str, Any]]
+    warehouse_clone_params: tuple[Callable[..., Any], list[Any], dict[str, Any]]
+    table: "Table"
+    query: "Select"
+    udf_fields: list[str]
+    batching: "BatchingStrategy"
+    processes: Optional[int]
+    is_generator: bool
+    cache: bool
+    rows_total: int
+
+
+class AbstractUDFDistributor(ABC):
+    @abstractmethod
+    def __init__(
+        self,
+        catalog: "Catalog",
+        table: "Table",
+        query: "Select",
+        udf_data: bytes,
+        batching: "BatchingStrategy",
+        workers: Union[bool, int],
+        processes: Union[bool, int],
+        udf_fields: list[str],
+        rows_total: int,
+        use_cache: bool,
+        is_generator: bool = False,
+        min_task_size: Optional[Union[str, int]] = None,
+    ) -> None: ...
+
+    @abstractmethod
+    def __call__(self) -> None: ...
+
+    @staticmethod
+    @abstractmethod
+    def run_worker() -> int: ...

{datachain-0.14.2 → datachain-0.14.3}/src/datachain/utils.py

@@ -286,15 +286,41 @@ def retry_with_backoff(retries=5, backoff_sec=1, errors=(Exception,)):
     return retry
 
 
-def determine_processes(parallel: Optional[Union[bool, int]]) -> Union[bool, int]:
+def determine_workers(
+    workers: Union[bool, int],
+    rows_total: Optional[int] = None,
+) -> Union[bool, int]:
+    """Determine the number of workers to use for distributed processing."""
+    if rows_total is not None and rows_total <= 1:
+        # Disable distributed processing if there is no rows or only one row.
+        return False
+    if (
+        workers is False
+        and os.environ.get("DATACHAIN_DISTRIBUTED")
+        and os.environ.get("DATACHAIN_SETTINGS_WORKERS")
+    ):
+        # Enable distributed processing by default if the module is available,
+        # and a default number of workers is provided.
+        workers = int(os.environ["DATACHAIN_SETTINGS_WORKERS"])
+    if not workers or workers <= 0:
+        return False
+    return workers
+
+
+def determine_processes(
+    parallel: Optional[Union[bool, int]] = None,
+    rows_total: Optional[int] = None,
+) -> Union[bool, int]:
+    """Determine the number of processes to use for parallel processing."""
+    if rows_total is not None and rows_total <= 1:
+        # Disable parallel processing if there is no rows or only one row.
+        return False
     if parallel is None and os.environ.get("DATACHAIN_SETTINGS_PARALLEL") is not None:
         parallel = int(os.environ["DATACHAIN_SETTINGS_PARALLEL"])
-    if parallel is None or parallel is False:
+    if parallel is None or parallel is False or parallel == 0:
         return False
     if parallel is True:
         return True
-    if parallel == 0:
-        return False
     if parallel < 0:
         return True
     return parallel

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

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

{datachain-0.14.2 → datachain-0.14.3}/tests/unit/lib/test_datachain.py

@@ -4,7 +4,7 @@ import math
 import os
 import re
 from collections.abc import Generator, Iterator
-from unittest.mock import ANY
+from unittest.mock import ANY, patch
 
 import numpy as np
 import pandas as pd
@@ -26,6 +26,7 @@ from datachain.lib.signal_schema import (
     SignalResolvingTypeError,
     SignalSchema,
 )
+from datachain.lib.udf import UDFAdapter
 from datachain.lib.udf_signature import UdfSignatureError
 from datachain.lib.utils import DataChainColumnError, DataChainParamsError
 from datachain.sql.types import Float, Int64, String
@@ -270,6 +271,17 @@ def test_read_record_empty_chain_without_schema(test_session):
     )
 
 
+def test_empty_chain_skip_udf_run(test_session):
+    # Test that UDF is not called for empty chain
+    with patch.object(UDFAdapter, "run") as mock_udf_run:
+        (
+            dc.read_records([], schema={"val": int}, session=test_session)
+            .map(lambda val: val * 2, params="val", output={"val2": int})
+            .exec()
+        )
+        mock_udf_run.assert_not_called()
+
+
 def test_datasets(test_session):
     ds = dc.datasets(session=test_session)
     datasets = [d for d in ds.collect("dataset") if d.name == "fibonacci"]

{datachain-0.14.2 → datachain-0.14.3}/tests/unit/test_catalog_loader.py

@@ -5,8 +5,8 @@ import pytest
 
 from datachain.catalog.loader import (
     get_catalog,
-    get_distributed_class,
     get_metastore,
+    get_udf_distributor_class,
     get_warehouse,
 )
 from datachain.data_storage.sqlite import (
@@ -44,7 +44,7 @@ def test_get_metastore(sqlite_db):
 def test_get_metastore_in_memory():
     if os.environ.get("DATACHAIN_METASTORE"):
         with pytest.raises(RuntimeError):
-            metastore = get_metastore(in_memory=True)
+            get_metastore(in_memory=True)
     else:
         metastore = get_metastore(in_memory=True)
         assert isinstance(metastore, SQLiteMetastore)
@@ -71,7 +71,7 @@ def test_get_warehouse(sqlite_db):
 def test_get_warehouse_in_memory():
     if os.environ.get("DATACHAIN_WAREHOUSE"):
         with pytest.raises(RuntimeError):
-            warehouse = get_warehouse(in_memory=True)
+            get_warehouse(in_memory=True)
     else:
         warehouse = get_warehouse(in_memory=True)
         assert isinstance(warehouse, SQLiteWarehouse)
@@ -80,38 +80,31 @@ def test_get_warehouse_in_memory():
 
 
 def test_get_distributed_class():
-    distributed_args = {"foo": "bar", "baz": "37", "empty": ""}
     env = {
         "DATACHAIN_DISTRIBUTED": "tests.unit.test_catalog_loader.DistributedClass",
-        "DATACHAIN_DISTRIBUTED_ARG_FOO": "bar",
-        "DATACHAIN_DISTRIBUTED_ARG_BAZ": "37",
-        "DATACHAIN_DISTRIBUTED_ARG_EMPTY": "",
     }
 
     with patch.dict(os.environ, env):
-        distributed = get_distributed_class()
-        assert distributed
-        assert isinstance(distributed, DistributedClass)
-        assert distributed.kwargs == distributed_args
+        assert get_udf_distributor_class() == DistributedClass
 
     with patch.dict(os.environ, {"DATACHAIN_DISTRIBUTED": ""}):
         with pytest.raises(
             RuntimeError, match="DATACHAIN_DISTRIBUTED import path is required"
         ):
-            get_distributed_class()
+            get_udf_distributor_class()
 
     with patch.dict(
         os.environ,
         {"DATACHAIN_DISTRIBUTED": "tests.unit.test_catalog_loader.NonExistent"},
     ):
         with pytest.raises(AttributeError, match="has no attribute 'NonExistent'"):
-            get_distributed_class()
+            get_udf_distributor_class()
 
     with patch.dict(os.environ, {"DATACHAIN_DISTRIBUTED": "DistributionClass"}):
         with pytest.raises(
             RuntimeError, match="Invalid DATACHAIN_DISTRIBUTED import path"
         ):
-            get_distributed_class()
+            get_udf_distributor_class()
 
 
 def test_get_catalog(sqlite_db):

{datachain-0.14.2 → datachain-0.14.3}/tests/unit/test_utils.py

@@ -5,6 +5,7 @@ import pytest
 from datachain.utils import (
     datachain_paths_join,
     determine_processes,
+    determine_workers,
     nested_dict_path_set,
     retry_with_backoff,
     row_to_nested_dict,
@@ -141,21 +142,56 @@ def test_retry_with_backoff():
 
 
 @pytest.mark.parametrize(
-    "parallel,settings,expected",
+    "workers,rows_total,settings,expected",
     (
-        (None, None, False),
-        (None, "-1", True),
-        (None, "0", False),
-        (None, "5", 5),
-        (-1, "5", True),
-        (0, "5", False),
-        (10, "5", 10),
+        (False, None, None, False),
+        (False, None, "-1", False),
+        (False, None, "0", False),
+        (False, None, "5", 5),
+        (-1, None, "5", False),
+        (0, None, "5", False),
+        (10, None, "5", 10),
+        (True, None, None, True),
+        (True, None, "5", True),
+        (10, 0, None, False),
+        (10, 1, None, False),
+        (10, 2, None, 10),
+        (True, 0, None, False),
+        (True, 1, None, False),
+        (True, 2, None, True),
     ),
 )
-def test_determine_processes(parallel, settings, expected):
+def test_determine_workers(workers, rows_total, settings, expected):
+    if settings is not None:
+        os.environ["DATACHAIN_DISTRIBUTED"] = "some_defined_value"
+        os.environ["DATACHAIN_SETTINGS_WORKERS"] = settings
+    assert determine_workers(workers, rows_total=rows_total) == expected
+
+
+@pytest.mark.parametrize(
+    "parallel,rows_total,settings,expected",
+    (
+        (None, None, None, False),
+        (None, None, "-1", True),
+        (None, None, "0", False),
+        (None, None, "5", 5),
+        (-1, None, "5", True),
+        (0, None, "5", False),
+        (10, None, "5", 10),
+        (True, None, None, True),
+        (True, None, "5", True),
+        (10, 0, None, False),
+        (10, 1, None, False),
+        (10, 2, None, 10),
+        (True, 0, None, False),
+        (True, 1, None, False),
+        (True, 2, None, True),
+    ),
+)
+def test_determine_processes(parallel, rows_total, settings, expected):
     if settings is not None:
         os.environ["DATACHAIN_SETTINGS_PARALLEL"] = settings
-    assert determine_processes(parallel) == expected
+    assert determine_processes(parallel, rows_total=rows_total) == expected
 
 
 @pytest.mark.parametrize(

datachain-0.14.2/src/datachain/query/udf.py

@@ -1,20 +0,0 @@
-from typing import TYPE_CHECKING, Any, Callable, Optional, TypedDict
-
-if TYPE_CHECKING:
-    from sqlalchemy import Select, Table
-
-    from datachain.query.batch import BatchingStrategy
-
-
-class UdfInfo(TypedDict):
-    udf_data: bytes
-    catalog_init: dict[str, Any]
-    metastore_clone_params: tuple[Callable[..., Any], list[Any], dict[str, Any]]
-    warehouse_clone_params: tuple[Callable[..., Any], list[Any], dict[str, Any]]
-    table: "Table"
-    query: "Select"
-    udf_fields: list[str]
-    batching: "BatchingStrategy"
-    processes: Optional[int]
-    is_generator: bool
-    cache: bool