Push document-id filter into chunk search query

CHANGELOG.md

@@ -11,6 +11,7 @@
 ### Fixed
 
 - **Chat TUI now renders citations again.** After the 0.42.1 flattening of skill state `citations` to `list[str]`, the TUI still indexed `citations[-1]` and iterated the resulting chunk-id string character-by-character, so no citations resolved through `citation_index` and the citation panel stayed empty. Fixed by iterating `state.citations` directly.
+- **`search(..., filter=...)` no longer silently under-returns.** The filter path used to materialize LanceDB's top-N window, filter to matching `document_id`s in pandas, and `head(limit)`. When matching chunks lived outside that top-N window (selective filters, broad queries), the caller got fewer than `limit` results even though plenty of matching chunks existed in the index. The document filter is now pushed down into the chunk query as `document_id IN (...)` so `.limit(limit)` applies to matching chunks directly. Behavior change: searches that previously under-returned will start returning the requested count.
 
 ## [0.42.1] - 2026-04-22
 

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

@@ -236,24 +236,25 @@ class ChunkRepository:
         """
         if not query.strip():
             return []
-        filtered_doc_ids = None
+
+        chunk_filter: str | None = 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.
+            # Translate the document-level filter into a chunk-level
+            # document_id IN (...) clause so LanceDB can combine it with
+            # limit. The previous two-step pattern (materialize top-N,
+            # filter in pandas, head(limit)) silently under-returned
+            # whenever the top-N window lacked `limit` matching chunks.
             docs_df = await (
                 self.store.documents_table.query()
                 .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"]
+            id_list = ", ".join(f"'{d}'" for d in docs_df["id"])
+            chunk_filter = f"document_id IN ({id_list})"
 
-        # Prepare search query based on search type
         if search_type == "vector":
             query_embedding = await self.embedder.embed_query(query)
             results = (
@@ -262,17 +263,13 @@ class ChunkRepository:
                 .column("vector")
                 .refine_factor(self.store._config.search.vector_refine_factor)
             )
-
         elif search_type == "fts":
             results = self.store.chunks_table.query().nearest_to_text(
                 query, columns="content_fts"
             )
-
         else:  # hybrid (default)
             query_embedding = await self.embedder.embed_query(query)
-            # Create RRF reranker
             reranker = RRFReranker()
-            # Perform native hybrid search with RRF reranking
             results = (
                 self.store.chunks_table.query()
                 .nearest_to(query_embedding)
@@ -282,15 +279,8 @@ class ChunkRepository:
                 .rerank(reranker)
             )
 
-        # Apply filtering if needed (common for all search types)
-        if filtered_doc_ids is not None:
-            chunks_df = await results.to_pandas()
-            filtered_chunks_df = chunks_df.loc[
-                chunks_df["document_id"].isin(filtered_doc_ids)
-            ].head(limit)
-            return await self._process_search_results(filtered_chunks_df)
-
-        # No filtering needed, apply limit and return
+        if chunk_filter is not None:
+            results = results.where(chunk_filter)
         results = results.limit(limit)
         return await self._process_search_results(results)
 

tests/test_filter.py

@@ -179,3 +179,45 @@ async def test_search_filter_with_all_search_types(temp_db_path):
         for result in results:
             assert result.document_uri is not None
             assert "other.com" in result.document_uri
+
+
+@pytest.mark.vcr()
+async def test_search_with_filter_returns_full_limit(temp_db_path):
+    """Regression: filter + limit must return up to `limit` matching chunks
+    even when non-matching chunks would dominate the top-N window.
+
+    Previously the filter path materialized LanceDB's default top-N window
+    (~10), filtered to matching document_ids in pandas, then took `head(limit)`.
+    If the top-N window was dominated by non-matching chunks, the caller got
+    silently fewer results than requested — even when plenty of matching
+    chunks existed further down the ranking. This test puts the target
+    document behind many distractor documents and asserts we still get the
+    requested count back.
+    """
+    async with HaikuRAG(db_path=temp_db_path, create=True) as client:
+        for i in range(12):
+            await client.create_document(
+                content=(
+                    "machine learning neural network deep learning model "
+                    "machine learning neural network deep learning model "
+                    "machine learning neural network deep learning model"
+                ),
+                uri=f"https://distractor.com/doc{i}.html",
+                title=f"Distractor {i}",
+            )
+
+        await client.create_document(
+            content="one passing mention of machine learning here",
+            uri="https://target.com/one.html",
+            title="Target One",
+        )
+
+        results = await client.search(
+            "machine learning",
+            limit=5,
+            search_type="fts",
+            filter="uri LIKE '%target.com%'",
+        )
+
+        assert len(results) == 1
+        assert results[0].document_uri == "https://target.com/one.html"