datachain 0.24.2__tar.gz → 0.24.4__tar.gz

{datachain-0.24.2 → datachain-0.24.4}/.github/workflows/tests-studio.yml

@@ -98,7 +98,7 @@ jobs:
       - name: Run tests
         # Generate `.test_durations` file with `pytest --store-durations --durations-path ../.github/.test_durations ...`
         run: >
-          DATACHAIN_METASTORE_ARG_PROJECT=john
+          DATACHAIN_METASTORE_ARG_USERNAME=john
           PYTHONPATH="$(pwd)/..:${PYTHONPATH}"
           pytest
           --config-file=pyproject.toml -rs

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

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

{datachain-0.24.2 → datachain-0.24.4}/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,
@@ -1111,6 +1120,14 @@ class Catalog:
         pull_dataset: bool = False,
         update: bool = False,
     ) -> DatasetRecord:
+        # Intentionally ignore update flag is version is provided. Here only exact
+        # version can be provided and update then doesn't make sense.
+        # It corresponds to a query like this for example:
+        #
+        #    dc.read_dataset("some.remote.dataset", version="1.0.0", update=True)
+        if version:
+            update = False
+
         if self.metastore.is_local_dataset(namespace_name) or not update:
             try:
                 project = self.metastore.get_project(project_name, namespace_name)
@@ -1124,7 +1141,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.2 → datachain-0.24.4}/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.2 → datachain-0.24.4}/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.2 → datachain-0.24.4}/src/datachain/dataset.py

@@ -32,7 +32,7 @@ QUERY_DATASET_PREFIX = "ds_query_"
 LISTING_PREFIX = "lst__"
 
 DEFAULT_DATASET_VERSION = "1.0.0"
-DATASET_NAME_RESERVED_CHARS = ["."]
+DATASET_NAME_RESERVED_CHARS = [".", "@"]
 DATASET_NAME_REPLACEMENT_CHAR = "_"
 
 

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

