datachain 0.24.1__tar.gz → 0.24.3__tar.gz

{datachain-0.24.1 → datachain-0.24.3}/.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.1 → datachain-0.24.3}/PKG-INFO

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

{datachain-0.24.1 → datachain-0.24.3}/src/datachain/catalog/catalog.py

@@ -1098,9 +1098,18 @@ class Catalog:
     ) -> DatasetRecord:
         from datachain.lib.listing import is_listing_dataset
 
+        project = project or self.metastore.default_project
+
         if is_listing_dataset(name):
             project = self.metastore.listing_project
-        return self.metastore.get_dataset(name, project.id if project else None)
+
+        try:
+            return self.metastore.get_dataset(name, project.id if project else None)
+        except DatasetNotFoundError:
+            raise DatasetNotFoundError(
+                f"Dataset {name} not found in namespace {project.namespace.name}"
+                f" and project {project.name}"
+            ) from None
 
     def get_dataset_with_remote_fallback(
         self,
@@ -1124,7 +1133,7 @@ class Catalog:
             raise DatasetNotFoundError(
                 f"Dataset {name}"
                 + (f" version {version} " if version else " ")
-                + "not found"
+                + f"not found in namespace {namespace_name} and project {project_name}"
             )
 
         if pull_dataset:

{datachain-0.24.1 → datachain-0.24.3}/src/datachain/data_storage/metastore.py

@@ -1194,14 +1194,16 @@ class AbstractDBMetastore(AbstractMetastore):
         Gets a single dataset in project by dataset name.
         """
         project_id = project_id or self.default_project.id
+
         d = self._datasets
         query = self._base_dataset_query()
         query = query.where(d.c.name == name, d.c.project_id == project_id)  # type: ignore [attr-defined]
         ds = self._parse_dataset(self.db.execute(query, conn=conn))
         if not ds:
             raise DatasetNotFoundError(
-                f"Dataset {name} not found in project {project_id}"
+                f"Dataset {name} not found in project with id {project_id}"
             )
+
         return ds
 
     def remove_dataset_version(

{datachain-0.24.1 → datachain-0.24.3}/src/datachain/data_storage/sqlite.py

@@ -774,7 +774,15 @@ class SQLiteWarehouse(AbstractWarehouse):
         query: Select,
         progress_cb: Optional[Callable[[int], None]] = None,
     ) -> None:
-        if len(query._group_by_clause) > 0:
+        col_id = (
+            query.selected_columns.sys__id
+            if "sys__id" in query.selected_columns
+            else None
+        )
+
+        # If there is no sys__id column, we cannot copy the table in batches,
+        # and we need to copy all rows at once. Same if there is a group by clause.
+        if col_id is None or len(query._group_by_clause) > 0:
             select_q = query.with_only_columns(
                 *[c for c in query.selected_columns if c.name != "sys__id"]
             )
@@ -782,12 +790,7 @@ class SQLiteWarehouse(AbstractWarehouse):
             self.db.execute(q)
             return
 
-        if "sys__id" in query.selected_columns:
-            col_id = query.selected_columns.sys__id
-        else:
-            col_id = sqlalchemy.column("sys__id")
         select_ids = query.with_only_columns(col_id)
-
         ids = self.db.execute(select_ids).fetchall()
 
         select_q = (

{datachain-0.24.1 → datachain-0.24.3}/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.1 → datachain-0.24.3}/src/datachain/listing.py

@@ -65,10 +65,17 @@ class Listing:
 
     @cached_property
     def dataset(self) -> "DatasetRecord":
+        from datachain.error import DatasetNotFoundError
+
         assert self.dataset_name
-        return self.metastore.get_dataset(
-            self.dataset_name, self.metastore.listing_project.id
-        )
+        project = self.metastore.listing_project
+        try:
+            return self.metastore.get_dataset(self.dataset_name, project.id)
+        except DatasetNotFoundError:
+            raise DatasetNotFoundError(
+                f"Dataset {self.dataset_name} not found in namespace"
+                f" {project.namespace.name} and project {project.name}"
+            ) from None
 
     @cached_property
     def dataset_rows(self):

{datachain-0.24.1 → datachain-0.24.3}/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.1 → datachain-0.24.3}/src/datachain.egg-info/PKG-INFO

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

{datachain-0.24.1 → datachain-0.24.3}/tests/func/test_datachain.py

@@ -236,6 +236,22 @@ def test_read_storage_dependencies(cloud_test_catalog, cloud_type):
     assert dependencies[0].name == dep_name
 
 
+def test_persist_after_mutate(test_session):
+    chain = (
+        dc.read_values(fib=[1, 1, 2, 3, 5, 8, 13, 21], session=test_session)
+        .map(mod3=lambda fib: fib % 3, output=int)
+        .group_by(
+            cnt=dc.func.count(),
+            partition_by="mod3",
+        )
+        .mutate(x=1)
+        .persist()
+    )
+
+    assert chain.count() == 3
+    assert set(chain.to_values("mod3")) == {0, 1, 2}
+
+
 def test_persist_not_affects_dependencies(tmp_dir, test_session):
     for i in range(4):
         (tmp_dir / f"file{i}.txt").write_text(f"file{i}")
@@ -2324,3 +2340,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.1 → datachain-0.24.3}/tests/func/test_delta.py

@@ -248,6 +248,9 @@ def test_delta_update_check_num_calls(test_session, tmp_dir, tmp_path, capsys):
 
 
 def test_delta_update_no_diff(test_session, tmp_dir, tmp_path):
+    catalog = test_session.catalog
+    default_namespace_name = catalog.metastore.default_namespace_name
+    default_project_name = catalog.metastore.default_project_name
     ds_name = "delta_ds"
     path = tmp_dir.as_uri()
     tmp_dir = tmp_dir / "images"
@@ -296,7 +299,10 @@ def test_delta_update_no_diff(test_session, tmp_dir, tmp_path):
     with pytest.raises(DatasetNotFoundError) as exc_info:
         dc.read_dataset(ds_name, version="1.0.1")
 
-    assert str(exc_info.value) == f"Dataset {ds_name} version 1.0.1 not found"
+    assert str(exc_info.value) == (
+        f"Dataset {ds_name} version 1.0.1 not found in namespace "
+        f"{default_namespace_name} and project {default_project_name}"
+    )
 
 
 @pytest.fixture

{datachain-0.24.1 → datachain-0.24.3}/tests/unit/lib/test_datachain.py

@@ -16,6 +16,7 @@ from pydantic import BaseModel
 
 import datachain as dc
 from datachain import Column
+from datachain.data_storage import AbstractMetastore
 from datachain.error import (
     DatasetInvalidVersionError,
     DatasetNotFoundError,
@@ -3428,6 +3429,29 @@ def test_save_to_non_default_namespace_and_project(
         dc.read_dataset(name="fibonacci")
 
 
+def test_dataset_not_found_in_default_project(test_session):
+    metastore = test_session.catalog.metastore
+    with pytest.raises(DatasetNotFoundError) as excinfo:
+        with patch.object(AbstractMetastore, "is_local_dataset", return_value=True):
+            dc.read_dataset("fibonacci")
+    assert str(excinfo.value) == (
+        f"Dataset fibonacci not found in namespace {metastore.default_namespace_name}"
+        f" and project {metastore.default_project_name}"
+    )
+
+
+@pytest.mark.parametrize("project_created", (True, False))
+def test_dataset_not_found_in_non_default_project(test_session, project_created):
+    if project_created:
+        dc.create_project("dev", "numbers")
+    with pytest.raises(DatasetNotFoundError) as excinfo:
+        with patch.object(AbstractMetastore, "is_local_dataset", return_value=True):
+            dc.read_dataset("dev.numbers.fibonacci")
+    assert str(excinfo.value) == (
+        "Dataset fibonacci not found in namespace dev and project numbers"
+    )
+
+
 @pytest.mark.parametrize("use_settings", (True, False))
 @pytest.mark.parametrize("project_created_upfront", (True, False))
 def test_save_specify_only_non_default_project(