[NOTICKET][doc] add sheet-level leg and RRF voting for tabular retrieval

src/rag/retrievers/schema.py

@@ -214,32 +214,82 @@ class SchemaRetriever(BaseRetriever):
             for row in rows
         ]
 
-    @staticmethod
-    def _chunk_key(r: RetrievalResult) -> tuple:
-        """Stable identity for dedup/RRF.
-
-        Includes filename, sheet_name, and chunk_level so that column chunks
-        and sheet chunks for the same file/sheet don't collide, and column
-        chunks with the same name across different files (e.g. `id` in two CSVs)
-        are kept distinct.
+    def _rank_tabular_sheets(
+        self,
+        sheet_results: list[RetrievalResult],
+        column_results: list[RetrievalResult],
+        top_k: int,
+        k_rrf: int = 60,
+    ) -> list[RetrievalResult]:
+        """Rank tabular sheets by RRF across two voting legs:
+          L1 (primary): sheet-chunk cosine score
+          L2 (vote):    best column-chunk position per (doc_id, sheet_name)
+
+        Returns top-k sheet-level RetrievalResults. The full column list of
+        each sheet is already in the sheet chunk's data.column_names from
+        ingestion, so downstream tabular_executor can read full sheet context.
+
+        For sheets surfaced by column votes but missing a sheet chunk (rare —
+        ingestion always creates one), a minimal stub is returned and
+        tabular_executor falls back to reading columns from the parquet.
         """
-        d = r.metadata.get("data", {})
-        return (
-            d.get("table_name"),
-            d.get("column_name"),
-            d.get("filename"),
-            d.get("sheet_name"),
-            r.metadata.get("chunk_level"),
-        )
+        # L1: sheets indexed by (doc_id, sheet_name) from sheet chunks
+        sheet_index: dict[tuple, RetrievalResult] = {}
+        sheet_ranked: list[tuple] = []
+        for r in sheet_results:
+            d = r.metadata.get("data", {})
+            key = (d.get("document_id"), d.get("sheet_name"))
+            if key[0] and key not in sheet_index:
+                sheet_index[key] = r
+                sheet_ranked.append(key)
+
+        # L2: sheets ranked by first-appearance in column-chunk results
+        col_sheet_ranked: list[tuple] = []
+        seen: set[tuple] = set()
+        for r in column_results:
+            d = r.metadata.get("data", {})
+            key = (d.get("document_id"), d.get("sheet_name"))
+            if key[0] and key not in seen:
+                col_sheet_ranked.append(key)
+                seen.add(key)
+
+        # RRF over (doc_id, sheet_name) across the two legs
+        rrf_scores: dict[tuple, float] = {}
+        for ranked_list in [sheet_ranked, col_sheet_ranked]:
+            for rank, key in enumerate(ranked_list):
+                rrf_scores[key] = rrf_scores.get(key, 0.0) + 1.0 / (k_rrf + rank + 1)
+
+        top_sheets = sorted(rrf_scores, key=lambda k: rrf_scores[k], reverse=True)[:top_k]
 
-    def _dedup(self, results: list[RetrievalResult]) -> list[RetrievalResult]:
-        """Deduplicate by chunk identity, keeping highest score per unique key."""
-        seen: dict[tuple, RetrievalResult] = {}
-        for r in results:
-            key = self._chunk_key(r)
-            if key not in seen or r.score > seen[key].score:
-                seen[key] = r
-        return sorted(seen.values(), key=lambda r: r.score, reverse=True)
+        results: list[RetrievalResult] = []
+        for key in top_sheets:
+            if key in sheet_index:
+                r = sheet_index[key]
+                r.score = rrf_scores[key]
+                results.append(r)
+            else:
+                # Surfaced by column votes only — build stub from a representative
+                # column result so tabular_executor can group correctly.
+                doc_id, sheet_name = key
+                rep = next(
+                    (r for r in column_results
+                     if r.metadata.get("data", {}).get("document_id") == doc_id
+                     and r.metadata.get("data", {}).get("sheet_name") == sheet_name),
+                    None,
+                )
+                if rep is None:
+                    continue
+                stub_data = dict(rep.metadata.get("data", {}))
+                stub_data.pop("column_name", None)
+                stub_data.pop("column_type", None)
+                results.append(RetrievalResult(
+                    content=f"Sheet: {stub_data.get('filename', '')}"
+                            + (f" / sheet: {sheet_name}" if sheet_name else ""),
+                    metadata={**rep.metadata, "data": stub_data, "chunk_level": "sheet"},
+                    score=rrf_scores[key],
+                    source_type="document",
+                ))
+        return results
 
     def _rank_db_tables(
         self,
@@ -325,8 +375,11 @@ class SchemaRetriever(BaseRetriever):
         ranked table set via _fetch_full_schema — the column chunks returned
         here are intentionally NOT used as the schema source, only for voting.
 
-        Tabular (CSV/XLSX) chunks remain at column/sheet level since they have
-        no table-level chunks.
+        Tabular (CSV/XLSX) sheets are ranked via RRF across two legs:
+          L1: sheet-chunk cosine
+          L2: column-chunk votes (best position per sheet)
+        Returns sheet-level RetrievalResults so tabular_executor receives
+        full sheet context (all columns) rather than fragmented column hits.
         """
         embedding = await self._embed_query(query)
         db_col_results, db_tbl_results, tabular_results, fts_results, sheet_results = await asyncio.gather(
@@ -338,17 +391,18 @@ class SchemaRetriever(BaseRetriever):
         )
 
         db_ranked = self._rank_db_tables(db_tbl_results, db_col_results, fts_results, top_k=k)
-        tabular_final = self._dedup(tabular_results + sheet_results)[:k]
+        tabular_ranked = self._rank_tabular_sheets(sheet_results, tabular_results, top_k=k)
 
-        results = db_ranked + tabular_final
+        results = db_ranked + tabular_ranked
         logger.info(
             "schema retrieval",
             count=len(results),
             db_tables_ranked=len(db_ranked),
             db_cols=len(db_col_results),
             db_tables=len(db_tbl_results),
-            tabular=len(tabular_results),
+            tabular_cols=len(tabular_results),
             tabular_sheets=len(sheet_results),
+            tabular_ranked=len(tabular_ranked),
             fts=len(fts_results),
         )
         return results