Update search_utils.py

search_utils.py

@@ -40,11 +40,11 @@ class MetadataManager:
         self.shard_dir.mkdir(parents=True, exist_ok=True)
 
     def _unzip_if_needed(self):
-        """Handle ZIP extraction with validation and retries"""
+        """Handle ZIP extraction with nested directory handling"""
         zip_path = Path("metadata_shards.zip")
         
-        # Check if we need to unzip
-        if not any(self.shard_dir.glob("*.parquet")):
+        # Check if we need to unzip by looking for parquet files in any subdirectory
+        if not any(self.shard_dir.rglob("*.parquet")):
             logger.info("No parquet files found, checking for zip archive")
             
             if not zip_path.exists():
@@ -53,28 +53,63 @@ class MetadataManager:
             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")
-                        
+                    # Check for nested directory structure in zip
+                    zip_root = self._get_zip_root(zip_ref)
+                    
+                    # Extract while preserving structure
                     zip_ref.extractall(self.shard_dir)
-                    logger.info(f"Extracted {len(zip_files)} files")
                     
-                    # Verify extraction succeeded
-                    if not any(self.shard_dir.glob("*.parquet")):
+                    # Handle nested directory if exists
+                    if zip_root:
+                        nested_dir = self.shard_dir / zip_root
+                        if nested_dir.exists():
+                            # Move files up from nested directory
+                            self._flatten_directory(nested_dir, self.shard_dir)
+                            nested_dir.rmdir()
+
+                    # Verify extraction
+                    parquet_files = list(self.shard_dir.rglob("*.parquet"))
+                    if not parquet_files:
                         raise RuntimeError("Extraction completed but no parquet files found")
                         
+                    logger.info(f"Found {len(parquet_files)} parquet files after extraction")
+                    
             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()
+                self._clean_failed_extraction()
                 raise
-                
-        else:
-            logger.info("Parquet files already exist in cache directory")
+
+    def _get_zip_root(self, zip_ref):
+        """Identify common root directory in zip file"""
+        try:
+            first_file = zip_ref.namelist()[0]
+            if '/' in first_file:
+                return first_file.split('/')[0]
+            return ""
+        except Exception as e:
+            logger.warning(f"Error detecting zip root: {str(e)}")
+            return ""
+
+    def _flatten_directory(self, src_dir, dest_dir):
+        """Move files from nested directory to destination"""
+        for item in src_dir.iterdir():
+            if item.is_dir():
+                self._flatten_directory(item, dest_dir)
+                item.rmdir()
+            else:
+                target = dest_dir / item.name
+                if target.exists():
+                    target.unlink()
+                item.rename(target)
+
+    def _clean_failed_extraction(self):
+        """Remove any extracted files after failed attempt"""
+        logger.info("Cleaning up failed extraction")
+        for item in self.shard_dir.iterdir():
+            if item.is_dir():
+                shutil.rmtree(item)
+            else:
+                item.unlink()
 
     def _build_shard_map(self):
         """Create validated index range to shard mapping"""