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