import os
import re
import time
import threading
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
import streamlit as st
import duckdb
import pandas as pd
# ── Config ─────────────────────────────────────────────────────────────────────
HF_TOKEN = os.environ.get("HF_TOKEN", "")
HF_DATASET = "hf://datasets/styhero/styhero_dataset/data"
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
_L1 = os.path.join(BASE_DIR, "mobile1_index.parquet")
_L2 = os.path.join(BASE_DIR, "mobile2_index.parquet")
MOBILE1_IDX = _L1 if os.path.exists(_L1) else "hf://datasets/styhero/styhero_dataset/data/mobile1_index.parquet"
MOBILE2_IDX = _L2 if os.path.exists(_L2) else "hf://datasets/styhero/styhero_dataset/data/mobile2_index.parquet"
# Allowlist for column names that go into f-strings (can't be SQL params)
_ALLOWED_COLS = {"mobile1", "mobile2"}
# Mobile number validation: digits only, 10–15 chars (covers Indian + intl. format)
_MOBILE_RE = re.compile(r"^\d{10,15}$")
# ── Connection pool (per-thread reuse) ────────────────────────────────────────
_local = threading.local()
def get_con():
"""Reuse one DuckDB connection per thread. Connection setup costs 50-150ms
so reusing it across queries in the same request is a big win."""
con = getattr(_local, "con", None)
if con is not None:
return con
con = duckdb.connect()
con.execute("INSTALL httpfs; LOAD httpfs;")
con.execute("CREATE SECRET IF NOT EXISTS hf_secret (TYPE HUGGINGFACE, TOKEN ?)", [HF_TOKEN])
con.execute("SET threads=4;")
con.execute("SET enable_progress_bar=false;")
con.execute("SET enable_object_cache=true;") # cache parquet metadata across queries
con.execute("SET http_keep_alive=true;")
con.execute("SET http_retries=3;")
_local.con = con
return con
# ── Helpers ────────────────────────────────────────────────────────────────────
def fmt_ms(ms):
return f"{ms/1000:.2f}s" if ms >= 1000 else f"{ms:.0f}ms"
def clean(val):
if val is None or (isinstance(val, float) and pd.isna(val)):
return ""
s = str(val).strip()
if s.lower() in ("nan", "none", ""):
return ""
if s.endswith(".0"):
try:
s = str(int(float(s)))
except Exception:
pass
return ", ".join(p.strip() for p in s.split("!") if p.strip())
def sanitize_df(df):
df = df.copy()
for col in ["name", "fname", "email", "id", "address", "circle"]:
if col in df.columns:
df[col] = df[col].apply(clean)
return df
# ── Queries ────────────────────────────────────────────────────────────────────
def query_both_indexes(target):
"""Combined UNION ALL query over both index files in a single round trip.
Parameterized to prevent SQL injection."""
con = get_con()
t0 = time.perf_counter()
df = con.execute("""
SELECT mobile1, mobile2, row_group_id, row_start, row_end, filename, 'mobile1' AS src
FROM read_parquet(?)
WHERE mobile1 = ?
UNION ALL
SELECT mobile1, mobile2, row_group_id, row_start, row_end, filename, 'mobile2' AS src
FROM read_parquet(?)
WHERE mobile2 = ?
""", [MOBILE1_IDX, target, MOBILE2_IDX, target]).fetchdf()
return df, (time.perf_counter() - t0) * 1000
def _fetch_file_batch(filename, ranges, target):
"""Fetch all matching row-groups for a single parquet file in one query.
`ranges` is a list of (row_start, row_end, search_col) tuples."""
# Validate search columns against allowlist before string interpolation
for _, _, col in ranges:
if col not in _ALLOWED_COLS:
raise ValueError(f"Invalid search column: {col!r}")
con = get_con()
hf_path = f"{HF_DATASET}/{filename.split('/')[-1]}"
# Build OR conditions; row_start/row_end are ints (we cast on input), col is allowlisted
or_parts = []
params = [hf_path]
for rs, re_, col in ranges:
or_parts.append(f"(file_row_number BETWEEN {int(rs)} AND {int(re_)} AND {col} = ?)")
params.append(target)
sql = f"""
SELECT mobile1, mobile2, email, id, name, fname, address, circle
FROM read_parquet(?, file_row_number=true)
WHERE {' OR '.join(or_parts)}
"""
t0 = time.perf_counter()
df = con.execute(sql, params).fetchdf()
return sanitize_df(df), (time.perf_counter() - t0) * 1000
# ── Cached search ─────────────────────────────────────────────────────────────
@st.cache_data(ttl=300, max_entries=200, show_spinner=False)
def search(target: str):
"""Top-level cached search. Returns (combined_df, idx_ms, fetch_ms, errors).
TTL kept short because results contain PII — don't want them sitting in
memory indefinitely."""
errors = []
t_start = time.perf_counter()
idx_df, idx_ms = query_both_indexes(target)
if idx_df.empty:
return pd.DataFrame(), idx_ms, 0.0, errors
# Group row ranges by file so we make 1 HTTP query per file, not per row-group
by_file = defaultdict(list)
for _, rg in idx_df[["filename", "row_start", "row_end", "src"]].drop_duplicates().iterrows():
by_file[str(rg["filename"])].append(
(int(rg["row_start"]), int(rg["row_end"]), str(rg["src"]))
)
all_dfs = []
if by_file:
with ThreadPoolExecutor(max_workers=min(8, len(by_file))) as ex:
futures = {
ex.submit(_fetch_file_batch, fname, ranges, target): fname
for fname, ranges in by_file.items()
}
for f in as_completed(futures):
try:
df, _ = f.result()
if not df.empty:
all_dfs.append(df)
except Exception as e:
errors.append(str(e))
fetch_ms = (time.perf_counter() - t_start) * 1000
if not all_dfs:
return pd.DataFrame(), idx_ms, fetch_ms, errors
combined = pd.concat(all_dfs, ignore_index=True).drop_duplicates()
return combined, idx_ms, fetch_ms, errors
# ── SVG icons ──────────────────────────────────────────────────────────────────
def ico(path, color="#8b8bff", size=14):
return f''
ICO_PHONE = ico('')
ICO_EMAIL = ico('')
ICO_ID = ico('')
ICO_PIN = ico('')
ICO_CIRCLE = ico('')
ICO_PERSON = ico('')
# ── Result card ────────────────────────────────────────────────────────────────
def field_html(ico_svg, label, value, full=False):
if not value:
return ""
span = 'style="grid-column:1/-1"' if full else ""
return f"""
{ico_svg}{label}
{value}
"""
def show_card(row, idx):
name = clean(row.get("name")) or "Unknown"
fname = clean(row.get("fname")) or ""
mobile1 = clean(row.get("mobile1")) or ""
mobile2 = clean(row.get("mobile2")) or ""
email = clean(row.get("email")) or ""
rid = clean(row.get("id")) or ""
address = clean(row.get("address")) or ""
circle = clean(row.get("circle")) or ""
copy_text = "\n".join(filter(None, [
f"Name: {name}",
f"Father: {fname}" if fname else "",
f"Mobile1: {mobile1}" if mobile1 else "",
f"Mobile2: {mobile2}" if mobile2 else "",
f"Email: {email}" if email else "",
f"ID: {rid}" if rid else "",
f"Address: {address}" if address else "",
f"Circle: {circle}" if circle else "",
]))
copy_safe = copy_text.replace("\\","\\\\").replace("`","\\`").replace("${","\\${")
initials = (name[0] if name and name != "Unknown" else "?").upper()
card = f"""
""", unsafe_allow_html=True)
# ── Search form ─────────────────────────────────────────────────────────────────
with st.form("search_form"):
number = st.text_input(
"Mobile number",
placeholder="Enter mobile number — e.g. 919876543210",
label_visibility="collapsed",
)
submitted = st.form_submit_button("Search", use_container_width=True, type="primary")
if not submitted or not number.strip():
st.stop()
target = number.strip()
# Input validation — reject anything that isn't 10-15 digits
if not _MOBILE_RE.match(target):
st.error("Please enter a valid mobile number (10–15 digits, numbers only).")
st.stop()
# ── Search with progress ────────────────────────────────────────────────────────
with st.status("Searching...", expanded=False) as status:
combined, idx_ms, fetch_ms, errors = search(target)
if combined.empty:
status.update(label="No records found", state="error", expanded=False)
for err in errors:
st.error(err)
st.warning(f"No records found for **{target}**")
st.stop()
status.update(label=f"Done — {fmt_ms(fetch_ms)}", state="complete", expanded=False)
for err in errors:
st.error(err)
# ── Results ─────────────────────────────────────────────────────────────────────
st.markdown(f"""
{len(combined)} record(s)
{fmt_ms(idx_ms)} index · {fmt_ms(fetch_ms)} total
""", unsafe_allow_html=True)
for i, (_, row) in enumerate(combined.iterrows()):
show_card(row, i)