refactoring table naming methods

Author: ilongin

src/datachain/catalog/catalog.py

@@ -25,6 +25,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,
@@ -2096,22 +2097,15 @@ def index(
         ):
             pass
 
-    def _cleanup_udf_tables(self, prefix: str, suffix: str) -> None:
-        """Remove all UDF tables matching a prefix and optional suffix."""
-        tables = self.warehouse.db.list_tables(prefix=prefix)
-        tables = [t for t in tables if t.endswith(suffix)]
-        if tables:
-            logger.info("Removing %d UDF tables: %s", len(tables), tables)
-            self.warehouse.cleanup_tables(tables)
-
     def cleanup_checkpoints(self, ttl_seconds: int | None = None) -> int:
         """Clean up outdated checkpoints and their associated UDF tables.
 
         Algorithm:
         1. Find inactive jobs (all checkpoints older than TTL) — single query
-        2. For each inactive job: remove output/partial tables and checkpoints
+        2. Collect output/partial table names from checkpoints and remove them
         3. For run groups where all member jobs are inactive: also remove
-           shared input tables
+           shared input tables (prefix-based, since input hash is not on checkpoint)
+        4. Soft-delete checkpoint metadata
         """
         if ttl_seconds is None:
             ttl_seconds = CHECKPOINTS_TTL
@@ -2127,24 +2121,37 @@ def cleanup_checkpoints(self, ttl_seconds: int | None = None) -> int:
         inactive_job_ids = [job.id for job in inactive_jobs]
         run_group_ids = {job.run_group_id for job in inactive_jobs if job.run_group_id}
 
+        checkpoints = list(self.metastore.list_checkpoints(job_id=inactive_job_ids))
+
         logger.info(
-            "Cleaning %d inactive jobs across %d run groups",
+            "Cleaning %d inactive jobs across %d run groups (%d checkpoints)",
             len(inactive_jobs),
             len(run_group_ids),
+            len(checkpoints),
         )
 
-        for job in inactive_jobs:
-            self._cleanup_udf_tables(f"udf_{job.id}_", "_output")
-            self._cleanup_udf_tables(f"udf_{job.id}_", "_output_partial")
+        # Remove output/partial tables using exact names from checkpoints
+        tables = [ch.table_name for ch in checkpoints]
+        if tables:
+            logger.info("Removing %d UDF output tables: %s", len(tables), tables)
+            self.warehouse.cleanup_tables(tables)
 
         # Shared input tables — only when entire run group is inactive
         for group_id in run_group_ids:
             if not self.metastore.has_active_checkpoints_in_run_group(
                 group_id, ttl_threshold
             ):
-                self._cleanup_udf_tables(f"udf_{group_id}_", "_input")
+                input_tables = self.warehouse.db.list_tables(
+                    pattern=Checkpoint.input_table_pattern(group_id)
+                )
+                if input_tables:
+                    logger.info(
+                        "Removing %d shared input tables: %s",
+                        len(input_tables),
+                        input_tables,
+                    )
+                    self.warehouse.cleanup_tables(input_tables)
 
