datachain 0.28.1__tar.gz → 0.28.2__tar.gz

{datachain-0.28.1 → datachain-0.28.2}/.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.12.5'
+    rev: 'v0.12.7'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.28.1 → datachain-0.28.2}/PKG-INFO

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

{datachain-0.28.1 → datachain-0.28.2}/src/datachain/lib/dc/datachain.py

@@ -324,6 +324,7 @@ class DataChain:
         sys: Optional[bool] = None,
         namespace: Optional[str] = None,
         project: Optional[str] = None,
+        batch_rows: Optional[int] = None,
     ) -> "Self":
         """Change settings for chain.
 
@@ -331,22 +332,24 @@ class DataChain:
         It returns chain, so, it can be chained later with next operation.
 
         Parameters:
-            cache : data caching (default=False)
+            cache : data caching. (default=False)
             parallel : number of thread for processors. True is a special value to
-                enable all available CPUs (default=1)
+                enable all available CPUs. (default=1)
             workers : number of distributed workers. Only for Studio mode. (default=1)
-            min_task_size : minimum number of tasks (default=1)
-            prefetch: number of workers to use for downloading files in advance.
+            min_task_size : minimum number of tasks. (default=1)
+            prefetch : number of workers to use for downloading files in advance.
                       This is enabled by default and uses 2 workers.
                       To disable prefetching, set it to 0.
-            namespace: namespace name.
-            project: project name.
+            namespace : namespace name.
+            project : project name.
+            batch_rows : row limit per insert to balance speed and memory usage.
+                      (default=2000)
 
         Example:
             ```py
             chain = (
                 chain
-                .settings(cache=True, parallel=8)
+                .settings(cache=True, parallel=8, batch_rows=300)
                 .map(laion=process_webdataset(spec=WDSLaion), params="file")
             )
             ```
@@ -356,7 +359,14 @@ class DataChain:
         settings = copy.copy(self._settings)
         settings.add(
             Settings(
-                cache, parallel, workers, min_task_size, prefetch, namespace, project
+                cache,
+                parallel,
+                workers,
+                min_task_size,
+                prefetch,
+                namespace,
+                project,
+                batch_rows,
             )
         )
         return self._evolve(settings=settings, _sys=sys)
@@ -711,7 +721,7 @@ class DataChain:
 
         return self._evolve(
             query=self._query.add_signals(
-                udf_obj.to_udf_wrapper(),
+                udf_obj.to_udf_wrapper(self._settings.batch_rows),
                 **self._settings.to_dict(),
             ),
             signal_schema=self.signals_schema | udf_obj.output,
@@ -749,7 +759,7 @@ class DataChain:
             udf_obj.prefetch = prefetch
         return self._evolve(
             query=self._query.generate(
-                udf_obj.to_udf_wrapper(),
+                udf_obj.to_udf_wrapper(self._settings.batch_rows),
                 **self._settings.to_dict(),
             ),
             signal_schema=udf_obj.output,
@@ -885,7 +895,7 @@ class DataChain:
         udf_obj = self._udf_to_obj(Aggregator, func, params, output, signal_map)
         return self._evolve(
             query=self._query.generate(
-                udf_obj.to_udf_wrapper(),
+                udf_obj.to_udf_wrapper(self._settings.batch_rows),
                 partition_by=processed_partition_by,
                 **self._settings.to_dict(),
             ),
@@ -917,11 +927,24 @@ class DataChain:
             )
             chain.save("new_dataset")
             ```
+
+        .. deprecated:: 0.29.0
+            This method is deprecated and will be removed in a future version.
+            Use `agg()` instead, which provides the similar functionality.
         """
