Add order to SearchResult, add ChunkRepository.get_chunks_in_range(), to use them in _expand_with_chunks to fetch only nearby chunks

CHANGELOG.md

@@ -4,10 +4,14 @@
 ### Changed
 
 - **Dependency updates**: lancedb 0.30.2, pydantic-ai-slim ≥1.77.0, docling ≥2.84.0, docling-core ≥2.71.0, haiku.skills ≥0.13.0, cachetools ≥7.0.5, pydantic-monty ≥0.0.9, cohere ≥5.21.1, textual ≥8.2.1, ty ≥0.0.28, ruff ≥0.15.9
+- **Search result model**: `SearchResult` now includes `order` field propagated from chunk order
 
 ### Fixed
 
 - **Type checking**: Fix 37 new ty 0.0.28 diagnostics with proper None guards, assertions, and specific ignore codes
+- **Search performance**: Avoid loading full document blobs (docling_document, content) during search — use column projection to fetch only needed metadata (id, uri, title, metadata)
+- **Context expansion performance**: Load only docling columns during expand_context (skip content blob), and only when doc_item_refs exist
+- **Chunk expansion performance**: Fetch only chunks in the needed order range during context expansion instead of all chunks for a document
 
 ## [0.36.3] - 2026-04-01
 

haiku_rag_slim/haiku/rag/client.py

@@ -1400,32 +1400,40 @@ class HaikuRAG:
         radius: int,
     ) -> list[SearchResult]:
         """Expand results using chunk-based adjacency."""
-        all_chunks = await self.chunk_repository.get_by_document_id(doc_id)
-        if not all_chunks:
-            return results
-
-        content_to_chunk = {c.content: c for c in all_chunks}
-        chunk_by_order = {c.order: c for c in all_chunks}
-        min_order, max_order = min(chunk_by_order.keys()), max(chunk_by_order.keys())
-
-        # Build ranges
+        # Build ranges from result orders
         ranges: list[tuple[int, int, SearchResult]] = []
         passthrough: list[SearchResult] = []
 
         for result in results:
-            chunk = content_to_chunk.get(result.content)
-            if chunk is None:
+            if result.chunk_id is None:
                 passthrough.append(result)
                 continue
-            start = max(min_order, chunk.order - radius)
-            end = min(max_order, chunk.order + radius)
+            start = result.order - radius
+            end = result.order + radius
             ranges.append((start, end, result))
 
+        if not ranges:
+            return results
+
+        # Compute the full order range needed and fetch only those chunks
+        all_starts = [s for s, _, _ in ranges]
+        all_ends = [e for _, e, _ in ranges]
+        range_min = min(all_starts)
+        range_max = max(all_ends)
+
+        chunks_in_range = await self.chunk_repository.get_chunks_in_range(
+            doc_id, range_min, range_max
+        )
+        if not chunks_in_range:
+            return results
+
+        chunk_by_order = {c.order: c for c in chunks_in_range}
+
         # Merge and build results
         final_results: list[SearchResult] = []
         for min_idx, max_idx, original_results in self._merge_ranges(ranges):
             # Collect chunks in order
-            chunks_in_range = [
+            merged_chunks = [
                 chunk_by_order[o]
                 for o in range(min_idx, max_idx + 1)
                 if o in chunk_by_order
@@ -1433,7 +1441,7 @@ class HaikuRAG:
             first = original_results[0]
             final_results.append(
                 SearchResult(
-                    content="".join(c.content for c in chunks_in_range),
+                    content="".join(c.content for c in merged_chunks),
                     score=max(r.score for r in original_results),
                     chunk_id=first.chunk_id,
                     document_id=first.document_id,

haiku_rag_slim/haiku/rag/store/models/chunk.py

@@ -117,6 +117,7 @@ class SearchResult(BaseModel):
     document_id: str | None = None
     document_uri: str | None = None
     document_title: str | None = None
+    order: int = 0
     doc_item_refs: list[str] = []
     page_numbers: list[int] = []
     headings: list[str] | None = None
@@ -137,6 +138,7 @@ class SearchResult(BaseModel):
             document_id=chunk.document_id,
             document_uri=chunk.document_uri,
             document_title=chunk.document_title,
+            order=chunk.order,
             doc_item_refs=meta.doc_item_refs,
             page_numbers=meta.page_numbers,
             headings=meta.headings,

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

@@ -363,22 +363,68 @@ class ChunkRepository:
         )
         return len(df)
 
+    async def get_chunks_in_range(
+        self, document_id: str, min_order: int, max_order: int
+    ) -> list[Chunk]:
+        """Get chunks for a document within an order range.
+
+        Args:
+            document_id: The document ID to get chunks for.
+            min_order: Minimum order value (inclusive).
+            max_order: Maximum order value (inclusive).
+
+        Returns:
+            List of chunks within the order range.
+        """
+        where = (
+            f"document_id = '{document_id}'"
+            f" AND `order` >= {min_order}"
+            f" AND `order` <= {max_order}"
+        )
+        results = list(
+            self.store.chunks_table.search()
+            .where(where)
+            .to_pydantic(self.store.ChunkRecord)
+        )
+        return [
+            Chunk(
+                id=rec.id,
+                document_id=rec.document_id,
+                content=rec.content,
+                metadata=json.loads(rec.metadata),
+                order=rec.order,
+            )
+            for rec in results
+        ]
+
     async def get_adjacent_chunks(self, chunk: Chunk, num_adjacent: int) -> list[Chunk]:
         """Get adjacent chunks before and after the given chunk within the same document."""
         assert chunk.document_id, "Document id is required for adjacent chunk finding"
 
-        chunk_order = chunk.order
+        min_order = chunk.order - num_adjacent
+        max_order = chunk.order + num_adjacent
 
-        # Fetch chunks for the same document and filter by order proximity
-        all_chunks = await self.get_by_document_id(chunk.document_id)
-
-        adjacent_chunks: list[Chunk] = []
-        for c in all_chunks:
-            c_order = c.order
-            if c.id != chunk.id and abs(c_order - chunk_order) <= num_adjacent:
-                adjacent_chunks.append(c)
-
-        return adjacent_chunks
+        where = (
+            f"document_id = '{chunk.document_id}'"
+            f" AND `order` >= {min_order}"
+            f" AND `order` <= {max_order}"
+            f" AND id != '{chunk.id}'"
+        )
+        results = list(
+            self.store.chunks_table.search()
+            .where(where)
+            .to_pydantic(self.store.ChunkRecord)
+        )
+        return [
+            Chunk(
+                id=rec.id,
+                document_id=rec.document_id,
+                content=rec.content,
+                metadata=json.loads(rec.metadata),
+                order=rec.order,
+            )
+            for rec in results
+        ]
 
     async def _process_search_results(
         self, query_result: "pd.DataFrame | LanceQueryBuilder"