-        checkpoints = list(self.metastore.list_checkpoints(job_id=inactive_job_ids))
         self.metastore.remove_checkpoints([ch.id for ch in checkpoints])
 
         logger.info(

src/datachain/checkpoint.py

@@ -32,6 +32,33 @@ class Checkpoint:
     created_at: datetime
     status: int = CheckpointStatus.ACTIVE
 
+    @staticmethod
+    def output_table_name(job_id: str, _hash: str) -> str:
+        """Final UDF output table. Job-specific, created when UDF completes."""
+        return f"udf_{job_id}_{_hash}_output"
+
+    @staticmethod
+    def partial_output_table_name(job_id: str, _hash: str) -> str:
+        """Partial UDF output table. Temporary, renamed to final on completion."""
+        return f"udf_{job_id}_{_hash}_output_partial"
+
+    @staticmethod
+    def input_table_name(group_id: str, _hash: str) -> str:
+        """Shared UDF input table. Scoped to run group, reused across jobs."""
+        return f"udf_{group_id}_{_hash}_input"
+
+    @staticmethod
+    def input_table_pattern(group_id: str) -> str:
+        """LIKE pattern for finding all input tables in a run group."""
+        return f"udf_{group_id}_%_input"
+
+    @property
+    def table_name(self) -> str:
+        """UDF output table name associated with this checkpoint."""
+        if self.partial:
+            return self.partial_output_table_name(self.job_id, self.hash)
+        return self.output_table_name(self.job_id, self.hash)
+
     @classmethod
     def parse(
         cls,

src/datachain/data_storage/db_engine.py

@@ -123,15 +123,15 @@ def has_table(self, name: str) -> bool:
         return sa.inspect(self.engine).has_table(name)
 
     @abstractmethod
-    def list_tables(self, prefix: str = "") -> list[str]:
+    def list_tables(self, pattern: str = "") -> list[str]:
         """
-        List all table names, optionally filtered by prefix.
+        List all table names, optionally filtered by a SQL LIKE pattern.
 
         Args:
-            prefix: Optional prefix to filter table names
+            pattern: SQL LIKE pattern to filter table names (e.g. 'udf_%')
 
         Returns:
-            List of table names matching the prefix
+            List of table names matching the pattern
         """
 
     @abstractmethod

src/datachain/data_storage/sqlite.py

@@ -285,8 +285,8 @@ def execute_str(self, sql: str, parameters=None) -> sqlite3.Cursor:
             return self.db.execute(sql)
         return self.db.execute(sql, parameters)
 
-    def list_tables(self, prefix: str = "") -> list[str]:
-        """List all table names, optionally filtered by prefix."""
+    def list_tables(self, pattern: str = "") -> list[str]:
+        """List all table names, optionally filtered by a SQL LIKE pattern."""
         sqlite_master = sqlalchemy.table(
             "sqlite_master",
             sqlalchemy.column("type"),
@@ -295,8 +295,8 @@ def list_tables(self, prefix: str = "") -> list[str]:
         query = sqlalchemy.select(sqlite_master.c.name).where(
             sqlite_master.c.type == "table"
         )
-        if prefix:
-            query = query.where(sqlite_master.c.name.like(f"{prefix}%"))
+        if pattern:
+            query = query.where(sqlite_master.c.name.like(pattern))
         result = self.execute(query)
         return [row[0] for row in result.fetchall()]
 

src/datachain/query/dataset.py

@@ -101,18 +101,6 @@
 T = TypeVar("T", bound="DatasetQuery")
 
 
-def udf_input_table_name(run_group_id: str, _hash: str) -> str:
-    return f"udf_{run_group_id}_{_hash}_input"
-
-
-def udf_output_table_name(job_id: str, _hash: str) -> str:
-    return f"udf_{job_id}_{_hash}_output"
-
-
-def udf_partial_output_table_name(job_id: str, _hash: str) -> str:
-    return f"udf_{job_id}_{_hash}_output_partial"
-
-
 def detach(
     method: "Callable[Concatenate[T, P], T]",
 ) -> "Callable[Concatenate[T, P], T]":
@@ -970,7 +958,7 @@ def get_or_create_input_table(self, query: Select, _hash: str) -> "Table":
         Returns the input table.
         """
         group_id = self.job.run_group_id or self.job.id
-        input_table_name = udf_input_table_name(group_id, _hash)
+        input_table_name = Checkpoint.input_table_name(group_id, _hash)
 
         # Check if input table already exists (created by ancestor job)
         if self.warehouse.db.has_table(input_table_name):
@@ -1063,10 +1051,10 @@ def _skip_udf(
             checkpoint.job_id,
         )
         existing_output_table = self.warehouse.get_table(
-            udf_output_table_name(checkpoint.job_id, checkpoint.hash)
+            Checkpoint.output_table_name(checkpoint.job_id, checkpoint.hash)
         )
         output_table = self.warehouse.create_table_from_query(
-            udf_output_table_name(self.job.id, checkpoint.hash),
+            Checkpoint.output_table_name(self.job.id, checkpoint.hash),
             sa.select(existing_output_table),
             create_fn=self.create_output_table,
         )
@@ -1140,7 +1128,7 @@ def _run_from_scratch(
         input_table = self.get_or_create_input_table(query, hash_input)
 
         partial_output_table = self.create_output_table(
-            udf_partial_output_table_name(self.job.id, partial_hash),
+            Checkpoint.partial_output_table_name(self.job.id, partial_hash),
         )
 
         if self.partition_by is not None:
@@ -1151,7 +1139,7 @@ def _run_from_scratch(
         self.populate_udf_output_table(partial_output_table, input_query)
 
         output_table = self.warehouse.rename_table(
-            partial_output_table, udf_output_table_name(self.job.id, hash_output)
+            partial_output_table, Checkpoint.output_table_name(self.job.id, hash_output)
         )
 
         if partial_checkpoint:
@@ -1218,7 +1206,7 @@ def _continue_udf(
 
         try:
             parent_partial_table = self.warehouse.get_table(
-                udf_partial_output_table_name(
+                Checkpoint.partial_output_table_name(
                     self.job.rerun_from_job_id, checkpoint.hash
                 )
             )
@@ -1246,7 +1234,9 @@ def _continue_udf(
                 len(incomplete_input_ids),
             )
 
-        partial_table_name = udf_partial_output_table_name(self.job.id, checkpoint.hash)
+        partial_table_name = Checkpoint.partial_output_table_name(
+            self.job.id, checkpoint.hash
+        )
         if incomplete_input_ids:
             # Filter out incomplete inputs - they will be re-processed
             filtered_query = sa.select(parent_partial_table).where(
@@ -1288,7 +1278,7 @@ def _continue_udf(
         )
 
         output_table = self.warehouse.rename_table(
-            partial_table, udf_output_table_name(self.job.id, hash_output)
+            partial_table, Checkpoint.output_table_name(self.job.id, hash_output)
         )
 
         self.metastore.remove_checkpoints([partial_checkpoint.id])

tests/conftest.py

@@ -220,7 +220,9 @@ def cleanup_udf_tables(warehouse):
     UDF tables are shared across jobs and persist after chain finishes,
     so we need to clean them up after each test to prevent interference.
     """
-    for table_name in warehouse.db.list_tables(prefix=warehouse.UDF_TABLE_NAME_PREFIX):
+    for table_name in warehouse.db.list_tables(
+        pattern=f"{warehouse.UDF_TABLE_NAME_PREFIX}%"
+    ):
         table = warehouse.db.get_table(table_name)
         warehouse.db.drop_table(table, if_exists=True)
 

tests/func/checkpoints/test_checkpoint_cleanup.py

@@ -24,7 +24,7 @@ def test_cleanup_checkpoints_with_ttl(test_session, nums_dataset):
     job_id = test_session.get_or_create_job().id
 
     assert len(list(metastore.list_checkpoints(job_id))) == 4
-    assert len(warehouse.db.list_tables(prefix="udf_")) > 0
+    assert len(warehouse.db.list_tables(pattern="udf_%")) > 0
 
     # Make all checkpoints older than the default TTL (4h)
     old_time = datetime.now(timezone.utc) - timedelta(hours=5)
@@ -33,7 +33,7 @@ def test_cleanup_checkpoints_with_ttl(test_session, nums_dataset):
     catalog.cleanup_checkpoints()
 
     assert len(list(metastore.list_checkpoints(job_id))) == 0
-    assert len(warehouse.db.list_tables(prefix=f"udf_{job_id}_")) == 0
+    assert len(warehouse.db.list_tables(pattern=f"udf_{job_id}_%")) == 0
 
 
 def test_cleanup_checkpoints_with_custom_ttl(test_session, nums_dataset):
@@ -101,7 +101,7 @@ def test_cleanup_does_not_remove_unrelated_tables(test_session, nums_dataset):
     catalog.cleanup_checkpoints()
 
     assert len(list(metastore.list_checkpoints(job_id))) == 0
-    assert len(warehouse.db.list_tables(prefix=f"udf_{job_id}_")) == 0
+    assert len(warehouse.db.list_tables(pattern=f"udf_{job_id}_%")) == 0
     assert warehouse.db.has_table(fake_table_name)
 
     warehouse.cleanup_tables([fake_table_name])
@@ -194,9 +194,6 @@ def test_cleanup_preserves_input_tables_when_run_group_active(
     warehouse.create_udf_table(name=input_table)
     assert warehouse.db.has_table(input_table)
 
-    output_table = f"udf_{job1_id}_somehash_output"
-    warehouse.create_udf_table(name=output_table)
-
     # Create a second job in the same run group with active checkpoints
     reset_session_job_state()
     job2_id = metastore.create_job(
@@ -208,6 +205,14 @@ def test_cleanup_preserves_input_tables_when_run_group_active(
     monkeypatch.setenv("DATACHAIN_JOB_ID", job2_id)
     chain.map(tripled=lambda num: num * 3, output=int).save("nums_tripled")
 
+    # Record job1's UDF output tables before cleanup
+    job1_output_tables = [
+        t
+        for t in warehouse.db.list_tables(pattern=f"udf_{job1_id}_%")
+        if t.endswith("_output")
+    ]
+    assert len(job1_output_tables) > 0
+
     # Make only job 1's checkpoints outdated
     old_time = datetime.now(timezone.utc) - timedelta(hours=5)
     for ch in metastore.list_checkpoints(job1_id):
@@ -219,7 +224,8 @@ def test_cleanup_preserves_input_tables_when_run_group_active(
 
     catalog.cleanup_checkpoints()
 
-    assert not warehouse.db.has_table(output_table)
+    # Job 1's output tables removed
+    assert all(not warehouse.db.has_table(t) for t in job1_output_tables)
     assert len(list(metastore.list_checkpoints(job1_id))) == 0
     assert len(list(metastore.list_checkpoints(job2_id))) > 0
     # Shared input table preserved (run group still active)

tests/func/checkpoints/test_checkpoint_recovery.py

@@ -586,7 +586,7 @@ def mapper(num) -> int:
     catalog = test_session.catalog
 
     # Drop all UDF output tables from first run
-    for table_name in catalog.warehouse.db.list_tables(prefix="udf_"):
+    for table_name in catalog.warehouse.db.list_tables(pattern="udf_%"):
         if "_output" in table_name and "_partial" not in table_name:
             table = catalog.warehouse.db.get_table(table_name)
             catalog.warehouse.db.drop_table(table, if_exists=True)
@@ -622,7 +622,7 @@ def mapper(num) -> int:
     test_session.get_or_create_job()
 
     # Drop all partial output tables from first run
-    for table_name in catalog.warehouse.db.list_tables(prefix="udf_"):
+    for table_name in catalog.warehouse.db.list_tables(pattern="udf_%"):
         if "_partial" in table_name:
             table = catalog.warehouse.db.get_table(table_name)
             catalog.warehouse.db.drop_table(table, if_exists=True)

tests/unit/test_data_storage.py

@@ -95,8 +95,8 @@ def test_list_tables(catalog):
         tables_after = db.list_tables()
         assert "test_list_tables_abc" in tables_after
 
-        # Test with prefix filter
-        filtered = db.list_tables(prefix="test_list_tables")
+        # Test with pattern filter
+        filtered = db.list_tables(pattern="test_list_tables%")
         assert "test_list_tables_abc" in filtered
     finally:
         db.drop_table(table)