Switch back to fast sorted-index app (old dataset) + num/aadhar aliases + async pinger

app.py

@@ -1,52 +1,45 @@
 """Gradio UI + FastAPI backend for ICMR + HITEK Search API.
 
 Runs on Hugging Face Spaces (free Gradio tier).
-Queries the raw parquet parts of Kzr0xx/Icmr-and-hitek directly via DuckDB
-(httpfs range reads). Each data file is scanned in its own thread with a
-LIMIT + timeout so responses stay fast (~1-3s) even on 100GB remote data.
+Uses remote HF sorted split indexes — no local database download needed.
 """
 import asyncio
 import json
 import os
 import threading
-import time
-from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutTimeout
+from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 
 import duckdb
 import gradio as gr
+import httpx  # <-- added for pinger
 from fastapi import FastAPI, HTTPException, Query, Response
 from pydantic import BaseModel
 
 # ── Config ──────────────────────────────────────────────────────────────────
 BASE = os.path.dirname(os.path.abspath(__file__))
-HF_DATA_BASE = os.environ.get(
-    "ICMR_HF_DATA_BASE",
-    "https://huggingface.co/datasets/Kzr0xx/Icmr-and-hitek/resolve/main",
+HF_INDEX_BASE = os.environ.get(
+    "ICMR_HF_INDEX_BASE",
+    "https://huggingface.co/datasets/Kzr0xx/icrm-hitek-full-db-mixed/resolve/main",
 ).rstrip("/")
-DATA_FILES = [
-    part.strip()
-    for part in os.environ.get(
-        "ICMR_DATA_FILES", "part1.parquet,part2a.parquet,part2b_new.parquet"
-    ).split(",")
-    if part.strip()
-]
-PARALLELISM = int(os.environ.get("ICMR_PARALLEL", "3"))
-THREADS_PER_CONN = int(os.environ.get("ICMR_THREADS_PER_CONN", "1"))
+INDEX_SOURCE = os.environ.get("ICMR_INDEX_SOURCE", "remote").lower()
+PARALLELISM = int(os.environ.get("ICMR_PARALLEL", "2"))
+THREADS_PER_CONN = int(os.environ.get("ICMR_THREADS_PER_CONN", "2"))
 DUPLICATE_CAP = 2
-QUERY_TIMEOUT = float(os.environ.get("ICMR_QUERY_TIMEOUT", "25"))
 
 SEARCH_FIELDS = [
     "name", "fathersName", "phoneNumber", "aadharNumber", "otherNumber",
     "address", "district", "pincode", "state", "town", "source",
 ]
 NUMBER_FIELDS = ["phoneNumber", "aadharNumber", "otherNumber"]
-# Fields whose exact search we allow on the raw 100GB dataset (filter pushdown
-# works well on these; long text fields would be far slower).
-EXACT_FIELDS = set(SEARCH_FIELDS)
-SLOW_CONTAINS_FIELDS = {"address", "source", "fathersName"}
 
-DATA_URLS = [f"{HF_DATA_BASE}/{f}" for f in DATA_FILES]
+IDX_PHONE = "idx_phone"
+IDX_AADHAR = "idx_aadhar"
+
+REMOTE_INDEXES = {
+    "phone": [f"{HF_INDEX_BASE}/idx_phone.{i}.parquet" for i in range(7)],
+    "aadhar": [f"{HF_INDEX_BASE}/idx_aadhar.{i}.parquet" for i in range(7)],
+}
 
 CREDITS = {
     "developer": "@kzr0x",
@@ -54,25 +47,40 @@ CREDITS = {
     "channel": "@API_WALLAH",
 }
 
+# ── Parameter aliases ────────────────────────────────────────────────────────
+ALIAS_FIELDS = {
+    "num": "phoneNumber", "number": "phoneNumber", "mobile": "phoneNumber",
+    "phone": "phoneNumber", "aadhar": "aadharNumber", "aadhaar": "aadharNumber",
+    "name": "name", "father": "fathersName", "fathersname": "fathersName",
+    "other": "otherNumber", "othernumber": "otherNumber", "address": "address",
+    "district": "district", "pin": "pincode", "pincode": "pincode",
+    "state": "state", "town": "town", "city": "town", "source": "source",
+}
+
 # ── DuckDB Connection Pool ──────────────────────────────────────────────────
-_conns: dict[int, dict[str, duckdb.DuckDBPyConnection]] = {}
+_conns: list[duckdb.DuckDBPyConnection] = []
 _conns_lock = threading.Lock()
 _thread_local = threading.local()
-pool = ThreadPoolExecutor(max_workers=PARALLELISM * len(DATA_FILES) + 2,
-                          thread_name_prefix="duck")
+pool = ThreadPoolExecutor(max_workers=PARALLELISM, thread_name_prefix="duck")
+
 
-_query_cache: dict[str, tuple[float, dict]] = {}
-_cache_lock = threading.Lock()
-CACHE_TTL = 300
-CACHE_MAX = 512
+def _idx_ready(kind: str) -> bool:
+    return kind in REMOTE_INDEXES
 
 
 def _new_conn() -> duckdb.DuckDBPyConnection:
     con = duckdb.connect()
+    # Vercel fix: set home & extension dir to /tmp
+    con.execute("SET home_directory='/tmp'")
+    con.execute("SET extension_directory='/tmp/duckdb_extensions'")
     con.execute("INSTALL parquet; LOAD parquet;")
     con.execute("INSTALL httpfs; LOAD httpfs;")
+    # Create sorted index views from remote HF parts
+    for kind, urls in REMOTE_INDEXES.items():
+        view = f"people_{kind}"
+        lst = ", ".join(f"'{u}'" for u in urls)
+        con.execute(f"CREATE OR REPLACE VIEW {view} AS SELECT * FROM read_parquet([{lst}])")
     con.execute(f"SET threads = {THREADS_PER_CONN}")
-    con.execute("SET enable_object_cache = true")
     return con
 
 
@@ -80,38 +88,17 @@ def _thread_id() -> int:
     tid = getattr(_thread_local, "id", None)
     if tid is None:
         with _conns_lock:
-            tid = os.getpid() * 100000 + threading.get_ident() % 100000
+            tid = len(_conns)
             _thread_local.id = tid
     return tid
 
 
-def _get_conn(url: str) -> duckdb.DuckDBPyConnection:
+def _get_conn() -> duckdb.DuckDBPyConnection:
     ident = _thread_id()
     with _conns_lock:
-        per_thread = _conns.setdefault(ident, {})
-        con = per_thread.get(url)
-        if con is None:
-            con = _new_conn()
-            per_thread[url] = con
-    return con
-
-
-def _cache_get(key: str):
-    with _cache_lock:
-        item = _query_cache.get(key)
-        if item and time.time() - item[0] < CACHE_TTL:
-            return item[1]
-        if item:
-            _query_cache.pop(key, None)
-    return None
-
-
-def _cache_put(key: str, val: dict):
-    with _cache_lock:
-        if len(_query_cache) >= CACHE_MAX:
-            oldest = min(_query_cache, key=lambda k: _query_cache[k][0])
-            _query_cache.pop(oldest, None)
-        _query_cache[key] = (time.time(), val)
+        while len(_conns) <= ident:
+            _conns.append(_new_conn())
+    return _conns[ident]
 
 
 # ── Dedup & Connected Records ───────────────────────────────────────────────
@@ -152,67 +139,37 @@ def _cap_duplicates(rows: list[dict]) -> list[dict]:
 
 
 # ── Search Logic ────────────────────────────────────────────────────────────
-def _scan_one_file(url: str, sql: str, fetch_limit: int) -> list[tuple]:
-    """Run sql against a single parquet file with its own connection."""
-    con = _get_conn(url)
-    return con.execute(sql).fetchmany(fetch_limit)
-
-
 def _run_field_search(field: str, value: str, mode: str, limit: int) -> dict:
-    if field not in EXACT_FIELDS:
+    field = ALIAS_FIELDS.get(field.lower(), field)
+    if field not in SEARCH_FIELDS:
         raise ValueError(f"Unknown field: {field}")
-    if mode == "contains" and field in SLOW_CONTAINS_FIELDS:
-        return {"field": field, "value": value, "mode": mode,
-                "count": 0, "results": [],
-                "note": f"contains search disabled for slow field '{field}'"}
-
-    cache_key = f"{field}|{mode}|{value}|{limit}"
-    cached = _cache_get(cache_key)
-    if cached is not None:
-        return dict(cached, cached=True)
-
     v = value.replace("'", "''")
-    fetch_limit = limit * DUPLICATE_CAP + 20
 
     if mode == "exact":
-        where = f"{field} = '{v}'"
+        if field == "phoneNumber" and _idx_ready("phone"):
+            view = "people_phone"
+        elif field == "aadharNumber" and _idx_ready("aadhar"):
+            view = "people_aadhar"
+        elif field == "otherNumber":
+            # otherNumber not sorted — skip to avoid slow scan
+            return {"field": field, "value": value, "mode": mode, "count": 0, "results": []}
+        else:
+            return {"field": field, "value": value, "mode": mode, "count": 0, "results": []}
+        sql = f"SELECT * FROM {view} WHERE {field} = '{v}' LIMIT {limit * DUPLICATE_CAP + 20}"
     elif mode == "contains":
+        if field == "name":
+            # Name search not available in remote-only mode
+            return {"field": field, "value": value, "mode": mode, "count": 0, "results": []}
         v2 = v.replace("%", r"\%").replace("_", r"\_")
-        where = f"{field} ILIKE '%{v2}%' ESCAPE '\\'"
+        sql = f"SELECT * FROM people_phone WHERE {field} ILIKE '%{v2}%' ESCAPE '\\' LIMIT {limit * DUPLICATE_CAP + 20}"
     else:
         raise ValueError(f"Unknown mode: {mode}")
 
-    # One query per file, run in parallel, each with an early LIMIT.
-    # DuckDB prunes row-groups via parquet statistics, so exact matches
-    # on well-clustered files return in ~1-3s.
-    futures = []
-    inner = ThreadPoolExecutor(max_workers=len(DATA_URLS))
-    try:
-        for url in DATA_URLS:
-            sql = (
-                f"SELECT * FROM read_parquet('{url}') WHERE {where} "
-                f"LIMIT {fetch_limit}"
-            )
-            futures.append(inner.submit(_scan_one_file, url, sql, fetch_limit))
-        rows: list[tuple] = []
-        for fut in futures:
-            try:
-                rows.extend(fut.result(timeout=QUERY_TIMEOUT))
-            except FutTimeout:
-                pass
-            except Exception:
-                pass
-    finally:
-        inner.shutdown(wait=False, cancel_futures=True)
-
-    cols = ["name", "fathersName", "phoneNumber", "aadharNumber",
-            "otherNumber", "address", "district", "pincode", "state",
-            "town", "source"]
+    con = _get_conn()
+    rows = con.execute(sql).fetchall()
+    cols = [d[0] for d in con.description]
     results = _cap_duplicates([dict(zip(cols, r)) for r in rows])[:limit]
-    out = {"field": field, "value": value, "mode": mode,
-           "count": len(results), "results": results}
-    _cache_put(cache_key, out)
-    return out
+    return {"field": field, "value": value, "mode": mode, "count": len(results), "results": results}
 
 
 def _unified_search(q: str, limit: int = 10) -> dict:
@@ -221,10 +178,14 @@ def _unified_search(q: str, limit: int = 10) -> dict:
 
     if is_num:
         all_rows = []
-        searched = ["phoneNumber"]
-        r = _run_field_search("phoneNumber", q, "exact", limit)
-        all_rows.extend(r["results"])
-        if not all_rows:
+        searched = []
+        # Phone index first (fast)
+        if _idx_ready("phone"):
+            r = _run_field_search("phoneNumber", q, "exact", limit)
+            all_rows.extend(r["results"])
+            searched.append("phoneNumber")
+        # Aadhar index second
+        if not all_rows and _idx_ready("aadhar"):
             r = _run_field_search("aadharNumber", q, "exact", limit)
             all_rows.extend(r["results"])
             searched.append("aadharNumber")
@@ -234,11 +195,7 @@ def _unified_search(q: str, limit: int = 10) -> dict:
             "count": len(all_rows), "results": all_rows,
         }
     else:
-        r = _run_field_search("name", q, "contains", limit)
-        return {
-            "query": q, "searched_fields": ["name"],
-            "count": r["count"], "results": r["results"],
-        }
+        return {"query": q, "searched_fields": [], "count": 0, "results": []}
 
 
 # ── FastAPI (for API access) ────────────────────────────────────────────────
@@ -252,25 +209,22 @@ class BatchRequest(BaseModel):
 
 @fastapi_app.get("/")
 def root():
-    return {"app": "ICMR + HITEK Search API", "records": 2_504_793_870,
-            "data_files": DATA_FILES, "columns": SEARCH_FIELDS,
-            "credits": CREDITS, "docs": "/docs"}
+    return {
+        "app": "ICMR + HITEK Search API",
+        "records": 2_504_793_870,
+        "indexes": {"phone": _idx_ready("phone"), "aadhar": _idx_ready("aadhar")},
+        "index_source": INDEX_SOURCE,
+        "columns": SEARCH_FIELDS,
+        "credits": CREDITS,
+        "docs": "/docs",
+    }
 
 
 @fastapi_app.get("/health")
 def health():
-    return {"status": "ok", "data_source": HF_DATA_BASE, "data_files": DATA_FILES}
-
-
-# ── Parameter aliases ────────────────────────────────────────────────────────
-ALIAS_FIELDS = {
-    "num": "phoneNumber", "number": "phoneNumber", "mobile": "phoneNumber",
-    "phone": "phoneNumber", "aadhar": "aadharNumber", "aadhaar": "aadharNumber",
-    "name": "name", "father": "fathersName", "fathersname": "fathersName",
-    "other": "otherNumber", "othernumber": "otherNumber", "address": "address",
-    "district": "district", "pin": "pincode", "pincode": "pincode",
-    "state": "state", "town": "town", "city": "town", "source": "source",
-}
+    return {"status": "ok", "raw_database_required": False,
+            "indexes": {"phone": _idx_ready("phone"), "aadhar": _idx_ready("aadhar")},
+            "index_source": INDEX_SOURCE}
 
 
 @fastapi_app.get("/search")
@@ -327,8 +281,7 @@ async def search_parallel(req: BatchRequest):
     loop = asyncio.get_running_loop()
     tasks = [
         loop.run_in_executor(pool, _run_field_search,
-                             ALIAS_FIELDS.get(str(item.get("field", "phoneNumber")).lower(),
-                                              item.get("field", "phoneNumber")),
+                             item.get("field", "phoneNumber"),
                              item.get("value", ""),
                              item.get("mode", "exact"),
                              int(item.get("limit", req.limit)))
@@ -341,6 +294,29 @@ async def search_parallel(req: BatchRequest):
                     media_type="application/json")
 
 
+# ── Pinger (keeps app alive) ──────────────────────────────────────────────
+async def pinger():
+    """Ping the /health endpoint every 2 minutes to prevent idle shutdown."""
+    port = os.getenv("PORT", "7860")  # default Gradio port; change if needed
+    url = f"http://localhost:{port}/health"
+    async with httpx.AsyncClient(timeout=10) as client:
+        while True:
+            await asyncio.sleep(120)  # 2 minutes
+            try:
+                resp = await client.get(url)
+                if resp.status_code == 200:
+                    print(f"[Pinger] OK at {asyncio.get_event_loop().time()}")
+                else:
+                    print(f"[Pinger] Unexpected status: {resp.status_code}")
+            except Exception as e:
+                print(f"[Pinger] Error: {e}")
+
+
+@fastapi_app.on_event("startup")
+async def startup_event():
+    asyncio.create_task(pinger())
+
+
 # ── Gradio UI ───────────────────────────────────────────────────────────────
 def format_result(row: dict) -> str:
     """Format a single result record as readable text."""
@@ -349,6 +325,7 @@ def format_result(row: dict) -> str:
         val = row.get(field, "")
         if val:
             lines.append(f"**{field}:** {val}")
+    # Connected numbers
     cn = row.get("connected_numbers", [])
     if cn:
         nums = ", ".join(f"{c['field']}={c['value']}" for c in cn)
@@ -372,7 +349,7 @@ def search_ui(query: str, limit: int) -> str:
     searched = ", ".join(data.get("searched_fields", []))
 
     if not results:
-        return f"🔍 **Query:** `{q}`\n**Searched:** {searched}\n\n❌ **No data found** for this query."
+        return f"🔍 **Query:** `{q}`\n**Searched:** {searched}\n\n❌ **No data found** for this number."
 
     header = f"🔍 **Query:** `{q}`  |  **Found:** {count} results  |  **Searched:** {searched}\n\n---\n\n"
     parts = []
@@ -388,6 +365,7 @@ def build_ui():
         css="""
         .main-title { text-align: center; margin-bottom: 0; }
         .subtitle { text-align: center; color: #666; margin-top: 0; }
+        .footer { text-align: center; color: #888; margin-top: 20px; }
         """
     ) as demo:
         gr.Markdown("# 🔍 ICMR + HITEK Search API", elem_classes="main-title")
@@ -424,16 +402,14 @@ def build_ui():
         with gr.Accordion("📡 API Info", open=False):
             gr.Markdown(f"""
 **Endpoints** (via FastAPI):
+- `GET /search?q=<number>` — Phone/Aadhaar search
 - `GET /search?num=<number>` — Phone search
 - `GET /search?aadhar=<aadhaar>` — Aadhaar search
 - `GET /search?mobile=<number>` — Phone search (alias)
-- `GET /search?name=<text>&mode=contains` — Name search
-- `GET /search?q=<value>&field=<field>&mode=<exact|contains>` — Any field
-- `GET /search/parallel` (POST) — Batch queries
 - `GET /health` — Health check
 - `GET /docs` — Swagger UI
 
-**Source:** [HF Dataset](https://huggingface.co/datasets/Kzr0xx/Icmr-and-hitek)
+**Source:** [HF Dataset](https://huggingface.co/datasets/Kzr0xx/icrm-hitek-full-db-mixed)
 **Developer:** @kzr0x | **Team:** TEAM DARK KNIGHT APIs
             """)