Upload app.py with huggingface_hub

app.py

@@ -36,13 +36,18 @@ class UserIdIndex:
     def build(self, arrow_dataset):
         logger.info("Building user_id index from parquet metadata...")
         t0 = time.time()
+        schema = arrow_dataset.schema
+        uid_idx = schema.get_field_index("user_id")
+        if uid_idx < 0:
+            logger.error("user_id column not found in schema!")
+            return
         for frag in arrow_dataset.get_fragments():
             meta = frag.metadata
             path = frag.path
             for rg_idx in range(meta.num_row_groups):
-                col = meta.row_group(rg_idx).column(0)
+                col = meta.row_group(rg_idx).column(uid_idx)
                 s = col.statistics
-                if s:
+                if s and s.min is not None and s.max is not None:
                     self.entries.append((int(s.min), int(s.max), path, rg_idx))
         self.entries.sort(key=lambda x: x[0])
         self.keys = [e[0] for e in self.entries]
@@ -93,6 +98,8 @@ def init_dataset():
     stats["startup_time"] = time.time()
 
 
+ALL_COLS = ["user_id", "username", "first_name", "last_name", "phone", "email", "status", "linked_id", "linked_name", "linked_handle"]
+
 def query_user(user_id: int):
     stats["queries"] += 1
     t0 = time.time()
@@ -102,19 +109,15 @@ def query_user(user_id: int):
         path, rg_idx = result
         try:
             pf = pq.ParquetFile(path)
-            tbl = pf.read_row_group(rg_idx, columns=[f.name for f in pf.schema_arrow])
-            mask = pc.equal(pc.field("user_id"), user_id)
-            tbl = tbl.filter(mask)
+            tbl = pf.read_row_group(rg_idx, columns=ALL_COLS)
+            tbl = tbl.filter(pc.equal(pc.field("user_id"), user_id))
             if len(tbl):
                 elapsed = time.time() - t0
                 return tbl.to_pylist()[0], elapsed
         except Exception as e:
             logger.warning(f"Index lookup failed for {user_id}: {e}")
 
-    tbl = dataset.to_table(
-        filter=pc.equal(pc.field("user_id"), user_id),
-        columns=["user_id", "username", "first_name", "last_name", "phone", "email", "status", "linked_id", "linked_name", "linked_handle"]
-    )
+    tbl = dataset.to_table(filter=pc.equal(pc.field("user_id"), user_id), columns=ALL_COLS)
     elapsed = time.time() - t0
     if len(tbl) == 0:
         return None, elapsed