Added update_versions argument to DataChain.save() (#1100)

Added `update_versions` argument to `DataChain.save()`

Author: ilongin

src/datachain/catalog/catalog.py

@@ -779,6 +779,7 @@ def create_dataset(
         uuid: Optional[str] = None,
         description: Optional[str] = None,
         attrs: Optional[list[str]] = None,
+        update_version: Optional[str] = "patch",
     ) -> "DatasetRecord":
         """
         Creates new dataset of a specific version.
@@ -795,6 +796,11 @@ def create_dataset(
         try:
             dataset = self.get_dataset(name)
             default_version = dataset.next_version_patch
+            if update_version == "major":
+                default_version = dataset.next_version_major
+            if update_version == "minor":
+                default_version = dataset.next_version_minor
+
             if (description or attrs) and (
                 dataset.description != description or dataset.attrs != attrs
             ):

src/datachain/lib/dc/datachain.py

@@ -461,6 +461,7 @@ def save(  # type: ignore[override]
         version: Optional[str] = None,
         description: Optional[str] = None,
         attrs: Optional[list[str]] = None,
+        update_version: Optional[str] = "patch",
         **kwargs,
     ) -> "Self":
         """Save to a Dataset. It returns the chain itself.
@@ -472,10 +473,22 @@ def save(  # type: ignore[override]
             description : description of a dataset.
             attrs : attributes of a dataset. They can be without value, e.g "NLP",
                 or with a value, e.g "location=US".
+            update_version: which part of the dataset version to automatically increase.
+                Available values: `major`, `minor` or `patch`. Default is `patch`.
         """
         if version is not None:
             semver.validate(version)
 
+        if update_version is not None and update_version not in [
+            "patch",
+            "major",
+            "minor",
+        ]:
+            raise ValueError(
+                "update_version can have one of the following values: major, minor or"
+                " patch"
+            )
+
         schema = self.signals_schema.clone_without_sys_signals().serialize()
         return self._evolve(
             query=self._query.save(
@@ -484,6 +497,7 @@ def save(  # type: ignore[override]
                 description=description,
                 attrs=attrs,
                 feature_schema=schema,
+                update_version=update_version,
                 **kwargs,
             )
         )

src/datachain/lib/dc/storage.py

@@ -5,7 +5,6 @@
     Union,
 )
 
-from datachain.error import DatasetNotFoundError
 from datachain.lib.file import (
     FileType,
     get_file_type,
@@ -132,11 +131,6 @@ def read_storage(
 
             def lst_fn(ds_name, lst_uri):
                 # disable prefetch for listing, as it pre-downloads all files
-                try:
-                    version = catalog.get_dataset(ds_name).next_version_major
-                except DatasetNotFoundError:
-                    version = None
-
                 (
                     read_records(
                         DataChain.DEFAULT_FILE_RECORD,
@@ -150,7 +144,7 @@ def lst_fn(ds_name, lst_uri):
                         output={f"{column}": file_type},
                     )
                     # for internal listing datasets, we always bump major version
-                    .save(ds_name, listing=True, version=version)
+                    .save(ds_name, listing=True, update_version="major")
                 )
 
             dc._query.set_listing_fn(

src/datachain/query/dataset.py

@@ -1689,6 +1689,7 @@ def save(
         feature_schema: Optional[dict] = None,
         description: Optional[str] = None,
         attrs: Optional[list[str]] = None,
+        update_version: Optional[str] = "patch",
         **kwargs,
     ) -> "Self":
         """Save the query as a dataset."""
@@ -1723,6 +1724,7 @@ def save(
                 columns=columns,
                 description=description,
                 attrs=attrs,
+                update_version=update_version,
                 **kwargs,
             )
             version = version or dataset.latest_version

tests/unit/lib/test_datachain.py

@@ -3173,18 +3173,41 @@ def test_delete_dataset_from_studio_not_found(
     assert str(exc_info.value) == error_message
 
 
+@pytest.mark.parametrize(
+    "update_version,versions",
+    [
+        ("patch", ["1.0.0", "1.0.1", "1.0.2"]),
+        ("minor", ["1.0.0", "1.1.0", "1.2.0"]),
+        ("major", ["1.0.0", "2.0.0", "3.0.0"]),
+    ],
+)
+def test_update_versions(test_session, update_version, versions):
+    ds_name = "fibonacci"
+    chain = dc.read_values(fib=[1, 1, 2, 3, 5, 8], session=test_session)
+    chain.save(ds_name, update_version=update_version)
+    chain.save(ds_name, update_version=update_version)
+    chain.save(ds_name, update_version=update_version)
+    assert sorted(
+        [
+            ds.version
+            for ds in dc.datasets(column="dataset", session=test_session).collect(
+                "dataset"
+            )
+        ]
+    ) == sorted(versions)
+
+
 def test_update_versions_mix_major_minor_patch(test_session):
     ds_name = "fibonacci"
     chain = dc.read_values(fib=[1, 1, 2, 3, 5, 8], session=test_session)
     chain.save(ds_name)
+    chain.save(ds_name, update_version="patch")
+    chain.save(ds_name, update_version="minor")
+    chain.save(ds_name, update_version="major")
+    chain.save(ds_name, update_version="minor")
+    chain.save(ds_name, update_version="patch")
     chain.save(ds_name)
-    chain.save(ds_name, version="1.1.0")
-    chain.save(ds_name)
-    chain.save(ds_name, version="2.0.0")
-    chain.save(ds_name)
-    chain.save(ds_name, version="2.1.0")
-    chain.save(ds_name)
-    chain.save(ds_name)
+    chain.save(ds_name, version="3.0.0")
     assert sorted(
         [
             ds.version
@@ -3197,16 +3220,27 @@ def test_update_versions_mix_major_minor_patch(test_session):
             "1.0.0",
             "1.0.1",
             "1.1.0",
-            "1.1.1",
             "2.0.0",
-            "2.0.1",
             "2.1.0",
             "2.1.1",
             "2.1.2",
+            "3.0.0",
         ]
     )
 
 
+def test_update_versions_wrong_value(test_session):
+    ds_name = "fibonacci"
+    chain = dc.read_values(fib=[1, 1, 2, 3, 5, 8], session=test_session)
+    chain.save(ds_name)
+    with pytest.raises(ValueError) as excinfo:
+        chain.save(ds_name, update_version="wrong")
+
+    assert str(excinfo.value) == (
+        "update_version can have one of the following values: major, minor or patch"
+    )
+
+
 def test_from_dataset_version_int_backward_compatible(test_session):
     ds_name = "numbers"
     dc.read_values(nums=[1], session=test_session).save(ds_name, version="1.0.0")