UDF Checkpoints cleanup

noxfile.py

@@ -38,11 +38,9 @@ def bench(session: nox.Session) -> None:
 def tests(session: nox.Session) -> None:
     session.install(".[tests]")
     env = {"COVERAGE_FILE": f".coverage.{session.python}"}
-    if session.python in ("3.12", "3.13"):
-        # improve performance of tests in Python>=3.12 when used with coverage
-        # https://github.com/nedbat/coveragepy/issues/1665
-        # https://github.com/python/cpython/issues/107674
-        env["COVERAGE_CORE"] = "sysmon"
+    # Note: Previously used COVERAGE_CORE=sysmon for Python 3.12/3.13 performance,
+    # but sysmon doesn't support branch coverage in those versions.
+    # Removed to avoid: "Can't use core=sysmon: sys.monitoring can't measure branches"
     session.run(
         "pytest",
         "--cov",

pyproject.toml

@@ -108,6 +108,7 @@ tests = [
   "pytest-asyncio",
   "pytest-sugar>=0.9.6",
   "pytest-cov>=4.1.0",
+  "coverage>=7.11.1",
   "pytest-mock>=3.12.0",
   "pytest-servers[all]>=0.5.9",
   "pytest-benchmark[histogram]",

src/datachain/catalog/catalog.py

@@ -12,6 +12,7 @@
 from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
 from copy import copy
 from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
 from functools import cached_property, reduce
 from threading import Thread
 from typing import IO, TYPE_CHECKING, Any, NoReturn
@@ -22,6 +23,7 @@
 from tqdm.auto import tqdm
 
 from datachain.cache import Cache
+from datachain.checkpoint import Checkpoint
 from datachain.client import Client
 from datachain.dataset import (
     DATASET_PREFIX,
@@ -2039,3 +2041,66 @@ def index(
             client_config=client_config or self.client_config,
             only_index=True,
         )
+
+    def _remove_checkpoint(self, checkpoint: Checkpoint) -> None:
+        """
+        Remove a checkpoint and its associated job-specific UDF tables.
+
+        Since tables are now job-scoped, this removes only the tables
+        belonging to this specific checkpoint's job.
+
+        Args:
+            checkpoint: The checkpoint object to remove.
+        """
+        # Remove the checkpoint from metastore first
+        self.metastore.remove_checkpoint(checkpoint)
+
+        # Remove job-specific tables for this checkpoint
+        # Table patterns: udf_{job_id}_{hash}_{suffix}
+        # where suffix can be: input, output, output_partial, processed
+        job_id_sanitized = checkpoint.job_id.replace("-", "")
+        table_prefix = f"udf_{job_id_sanitized}_{checkpoint.hash}_"
+        matching_tables = self.warehouse.db.list_tables(prefix=table_prefix)
+
+        if matching_tables:
+            self.warehouse.cleanup_tables(matching_tables)
+
+    def cleanup_checkpoints(self, ttl_seconds: int | None = None) -> None:
+        """
+        Clean up outdated checkpoints and their associated UDF tables.
+
+        Uses optimized branch pruning: removes outdated checkpoints if no
+        descendants have active (non-outdated) checkpoints that depend on them.
+
+        This prevents accumulation of checkpoints while ensuring that ancestor
+        tables are preserved when descendants still need them.
+
+        Args:
+            ttl_seconds: Time-to-live in seconds. Checkpoints older than this
+                        are considered outdated. If None, uses CHECKPOINT_TTL
+                        environment variable or default.
+        """
+        if ttl_seconds is None:
+            ttl_seconds = int(os.environ.get("CHECKPOINT_TTL", str(TTL_INT)))
+
+        ttl_threshold = datetime.now(timezone.utc) - timedelta(seconds=ttl_seconds)
+
+        # Cache descendant check results per job_id to avoid redundant checks
+        has_active_descendants_cache: dict[str, bool] = {}
+
+        # For each outdated checkpoint, check if it's safe to remove
+        for ch in self.metastore.list_checkpoints(created_before=ttl_threshold):
+            # Check once per job_id if descendants have active checkpoints (cached)
+            if ch.job_id not in has_active_descendants_cache:
+                has_active_descendants_cache[ch.job_id] = any(
+                    list(
+                        self.metastore.list_checkpoints(
+                            desc_id, created_after=ttl_threshold
+                        )
+                    )
+                    for desc_id in self.metastore.get_descendant_job_ids(ch.job_id)
+                )
+
+            # If no active descendants, remove the checkpoint
+            if not has_active_descendants_cache[ch.job_id]:
+                self._remove_checkpoint(ch)

src/datachain/cli/commands/misc.py

@@ -12,12 +12,18 @@ def clear_cache(catalog: "Catalog"):
 
 def garbage_collect(catalog: "Catalog"):
     temp_tables = catalog.get_temp_table_names()
-    if not temp_tables:
-        print("Nothing to clean up.")
-    else:
-        print(f"Garbage collecting {len(temp_tables)} tables.")
+    has_tables = bool(temp_tables)
+
+    if has_tables:
+        print(f"Garbage collecting {len(temp_tables)} temporary tables.")
         catalog.cleanup_tables(temp_tables)
 
+    print("Cleaning up outdated checkpoints.")
+    catalog.cleanup_checkpoints()
+
+    if not has_tables:
+        print("No temporary tables to clean up.")
+
 
 def completion(shell: str) -> str:
     from datachain.cli import get_parser

src/datachain/data_storage/db_engine.py

@@ -80,12 +80,19 @@ def execute(
     ) -> Iterator[tuple[Any, ...]]: ...
 
     def get_table(self, name: str) -> "Table":
+        from datachain.error import TableMissingError
+
         table = self.metadata.tables.get(name)
         if table is None:
-            sa.Table(name, self.metadata, autoload_with=self.engine)
-            # ^^^ This table may not be correctly initialised on some dialects
-            # Grab it from metadata instead.
-            table = self.metadata.tables[name]
+            try:
+                sa.Table(name, self.metadata, autoload_with=self.engine)
+                # ^^^ This table may not be correctly initialised on some dialects
+                # Grab it from metadata instead.
+                table = self.metadata.tables.get(name)
+                if table is None:
+                    raise TableMissingError(f"Table '{name}' not found")
+            except (KeyError, sa.exc.NoSuchTableError) as e:
+                raise TableMissingError(f"Table '{name}' not found") from e
         return table
 
     @abstractmethod
@@ -111,6 +118,16 @@ def has_table(self, name: str) -> bool:
         """
         return sa.inspect(self.engine).has_table(name)
 
+    def list_tables(self, prefix: str = "") -> list[str]:
+        """
+        Return a list of table names that start with the given prefix.
+        If no prefix is provided, returns all table names.
+        """
+        all_tables = sa.inspect(self.engine).get_table_names()
+        if not prefix:
+            return all_tables
+        return [table for table in all_tables if table.startswith(prefix)]
+
     @abstractmethod
     def create_table(
         self,

src/datachain/data_storage/metastore.py

@@ -21,6 +21,7 @@
     Table,
     Text,
     UniqueConstraint,
+    cast,
     desc,
     literal,
     select,
@@ -422,6 +423,20 @@ def create_job(
     def get_job(self, job_id: str) -> Job | None:
         """Returns the job with the given ID."""
 
+    @abstractmethod
+    def get_ancestor_job_ids(self, job_id: str, conn=None) -> list[str]:
+        """
+        Returns list of ancestor job IDs in order from parent to root.
+        Uses recursive CTE to get all ancestors in a single query.
+        """
+
+    @abstractmethod
+    def get_descendant_job_ids(self, job_id: str, conn=None) -> list[str]:
+        """
+        Returns list of descendant job IDs (children, grandchildren, etc.).
+        Uses recursive CTE to get all descendants in a single query.
+        """
+
     @abstractmethod
     def update_job(
         self,
@@ -457,8 +472,24 @@ def get_last_job_by_name(self, name: str, conn=None) -> "Job | None":
     #
 
     @abstractmethod
-    def list_checkpoints(self, job_id: str, conn=None) -> Iterator[Checkpoint]:
-        """Returns all checkpoints related to some job"""
+    def list_checkpoints(
+        self,
+        job_id: str | None = None,
+        created_after: datetime | None = None,
+        created_before: datetime | None = None,
+        conn=None,
+    ) -> Iterator[Checkpoint]:
+        """
+        List checkpoints by job id, or all checkpoints if job_id is None.
+
+        Args:
+            job_id: Filter by job ID. If None, lists all checkpoints.
+            created_after: Filter by creation date. If provided, only returns
+                          checkpoints created after this timestamp.
+            created_before: Filter by creation date. If provided, only returns
+                           checkpoints created before this timestamp.
+            conn: Database connection to use.
+        """
 
     @abstractmethod
     def get_last_checkpoint(self, job_id: str, conn=None) -> Checkpoint | None:
@@ -485,6 +516,12 @@ def create_checkpoint(
     ) -> Checkpoint:
         """Creates new checkpoint"""
 
+    @abstractmethod
+    def remove_checkpoint(
+        self, checkpoint: Checkpoint, conn: Any | None = None
+    ) -> None:
+        """Removes a checkpoint by checkpoint object"""
+
 
 class AbstractDBMetastore(AbstractMetastore):
     """
@@ -1746,6 +1783,73 @@ def get_job(self, job_id: str, conn=None) -> Job | None:
             return None
         return self._parse_job(results[0])
 
+    def get_ancestor_job_ids(self, job_id: str, conn=None) -> list[str]:
+        # Use recursive CTE to walk up the parent chain
+        # Format: WITH RECURSIVE ancestors(id, parent_job_id) AS (...)
+        ancestors_cte = (
+            select(
+                self._jobs.c.id.label("id"),
+                self._jobs.c.parent_job_id.label("parent_job_id"),
+            )
+            .where(self._jobs.c.id == job_id)
+            .cte(name="ancestors", recursive=True)
+        )
+
+        # Recursive part: join with parent jobs
+        ancestors_recursive = ancestors_cte.union_all(
+            select(
+                self._jobs.c.id.label("id"),
+                self._jobs.c.parent_job_id.label("parent_job_id"),
+            ).select_from(
+                self._jobs.join(
+                    ancestors_cte,
+                    self._jobs.c.id
+                    == cast(ancestors_cte.c.parent_job_id, self._jobs.c.id.type),
+                )
+            )
+        )
+
+        # Select all ancestor IDs except the starting job itself
+        query = select(ancestors_recursive.c.id).where(
+            ancestors_recursive.c.id != job_id
+        )
+
+        results = list(self.db.execute(query, conn=conn))
+        return [str(row[0]) for row in results]
+
+    def get_descendant_job_ids(self, job_id: str, conn=None) -> list[str]:
+        # Use recursive CTE to walk down the child chain
+        descendants_cte = (
+            select(
+                self._jobs.c.id.label("id"),
+                self._jobs.c.parent_job_id.label("parent_job_id"),
+            )
+            .where(self._jobs.c.id == job_id)
+            .cte(name="descendants", recursive=True)
+        )
+
+        # Recursive part: join with child jobs
+        descendants_recursive = descendants_cte.union_all(
+            select(
+                self._jobs.c.id.label("id"),
+                self._jobs.c.parent_job_id.label("parent_job_id"),
+            ).select_from(
+                self._jobs.join(
+                    descendants_cte,
+                    cast(self._jobs.c.parent_job_id, self._jobs.c.id.type)
+                    == descendants_cte.c.id,
+                )
+            )
+        )
+
+        # Select all descendant IDs except the starting job itself
+        query = select(descendants_recursive.c.id).where(
+            descendants_recursive.c.id != job_id
+        )
+
+        results = list(self.db.execute(query, conn=conn))
+        return [str(row[0]) for row in results]
+
     def update_job(
         self,
         job_id: str,
@@ -1833,7 +1937,7 @@ def _checkpoints_columns() -> "list[SchemaItem]":
             Column("hash", Text, nullable=False),
             Column("partial", Boolean, default=False),
             Column("created_at", DateTime(timezone=True), nullable=False),
-            UniqueConstraint("job_id", "hash"),
+            UniqueConstraint("job_id", "hash", "partial"),
         ]
 
     @cached_property
@@ -1872,24 +1976,48 @@ def create_checkpoint(
         conn: Any | None = None,
     ) -> Checkpoint:
         """
-        Creates a new job query step.
+        Creates a new checkpoint or returns existing one if already exists.
+        This is idempotent - calling it multiple times with the same job_id and hash
+        will not create duplicates.
         """
+        # First check if checkpoint already exists
+        existing = self.find_checkpoint(job_id, _hash, partial=partial, conn=conn)
+        if existing:
+            return existing
+
         checkpoint_id = str(uuid4())
-        self.db.execute(
-            self._checkpoints_insert().values(
-                id=checkpoint_id,
-                job_id=job_id,
-                hash=_hash,
-                partial=partial,
-                created_at=datetime.now(timezone.utc),
-            ),
-            conn=conn,
+        query = self._checkpoints_insert().values(
+            id=checkpoint_id,
+            job_id=job_id,
+            hash=_hash,
+            partial=partial,
+            created_at=datetime.now(timezone.utc),
         )
-        return self.get_checkpoint_by_id(checkpoint_id)
 
-    def list_checkpoints(self, job_id: str, conn=None) -> Iterator[Checkpoint]:
-        """List checkpoints by job id."""
-        query = self._checkpoints_query().where(self._checkpoints.c.job_id == job_id)
+        # Use on_conflict_do_nothing to handle race conditions
+        if hasattr(query, "on_conflict_do_nothing"):
+            query = query.on_conflict_do_nothing(
+                index_elements=["job_id", "hash", "partial"]
+            )
+
+        self.db.execute(query, conn=conn)
+
+        return self.find_checkpoint(job_id, _hash, partial=partial, conn=conn)  # type: ignore[return-value]
+
+    def list_checkpoints(
+        self,
+        job_id: str | None = None,
+        created_after: datetime | None = None,
+        created_before: datetime | None = None,
+        conn=None,
+    ) -> Iterator[Checkpoint]:
+        query = self._checkpoints_query()
+        if job_id is not None:
+            query = query.where(self._checkpoints.c.job_id == job_id)
+        if created_after is not None:
+            query = query.where(self._checkpoints.c.created_at >= created_after)
+        if created_before is not None:
+            query = query.where(self._checkpoints.c.created_at < created_before)
         rows = list(self.db.execute(query, conn=conn))
 
         yield from [self.checkpoint_class.parse(*r) for r in rows]
@@ -1929,3 +2057,13 @@ def get_last_checkpoint(self, job_id: str, conn=None) -> Checkpoint | None:
         if not rows:
             return None
         return self.checkpoint_class.parse(*rows[0])
+
+    def remove_checkpoint(
+        self, checkpoint: Checkpoint, conn: Any | None = None
+    ) -> None:
+        """Removes a checkpoint by checkpoint object"""
+        ch = self._checkpoints
+        self.db.execute(
+            self._checkpoints_delete().where(ch.c.id == checkpoint.id),
+            conn=conn,
+        )