+        import warnings
+
+        warnings.warn(
+            "batch_map() is deprecated and will be removed in a future version. "
+            "Use agg() instead, which provides the similar functionality.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         udf_obj = self._udf_to_obj(BatchMapper, func, params, output, signal_map)
+
         return self._evolve(
             query=self._query.add_signals(
-                udf_obj.to_udf_wrapper(batch),
+                udf_obj.to_udf_wrapper(self._settings.batch_rows, batch=batch),
                 **self._settings.to_dict(),
             ),
             signal_schema=self.signals_schema | udf_obj.output,
@@ -2340,7 +2363,7 @@ class DataChain:
     def setup(self, **kwargs) -> "Self":
         """Setup variables to pass to UDF functions.
 
-        Use before running map/gen/agg/batch_map to save an object and pass it as an
+        Use before running map/gen/agg to save an object and pass it as an
         argument to the UDF.
 
         The value must be a callable (a `lambda: <value>` syntax can be used to quickly

{datachain-0.28.1 → datachain-0.28.2}/src/datachain/lib/dc/records.py

@@ -15,6 +15,8 @@ if TYPE_CHECKING:
 
     P = ParamSpec("P")
 
+READ_RECORDS_BATCH_SIZE = 10000
+
 
 def read_records(
     to_insert: Optional[Union[dict, Iterable[dict]]],
@@ -41,7 +43,7 @@ def read_records(
     Notes:
         This call blocks until all records are inserted.
     """
-    from datachain.query.dataset import INSERT_BATCH_SIZE, adjust_outputs, get_col_types
+    from datachain.query.dataset import adjust_outputs, get_col_types
     from datachain.sql.types import SQLType
     from datachain.utils import batched
 
@@ -94,7 +96,7 @@ def read_records(
         {c.name: c.type for c in columns if isinstance(c.type, SQLType)},
     )
     records = (adjust_outputs(warehouse, record, col_types) for record in to_insert)
-    for chunk in batched(records, INSERT_BATCH_SIZE):
+    for chunk in batched(records, READ_RECORDS_BATCH_SIZE):
         warehouse.insert_rows(table, chunk)
     warehouse.insert_rows_done(table)
     return read_dataset(name=dsr.full_name, session=session, settings=settings)

{datachain-0.28.1 → datachain-0.28.2}/src/datachain/lib/settings.py

@@ -1,4 +1,5 @@
 from datachain.lib.utils import DataChainParamsError
+from datachain.utils import DEFAULT_CHUNK_ROWS
 
 
 class SettingsError(DataChainParamsError):
@@ -16,6 +17,7 @@ class Settings:
         prefetch=None,
         namespace=None,
         project=None,
+        batch_rows=None,
     ):
         self._cache = cache
         self.parallel = parallel
@@ -24,6 +26,7 @@ class Settings:
         self.prefetch = prefetch
         self.namespace = namespace
         self.project = project
+        self._chunk_rows = batch_rows
 
         if not isinstance(cache, bool) and cache is not None:
             raise SettingsError(
@@ -53,6 +56,18 @@ class Settings:
                 f", {min_task_size.__class__.__name__} was given"
             )
 
+        if batch_rows is not None and not isinstance(batch_rows, int):
+            raise SettingsError(
+                "'batch_rows' argument must be int or None"
+                f", {batch_rows.__class__.__name__} was given"
+            )
+
+        if batch_rows is not None and batch_rows <= 0:
+            raise SettingsError(
+                "'batch_rows' argument must be positive integer"
+                f", {batch_rows} was given"
+            )
+
     @property
     def cache(self):
         return self._cache if self._cache is not None else False
@@ -61,6 +76,10 @@ class Settings:
     def workers(self):
         return self._workers if self._workers is not None else False
 
+    @property
+    def batch_rows(self):
+        return self._chunk_rows if self._chunk_rows is not None else DEFAULT_CHUNK_ROWS
+
     def to_dict(self):
         res = {}
         if self._cache is not None:
@@ -75,6 +94,8 @@ class Settings:
             res["namespace"] = self.namespace
         if self.project is not None:
             res["project"] = self.project
+        if self._chunk_rows is not None:
+            res["batch_rows"] = self._chunk_rows
         return res
 
     def add(self, settings: "Settings"):
@@ -86,3 +107,5 @@ class Settings:
         self.project = settings.project or self.project
         if settings.prefetch is not None:
             self.prefetch = settings.prefetch
+        if settings._chunk_rows is not None:
+            self._chunk_rows = settings._chunk_rows

{datachain-0.28.1 → datachain-0.28.2}/src/datachain/lib/udf.py

@@ -62,19 +62,21 @@ class UDFProperties:
         return self.udf.get_batching(use_partitioning)
 
     @property
-    def batch(self):
-        return self.udf.batch
+    def batch_rows(self):
+        return self.udf.batch_rows
 
 
 @attrs.define(slots=False)
 class UDFAdapter:
     inner: "UDFBase"
     output: UDFOutputSpec
+    batch_rows: Optional[int] = None
     batch: int = 1
 
     def get_batching(self, use_partitioning: bool = False) -> BatchingStrategy:
         if use_partitioning:
             return Partition()
+
         if self.batch == 1:
             return NoBatching()
         if self.batch > 1:
@@ -233,10 +235,15 @@ class UDFBase(AbstractUDF):
     def signal_names(self) -> Iterable[str]:
         return self.output.to_udf_spec().keys()
 
-    def to_udf_wrapper(self, batch: int = 1) -> UDFAdapter:
+    def to_udf_wrapper(
+        self,
+        batch_rows: Optional[int] = None,
+        batch: int = 1,
+    ) -> UDFAdapter:
         return UDFAdapter(
             self,
             self.output.to_udf_spec(),
+            batch_rows,
             batch,
         )
 
@@ -418,11 +425,27 @@ class Mapper(UDFBase):
 
 
 class BatchMapper(UDFBase):
-    """Inherit from this class to pass to `DataChain.batch_map()`."""
+    """Inherit from this class to pass to `DataChain.batch_map()`.
+
+    .. deprecated:: 0.29.0
+        This class is deprecated and will be removed in a future version.
+        Use `Aggregator` instead, which provides the similar functionality.
+    """
 
     is_input_batched = True
     is_output_batched = True
 
+    def __init__(self):
+        import warnings
+
+        warnings.warn(
+            "BatchMapper is deprecated and will be removed in a future version. "
+            "Use Aggregator instead, which provides the similar functionality.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        super().__init__()
+
     def run(
         self,
         udf_fields: Sequence[str],

{datachain-0.28.1 → datachain-0.28.2}/src/datachain/query/dataset.py

@@ -333,32 +333,24 @@ def process_udf_outputs(
     udf_table: "Table",
     udf_results: Iterator[Iterable["UDFResult"]],
     udf: "UDFAdapter",
-    batch_size: int = INSERT_BATCH_SIZE,
     cb: Callback = DEFAULT_CALLBACK,
 ) -> None:
-    import psutil
-
-    rows: list[UDFResult] = []
     # Optimization: Compute row types once, rather than for every row.
     udf_col_types = get_col_types(warehouse, udf.output)
+    batch_rows = udf.batch_rows or INSERT_BATCH_SIZE
 
-    for udf_output in udf_results:
-        if not udf_output:
-            continue
-        with safe_closing(udf_output):
-            for row in udf_output:
-                cb.relative_update()
-                rows.append(adjust_outputs(warehouse, row, udf_col_types))
-                if len(rows) >= batch_size or (
-                    len(rows) % 10 == 0 and psutil.virtual_memory().percent > 80
-                ):
-                    for row_chunk in batched(rows, batch_size):
-                        warehouse.insert_rows(udf_table, row_chunk)
-                    rows.clear()
+    def _insert_rows():
+        for udf_output in udf_results:
+            if not udf_output:
+                continue
+
+            with safe_closing(udf_output):
+                for row in udf_output:
+                    cb.relative_update()
+                    yield adjust_outputs(warehouse, row, udf_col_types)
 
-    if rows:
-        for row_chunk in batched(rows, batch_size):
-            warehouse.insert_rows(udf_table, row_chunk)
+    for row_chunk in batched(_insert_rows(), batch_rows):
+        warehouse.insert_rows(udf_table, row_chunk)
 
     warehouse.insert_rows_done(udf_table)
 
@@ -401,6 +393,7 @@ class UDFStep(Step, ABC):
     min_task_size: Optional[int] = None
     is_generator = False
     cache: bool = False
+    batch_rows: Optional[int] = None
 
     @abstractmethod
     def create_udf_table(self, query: Select) -> "Table":
@@ -602,6 +595,7 @@ class UDFStep(Step, ABC):
                 parallel=self.parallel,
                 workers=self.workers,
                 min_task_size=self.min_task_size,
+                batch_rows=self.batch_rows,
             )
         return self.__class__(self.udf, self.catalog)
 
@@ -1633,6 +1627,7 @@ class DatasetQuery:
         min_task_size: Optional[int] = None,
         partition_by: Optional[PartitionByType] = None,
         cache: bool = False,
+        batch_rows: Optional[int] = None,
     ) -> "Self":
         """
         Adds one or more signals based on the results from the provided UDF.
@@ -1658,6 +1653,7 @@ class DatasetQuery:
                 workers=workers,
                 min_task_size=min_task_size,
                 cache=cache,
+                batch_rows=batch_rows,
             )
         )
         return query
@@ -1679,6 +1675,7 @@ class DatasetQuery:
         namespace: Optional[str] = None,
         project: Optional[str] = None,
         cache: bool = False,
+        batch_rows: Optional[int] = None,
     ) -> "Self":
         query = self.clone()
         steps = query.steps
@@ -1691,6 +1688,7 @@ class DatasetQuery:
                 workers=workers,
                 min_task_size=min_task_size,
                 cache=cache,
+                batch_rows=batch_rows,
             )
         )
         return query

{datachain-0.28.1 → datachain-0.28.2}/src/datachain/utils.py

@@ -11,7 +11,6 @@ import time
 from collections.abc import Iterable, Iterator, Sequence
 from contextlib import contextmanager
 from datetime import date, datetime, timezone
-from itertools import chain, islice
 from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
 from uuid import UUID
 
@@ -26,6 +25,8 @@ if TYPE_CHECKING:
     from typing_extensions import Self
 
 
+DEFAULT_CHUNK_ROWS = 2000
+
 logger = logging.getLogger("datachain")
 
 NUL = b"\0"
@@ -225,30 +226,44 @@ def get_envs_by_prefix(prefix: str) -> dict[str, str]:
 _T_co = TypeVar("_T_co", covariant=True)
 
 
-def batched(iterable: Iterable[_T_co], n: int) -> Iterator[tuple[_T_co, ...]]:
-    """Batch data into tuples of length n. The last batch may be shorter."""
-    # Based on: https://docs.python.org/3/library/itertools.html#itertools-recipes
-    # batched('ABCDEFG', 3) --> ABC DEF G
-    if n < 1:
-        raise ValueError("Batch size must be at least one")
-    it = iter(iterable)
-    while batch := tuple(islice(it, n)):
+def _dynamic_batched_core(
+    iterable: Iterable[_T_co],
+    batch_rows: int,
+) -> Iterator[list[_T_co]]:
+    """Core batching logic that yields lists."""
+
+    batch: list[_T_co] = []
+
+    for item in iterable:
+        # Check if adding this item would exceed limits
+        if len(batch) >= batch_rows and batch:  # Yield current batch if we have one
+            yield batch
+            batch = []
+
+        batch.append(item)
+
+    # Yield any remaining items
+    if batch:
         yield batch
 
 
-def batched_it(iterable: Iterable[_T_co], n: int) -> Iterator[Iterator[_T_co]]:
-    """Batch data into iterators of length n. The last batch may be shorter."""
-    # batched('ABCDEFG', 3) --> ABC DEF G
-    if n < 1:
-        raise ValueError("Batch size must be at least one")
-    it = iter(iterable)
-    while True:
-        chunk_it = islice(it, n)
-        try:
-            first_el = next(chunk_it)
-        except StopIteration:
-            return
-        yield chain((first_el,), chunk_it)
+def batched(iterable: Iterable[_T_co], batch_rows: int) -> Iterator[tuple[_T_co, ...]]:
+    """
+    Batch data into tuples of length batch_rows .
+    The last batch may be shorter.
+    """
+    yield from (tuple(batch) for batch in _dynamic_batched_core(iterable, batch_rows))
+
+
+def batched_it(
+    iterable: Iterable[_T_co],
+    batch_rows: int = DEFAULT_CHUNK_ROWS,
+) -> Iterator[Iterator[_T_co]]:
+    """
+    Batch data into iterators with dynamic sizing
+    based on row count and memory usage.
+    """
+    yield from (iter(batch) for batch in _dynamic_batched_core(iterable, batch_rows))
 
 
 def flatten(items):

{datachain-0.28.1 → datachain-0.28.2}/src/datachain.egg-info/PKG-INFO

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

{datachain-0.28.1 → datachain-0.28.2}/src/datachain.egg-info/SOURCES.txt

@@ -379,6 +379,7 @@ tests/unit/lib/test_partition_by.py
 tests/unit/lib/test_project.py
 tests/unit/lib/test_python_to_sql.py
 tests/unit/lib/test_schema.py
+tests/unit/lib/test_settings.py
 tests/unit/lib/test_signal_schema.py
 tests/unit/lib/test_sql_to_python.py
 tests/unit/lib/test_text.py

{datachain-0.28.1 → datachain-0.28.2}/tests/func/test_datachain.py

@@ -2420,3 +2420,22 @@ def test_agg_sample(catalog_tmpfile, parallel, sample):
     records = list(ds.to_records())
     assert len(records) == sample
     assert all(row["count"] == 1 for row in records)
+
+
+def test_batch_for_map(test_session):
+    # Create a chain with batch settings
+    chain = dc.read_values(x=list(range(100)), session=test_session)
+    chain_with_settings = chain.settings(batch_rows=15)
+
+    def add_one(x):
+        return x + 1
+
+    result = chain_with_settings.map(add_one, output={"result": int})
+
+    results = [
+        r[0] for r in result.to_iter("result")
+    ]  # Access the first element of each tuple
+
+    assert len(results) == 100
+
+    assert set(results) == set(range(1, 101))

{datachain-0.28.1 → datachain-0.28.2}/tests/func/test_hf.py

@@ -119,6 +119,7 @@ def test_hf_image(tmp_path):
     assert row.image.img == image_to_bytes(img)
 
 
+@pytest.mark.skip("fails with 'NotImplementedError', need to investigate")
 @require_torchcodec
 def test_hf_audio(tmp_path):
     # See https://stackoverflow.com/questions/66191480/how-to-convert-a-numpy-array-to-a-mp3-file

datachain-0.28.2/tests/unit/lib/test_settings.py

@@ -0,0 +1,61 @@
+import pytest
+
+from datachain.lib.settings import Settings, SettingsError
+
+
+def test_settings_defaults_and_custom():
+    """Test Settings class with default and custom batch parameters."""
+    # Default values
+    settings = Settings()
+    assert settings.batch_rows == 2000
+
+    # Custom values
+    settings = Settings(batch_rows=500)
+    assert settings.batch_rows == 500
+
+    # to_dict method
+    d = settings.to_dict()
+    assert d["batch_rows"] == 500
+
+    # Chaining
+    s2 = settings
+    s3 = s2
+    assert s3.batch_rows == 500
+
+
+def test_settings_validation():
+    # Valid
+    settings = Settings(batch_rows=100)
+    assert settings.batch_rows == 100
+
+    # Invalid batch_rows
+    with pytest.raises(SettingsError):
+        Settings(batch_rows="invalid")
+
+    # Zero batch_rows
+    with pytest.raises(SettingsError):
+        Settings(batch_rows=0)
+
+
+def test_settings_add():
+    """Test Settings.add() method with batch_rows."""
+    # Create base settings
+    base_settings = Settings(batch_rows=1000)
+
+    # Create settings to add
+    add_settings = Settings(batch_rows=2000)
+
+    # Add settings
+    base_settings.add(add_settings)
+
+    # Verify that values from add_settings override base_settings
+    assert base_settings.batch_rows == 2000
+
+    # Test with None values (should not override)
+    base_settings = Settings(batch_rows=1000)
+    none_settings = Settings(batch_rows=None)
+
+    base_settings.add(none_settings)
+
+    # Verify that None values don't override existing values
+    assert base_settings.batch_rows == 1000

{datachain-0.28.1 → datachain-0.28.2}/tests/unit/test_utils.py

@@ -3,6 +3,7 @@ import os
 import pytest
 
 from datachain.utils import (
+    batched,
     datachain_paths_join,
     determine_processes,
     determine_workers,
@@ -253,3 +254,21 @@ def test_nested_dict_path_set(data, path, value, expected):
 )
 def test_row_to_nested_dict(headers, row, expected):
     assert row_to_nested_dict(headers, row) == expected
+
+
+def test_batched_basic():
+    """Test basic batching functionality."""
+    data = list(range(10))
+    batches = list(batched(data, 3))
+    assert batches == [(0, 1, 2), (3, 4, 5), (6, 7, 8), (9,)]
+
+
+def test_batched_row_limit():
+    """Test dynamic batching with row count limit."""
+    data = list(range(15))
+    batches = list(batched(data, batch_rows=4))
+    assert len(batches) == 4  # 15 items / 4 max = 4 batches
+    assert batches[0] == (0, 1, 2, 3)
+    assert batches[1] == (4, 5, 6, 7)
+    assert batches[2] == (8, 9, 10, 11)
+    assert batches[3] == (12, 13, 14)