feat: optimize LanceDB compaction and FTS indexing

Jay-ju · Jay-ju · commit d6d59ce95783 · 2026-04-10T14:27:55.000+08:00
- Refactored LanceGraphDB to use native ds.optimize() for compaction and snapshot pruning.
- Automated FTS index creation in _ensure_schema.
- Added configurable cleanup_older_than_days property.
- Added pytest case to verify FTS index effectiveness and optimization.
diff --git a/src/memos/configs/graph_db.py b/src/memos/configs/graph_db.py
@@ -252,6 +252,15 @@ class LanceGraphDBConfig(BaseConfig):
         description="Logical user or tenant ID for data isolation",
     )
     embedding_dimension: int = Field(default=768, description="Dimension of vector embedding")
+    compaction_version_threshold: int = Field(
+        default=500, description="Number of new versions to accumulate before triggering compaction"
+    )
+    compaction_interval_mins: int = Field(
+        default=30, description="Fallback interval in minutes to check and run compaction"
+    )
+    cleanup_older_than_days: int = Field(
+        default=7, description="Number of days to keep old versions before pruning"
+    )
 
 
 class GraphDBConfigFactory(BaseModel):
diff --git a/src/memos/graph_dbs/lance.py b/src/memos/graph_dbs/lance.py
@@ -2,9 +2,11 @@
 
 import json
 import os
+import threading
+import time
 import uuid
 
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any
 
 
@@ -35,11 +37,123 @@ def __init__(self, config: LanceGraphDBConfig):
         self.user_name = config.user_name or "default"
         self.dim = config.embedding_dimension
 
+        # Compaction settings
+        self.compaction_version_threshold = config.compaction_version_threshold
+        self.compaction_interval_mins = config.compaction_interval_mins
+        self.cleanup_older_than_days = config.cleanup_older_than_days
+
         self.memories_uri = os.path.join(self.uri, "memories")
         self.edges_uri = os.path.join(self.uri, "edges")
 
         self._init_schema()
 
+        # Start LanceDB background optimizer thread
+        self._last_compact_versions = {
+            "memories": self._get_memories_table().version,
+            "edges": self._get_edges_table().version,
+        }
+        self._optimizer_thread = threading.Thread(
+            target=self._db_optimizer_loop,
+            daemon=True,
+            name="lancedb-optimizer",
+        )
+        self._optimizer_thread.start()
+
+    def _db_optimizer_loop(self):
+        """Background loop to periodically trigger table optimization."""
+        import schedule
+
+        schedule.every(self.compaction_interval_mins).minutes.do(self._force_optimize)
+
+        logger.info(
+            f"Started LanceDB optimizer thread. Compaction interval: {self.compaction_interval_mins}m, "
+            f"Version threshold: {self.compaction_version_threshold}"
+        )
+
+        while True:
+            try:
+                # 1. Check version threshold
+                self._check_and_trigger_compaction()
+
+                # 2. Run scheduled fallback compaction
+                schedule.run_pending()
+            except Exception as e:
+                logger.error(f"Error in LanceDB optimizer loop: {e}", stack_info=True)
+
+            time.sleep(5)  # Avoid busy waiting
+
+    def _check_and_trigger_compaction(self):
+        """Trigger compaction if any table's version diff exceeds the threshold."""
+        try:
+            memories_ds = self._get_memories_table()
+            if (
+                memories_ds.version - self._last_compact_versions["memories"]
+                > self.compaction_version_threshold
+            ):
+                self._optimize_table("memories", memories_ds)
+
+            edges_ds = self._get_edges_table()
+            if (
+                edges_ds.version - self._last_compact_versions["edges"]
+                > self.compaction_version_threshold
+            ):
+                self._optimize_table("edges", edges_ds)
+        except Exception as e:
+            logger.error(f"Failed to check compaction versions: {e}")
+
+    def _optimize_table(self, table_name: str, ds):
+        """Helper method to optimize a specific LanceDB table."""
+        try:
+            current_version = ds.version
+            last_version = self._last_compact_versions[table_name]
+
+            if current_version > last_version:
+                logger.info(
+                    f"Triggering LanceDB optimization for '{table_name}'. "
+                    f"Current version: {current_version}, Last compacted: {last_version}"
+                )
+
+                stats = ds.optimize(cleanup_older_than=timedelta(days=self.cleanup_older_than_days))
+
+                stats_msg = ""
+                if stats:
+                    compaction = getattr(stats, "compaction", None)
+                    if compaction:
+                        stats_msg += (
+                            f" | Compaction: "
+                            f"-{getattr(compaction, 'fragments_removed', 0)}/"
+                            f"+{getattr(compaction, 'fragments_added', 0)} fragments, "
+                            f"-{getattr(compaction, 'files_removed', 0)}/"
+                            f"+{getattr(compaction, 'files_added', 0)} files"
+                        )
+
+                    prune = getattr(stats, "prune", None)
+                    if prune:
+                        stats_msg += (
+                            f" | Prune: -{getattr(prune, 'bytes_removed', 0)} bytes, "
+                            f"-{getattr(prune, 'old_versions_removed', 0)} versions"
+                        )
+
+                # Reload the table to get the updated version after optimization
+                if table_name == "memories":
+                    ds = self._get_memories_table()
+                elif table_name == "edges":
+                    ds = self._get_edges_table()
+
+                self._last_compact_versions[table_name] = ds.version
+                logger.info(
+                    f"LanceDB '{table_name}' optimization completed successfully. "
+                    f"New version: {self._last_compact_versions[table_name]}{stats_msg}"
+                )
+        except Exception as e:
+            logger.error(f"LanceDB '{table_name}' optimization failed: {e}")
+
+    def _force_optimize(self):
+        # Optimize Memories Table
+        self._optimize_table("memories", self._get_memories_table())
+        # Optimize Edges Table
+        self._optimize_table("edges", self._get_edges_table())
+
     def _init_schema(self):
         import lancedb
         import pyarrow as pa
