From 576ac9cecf6735d35c10bc8a1c0defe89a286f8e Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 4 Nov 2025 09:37:39 +0200 Subject: [PATCH 1/4] Perform join when filtering using pandas DataFrames instead ofthe slower WHERE IN clause --- src/haiku/rag/store/repositories/chunk.py | 132 ++++++++++++---------- 1 file changed, 71 insertions(+), 61 deletions(-) diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index 397f4117..ab60e584 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -249,64 +249,51 @@ class ChunkRepository: """ if not query.strip(): return [] - - chunk_where_clause = None + filtered_docs_df = 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) - - if not matching_doc_ids: - return [] - - # Build WHERE clause for chunks table - # Use IN clause with document IDs - id_list = "', '".join(matching_doc_ids) - chunk_where_clause = f"document_id IN ('{id_list}')" + filtered_docs_df = ( + self.store.documents_table.search() + .select(["id"]) + .where(filter) + .to_pandas() + ) + # Prepare search query based on search type 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_docs_df is not None: + chunks_df = results.to_pandas() + filtered_chunks_df = chunks_df.loc[ + chunks_df["document_id"].isin(filtered_docs_df["id"]) + ].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 +351,62 @@ class ChunkRepository: return adjacent_chunks - def _get_filtered_document_ids(self, filter: str) -> list[str]: - """Query documents table with filter and return matching document IDs.""" - filtered_docs = ( - self.store.documents_table.search() - .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]]: """Process search results into chunks with document info and scores.""" chunks_with_scores = [] - # Get both arrow and pydantic results to access scores - arrow_result = query_result.to_arrow() - pydantic_results = list(query_result.to_pydantic(self.store.ChunkRecord)) + # Handle pandas DataFrame (from filtered results) + import pandas as pd - # Extract scores from arrow result based on search type - scores = [] - column_names = arrow_result.column_names + if isinstance(query_result, pd.DataFrame): + # DataFrame already contains the data we need + pydantic_results = [] + for _, row in query_result.iterrows(): + chunk_record = 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, + ) + pydantic_results.append(chunk_record) - 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() + # Extract scores from DataFrame columns + scores = [] + if "_distance" in query_result.columns: + # Vector search - distance (lower is better, convert to similarity) + distances = query_result["_distance"].tolist() + scores = [max(0.0, 1.0 / (1.0 + dist)) for dist in distances] + elif "_relevance_score" in query_result.columns: + # Hybrid search - relevance score (higher is better) + scores = query_result["_relevance_score"].tolist() + elif "_score" in query_result.columns: + # FTS search - score (higher is better) + scores = query_result["_score"].tolist() + else: + raise ValueError("Unknown search result format, cannot extract scores") else: - raise ValueError("Unknown search result format, cannot extract scores") + # Handle LanceDB query result (original behavior) + # Get both arrow and pydantic results to access scores + arrow_result = query_result.to_arrow() + pydantic_results = list(query_result.to_pydantic(self.store.ChunkRecord)) + + # Extract scores from arrow result based on search type + scores = [] + column_names = arrow_result.column_names + + 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") # Collect all unique document IDs for batch lookup document_ids = list(set(chunk.document_id for chunk in pydantic_results)) From cd2a23272e2b0cff5a808fe45d62978d78f56c5d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 4 Nov 2025 09:48:00 +0200 Subject: [PATCH 2/4] Simplify _process_search_results --- src/haiku/rag/store/repositories/chunk.py | 111 ++++++++++------------ 1 file changed, 48 insertions(+), 63 deletions(-) diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index ab60e584..c5b669b0 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/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 @@ -351,62 +355,51 @@ class ChunkRepository: return adjacent_chunks - async def _process_search_results(self, query_result) -> list[tuple[Chunk, float]]: - """Process search results into chunks with document info and scores.""" - chunks_with_scores = [] + async def _process_search_results( + self, query_result: "pd.DataFrame | LanceQueryBuilder" + ) -> list[tuple[Chunk, float]]: + """Process search results into chunks with document info and scores. - # Handle pandas DataFrame (from filtered results) + Args: + query_result: Either a pandas DataFrame or a LanceDB query result + """ import pandas as pd + def extract_scores(df: pd.DataFrame) -> list[float]: + """Extract scores from DataFrame columns based on search type.""" + if "_distance" in df.columns: + # Vector search - convert distance to similarity + return [max(0.0, 1.0 / (1.0 + d)) for d in df["_distance"].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") + + # Convert everything to DataFrame for uniform processing if isinstance(query_result, pd.DataFrame): - # DataFrame already contains the data we need - pydantic_results = [] - for _, row in query_result.iterrows(): - chunk_record = 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, - ) - pydantic_results.append(chunk_record) - - # Extract scores from DataFrame columns - scores = [] - if "_distance" in query_result.columns: - # Vector search - distance (lower is better, convert to similarity) - distances = query_result["_distance"].tolist() - scores = [max(0.0, 1.0 / (1.0 + dist)) for dist in distances] - elif "_relevance_score" in query_result.columns: - # Hybrid search - relevance score (higher is better) - scores = query_result["_relevance_score"].tolist() - elif "_score" in query_result.columns: - # FTS search - score (higher is better) - scores = query_result["_score"].tolist() - else: - raise ValueError("Unknown search result format, cannot extract scores") + df = query_result else: - # Handle LanceDB query result (original behavior) - # Get both arrow and pydantic results to access scores - arrow_result = query_result.to_arrow() - pydantic_results = list(query_result.to_pydantic(self.store.ChunkRecord)) + # Convert LanceDB query result to DataFrame + df = query_result.to_pandas() - # Extract scores from arrow result based on search type - scores = [] - column_names = arrow_result.column_names + # Extract scores + scores = extract_scores(df) - 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 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)) @@ -423,29 +416,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_title=doc_title, - document_meta=json.loads(doc_meta), + document_uri=doc.uri if doc else None, + document_title=doc.title if doc else None, + 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 From 5ffedf013b4126768d9d79dbd2329ab2a5405e6f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 4 Nov 2025 10:03:47 +0200 Subject: [PATCH 3/4] Minor optimizations --- src/haiku/rag/store/repositories/chunk.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index c5b669b0..467f20fb 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -253,17 +253,22 @@ class ChunkRepository: """ if not query.strip(): return [] - filtered_docs_df = None + filtered_doc_ids = 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. - filtered_docs_df = ( + docs_df = ( self.store.documents_table.search() .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"] # Prepare search query based on search type if search_type == "vector": @@ -288,10 +293,10 @@ class ChunkRepository: ) # Apply filtering if needed (common for all search types) - if filtered_docs_df is not None: + if filtered_doc_ids is not None: chunks_df = results.to_pandas() filtered_chunks_df = chunks_df.loc[ - chunks_df["document_id"].isin(filtered_docs_df["id"]) + chunks_df["document_id"].isin(filtered_doc_ids) ].head(limit) return await self._process_search_results(filtered_chunks_df) @@ -407,8 +412,9 @@ class ChunkRepository: # Batch fetch all documents at once documents_map = {} if document_ids: - # Create a WHERE clause for all document IDs - where_clause = " OR ".join(f"id = '{doc_id}'" for doc_id in document_ids) + # Use IN clause for efficient batch lookup + id_list = "', '".join(document_ids) + where_clause = f"id IN ('{id_list}')" doc_results = list( self.store.documents_table.search() .where(where_clause) From 4773b56f9d6d8eddde979795b140a0f7a2a5e7b7 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 4 Nov 2025 16:34:47 +0200 Subject: [PATCH 4/4] Improve performance of vector search score normalization Co-authored-by: Tres Seaver --- src/haiku/rag/store/repositories/chunk.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index 467f20fb..24f2bb68 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -374,7 +374,7 @@ class ChunkRepository: """Extract scores from DataFrame columns based on search type.""" if "_distance" in df.columns: # Vector search - convert distance to similarity - return [max(0.0, 1.0 / (1.0 + d)) for d in df["_distance"].tolist()] + 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()