datachain 0.24.0__tar.gz → 0.24.2__tar.gz

{datachain-0.24.0 → datachain-0.24.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.0'
+    rev: 'v0.12.1'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]
@@ -35,7 +35,7 @@ repos:
       - id: codespell
         additional_dependencies: ["tomli"]
   - repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks
-    rev: v2.14.0
+    rev: v2.15.0
     hooks:
       - id: pretty-format-toml
         args: [--autofix, --no-sort]

{datachain-0.24.0 → datachain-0.24.2}/PKG-INFO

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

{datachain-0.24.0 → datachain-0.24.2}/src/datachain/data_storage/warehouse.py

@@ -218,7 +218,7 @@ class AbstractWarehouse(ABC, Serializable):
         limit = query._limit
         paginated_query = query.limit(page_size)
 
-        offset = 0
+        offset = query._offset or 0
         num_yielded = 0
 
         # Ensure we're using a thread-local connection
@@ -234,13 +234,13 @@ class AbstractWarehouse(ABC, Serializable):
                 # Cursor results are not thread-safe, so we convert them to a list
                 results = list(wh.dataset_rows_select(paginated_query.offset(offset)))
 
-                processed = False
+                processed = 0
                 for row in results:
-                    processed = True
+                    processed += 1
                     yield row
                     num_yielded += 1
 
-                if not processed:
+                if processed < page_size:
                     break  # no more results
                 offset += page_size
 
@@ -343,6 +343,8 @@ class AbstractWarehouse(ABC, Serializable):
         if (id_col := get_query_id_column(query)) is None:
             raise RuntimeError("sys__id column not found in query")
 
+        query = query._clone().offset(None).limit(None).order_by(None)
+
         if is_batched:
             for batch in ids:
                 yield list(self.dataset_rows_select(query.where(id_col.in_(batch))))

{datachain-0.24.0 → datachain-0.24.2}/src/datachain/delta.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Union
 import datachain
 from datachain.dataset import DatasetDependency
 from datachain.error import DatasetNotFoundError
+from datachain.project import Project
 
 if TYPE_CHECKING:
     from typing_extensions import Concatenate, ParamSpec
@@ -50,15 +51,24 @@ def _append_steps(dc: "DataChain", other: "DataChain"):
 
 def _get_delta_chain(
     source_ds_name: str,
+    source_ds_project: Project,
     source_ds_version: str,
     source_ds_latest_version: str,
     on: Union[str, Sequence[str]],
     compare: Optional[Union[str, Sequence[str]]] = None,
 ) -> "DataChain":
     """Get delta chain for processing changes between versions."""
-    source_dc = datachain.read_dataset(source_ds_name, version=source_ds_version)
+    source_dc = datachain.read_dataset(
+        source_ds_name,
+        namespace=source_ds_project.namespace.name,
+        project=source_ds_project.name,
+        version=source_ds_version,
+    )
     source_dc_latest = datachain.read_dataset(
-        source_ds_name, version=source_ds_latest_version
+        source_ds_name,
+        namespace=source_ds_project.namespace.name,
+        project=source_ds_project.name,
+        version=source_ds_latest_version,
     )
 
     # Calculate diff between source versions