@@ -74,6 +188,34 @@ def _init_schema(self):
             self.db.create_table("memories", data=empty_table)
             logger.info("Created LanceDB table for memories.")
 
+            try:
+                ds = self.db.open_table("memories")
+
+                # Create vector index (aligned with memory-lancedb TS implementation)
+                import math
+
+                row_count = ds.count_rows()
+                if row_count > 256:  # LanceDB requires at least 256 rows to train vector index
+                    num_partitions = max(1, math.floor(math.sqrt(row_count)))
+                    ds.create_index(
+                        metric="cosine",
+                        vector_column_name="embedding",
+                        num_partitions=num_partitions,
+                    )
+                    logger.info(
+                        f"Created IVF_FLAT index for memories.embedding with metric=cosine, partitions={num_partitions}"
+                    )
+                else:
+                    logger.debug(
+                        f"Skipping vector index creation, not enough rows ({row_count} <= 256)"
+                    )
+
+                # Create full-text search index
+                ds.create_fts_index("memory", replace=True)
+                logger.info("Created FTS index for memories.memory")
+            except Exception as e:
+                logger.warning(f"Failed to create LanceDB indices: {e}")
+
         if "edges" not in table_names:
             edge_schema = pa.schema(
                 [
diff --git a/tests/graph_dbs/test_lance.py b/tests/graph_dbs/test_lance.py
@@ -129,5 +129,73 @@ def test_lance_graph_db():
         print("Test finished successfully in temporary directory!")
 
 
+def test_lance_compaction_and_fts_effectiveness():
+    """
+    Test the effectiveness of the LanceDB _optimize_table mechanism,
+    including compaction of small files and FTS index functionality.
+    """
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        db_uri = os.path.join(tmpdir, "test_lancedb_compaction")
+        # Use a low threshold to force triggering
+        config = LanceGraphDBConfig(
+            uri=db_uri, user_name="test_user", embedding_dimension=3, compaction_version_threshold=2
+        )
+        db = LanceGraphDB(config)
+
+        # 1. Insert multiple single nodes to create small fragments
+        print("\nInserting 5 separate fragments...")
+        for i in range(5):
+            node = {
+                "id": f"node_c_{i}",
+                "memory": f"Alice went to the magical forest number {i}",
+                "metadata": {"memory_type": "LongTermMemory", "status": "activated"},
+                "embedding": [0.1 * i, 0.2 * i, 0.3 * i],
+            }
+            db.add_nodes_batch([node])
+
+        import lance
+
+        ds = lance.dataset(os.path.join(db_uri, "memories.lance"))
+        fragments_before = len(ds.get_fragments())
+        print(f"Fragments BEFORE optimize: {fragments_before}")
+
+        # 2. Test FTS before optimization
+        try:
+            res_fts_before = db.search_by_fulltext(["magical"], top_k=10)
+            print(f"FTS hits BEFORE optimize: {len(res_fts_before)}")
+        except Exception as e:
+            print(f"FTS failed before optimize: {e}")
+
+        # 3. Force the internal optimizer
+        print("Forcing LanceDB optimizer...")
+        db._force_optimize()
+
+        ds = lance.dataset(os.path.join(db_uri, "memories.lance"))
+        fragments_after = len(ds.get_fragments())
+        print(f"Fragments AFTER optimize: {fragments_after}")
+
+        # 5. Verify FTS index effectiveness after optimization
+        res_fts_after = db.search_by_fulltext(["magical"], top_k=10)
+        assert len(res_fts_after) == 5, (
+            f"FTS should recall all 5 nodes, but got {len(res_fts_after)}"
+        )
+        print(f"FTS hits AFTER optimize: {len(res_fts_after)}")
+
+        # 6. Test prune/delete
+        db.delete_node("node_c_0")
+        db._force_optimize()
+
+        res_fts_deleted = db.search_by_fulltext(["magical"], top_k=10)
+        assert len(res_fts_deleted) == 4, (
+            f"FTS should recall 4 nodes after deletion, got {len(res_fts_deleted)}"
+        )
+        print(f"FTS hits AFTER deletion and optimize: {len(res_fts_deleted)}")
+
+        db.clear()
+
+
 if __name__ == "__main__":
     test_lance_graph_db()
+    test_lance_compaction_and_fts_effectiveness()
