Formatting

Author: Nitin Kanukolanu

redisvl/utils/migration.py

@@ -26,26 +26,26 @@
 
 class IndexMigrator:
     """Helper class to migrate indices to SVS-VAMANA with compression.
-    
+
     This class provides utilities to migrate existing FLAT or HNSW vector indices
     to SVS-VAMANA indices with compression, enabling significant memory savings.
-    
+
     Example:
         .. code-block:: python
-        
+
             from redisvl.index import SearchIndex
             from redisvl.utils import IndexMigrator
-            
+
             # Load existing index
             old_index = SearchIndex.from_existing("my_flat_index")
-            
+
             # Migrate to SVS-VAMANA with LVQ compression
             new_index = IndexMigrator.migrate_to_svs(
                 old_index,
                 compression="LVQ4x4",
                 batch_size=1000
             )
-            
+
             print(f"Migrated {new_index.info()['num_docs']} documents")
     """
 
@@ -60,10 +60,10 @@ def migrate_to_svs(
         progress_callback: Optional[Callable[[int, int], None]] = None,
     ) -> "SearchIndex":
         """Migrate an existing index to SVS-VAMANA with compression.
-        
+
         This method creates a new SVS-VAMANA index and copies all data from the
         old index in batches. The old index is not modified or deleted.
-        
+
         Args:
             old_index: The existing SearchIndex to migrate from.
             new_index_name: Name for the new index. If None, uses "{old_name}_svs".
@@ -74,20 +74,20 @@ def migrate_to_svs(
             batch_size: Number of documents to migrate per batch. Default: 1000.
             overwrite: Whether to overwrite the new index if it exists. Default: False.
             progress_callback: Optional callback function(current, total) for progress tracking.
-        
+
         Returns:
             SearchIndex: The new SVS-VAMANA index with migrated data.
-        
+
         Raises:
             RedisModuleVersionError: If Redis version doesn't support SVS-VAMANA.
             ValueError: If the old index has no vector fields or invalid parameters.
-        
+
         Example:
             .. code-block:: python
-            
+
                 def progress(current, total):
                     print(f"Migrated {current}/{total} documents")
-                
+
                 new_index = IndexMigrator.migrate_to_svs(
                     old_index,
                     compression="LVQ4x4",
@@ -106,35 +106,35 @@ def progress(current, total):
             raise RedisModuleVersionError.for_svs_vamana(
                 caps.redis_version, caps.search_version
             )
-        
+
         # Find vector fields in the old schema
         vector_fields = [
             (name, field)
             for name, field in old_index.schema.fields.items()
             if hasattr(field, "attrs") and hasattr(field.attrs, "algorithm")
         ]
-        
+
         if not vector_fields:
             raise ValueError("Old index has no vector fields to migrate")
-        
+
         # Create new schema based on old schema
         new_schema_dict = old_index.schema.to_dict()
-        
+
         # Update index name
         if new_index_name is None:
             new_index_name = f"{old_index.name}_svs"
         new_schema_dict["index"]["name"] = new_index_name
         new_schema_dict["index"]["prefix"] = new_index_name
-        
+
         # Update vector fields to use SVS-VAMANA
         for field_dict in new_schema_dict["fields"]:
             if field_dict["type"] == "vector":
                 attrs = field_dict.get("attrs", {})
                 dims = attrs.get("dims")
-                
+
                 if dims is None:
                     raise ValueError(f"Vector field '{field_dict['name']}' has no dims")
-                
+
                 # Use CompressionAdvisor if compression not specified
                 if compression is None:
                     config = CompressionAdvisor.recommend(
@@ -149,14 +149,14 @@ def progress(current, total):
                         f"CompressionAdvisor recommended: {compression} "
                         f"(reduce={reduce}) for {dims} dims"
                     )
-                
+
                 # Update to SVS-VAMANA
                 attrs["algorithm"] = "svs-vamana"
                 attrs["compression"] = compression
-                
+
                 if reduce is not None:
                     attrs["reduce"] = reduce
-                
+
                 # Set default SVS parameters if not present
                 if "graph_max_degree" not in attrs:
                     attrs["graph_max_degree"] = 40
@@ -168,31 +168,31 @@ def progress(current, total):
                 # Default is 10240, minimum is 1024 (DEFAULT_BLOCK_SIZE)
                 if "training_threshold" not in attrs:
                     attrs["training_threshold"] = 1024
-        
+
         # Create new index
         new_schema = IndexSchema.from_dict(new_schema_dict)
         new_index = SearchIndex(schema=new_schema, redis_client=old_index._redis_client)
         new_index.create(overwrite=overwrite)
-        
+
         logger.info(f"Created new SVS-VAMANA index: {new_index_name}")
-        
+
         # Get total document count
         old_info = old_index.info()
         total_docs = int(old_info.get("num_docs", 0))
-        
+
         if total_docs == 0:
             logger.warning("Old index has no documents to migrate")
             return new_index
-        
+
         logger.info(f"Migrating {total_docs} documents in batches of {batch_size}")
-        
+
         # Migrate data in batches using pagination
         migrated_count = 0
         query = FilterQuery(
             filter_expression=FilterExpression("*"),
             return_fields=list(old_index.schema.fields.keys()),
         )
-        
+
         for batch in old_index.paginate(query, page_size=batch_size):
             if batch:
                 # The 'id' field contains the full Redis key (e.g., "prefix:ulid")
@@ -223,7 +223,7 @@ def progress(current, total):
                     progress_callback(migrated_count, total_docs)
 
                 logger.debug(f"Migrated {migrated_count}/{total_docs} documents")
-        
+
         logger.info(f"Migration complete: {migrated_count} documents migrated")
 
         # Verify migration by checking Redis keys (not index count)
@@ -262,10 +262,10 @@ async def migrate_to_svs_async(
         progress_callback: Optional[Callable[[int, int], None]] = None,
     ) -> "AsyncSearchIndex":
         """Asynchronously migrate an existing index to SVS-VAMANA with compression.