@@ -21,6 +21,7 @@ from typing import (
 import orjson
 import sqlalchemy
 from pydantic import BaseModel
+from sqlalchemy.sql.elements import ColumnElement
 from tqdm import tqdm
 
 from datachain import semver
@@ -806,11 +807,35 @@ class DataChain:
             chain.save("new_dataset")
             ```
         """
+        # Convert string partition_by parameters to Column objects
+        processed_partition_by = partition_by
+        if partition_by is not None:
+            if isinstance(partition_by, (str, Function, ColumnElement)):
+                list_partition_by = [partition_by]
+            else:
+                list_partition_by = list(partition_by)
+
+            processed_partition_columns: list[ColumnElement] = []
+            for col in list_partition_by:
+                if isinstance(col, str):
+                    col_db_name = ColumnMeta.to_db_name(col)
+                    col_type = self.signals_schema.get_column_type(col_db_name)
+                    column = Column(col_db_name, python_to_sql(col_type))
+                    processed_partition_columns.append(column)
+                elif isinstance(col, Function):
+                    column = col.get_column(self.signals_schema)
+                    processed_partition_columns.append(column)
+                else:
+                    # Assume it's already a ColumnElement
+                    processed_partition_columns.append(col)
+
+            processed_partition_by = processed_partition_columns
+
         udf_obj = self._udf_to_obj(Aggregator, func, params, output, signal_map)
         return self._evolve(
             query=self._query.generate(
                 udf_obj.to_udf_wrapper(),
-                partition_by=partition_by,
+                partition_by=processed_partition_by,
                 **self._settings.to_dict(),
             ),
             signal_schema=udf_obj.output,

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

@@ -189,6 +189,7 @@ def read_dataset(
         namespace_name=namespace_name,
         version=version,  #  type: ignore[arg-type]
         session=session,
+        update=update,
     )
 
     signals_schema = SignalSchema({"sys": Sys})

{datachain-0.24.2 → datachain-0.24.4}/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.2 → datachain-0.24.4}/src/datachain/namespace.py

@@ -6,7 +6,7 @@ from typing import Any, Optional, TypeVar
 from datachain.error import InvalidNamespaceNameError
 
 N = TypeVar("N", bound="Namespace")
-NAMESPACE_NAME_RESERVED_CHARS = ["."]
+NAMESPACE_NAME_RESERVED_CHARS = [".", "@"]
 
 
 @dataclass(frozen=True)

{datachain-0.24.2 → datachain-0.24.4}/src/datachain/project.py

@@ -7,7 +7,7 @@ from datachain.error import InvalidProjectNameError
 from datachain.namespace import Namespace
 
 P = TypeVar("P", bound="Project")
-PROJECT_NAME_RESERVED_CHARS = ["."]
+PROJECT_NAME_RESERVED_CHARS = [".", "@"]
 
 
 @dataclass(frozen=True)

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

@@ -82,7 +82,10 @@ if TYPE_CHECKING:
 INSERT_BATCH_SIZE = 10000
 
 PartitionByType = Union[
-    Function, ColumnElement, Sequence[Union[Function, ColumnElement]]
+    str,
+    Function,
+    ColumnElement,
+    Sequence[Union[str, Function, ColumnElement]],
 ]
 JoinPredicateType = Union[str, ColumnClause, ColumnElement]
 DatasetDependencyType = tuple["DatasetRecord", str]
@@ -1142,6 +1145,7 @@ class DatasetQuery:
                     project_name=project_name,
                     version=version,
                     pull_dataset=True,
+                    update=update,
                 )
             )
 

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

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

{datachain-0.24.2 → datachain-0.24.4}/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}")

{datachain-0.24.2 → datachain-0.24.4}/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.2 → datachain-0.24.4}/tests/func/test_read_dataset_remote.py

@@ -362,8 +362,16 @@ def test_read_dataset_remote_update_flag(
     assert dc.datasets().to_values("version") == ["1.0.0"]
     assert ds1.to_values("version")[0] == "1.0.0"
 
+    # Read without update and version returns a cached version
+    ds1 = dc.read_dataset(
+        f"{REMOTE_NAMESPACE_NAME}.{REMOTE_PROJECT_NAME}.dogs",
+        session=test_session,
+    )
+    assert dc.datasets().to_values("version") == ["1.0.0"]
+    assert ds1.to_values("version")[0] == "1.0.0"
+
     # Second read with update=True with the exact version
-    # returns the same
+    # returns the same dataset version
     ds2 = dc.read_dataset(
         f"{REMOTE_NAMESPACE_NAME}.{REMOTE_PROJECT_NAME}.dogs",
         version="1.0.0",
@@ -385,9 +393,7 @@ def test_read_dataset_remote_update_flag(
     assert dc.datasets().to_values("version") == ["1.0.0"]
     assert ds3.to_values("version")[0] == "1.0.0"
 
-    # Finally, read with update=False even with version specifier
-    # that allows for newer version still bring the same version
-    # as the one already downloaded
+    # Finally, read with update=True brings the latest version
     ds4 = dc.read_dataset(
         f"{REMOTE_NAMESPACE_NAME}.{REMOTE_PROJECT_NAME}.dogs",
         version=">=1.0.0",
@@ -399,6 +405,45 @@ def test_read_dataset_remote_update_flag(
     assert dc.datasets().to_values("version") == ["1.0.0", "2.0.0"]
 
 
+@skip_if_not_sqlite
+def test_read_dataset_remote_update_flag_no_version(
+    studio_token,
+    test_session,
+    remote_dataset_multi_version,
+    mock_dataset_info_endpoint,
+    mock_export_endpoint_with_urls,
+    mock_export_status_completed,
+    mock_s3_parquet_download,
+    mock_dataset_rows_fetcher_status_check,
+    requests_mock,
+):
+    """Test read_dataset with update=True flag to force remote check."""
+
+    # Mock the Studio API responses
+    mock_dataset_info_endpoint(remote_dataset_multi_version)
+    mock_s3_parquet_download()
+
+    # First read - downloads version 1.0.0
+    ds1 = dc.read_dataset(
+        f"{REMOTE_NAMESPACE_NAME}.{REMOTE_PROJECT_NAME}.dogs",
+        version="1.0.0",
+        session=test_session,
+    )
+    assert dc.datasets().to_values("version") == ["1.0.0"]
+    assert ds1.to_values("version")[0] == "1.0.0"
+
+    # Read with update=True w/o version specifier also
+    # checks the most recent remote version and brings it
+    ds4 = dc.read_dataset(
+        f"{REMOTE_NAMESPACE_NAME}.{REMOTE_PROJECT_NAME}.dogs",
+        update=True,
+        session=test_session,
+    )
+
+    assert ds4.to_values("version")[0] == "2.0.0"
+    assert dc.datasets().to_values("version") == ["1.0.0", "2.0.0"]
+
+
 @skip_if_not_sqlite
 def test_read_dataset_remote_version_specifiers(
     studio_token,

{datachain-0.24.2 → datachain-0.24.4}/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(
@@ -3571,3 +3595,60 @@ def test_save_create_project_not_allowed(test_session, allow_create_project):
         dc.read_values(fib=[1, 1, 2, 3, 5, 8], session=test_session).save(
             "dev.numbers.fibonacci"
         )
+
+
+def test_agg_partition_by_string_notation(test_session):
+    """Test that agg method supports string notation for partition_by."""
+
+    class _ImageGroup(BaseModel):
+        name: str
+        size: int
+
+    def func(key, val) -> Iterator[tuple[File, _ImageGroup]]:
+        n = "-".join(key)
+        v = sum(val)
+        yield File(path=n), _ImageGroup(name=n, size=v)
+
+    keys = ["n1", "n2", "n1"]
+    values = [1, 5, 9]
+
+    # Test using string notation (NEW functionality)
+    ds = dc.read_values(key=keys, val=values, session=test_session).agg(
+        x=func,
+        partition_by="key",  # String notation instead of C("key")
+    )
+
+    assert ds.order_by("x_1.name").to_values("x_1.name") == ["n1-n1", "n2"]
+    assert ds.order_by("x_1.size").to_values("x_1.size") == [5, 10]
+
+
+def test_agg_partition_by_string_sequence(test_session):
+    """Test that agg method supports sequence of strings for partition_by."""
+
+    class _ImageGroup(BaseModel):
+        name: str
+        size: int
+
+    def func(key1, key2, val) -> Iterator[tuple[File, _ImageGroup]]:
+        n = f"{key1[0]}-{key2[0]}"
+        v = sum(val)
+        yield File(path=n), _ImageGroup(name=n, size=v)
+
+    key1_values = ["a", "a", "b"]
+    key2_values = ["x", "y", "x"]
+    values = [1, 5, 9]
+
+    # Test using sequence of strings (NEW functionality)
+    ds = dc.read_values(
+        key1=key1_values, key2=key2_values, val=values, session=test_session
+    ).agg(
+        x=func,
+        partition_by=["key1", "key2"],  # Sequence of strings
+    )
+
+    result_names = ds.order_by("x_1.name").to_values("x_1.name")
+    result_sizes = ds.order_by("x_1.size").to_values("x_1.size")
+
+    # Should have 3 partitions: (a,x), (a,y), (b,x)
+    assert len(result_names) == 3
+    assert len(result_sizes) == 3