@@ -67,12 +77,15 @@ def _get_delta_chain(
 
 def _get_retry_chain(
     name: str,
+    project: Project,
     latest_version: str,
     source_ds_name: str,
-    source_ds_latest_version: str,
+    source_ds_project: Project,
+    source_ds_version: str,
     on: Union[str, Sequence[str]],
     right_on: Optional[Union[str, Sequence[str]]],
     delta_retry: Optional[Union[bool, str]],
+    diff_chain: "DataChain",
 ) -> Optional["DataChain"]:
     """Get retry chain for processing error records and missing records."""
     # Import here to avoid circular import
@@ -81,35 +94,49 @@ def _get_retry_chain(
     retry_chain = None
 
     # Read the latest version of the result dataset for retry logic
-    result_dataset = datachain.read_dataset(name, version=latest_version)
-    source_dc_latest = datachain.read_dataset(
-        source_ds_name, version=source_ds_latest_version
+    result_dataset = datachain.read_dataset(
+        name,
+        namespace=project.namespace.name,
+        project=project.name,
+        version=latest_version,
+    )
+    source_dc = datachain.read_dataset(
+        source_ds_name,
+        namespace=source_ds_project.namespace.name,
+        project=source_ds_project.name,
+        version=source_ds_version,
     )
 
     # Handle error records if delta_retry is a string (column name)
     if isinstance(delta_retry, str):
         error_records = result_dataset.filter(C(delta_retry) != "")
-        error_source_records = source_dc_latest.merge(
+        error_source_records = source_dc.merge(
             error_records, on=on, right_on=right_on, inner=True
-        ).select(*list(source_dc_latest.signals_schema.values))
+        ).select(*list(source_dc.signals_schema.values))
         retry_chain = error_source_records
 
     # Handle missing records if delta_retry is True
     elif delta_retry is True:
-        missing_records = source_dc_latest.subtract(
-            result_dataset, on=on, right_on=right_on
-        )
+        missing_records = source_dc.subtract(result_dataset, on=on, right_on=right_on)
         retry_chain = missing_records
 
-    return retry_chain
+    # Subtract also diff chain since some items might be picked
+    # up by `delta=True` itself (e.g. records got modified AND are missing in the
+    # result dataset atm)
+    return retry_chain.subtract(diff_chain, on=on) if retry_chain else None
 
 
 def _get_source_info(
     name: str,
+    project: Project,
     latest_version: str,
     catalog,
 ) -> tuple[
-    Optional[str], Optional[str], Optional[str], Optional[list[DatasetDependency]]
+    Optional[str],
+    Optional[Project],
+    Optional[str],
+    Optional[str],
+    Optional[list[DatasetDependency]],
 ]:
     """Get source dataset information and dependencies.
 
@@ -118,23 +145,34 @@ def _get_source_info(
         Returns (None, None, None, None) if source dataset was removed.
     """
     dependencies = catalog.get_dataset_dependencies(
-        name, latest_version, indirect=False
+        name, latest_version, project=project, indirect=False
     )
 
     dep = dependencies[0]
     if not dep:
         # Starting dataset was removed, back off to normal dataset creation
-        return None, None, None, None
+        return None, None, None, None, None
 
+    source_ds_project = catalog.metastore.get_project(dep.project, dep.namespace)
     source_ds_name = dep.name
     source_ds_version = dep.version
-    source_ds_latest_version = catalog.get_dataset(source_ds_name).latest_version
-
-    return source_ds_name, source_ds_version, source_ds_latest_version, dependencies
+    source_ds_latest_version = catalog.get_dataset(
+        source_ds_name, project=source_ds_project
+    ).latest_version
+
+    return (
+        source_ds_name,
+        source_ds_project,
+        source_ds_version,
+        source_ds_latest_version,
+        dependencies,
+    )
 
 
 def delta_retry_update(
     dc: "DataChain",
+    namespace_name: str,
+    project_name: str,
     name: str,
     on: Union[str, Sequence[str]],
     right_on: Optional[Union[str, Sequence[str]]] = None,
@@ -173,11 +211,12 @@ def delta_retry_update(
     """
 
     catalog = dc.session.catalog
+    project = catalog.metastore.get_project(project_name, namespace_name)
     dc._query.apply_listing_pre_step()
 
     # Check if dataset exists
     try:
-        dataset = catalog.get_dataset(name)
+        dataset = catalog.get_dataset(name, project=project)
         latest_version = dataset.latest_version
     except DatasetNotFoundError:
         # First creation of result dataset
@@ -189,19 +228,29 @@ def delta_retry_update(
     retry_chain = None
     processing_chain = None
 
-    source_ds_name, source_ds_version, source_ds_latest_version, dependencies = (
-        _get_source_info(name, latest_version, catalog)
-    )
+    (
+        source_ds_name,
+        source_ds_project,
+        source_ds_version,
+        source_ds_latest_version,
+        dependencies,
+    ) = _get_source_info(name, project, latest_version, catalog)
 
     # If source_ds_name is None, starting dataset was removed
     if source_ds_name is None:
         return None, None, True
 
+    assert source_ds_project
     assert source_ds_version
     assert source_ds_latest_version
 
     diff_chain = _get_delta_chain(
-        source_ds_name, source_ds_version, source_ds_latest_version, on, compare
+        source_ds_name,
+        source_ds_project,
+        source_ds_version,
+        source_ds_latest_version,
+        on,
+        compare,
     )
 
     # Filter out removed dep
@@ -215,12 +264,15 @@ def delta_retry_update(
     if delta_retry:
         retry_chain = _get_retry_chain(
             name,
+            project,
             latest_version,
             source_ds_name,
-            source_ds_latest_version,
+            source_ds_project,
+            source_ds_version,
             on,
             right_on,
             delta_retry,
+            diff_chain,
         )
 
     # Combine delta and retry chains
@@ -236,7 +288,12 @@ def delta_retry_update(
     if processing_chain is None or (processing_chain and processing_chain.empty):
         return None, None, False
 
-    latest_dataset = datachain.read_dataset(name, version=latest_version)
+    latest_dataset = datachain.read_dataset(
+        name,
+        namespace=project.namespace.name,
+        project=project.name,
+        version=latest_version,
+    )
     compared_chain = latest_dataset.diff(
         processing_chain,
         on=right_on or on,

{datachain-0.24.0 → datachain-0.24.2}/src/datachain/lib/dc/datachain.py

@@ -598,6 +598,8 @@ class DataChain:
 
             result_ds, dependencies, has_changes = delta_retry_update(
                 self,
+                namespace_name,
+                project_name,
                 name,
                 on=self._delta_on,
                 right_on=self._delta_result_on,

{datachain-0.24.0 → datachain-0.24.2}/src/datachain/query/dataset.py

@@ -11,6 +11,7 @@ from collections.abc import Generator, Iterable, Iterator, Sequence
 from copy import copy
 from functools import wraps
 from secrets import token_hex
+from types import GeneratorType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -557,8 +558,8 @@ class UDFStep(Step, ABC):
         """
         assert self.partition_by is not None
 
-        if isinstance(self.partition_by, Sequence):
-            list_partition_by = self.partition_by
+        if isinstance(self.partition_by, (list, tuple, GeneratorType)):
+            list_partition_by = list(self.partition_by)
         else:
             list_partition_by = [self.partition_by]
 
@@ -575,7 +576,10 @@ class UDFStep(Step, ABC):
             f.dense_rank().over(order_by=partition_by).label(PARTITION_COLUMN_ID),
         ]
         self.catalog.warehouse.db.execute(
-            tbl.insert().from_select(cols, query.with_only_columns(*cols))
+            tbl.insert().from_select(
+                cols,
+                query.offset(None).limit(None).with_only_columns(*cols),
+            )
         )
 
         return tbl
@@ -601,13 +605,10 @@ class UDFStep(Step, ABC):
         if self.partition_by is not None:
             partition_tbl = self.create_partitions_table(query)
             temp_tables.append(partition_tbl.name)
-
-            subq = query.subquery()
-            query = (
-                sqlalchemy.select(*subq.c)
-                .outerjoin(partition_tbl, partition_tbl.c.sys__id == subq.c.sys__id)
-                .add_columns(*partition_columns())
-            )
+            query = query.outerjoin(
+                partition_tbl,
+                partition_tbl.c.sys__id == query.selected_columns.sys__id,
+            ).add_columns(*partition_columns())
 
         query, tables = self.process_input_query(query)
         temp_tables.extend(t.name for t in tables)

{datachain-0.24.0 → datachain-0.24.2}/src/datachain.egg-info/PKG-INFO

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

{datachain-0.24.0 → datachain-0.24.2}/tests/func/test_datachain.py

@@ -2324,3 +2324,72 @@ def test_agg(catalog_tmpfile, parallel):
         ],
         "parent",
     )
+
+
+@pytest.mark.parametrize("parallel", [1, 2])
+@pytest.mark.parametrize(
+    "offset,limit,files",
+    [
+        (None, 1000, [f"file{i:02d}" for i in range(100)]),
+        (None, 3, ["file00", "file01", "file02"]),
+        (0, 3, ["file00", "file01", "file02"]),
+        (97, 1000, ["file97", "file98", "file99"]),
+        (1, 2, ["file01", "file02"]),
+        (50, 3, ["file50", "file51", "file52"]),
+        (None, 0, []),
+        (50, 0, []),
+    ],
+)
+def test_agg_offset_limit(catalog_tmpfile, parallel, offset, limit, files):
+    def process(filename: list[str]) -> Iterator[tuple[str, int]]:
+        yield filename[0], len(filename)
+
+    ds = dc.read_values(
+        filename=[f"file{i:02d}" for i in range(100)],
+        value=list(range(100)),
+        session=catalog_tmpfile.session,
+    )
+    if offset is not None:
+        ds = ds.offset(offset)
+    if limit is not None:
+        ds = ds.limit(limit)
+    ds = (
+        ds.settings(parallel=parallel)
+        .agg(
+            process,
+            output={"filename": str, "count": int},
+            partition_by="filename",
+        )
+        .save("my-ds")
+    )
+
+    records = list(ds.to_records())
+    assert len(records) == len(files)
+    assert all(row["count"] == 1 for row in records)
+    assert sorted(row["filename"] for row in records) == sorted(files)
+
+
+@pytest.mark.parametrize("parallel", [1, 2])
+@pytest.mark.parametrize("sample", [0, 1, 3, 10, 50, 100])
+def test_agg_sample(catalog_tmpfile, parallel, sample):
+    def process(filename: list[str]) -> Iterator[tuple[str, int]]:
+        yield filename[0], len(filename)
+
+    ds = (
+        dc.read_values(
+            filename=[f"file{i:02d}" for i in range(100)],
+            session=catalog_tmpfile.session,
+        )
+        .sample(sample)
+        .settings(parallel=parallel)
+        .agg(
+            process,
+            output={"filename": str, "count": int},
+            partition_by="filename",
+        )
+        .save("my-ds")
+    )
+
+    records = list(ds.to_records())
+    assert len(records) == sample
+    assert all(row["count"] == 1 for row in records)

{datachain-0.24.0 → datachain-0.24.2}/tests/func/test_delta.py

@@ -14,15 +14,26 @@ from datachain.lib.file import File, ImageFile
 def _get_dependencies(catalog, name, version) -> list[tuple[str, str]]:
     return sorted(
         [
-            (d.name, d.version)
+            (f"{d.namespace}.{d.project}.{d.name}", d.version)
             for d in catalog.get_dataset_dependencies(name, version, indirect=False)
         ]
     )
 
 
-def test_delta_update_from_dataset(test_session, tmp_dir, tmp_path):
+@pytest.mark.parametrize("project", ("global.dev", ""))
+def test_delta_update_from_dataset(test_session, tmp_dir, tmp_path, project):
     catalog = test_session.catalog
-    starting_ds_name = "starting_ds"
+    default_namespace_name = catalog.metastore.default_namespace_name
+    default_project_name = catalog.metastore.default_project_name
+
+    if project:
+        starting_ds_name = f"{project}.starting_ds"
+        dependency_ds_name = starting_ds_name
+    else:
+        starting_ds_name = "starting_ds"
+        dependency_ds_name = (
+            f"{default_namespace_name}.{default_project_name}.{starting_ds_name}"
+        )
     ds_name = "delta_ds"
 
     images = [
@@ -55,12 +66,16 @@ def test_delta_update_from_dataset(test_session, tmp_dir, tmp_path):
     create_image_dataset(starting_ds_name, images[:2])
     # first version of delta dataset
     create_delta_dataset(ds_name)
-    assert _get_dependencies(catalog, ds_name, "1.0.0") == [(starting_ds_name, "1.0.0")]
+    assert _get_dependencies(catalog, ds_name, "1.0.0") == [
+        (dependency_ds_name, "1.0.0")
+    ]
     # second version of starting dataset
     create_image_dataset(starting_ds_name, images[2:])
     # second version of delta dataset
     create_delta_dataset(ds_name)
-    assert _get_dependencies(catalog, ds_name, "1.0.1") == [(starting_ds_name, "1.0.1")]
+    assert _get_dependencies(catalog, ds_name, "1.0.1") == [
+        (dependency_ds_name, "1.0.1")
+    ]
 
     assert (dc.read_dataset(ds_name, version="1.0.0").order_by("file.path")).to_values(
         "file.path"