Update search_utils.py

search_utils.py

@@ -21,60 +21,104 @@ logger = logging.getLogger("MetadataManager")
 
 class MetadataManager:
     def __init__(self):
-        self.shard_dir = Path("metadata_shards")
+        self.cache_dir = Path("unzipped_cache")
+        self.shard_dir = self.cache_dir / "metadata_shards"
         self.shard_map = {}
         self.loaded_shards = {}
         self.total_docs = 0
         
         logger.info("Initializing MetadataManager")
-        self._ensure_unzipped()
+        self._ensure_directories()
+        self._unzip_if_needed()
         self._build_shard_map()
         logger.info(f"Total documents indexed: {self.total_docs}")
         logger.info(f"Total shards found: {len(self.shard_map)}")
 
-    def _ensure_unzipped(self):
-        """Handle ZIP extraction without Streamlit elements"""
-        logger.info(f"Checking for shard directory: {self.shard_dir}")
-        if not self.shard_dir.exists():
-            zip_path = Path("metadata_shards.zip")
-            logger.info(f"Shard directory not found, looking for zip file: {zip_path}")
-            if zip_path.exists():
-                logger.info(f"Extracting from zip file: {zip_path}")
-                start_time = time.time()
+    def _ensure_directories(self):
+        """Create necessary directories if they don't exist"""
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.shard_dir.mkdir(parents=True, exist_ok=True)
+
+    def _unzip_if_needed(self):
+        """Handle ZIP extraction with validation and retries"""
+        zip_path = Path("metadata_shards.zip")
+        
+        # Check if we need to unzip
+        if not any(self.shard_dir.glob("*.parquet")):
+            logger.info("No parquet files found, checking for zip archive")
+            
+            if not zip_path.exists():
+                raise FileNotFoundError(f"Metadata ZIP file not found at {zip_path}")
+                
+            logger.info(f"Extracting {zip_path} to {self.shard_dir}")
+            try:
                 with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+                    # Validate zip contents before extraction
+                    zip_files = zip_ref.namelist()
+                    if not any(fname.endswith('.parquet') for fname in zip_files):
+                        raise ValueError("ZIP file contains no parquet files")
+                        
                     zip_ref.extractall(self.shard_dir)
-                logger.info(f"Extraction completed in {time.time() - start_time:.2f} seconds")
-            else:
-                error_msg = "Metadata ZIP file not found"
-                logger.error(error_msg)
-                raise FileNotFoundError(error_msg)
-        else:
-            logger.info("Shard directory exists, skipping extraction")
+                    logger.info(f"Extracted {len(zip_files)} files")
+                    
+                    # Verify extraction succeeded
+                    if not any(self.shard_dir.glob("*.parquet")):
+                        raise RuntimeError("Extraction completed but no parquet files found")
+                        
+            except Exception as e:
+                logger.error(f"Failed to extract zip file: {str(e)}")
+                # Clean up partial extraction
+                if any(self.shard_dir.iterdir()):
+                    for f in self.shard_dir.glob("*"):
+                        f.unlink()
+                raise
                 
+        else:
+            logger.info("Parquet files already exist in cache directory")
+
     def _build_shard_map(self):
-        """Create index range to shard mapping"""
+        """Create validated index range to shard mapping"""
         logger.info("Building shard map from parquet files")
-        self.total_docs = 0
-        shard_files = list(self.shard_dir.glob("*.parquet"))
-        logger.info(f"Found {len(shard_files)} parquet files")
+        parquet_files = list(self.shard_dir.glob("*.parquet"))
         
-        if not shard_files:
-            logger.warning("No parquet files found in shard directory")
+        if not parquet_files:
+            raise FileNotFoundError("No parquet files found after extraction")
             
-        for f in sorted(shard_files):
+        # Sort files by numerical order
+        parquet_files = sorted(parquet_files, key=lambda x: int(x.stem.split("_")[1]))
+        
+        # Track expected next index
+        expected_start = 0
+        
+        for f in parquet_files:
             try:
                 parts = f.stem.split("_")
-                if len(parts) < 3:
-                    logger.warning(f"Skipping file with invalid name format: {f}")
-                    continue
+                if len(parts) != 3:
+                    raise ValueError("Invalid filename format")
                     
                 start = int(parts[1])
                 end = int(parts[2])
+                
+                # Validate continuity
+                if start != expected_start:
+                    raise ValueError(f"Non-contiguous shard start: expected {expected_start}, got {start}")
+                    
+                # Validate range
+                if end <= start:
+                    raise ValueError(f"Invalid shard range: {start}-{end}")
+                    
                 self.shard_map[(start, end)] = f.name
-                self.total_docs = max(self.total_docs, end + 1)
+                self.total_docs = end + 1
+                expected_start = end + 1
+                
                 logger.debug(f"Mapped shard {f.name}: indices {start}-{end}")
+                
             except Exception as e:
-                logger.error(f"Error parsing shard filename {f}: {str(e)}")
+                logger.error(f"Error processing shard {f.name}: {str(e)}")
+                raise RuntimeError("Invalid shard structure") from e
+                
+        logger.info(f"Validated {len(self.shard_map)} continuous shards")
+        logger.info(f"Total document count: {self.total_docs}")
         
         # Log shard statistics
         logger.info(f"Shard map built with {len(self.shard_map)} shards")