datachain 0.24.3__tar.gz → 0.24.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.24.3
+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.3 → datachain-0.24.4}/src/datachain/catalog/catalog.py

@@ -1120,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)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.24.3
+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.3 → 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.3 → datachain-0.24.4}/tests/unit/lib/test_datachain.py

@@ -3595,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