improve db

data_loader.py

@@ -121,8 +121,11 @@ def refresh_data() -> List[Dict[str, Any]]:
     df["tags"] = df["tags"].apply(ensure_list_of_strings)
     df["language"] = df["language"].apply(ensure_list_of_strings)
 
-    # Convert 'features' column to string
-    df["features"] = df["features"].apply(lambda x: str(x) if x is not None else None)
+    # Ensure 'column_names' is a list
+    df["column_names"] = df["column_names"].apply(
+        lambda x: x if isinstance(x, list) else []
+    )
+
     df = df.astype({"hub_id": "string", "config_name": "string"})
 
     # save to parquet file with current date

main.py

@@ -20,6 +20,8 @@ logger = logging.getLogger(__name__)
 def get_db_connection():
     conn = sqlite3.connect("datasets.db")
     conn.row_factory = sqlite3.Row
+    conn.execute("PRAGMA journal_mode = WAL")
+    conn.execute("PRAGMA synchronous = NORMAL")
     return conn
 
 
@@ -31,17 +33,29 @@ def setup_database():
                  (hub_id TEXT PRIMARY KEY, 
                   likes INTEGER,
                   downloads INTEGER,
-                  tags TEXT,
+                  tags JSON,
                   created_at INTEGER,
                   last_modified INTEGER,
-                  license TEXT,
-                  language TEXT,
+                  license JSON,
+                  language JSON,
                   config_name TEXT,
-                  column_names TEXT,
-                  features TEXT)"""
+                  column_names JSON,
+                  features JSON)"""
+    )
+    c.execute(
+        """
+    CREATE INDEX IF NOT EXISTS idx_column_names 
+    ON datasets((json_each.value))
+    """
+    )
+    c.execute(
+        """
+    CREATE INDEX IF NOT EXISTS idx_downloads_likes
+    ON datasets(downloads DESC, likes DESC)
+    """
     )
-    c.execute("CREATE INDEX IF NOT EXISTS idx_column_names ON datasets (column_names)")
     conn.commit()
+    c.execute("ANALYZE")
     conn.close()
 
 
@@ -58,56 +72,46 @@ def serialize_numpy(obj):
     raise TypeError(f"Object of type {type(obj)} is not JSON serializable")
 
 
-def insert_data(conn, data):
-    c = conn.cursor()
-
-    created_at = data.get("created_at", 0)
-    if isinstance(created_at, Timestamp):
-        created_at = int(created_at.timestamp())
-
-    last_modified = data.get("last_modified", 0)
-    if isinstance(last_modified, Timestamp):
-        last_modified = int(last_modified.timestamp())
-
-    c.execute(
-        """
-        INSERT OR REPLACE INTO datasets 
-        (hub_id, likes, downloads, tags, created_at, last_modified, license, language, config_name, column_names, features) 
-        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
-    """,
-        (
-            data["hub_id"],
-            data.get("likes", 0),
-            data.get("downloads", 0),
-            json.dumps(data.get("tags", []), default=serialize_numpy),
-            created_at,
-            last_modified,
-            json.dumps(data.get("license", []), default=serialize_numpy),
-            json.dumps(data.get("language", []), default=serialize_numpy),
-            data.get("config_name", ""),
-            json.dumps(data.get("column_names", []), default=serialize_numpy),
-            json.dumps(data.get("features", []), default=serialize_numpy),
-        ),
-    )
-    conn.commit()
-
-
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Startup: Load data into the database
     setup_database()
     logger.info("Creating database connection")
     conn = get_db_connection()
     logger.info("Refreshing data")
     datasets = refresh_data()
 
