Merge pull request #127 from ggozad/feat/faster-filtering

Faster filtering through pandas

src/haiku/rag/store/repositories/chunk.py

@@ -4,6 +4,10 @@ import logging
 from typing import TYPE_CHECKING
 from uuid import uuid4
 
+if TYPE_CHECKING:
+    import pandas as pd
+    from lancedb.query import LanceQueryBuilder
+
 from lancedb.rerankers import RRFReranker
 
 from haiku.rag.store.engine import DocumentRecord, Store
@@ -249,64 +253,56 @@ class ChunkRepository:
         """
         if not query.strip():
             return []
-
+        filtered_doc_ids = None
-        chunk_where_clause = None
         if filter:
             # We perform filtering as a two-step process, first filtering documents, then
             # filtering chunks based on those document IDs.
             # This is because LanceDB does not support joins directly in search queries.
-            matching_doc_ids = self._get_filtered_document_ids(filter)
+            docs_df = (
-
+                self.store.documents_table.search()
-            if not matching_doc_ids:
+                .select(["id"])
+                .where(filter)
+                .to_pandas()
+            )
+            # Early exit if no documents match the filter
+            if docs_df.empty:
                 return []
+            # Keep as pandas Series for efficient vectorized operations
+            filtered_doc_ids = docs_df["id"]
 
-            # Build WHERE clause for chunks table
+        # Prepare search query based on search type
-            # Use IN clause with document IDs
-            id_list = "', '".join(matching_doc_ids)
-            chunk_where_clause = f"document_id IN ('{id_list}')"
-
         if search_type == "vector":
             query_embedding = await self.embedder.embed(query)
-
             results = self.store.chunks_table.search(
                 query_embedding, query_type="vector", vector_column_name="vector"
             )
 
-            if chunk_where_clause:
-                results = results.where(chunk_where_clause)
-
-            results = results.limit(limit)
-
-            return await self._process_search_results(results)
-
         elif search_type == "fts":
             results = self.store.chunks_table.search(query, query_type="fts")
 
-            if chunk_where_clause:
-                results = results.where(chunk_where_clause)
-
-            results = results.limit(limit)
-            return await self._process_search_results(results)
-
         else:  # hybrid (default)
             query_embedding = await self.embedder.embed(query)
-
             # Create RRF reranker
             reranker = RRFReranker()
-
             # Perform native hybrid search with RRF reranking
-            results = self.store.chunks_table.search(query_type="hybrid")
-
-            if chunk_where_clause:
-                results = results.where(chunk_where_clause)
-
             results = (
-                results.vector(query_embedding)
+                self.store.chunks_table.search(query_type="hybrid")
+                .vector(query_embedding)
                 .text(query)
                 .rerank(reranker)
-                .limit(limit)
             )
-            return await self._process_search_results(results)
+
+        # Apply filtering if needed (common for all search types)
+        if filtered_doc_ids is not None:
+            chunks_df = results.to_pandas()
+            filtered_chunks_df = chunks_df.loc[
+                chunks_df["document_id"].isin(filtered_doc_ids)
+            ].head(limit)
+            return await self._process_search_results(filtered_chunks_df)
+
+        # No filtering needed, apply limit and return
+        results = results.limit(limit)
+        return await self._process_search_results(results)
 
     async def get_by_document_id(self, document_id: str) -> list[Chunk]:
         """Get all chunks for a specific document."""
@@ -364,39 +360,51 @@ class ChunkRepository:
 
         return adjacent_chunks
 
-    def _get_filtered_document_ids(self, filter: str) -> list[str]:
+    async def _process_search_results(
-        """Query documents table with filter and return matching document IDs."""
+        self, query_result: "pd.DataFrame | LanceQueryBuilder"
-        filtered_docs = (
+    ) -> list[tuple[Chunk, float]]:
-            self.store.documents_table.search()
+        """Process search results into chunks with document info and scores.
-            .where(filter)
-            .to_pydantic(DocumentRecord)
-        )
-        return [doc.id for doc in filtered_docs]
 