-        
+
         This is the async version of migrate_to_svs(). See migrate_to_svs() for
         detailed documentation.
-        
+
         Args:
             old_index: The existing AsyncSearchIndex to migrate from.
             new_index_name: Name for the new index. If None, uses "{old_name}_svs".
@@ -274,16 +274,16 @@ async def migrate_to_svs_async(
             batch_size: Number of documents to migrate per batch. Default: 1000.
             overwrite: Whether to overwrite the new index if it exists. Default: False.
             progress_callback: Optional callback function(current, total) for progress.
-        
+
         Returns:
             AsyncSearchIndex: The new SVS-VAMANA index with migrated data.
-        
+
         Example:
             .. code-block:: python
-            
+
                 async def progress(current, total):
                     print(f"Migrated {current}/{total} documents")
-                
+
                 new_index = await IndexMigrator.migrate_to_svs_async(
                     old_index,
                     compression="LVQ4x4",
@@ -303,34 +303,34 @@ async def progress(current, total):
             raise RedisModuleVersionError.for_svs_vamana(
                 caps.redis_version, caps.search_version
             )
-        
+
         # Find vector fields
         vector_fields = [
             (name, field)
             for name, field in old_index.schema.fields.items()
             if hasattr(field, "attrs") and hasattr(field.attrs, "algorithm")
         ]
-        
+
         if not vector_fields:
             raise ValueError("Old index has no vector fields to migrate")
-        
+
         # Create new schema
         new_schema_dict = old_index.schema.to_dict()
-        
+
         if new_index_name is None:
             new_index_name = f"{old_index.name}_svs"
         new_schema_dict["index"]["name"] = new_index_name
         new_schema_dict["index"]["prefix"] = new_index_name
-        
+
         # Update vector fields
         for field_dict in new_schema_dict["fields"]:
             if field_dict["type"] == "vector":
                 attrs = field_dict.get("attrs", {})
                 dims = attrs.get("dims")
-                
+
                 if dims is None:
                     raise ValueError(f"Vector field '{field_dict['name']}' has no dims")
-                
+
                 if compression is None:
                     config = CompressionAdvisor.recommend(
                         dims=dims,
@@ -344,13 +344,13 @@ async def progress(current, total):
                         f"CompressionAdvisor recommended: {compression} "
                         f"(reduce={reduce}) for {dims} dims"
                     )
-                
+
                 attrs["algorithm"] = "svs-vamana"
                 attrs["compression"] = compression
-                
+
                 if reduce is not None:
                     attrs["reduce"] = reduce
-                
+
                 if "graph_max_degree" not in attrs:
                     attrs["graph_max_degree"] = 40
                 if "construction_window_size" not in attrs:
@@ -364,26 +364,26 @@ async def progress(current, total):
         new_schema = IndexSchema.from_dict(new_schema_dict)
         new_index = AsyncSearchIndex(schema=new_schema, redis_client=client)
         await new_index.create(overwrite=overwrite)
-        
+
         logger.info(f"Created new SVS-VAMANA index: {new_index_name}")
-        
+
         # Get total document count
         old_info = await old_index.info()
         total_docs = int(old_info.get("num_docs", 0))
-        
+
         if total_docs == 0:
             logger.warning("Old index has no documents to migrate")
             return new_index
-        
+
         logger.info(f"Migrating {total_docs} documents in batches of {batch_size}")
-        
+
         # Migrate data in batches
         migrated_count = 0
         query = FilterQuery(
             filter_expression=FilterExpression("*"),
             return_fields=list(old_index.schema.fields.keys()),
         )
-        
+
         async for batch in old_index.paginate(query, page_size=batch_size):
             if batch:
                 # Extract document IDs from full Redis keys
@@ -406,7 +406,7 @@ async def progress(current, total):
                     progress_callback(migrated_count, total_docs)
 
                 logger.debug(f"Migrated {migrated_count}/{total_docs} documents")
-        
+
         logger.info(f"Migration complete: {migrated_count} documents migrated")
 
         # Verify migration by checking Redis keys
@@ -430,4 +430,3 @@ async def progress(current, total):
                 )
 
         return new_index
-