danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Minor optimizations

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-11-04 10:03:47 +02:00
parent cd2a23272e
commit 5ffedf013b
No known key found for this signature in database
1 changed files with 12 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
src/haiku/rag/store/repositories/chunk.py
View file

@ -253,17 +253,22 @@ class ChunkRepository:
""" """
if not query.strip(): if not query.strip():
return [] return []
filtered_docs_df = None filtered_doc_ids = None
if filter: if filter:
# We perform filtering as a two-step process, first filtering documents, then # We perform filtering as a two-step process, first filtering documents, then
# filtering chunks based on those document IDs. # filtering chunks based on those document IDs.
# This is because LanceDB does not support joins directly in search queries. # This is because LanceDB does not support joins directly in search queries.
filtered_docs_df = ( docs_df = (
self.store.documents_table.search() self.store.documents_table.search()
.select(["id"]) .select(["id"])
.where(filter) .where(filter)
.to_pandas() .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 # Prepare search query based on search type
if search_type == "vector": if search_type == "vector":
@ -288,10 +293,10 @@ class ChunkRepository:
) )
# Apply filtering if needed (common for all search types) # 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() chunks_df = results.to_pandas()
filtered_chunks_df = chunks_df.loc[ 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) ].head(limit)
return await self._process_search_results(filtered_chunks_df) return await self._process_search_results(filtered_chunks_df)
@ -407,8 +412,9 @@ class ChunkRepository:
# Batch fetch all documents at once # Batch fetch all documents at once
documents_map = {} documents_map = {}
if document_ids: 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( doc_results = list(
self.store.documents_table.search() self.store.documents_table.search()
.where(where_clause) .where(where_clause)