-    async def _process_search_results(self, query_result) -> list[tuple[Chunk, float]]:
+        Args:
-        """Process search results into chunks with document info and scores."""
+            query_result: Either a pandas DataFrame or a LanceDB query result
-        chunks_with_scores = []
+        """
+        import pandas as pd
 
-        # Get both arrow and pydantic results to access scores
+        def extract_scores(df: pd.DataFrame) -> list[float]:
-        arrow_result = query_result.to_arrow()
+            """Extract scores from DataFrame columns based on search type."""
-        pydantic_results = list(query_result.to_pydantic(self.store.ChunkRecord))
+            if "_distance" in df.columns:
+                # Vector search - convert distance to similarity
+                return ((df["_distance"] + 1).rdiv(1)).clip(lower=0.0).tolist()
+            elif "_relevance_score" in df.columns:
+                # Hybrid search - relevance score (higher is better)
+                return df["_relevance_score"].tolist()
+            elif "_score" in df.columns:
+                # FTS search - score (higher is better)
+                return df["_score"].tolist()
+            else:
+                raise ValueError("Unknown search result format, cannot extract scores")
 
-        # Extract scores from arrow result based on search type
+        # Convert everything to DataFrame for uniform processing
-        scores = []
+        if isinstance(query_result, pd.DataFrame):
-        column_names = arrow_result.column_names
+            df = query_result
-
-        if "_distance" in column_names:
-            # Vector search - distance (lower is better, convert to similarity)
-            distances = arrow_result.column("_distance").to_pylist()
-            scores = [max(0.0, 1.0 / (1.0 + dist)) for dist in distances]
-        elif "_relevance_score" in column_names:
-            # Hybrid search - relevance score (higher is better)
-            scores = arrow_result.column("_relevance_score").to_pylist()
-        elif "_score" in column_names:
-            # FTS search - score (higher is better)
-            scores = arrow_result.column("_score").to_pylist()
         else:
-            raise ValueError("Unknown search result format, cannot extract scores")
+            # Convert LanceDB query result to DataFrame
+            df = query_result.to_pandas()
+
+        # Extract scores
+        scores = extract_scores(df)
+
+        # Convert DataFrame rows to ChunkRecords
+        pydantic_results = [
+            self.store.ChunkRecord(
+                id=str(row["id"]),
+                document_id=str(row["document_id"]),
+                content=str(row["content"]),
+                metadata=str(row["metadata"]),
+                order=int(row["order"]) if "order" in row else 0,
+            )
+            for _, row in df.iterrows()
+        ]
 
         # Collect all unique document IDs for batch lookup
         document_ids = list(set(chunk.document_id for chunk in pydantic_results))
@@ -404,8 +412,9 @@ class ChunkRepository:
         # Batch fetch all documents at once
         documents_map = {}
         if document_ids:
-            # Create a WHERE clause for all document IDs
+            # Use IN clause for efficient batch lookup
-            where_clause = " OR ".join(f"id = '{doc_id}'" for doc_id in document_ids)
+            id_list = "', '".join(document_ids)
+            where_clause = f"id IN ('{id_list}')"
             doc_results = list(
                 self.store.documents_table.search()
                 .where(where_clause)
@@ -413,29 +422,21 @@ class ChunkRepository:
             )
             documents_map = {doc.id: doc for doc in doc_results}
 
+        # Build final results with document info
+        chunks_with_scores = []
         for i, chunk_record in enumerate(pydantic_results):
-            # Get document info from pre-fetched map
             doc = documents_map.get(chunk_record.document_id)
-            doc_uri = doc.uri if doc else None
-            doc_title = doc.title if doc else None
-            doc_meta = doc.metadata if doc else "{}"
-
-            md = json.loads(chunk_record.metadata)
-
             chunk = Chunk(
                 id=chunk_record.id,
                 document_id=chunk_record.document_id,
                 content=chunk_record.content,
-                metadata=md,
+                metadata=json.loads(chunk_record.metadata),
                 order=chunk_record.order,
-                document_uri=doc_uri,
+                document_uri=doc.uri if doc else None,
-                document_title=doc_title,
+                document_title=doc.title if doc else None,
-                document_meta=json.loads(doc_meta),
+                document_meta=json.loads(doc.metadata if doc else "{}"),
             )
-
-            # Get score from arrow result
             score = scores[i] if i < len(scores) else 1.0
-
             chunks_with_scores.append((chunk, score))
 
         return chunks_with_scores