-    for data in datasets:
-        insert_data(conn, data)
+    c = conn.cursor()
+    c.executemany(
+        """
+        INSERT OR REPLACE INTO datasets 
+        (hub_id, likes, downloads, tags, created_at, last_modified, license, language, config_name, column_names, features) 
+        VALUES (?, ?, ?, json(?), ?, ?, json(?), json(?), ?, json(?), json(?))
+        """,
+        [
+            (
+                data["hub_id"],
+                data.get("likes", 0),
+                data.get("downloads", 0),
+                json.dumps(data.get("tags", []), default=serialize_numpy),
+                int(data["created_at"].timestamp())
+                if isinstance(data["created_at"], Timestamp)
+                else data.get("created_at", 0),
+                int(data["last_modified"].timestamp())
+                if isinstance(data["last_modified"], Timestamp)
+                else data.get("last_modified", 0),
+                json.dumps(data.get("license", []), default=serialize_numpy),
+                json.dumps(data.get("language", []), default=serialize_numpy),
+                data.get("config_name", ""),
+                json.dumps(data.get("column_names", []), default=serialize_numpy),
+                json.dumps(data.get("features", []), default=serialize_numpy),
+            )
+            for data in datasets
+        ],
+    )
+    conn.commit()
     conn.close()
     logger.info("Data refreshed")
     yield
-    # Shutdown: You can add any cleanup operations here if needed
-    # For example, closing database connections, clearing caches, etc.
 
 
 app = FastAPI(lifespan=lifespan)
@@ -140,43 +144,55 @@ async def search_datasets(
     try:
         if match_all:
             query = """
-            SELECT COUNT(*) as total FROM datasets
-            WHERE (SELECT COUNT(*) FROM json_each(column_names)
-                   WHERE value IN ({})) = ?
+            SELECT *, (
+                SELECT COUNT(*)
+                FROM json_each(column_names)
+                WHERE json_each.value IN ({})
+            ) as match_count
+            FROM datasets
+            WHERE match_count = ?
+            ORDER BY downloads DESC, likes DESC
+            LIMIT ? OFFSET ?
             """.format(",".join("?" * len(columns)))
-            c.execute(query, (*columns, len(columns)))
+            c.execute(query, (*columns, len(columns), page_size, offset))
         else:
             query = """
-            SELECT COUNT(*) as total FROM datasets
+            SELECT * FROM datasets
             WHERE EXISTS (
-                SELECT 1 FROM json_each(column_names)
-                WHERE value IN ({})
+                SELECT 1
+                FROM json_each(column_names)
+                WHERE json_each.value IN ({})
             )
+            ORDER BY downloads DESC, likes DESC
+            LIMIT ? OFFSET ?
             """.format(",".join("?" * len(columns)))
-            c.execute(query, columns)
+            c.execute(query, (*columns, page_size, offset))
 
-        total = c.fetchone()["total"]
+        results = [dict(row) for row in c.fetchall()]
 
+        # Get total count
         if match_all:
-            query = """
-            SELECT * FROM datasets
-            WHERE (SELECT COUNT(*) FROM json_each(column_names)
-                   WHERE value IN ({})) = ?
-            LIMIT ? OFFSET ?
+            count_query = """
+            SELECT COUNT(*) as total FROM datasets
+            WHERE (
+                SELECT COUNT(*)
+                FROM json_each(column_names)
+                WHERE json_each.value IN ({})
+            ) = ?
             """.format(",".join("?" * len(columns)))
-            c.execute(query, (*columns, len(columns), page_size, offset))
+            c.execute(count_query, (*columns, len(columns)))
         else:
-            query = """
-            SELECT * FROM datasets
+            count_query = """
+            SELECT COUNT(*) as total FROM datasets
             WHERE EXISTS (
-                SELECT 1 FROM json_each(column_names)
-                WHERE value IN ({})
+                SELECT 1
+                FROM json_each(column_names)
+                WHERE json_each.value IN ({})
             )
-            LIMIT ? OFFSET ?
             """.format(",".join("?" * len(columns)))
-            c.execute(query, (*columns, page_size, offset))
+            c.execute(count_query, columns)
 
-        results = [dict(row) for row in c.fetchall()]
+        total = c.fetchone()["total"]
 
         for result in results:
             result["tags"] = json.loads(result["tags"])