Ilongin/12872 dataset soft delete

src/datachain/catalog/catalog.py

@@ -29,6 +29,7 @@
     DATASET_PREFIX,
     DEFAULT_DATASET_VERSION,
     QUERY_DATASET_PREFIX,
+    REMOVED_VERSION_SUFFIX,
     DatasetDependency,
     DatasetListRecord,
     DatasetRecord,
@@ -48,7 +49,7 @@
     NamespaceNotFoundError,
     ProjectNotFoundError,
 )
-from datachain.lib.listing import get_listing
+from datachain.lib.listing import get_listing, is_listing_dataset
 from datachain.node import DirType, Node, NodeWithPath
 from datachain.nodes_thread_pool import NodesThreadPool
 from datachain.progress import tqdm
@@ -1094,21 +1095,52 @@ def update_dataset(self, dataset: DatasetRecord, **kwargs) -> DatasetRecord:
         self.warehouse.rename_dataset_tables(dataset, dataset_updated)
         return dataset_updated
 
-    def remove_dataset_version(
-        self, dataset: DatasetRecord, version: str, drop_rows: bool | None = True
-    ) -> None:
+    def remove_dataset_version(self, dataset: DatasetRecord, version: str) -> None:
         """
-        Deletes one single dataset version.
-        If it was last version, it removes dataset completely.
+        Remove a single dataset version.
+
+        For COMPLETE user-named versions this is a soft delete: rows table is
+        dropped, dependencies are preserved, and the version row stays with
+        status REMOVED so dependents can still render lineage.
+
+        For non-COMPLETE versions (CREATED/FAILED/STALE/REMOVING leftovers) and
+        for internal datasets (listing `lst__*` / `session_*` intermediates),
+        this is a hard delete: rows table dropped, dependencies removed,
+        version row deleted, and dataset row removed if it was the last
+        version.
         """
+        from datachain.query.session import Session
+
         if not dataset.has_version(version):
             return
+        v = dataset.get_version(version)
+        if v.status == DatasetStatus.REMOVED:
+            return
+
+        is_internal = is_listing_dataset(dataset.name) or dataset.name.startswith(
+            Session.DATASET_PREFIX
+        )
+        soft = v.status == DatasetStatus.COMPLETE and not is_internal
+
         self.metastore.update_dataset_version(
             dataset, version, status=DatasetStatus.REMOVING
         )
-        if drop_rows:
-            self.warehouse.drop_dataset_rows_table(dataset, version)
-        dataset = self.metastore.remove_dataset_version(dataset, version)
+        self.warehouse.drop_dataset_rows_table(dataset, version)
+
+        if soft:
+            # Rename the `version` column to free the (dataset_id, version)
+            # uniqueness slot so a future save can reclaim it. The original
+            # semver is still recoverable via DatasetVersion.display_version.
+            mangled_version = f"{version}{REMOVED_VERSION_SUFFIX}{v.id}"
+            self.metastore.update_dataset_version(
+                dataset,
+                version,
+                version=mangled_version,
+                status=DatasetStatus.REMOVED,
+                removed_at=datetime.now(timezone.utc),
+            )
+        else:
+            self.metastore.remove_dataset_version(dataset, version)
 
     def _remove_versions(self, pairs: Iterable[tuple[DatasetRecord, str]]) -> int:
         num_removed = 0

src/datachain/data_storage/metastore.py

@@ -318,9 +318,14 @@ def update_dataset(self, dataset: DatasetRecord, **kwargs) -> DatasetRecord:
 
     @abstractmethod
     def update_dataset_version(
-        self, dataset: DatasetRecord, version: str, **kwargs
+        self, dataset: DatasetRecord, dataset_version: str, **kwargs
     ) -> DatasetVersion:
-        """Updates dataset version fields."""
+        """Updates dataset version fields.
+
+        ``dataset_version`` is the lookup key (the current value of the row's
+        ``version`` column). Pass ``version=`` in ``kwargs`` to rename the
+        column itself.
+        """
 
     @abstractmethod
     def remove_dataset_version(
@@ -846,6 +851,7 @@ def _datasets_versions_columns(cls) -> list["SchemaItem"]:
             Column("schema", JSON, nullable=True),
             Column("job_id", Text, nullable=True),
             Column("content_hash", Text, nullable=True),
+            Column("removed_at", DateTime(timezone=True), nullable=True),
             UniqueConstraint("dataset_id", "version"),
         ]
 
@@ -1381,14 +1387,19 @@ def update_dataset(self, dataset: DatasetRecord, **kwargs) -> DatasetRecord:
         return result_ds
 
     def update_dataset_version(
-        self, dataset: DatasetRecord, version: str, **kwargs
+        self, dataset: DatasetRecord, dataset_version: str, **kwargs
     ) -> DatasetVersion:
-        """Updates dataset fields."""
+        """Updates dataset version fields.
+
+        ``dataset_version`` is the lookup key (the current value of the row's
+        ``version`` column). Pass ``version=`` in ``kwargs`` to rename the
+        column itself.
+        """
         logger.debug(
             "Metastore.update_dataset_version called for %s@%s: "
             "num_objects=%s, size=%s, preview_len=%s, all_fields=%s",
             dataset.name,
-            version,
+            dataset_version,
             kwargs.get("num_objects"),
             kwargs.get("size"),
             len(kwargs["preview"])
@@ -1440,16 +1451,16 @@ def update_dataset_version(
                 values[field] = value
                 version_values[field] = value
 
-        dataset_version = dataset.get_version(version)
+        version_obj = dataset.get_version(dataset_version)
 
         if not values:
-            return dataset_version
+            return version_obj
 
         logger.debug(
             "Writing to database for %s@%s: num_objects=%s, size=%s, "
             "preview_serialized=%s, fields_to_update=%s",
             dataset.name,
-            version,
+            dataset_version,
             values.get("num_objects"),
             values.get("size"),
             bool(values.get("preview")),
@@ -1459,21 +1470,21 @@ def update_dataset_version(
         dv = self._datasets_versions
         self.db.execute(
             self._datasets_versions_update()
-            .where(dv.c.dataset_id == dataset.id, dv.c.version == version)
+            .where(dv.c.dataset_id == dataset.id, dv.c.version == dataset_version)
             .values(values),
         )  # type: ignore [attr-defined]
 
-        dataset_version.update(**version_values)
+        version_obj.update(**version_values)
         logger.debug(
             "Dataset version updated successfully: %s@%s, "
             "final_num_objects=%s, final_size=%s, has_preview=%s",
             dataset.name,
-            version,
-            dataset_version.num_objects,
-            dataset_version.size,
-            bool(getattr(dataset_version, "_preview_data", None)),
+            dataset_version,
+            version_obj.num_objects,
+            version_obj.size,
+            bool(getattr(version_obj, "_preview_data", None)),
         )
-        return dataset_version
+        return version_obj
 
     def _parse_dataset(
         self,
@@ -2830,6 +2841,7 @@ def _get_dataset_version_for_job_ancestry_query(
                 self._projects.c.name == project_name,
                 self._dataset_version_jobs.c.job_id.in_(job_ancestry),
                 self._dataset_version_jobs.c.is_creator.is_(True),
+                self._datasets_versions.c.status != DatasetStatus.REMOVED,
             )
             .order_by(desc(self._dataset_version_jobs.c.created_at))
             .limit(1)

src/datachain/dataset.py

@@ -34,6 +34,10 @@
 DEFAULT_DATASET_VERSION = "1.0.0"
 DATASET_NAME_RESERVED_CHARS = [".", "@"]
 DATASET_NAME_REPLACEMENT_CHAR = "_"
+# Suffix appended to the `version` column on soft-delete so the
+# (dataset_id, version) uniqueness slot is freed for reuse. The semver
+# part before the suffix is the version originally claimed by the user.
+REMOVED_VERSION_SUFFIX = "~removed-"
 
 
 # StorageURI represents a normalized URI to a valid storage location
@@ -224,6 +228,13 @@ def parse(
 
         assert dataset_name is not None
 
+        # The `version` column on REMOVED tombstones carries a mangle suffix
+        # so the (dataset_id, version) slot is reusable. For lineage display
+        # we surface the original semver.
+        display_version = dataset_version
+        if display_version and REMOVED_VERSION_SUFFIX in display_version:
+            display_version = display_version.split(REMOVED_VERSION_SUFFIX, 1)[0]
+
         return cls(
             id,
             (
@@ -234,7 +245,7 @@ def parse(
             namespace_name,
             project_name,
             dataset_name,
-            dataset_version or None,  # type: ignore[arg-type]
+            display_version or None,  # type: ignore[arg-type]
             dataset_version_created_at,  # type: ignore[arg-type]
             [],
         )
@@ -264,6 +275,7 @@ class DatasetStatus:
     COMPLETE = 4
     STALE = 6
     REMOVING = 7
+    REMOVED = 8
 
 
 @dataclass
@@ -295,6 +307,7 @@ class DatasetVersion:
     query_script: str = ""
     job_id: str | None = None
     content_hash: str | None = None
+    removed_at: datetime | None = None
 
     @classmethod
     def parse(  # noqa: PLR0913
@@ -318,6 +331,7 @@ def parse(  # noqa: PLR0913
         query_script: str = "",
         job_id: str | None = None,
         content_hash: str | None = None,
+        removed_at: datetime | None = None,
         *,
         preview_loaded: bool = True,
     ):
@@ -346,12 +360,21 @@ def parse(  # noqa: PLR0913
             query_script=query_script,
             job_id=job_id,
             content_hash=content_hash,
+            removed_at=removed_at,
             _preview_loaded=preview_loaded,
         )
 
     @property
     def version_value(self) -> int:
-        return semver.value(self.version)
+        return semver.value(self.display_version)
+
+    @property
+    def display_version(self) -> str:
+        """Real semver version. For REMOVED tombstones, strips the internal
+        mangle suffix that frees the (dataset_id, version) slot."""
+        if REMOVED_VERSION_SUFFIX in self.version:
+            return self.version.split(REMOVED_VERSION_SUFFIX, 1)[0]
+        return self.version
 
     def __eq__(self, other):
         if not isinstance(other, DatasetVersion):
@@ -372,6 +395,7 @@ def is_final_status(self) -> bool:
             DatasetStatus.COMPLETE,
             DatasetStatus.STALE,
             DatasetStatus.REMOVING,
+            DatasetStatus.REMOVED,
         ]
 
     def update(self, **kwargs):
@@ -574,6 +598,7 @@ def parse(  # noqa: PLR0913
         version_schema: str | None = None,
         version_job_id: str | None = None,
         version_content_hash: str | None = None,
+        version_removed_at: datetime | None = None,
         *,
         versions_loaded: bool = True,
         preview_loaded: bool = True,
@@ -632,6 +657,7 @@ def parse(  # noqa: PLR0913
                 version_query_script or "",
                 version_job_id,
                 version_content_hash,
+                version_removed_at,
                 preview_loaded=preview_loaded,
             )
             versions_list = [dataset_version]
@@ -674,15 +700,14 @@ def has_version(self, version: str) -> bool:
     def is_valid_next_version(self, version: str) -> bool:
         """
         Checks if a number can be a valid next latest version for dataset.
-        The only rule is that it cannot be lower than current latest version
+        The only rule is that it cannot be lower than current latest live
+        version. REMOVED tombstones free their slot (see
+        REMOVED_VERSION_SUFFIX) so they don't participate in comparison.
         """
-        if not self.versions:
+        if not self.live_versions:
             return True
 
-        return not (
-            self.latest_version
-            and semver.value(self.latest_version) >= semver.value(version)
-        )
+        return semver.value(self.latest_version) < semver.value(version)
 
     def get_version(self, version: str) -> DatasetVersion:
         if not self.has_version(version):
@@ -730,7 +755,7 @@ def next_version_major(self) -> str:
         """
         Returns the next auto-incremented version if the major part is being bumped.
         """
-        if not self.versions:
+        if not self.live_versions:
             return "1.0.0"
 
         major, _, _ = semver.parse(self.latest_version)
@@ -741,7 +766,7 @@ def next_version_minor(self) -> str:
         """
         Returns the next auto-incremented version if the minor part is being bumped.
         """
-        if not self.versions:
+        if not self.live_versions:
             return "1.0.0"
 
         major, minor, _ = semver.parse(self.latest_version)
@@ -752,18 +777,23 @@ def next_version_patch(self) -> str:
         """
         Returns the next auto-incremented version if the patch part is being bumped.
         """
-        if not self.versions:
+        if not self.live_versions:
             return "1.0.0"
 
         major, minor, patch = semver.parse(self.latest_version)
         return semver.create(major, minor, patch + 1)
 
+    @property
+    def live_versions(self) -> list[DatasetVersion]:
+        """Versions excluding REMOVED tombstones."""
+        return [v for v in self.versions if v.status != DatasetStatus.REMOVED]
+
     @property
     def latest_version(self) -> str:
-        """Returns latest version of a dataset"""
-        if not self.versions:
+        """Latest user-visible version (skips REMOVED tombstones)."""
+        if not self.live_versions:
             raise DatasetVersionNotFoundError("Dataset has no versions")
-        return max(self.versions).version
+        return max(self.live_versions).version
 
     @property
     def latest_complete_version(self) -> str | None:
@@ -783,7 +813,9 @@ def latest_major_version(self, major: int) -> str | None:
         and we call `.latest_major_version(2)` it will return: "2.4.0".
         If no major version is find with input value, None will be returned
         """
-        versions = [v for v in self.versions if semver.parse(v.version)[0] == major]
+        versions = [
+            v for v in self.live_versions if semver.parse(v.version)[0] == major
+        ]
         if not versions:
             return None
         return max(versions).version
@@ -811,7 +843,7 @@ def latest_compatible_version(self, version_spec: str) -> str | None:
         # Convert dataset versions to packaging.Version objects
         # and filter compatible ones
         compatible_versions = []
-        for v in self.versions:
+        for v in self.live_versions:
             pkg_version = Version(v.version)
             if spec_set.contains(pkg_version):
                 compatible_versions.append(v)
@@ -943,7 +975,8 @@ def full_name(self) -> str:
         return f"{self.project.namespace.name}.{self.project.name}.{self.name}"
 
     def latest_version(self) -> DatasetListVersion:
-        return max(self.versions, key=lambda v: v.version_value)
+        live = [v for v in self.versions if v.status != DatasetStatus.REMOVED]
+        return max(live or self.versions, key=lambda v: v.version_value)
 
     @property
     def is_bucket_listing(self) -> bool:

src/datachain/delta.py

@@ -250,6 +250,12 @@ def _get_source_info(
         include_incomplete=False,
     )
 
+    # The version the dep points at may have been soft-deleted (REMOVED
+    # tombstone). Without a readable previous version we can't compute a
+    # diff; fall back to normal dataset creation, same as a missing dep.
+    if not source_ds.has_version(source_ds_dep.version):
+        return None, None, None, None
+
     return (
         source_ds.name,
         source_ds.project,