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

Use .select() projection to fetch only id/uri/title/metadata in search

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-07 11:04:40 +03:00
parent f2998ad0b5
commit 52712dbdb2
No known key found for this signature in database
1 changed files with 17 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -13,7 +13,7 @@ if TYPE_CHECKING:
from lancedb.rerankers import RRFReranker from lancedb.rerankers import RRFReranker
from haiku.rag.store.engine import DocumentRecord, Store from haiku.rag.store.engine import Store
from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.chunk import Chunk
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -321,17 +321,18 @@ class ChunkRepository:
results = list(query.to_pydantic(self.store.ChunkRecord)) results = list(query.to_pydantic(self.store.ChunkRecord))
# Get document info # Get document info (only metadata columns, skip content/docling blobs)
doc_results = list( doc_rows = list(
self.store.documents_table.search() self.store.documents_table.search()
.select(["id", "uri", "title", "metadata"])
.where(f"id = '{document_id}'") .where(f"id = '{document_id}'")
.limit(1) .limit(1)
.to_pydantic(DocumentRecord) .to_list()
) )
doc_uri = doc_results[0].uri if doc_results else None doc_uri = doc_rows[0]["uri"] if doc_rows else None
doc_title = doc_results[0].title if doc_results else None doc_title = doc_rows[0]["title"] if doc_rows else None
doc_meta = doc_results[0].metadata if doc_results else "{}" doc_meta = doc_rows[0].get("metadata", "{}") if doc_rows else "{}"
chunks: list[Chunk] = [] chunks: list[Chunk] = []
for rec in results: for rec in results:
@ -429,18 +430,18 @@ class ChunkRepository:
# Collect all unique document IDs for batch lookup # Collect all unique document IDs for batch lookup
document_ids = list(set(chunk.document_id for chunk in pydantic_results)) document_ids = list(set(chunk.document_id for chunk in pydantic_results))
# Batch fetch all documents at once # Batch fetch document metadata (skip content/docling blobs)
documents_map = {} documents_map: dict[str, dict] = {}
if document_ids: if document_ids:
# Use IN clause for efficient batch lookup
id_list = "', '".join(document_ids) id_list = "', '".join(document_ids)
where_clause = f"id IN ('{id_list}')" where_clause = f"id IN ('{id_list}')"
doc_results = list( doc_rows = list(
self.store.documents_table.search() self.store.documents_table.search()
.select(["id", "uri", "title", "metadata"])
.where(where_clause) .where(where_clause)
.to_pydantic(DocumentRecord) .to_list()
) )
documents_map = {doc.id: doc for doc in doc_results} documents_map = {str(row["id"]): row for row in doc_rows}
# Build final results with document info # Build final results with document info
chunks_with_scores = [] chunks_with_scores = []
@ -452,9 +453,9 @@ class ChunkRepository:
content=chunk_record.content, content=chunk_record.content,
metadata=json.loads(chunk_record.metadata), metadata=json.loads(chunk_record.metadata),
order=chunk_record.order, order=chunk_record.order,
document_uri=doc.uri if doc else None, document_uri=doc["uri"] if doc else None,
document_title=doc.title if doc else None, document_title=doc["title"] if doc else None,
document_meta=json.loads(doc.metadata if doc else "{}"), document_meta=json.loads(doc.get("metadata", "{}") if doc else "{}"),
) )
score = scores[i] if i < len(scores) else 1.0 score = scores[i] if i < len(scores) else 1.0
chunks_with_scores.append((chunk, score)) chunks_with_scores.append((chunk, score))