datachain 0.18.5__tar.gz → 0.18.6__tar.gz

{datachain-0.18.5/src/datachain.egg-info → datachain-0.18.6}/PKG-INFO

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

{datachain-0.18.5 → datachain-0.18.6}/src/datachain/query/dispatch.py

@@ -264,9 +264,10 @@ class UDFDispatcher:
             # Will be set to True when the input is exhausted
             input_finished = False
 
-            if not self.is_batching:
-                batch_size = self.input_batch_size(n_workers)
-                input_rows = batched(flatten(input_rows), batch_size)
+            input_rows = batched(
+                input_rows if self.is_batching else flatten(input_rows),
+                self.input_batch_size(n_workers),
+            )
 
             # Stop all workers after the input rows have finished processing
             input_data = chain(input_rows, [STOP_SIGNAL] * n_workers)

{datachain-0.18.5 → datachain-0.18.6/src/datachain.egg-info}/PKG-INFO

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

{datachain-0.18.5 → datachain-0.18.6}/tests/func/test_datachain.py

@@ -8,7 +8,7 @@ import re
 import uuid
 from collections.abc import Iterator
 from datetime import datetime, timedelta, timezone
-from pathlib import Path
+from pathlib import Path, PurePosixPath
 from unittest.mock import patch
 
 import numpy as np
@@ -2235,3 +2235,44 @@ def test_datachain_functional_after_exceptions(test_session):
     for _ in range(4):
         with pytest.raises(Exception, match="Test Error!"):
             chain.map(res=func).exec()
+
+
+@pytest.mark.parametrize("parallel", [1, 2])
+def test_agg(catalog_tmpfile, parallel):
+    from datachain import func
+
+    session = catalog_tmpfile.session
+
+    def process(files: list[str]) -> Iterator[tuple[str, int]]:
+        yield str(PurePosixPath(files[0]).parent), len(files)
+
+    ds = (
+        dc.read_values(
+            filename=(
+                "cats/cat1",
+                "cats/cat2",
+                "dogs/dog1",
+                "dogs/dog2",
+                "dogs/dog3",
+                "dogs/others/dog4",
+            ),
+            session=session,
+        )
+        .settings(parallel=parallel)
+        .agg(
+            process,
+            params=["filename"],
+            output={"parent": str, "count": int},
+            partition_by=func.path.parent("filename"),
+        )
+        .save("my-ds")
+    )
+
+    assert sorted_dicts(ds.to_records(), "parent") == sorted_dicts(
+        [
+            {"parent": "cats", "count": 2},
+            {"parent": "dogs", "count": 3},
+            {"parent": "dogs/others", "count": 1},
+        ],
+        "parent",
+    )

{datachain-0.18.5 → datachain-0.18.6}/tests/unit/lib/test_datachain.py

@@ -681,7 +681,7 @@ def test_agg_two_params(test_session):
     assert list(ds.order_by("x.cnt").collect("x.cnt")) == [7, 20]
 
 
-def test_agg_simple_iiterator(test_session):
+def test_agg_simple_iterator(test_session):
     def func(key, val) -> Iterator[tuple[File, str]]:
         for i in range(val):
             yield File(path=""), f"{key}_{i}"