Merge pull request #127 from ggozad/feat/faster-filtering
Faster filtering through pandas
This commit is contained in:
Yiorgis Gozadinos
GitHub
commit
bead05fffb
1 changed files with 80 additions and 79 deletions
|
|
@ -4,6 +4,10 @@ import logging
|
|||
from typing import TYPE_CHECKING
|
||||
from uuid import uuid4
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pandas as pd
|
||||
from lancedb.query import LanceQueryBuilder
|
||||
|
||||
from lancedb.rerankers import RRFReranker
|
||||
|
||||
from haiku.rag.store.engine import DocumentRecord, Store
|
||||
|
|
@ -249,64 +253,56 @@ class ChunkRepository:
|
|||
"""
|
||||
if not query.strip():
|
||||
return []
|
||||
|
||||
chunk_where_clause = None
|
||||
filtered_doc_ids = 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.
|
||||
matching_doc_ids = self._get_filtered_document_ids(filter)
|
||||
|
||||
if not matching_doc_ids:
|
||||
docs_df = (
|
||||
self.store.documents_table.search()
|
||||
.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"]
|
||||
|
||||
# Build WHERE clause for chunks table
|
||||
# Use IN clause with document IDs
|
||||
id_list = "', '".join(matching_doc_ids)
|
||||
chunk_where_clause = f"document_id IN ('{id_list}')"
|
||||
|
||||
# Prepare search query based on search type
|
||||
if search_type == "vector":
|
||||
query_embedding = await self.embedder.embed(query)
|
||||
|
||||
results = self.store.chunks_table.search(
|
||||
query_embedding, query_type="vector", vector_column_name="vector"
|
||||
)
|
||||
|
||||
if chunk_where_clause:
|
||||
results = results.where(chunk_where_clause)
|
||||
|
||||
results = results.limit(limit)
|
||||
|
||||
return await self._process_search_results(results)
|
||||
|
||||
elif search_type == "fts":
|
||||
results = self.store.chunks_table.search(query, query_type="fts")
|
||||
|
||||
if chunk_where_clause:
|
||||
results = results.where(chunk_where_clause)
|
||||
|
||||
results = results.limit(limit)
|
||||
return await self._process_search_results(results)
|
||||
|
||||
else: # hybrid (default)
|
||||
query_embedding = await self.embedder.embed(query)
|
||||
|
||||
# Create RRF reranker
|
||||
reranker = RRFReranker()
|
||||
|
||||
# Perform native hybrid search with RRF reranking
|
||||
results = self.store.chunks_table.search(query_type="hybrid")
|
||||
|
||||
if chunk_where_clause:
|
||||
results = results.where(chunk_where_clause)
|
||||
|
||||
results = (
|
||||
results.vector(query_embedding)
|
||||
self.store.chunks_table.search(query_type="hybrid")
|
||||
.vector(query_embedding)
|
||||
.text(query)
|
||||
.rerank(reranker)
|
||||
.limit(limit)
|
||||
)
|
||||
return await self._process_search_results(results)
|
||||
|
||||
# Apply filtering if needed (common for all search types)
|
||||
if filtered_doc_ids is not None:
|
||||
chunks_df = 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
|
||||
results = results.limit(limit)
|
||||
return await self._process_search_results(results)
|
||||
|
||||
async def get_by_document_id(self, document_id: str) -> list[Chunk]:
|
||||
"""Get all chunks for a specific document."""
|
||||
|
|
@ -364,39 +360,51 @@ class ChunkRepository:
|
|||
|
||||
return adjacent_chunks
|
||||
|
||||
def _get_filtered_document_ids(self, filter: str) -> list[str]:
|
||||
"""Query documents table with filter and return matching document IDs."""
|
||||
filtered_docs = (
|
||||
self.store.documents_table.search()
|
||||
.where(filter)
|
||||
.to_pydantic(DocumentRecord)
|
||||
)
|
||||
return [doc.id for doc in filtered_docs]
|
||||
async def _process_search_results(
|
||||
self, query_result: "pd.DataFrame | LanceQueryBuilder"
|
||||
) -> list[tuple[Chunk, float]]:
|
||||
"""Process search results into chunks with document info and scores.
|
||||
|
||||
async def _process_search_results(self, query_result) -> list[tuple[Chunk, float]]:
|
||||
"""Process search results into chunks with document info and scores."""
|
||||
chunks_with_scores = []
|
||||
Args:
|
||||
query_result: Either a pandas DataFrame or a LanceDB query result
|
||||
"""
|
||||
import pandas as pd
|
||||
|
||||
# Get both arrow and pydantic results to access scores
|
||||
arrow_result = query_result.to_arrow()
|
||||
pydantic_results = list(query_result.to_pydantic(self.store.ChunkRecord))
|
||||
def extract_scores(df: pd.DataFrame) -> list[float]:
|
||||
"""Extract scores from DataFrame columns based on search type."""
|
||||
if "_distance" in df.columns:
|
||||
# Vector search - convert distance to similarity
|
||||
return ((df["_distance"] + 1).rdiv(1)).clip(lower=0.0).tolist()
|
||||
elif "_relevance_score" in df.columns:
|
||||
# Hybrid search - relevance score (higher is better)
|
||||
return df["_relevance_score"].tolist()
|
||||
elif "_score" in df.columns:
|
||||
# FTS search - score (higher is better)
|
||||
return df["_score"].tolist()
|
||||
else:
|
||||
raise ValueError("Unknown search result format, cannot extract scores")
|
||||
|
||||
# Extract scores from arrow result based on search type
|
||||
scores = []
|
||||
column_names = arrow_result.column_names
|
||||
|
||||
if "_distance" in column_names:
|
||||
# Vector search - distance (lower is better, convert to similarity)
|
||||
distances = arrow_result.column("_distance").to_pylist()
|
||||
scores = [max(0.0, 1.0 / (1.0 + dist)) for dist in distances]
|
||||
elif "_relevance_score" in column_names:
|
||||
# Hybrid search - relevance score (higher is better)
|
||||
scores = arrow_result.column("_relevance_score").to_pylist()
|
||||
elif "_score" in column_names:
|
||||
# FTS search - score (higher is better)
|
||||
scores = arrow_result.column("_score").to_pylist()
|
||||
# Convert everything to DataFrame for uniform processing
|
||||
if isinstance(query_result, pd.DataFrame):
|
||||
df = query_result
|
||||
else:
|
||||
raise ValueError("Unknown search result format, cannot extract scores")
|
||||
# Convert LanceDB query result to DataFrame
|
||||
df = query_result.to_pandas()
|
||||
|
||||
# Extract scores
|
||||
scores = extract_scores(df)
|
||||
|
||||
# Convert DataFrame rows to ChunkRecords
|
||||
pydantic_results = [
|
||||
self.store.ChunkRecord(
|
||||
id=str(row["id"]),
|
||||
document_id=str(row["document_id"]),
|
||||
content=str(row["content"]),
|
||||
metadata=str(row["metadata"]),
|
||||
order=int(row["order"]) if "order" in row else 0,
|
||||
)
|
||||
for _, row in df.iterrows()
|
||||
]
|
||||
|
||||
# Collect all unique document IDs for batch lookup
|
||||
document_ids = list(set(chunk.document_id for chunk in pydantic_results))
|
||||
|
|
@ -404,8 +412,9 @@ class ChunkRepository:
|
|||
# Batch fetch all documents at once
|
||||
documents_map = {}
|
||||
if document_ids:
|
||||
# Create a WHERE clause for all document IDs
|
||||
where_clause = " OR ".join(f"id = '{doc_id}'" for doc_id in document_ids)
|
||||
# Use IN clause for efficient batch lookup
|
||||
id_list = "', '".join(document_ids)
|
||||
where_clause = f"id IN ('{id_list}')"
|
||||
doc_results = list(
|
||||
self.store.documents_table.search()
|
||||
.where(where_clause)
|
||||
|
|
@ -413,29 +422,21 @@ class ChunkRepository:
|
|||
)
|
||||
documents_map = {doc.id: doc for doc in doc_results}
|
||||
|
||||
# Build final results with document info
|
||||
chunks_with_scores = []
|
||||
for i, chunk_record in enumerate(pydantic_results):
|
||||
# Get document info from pre-fetched map
|
||||
doc = documents_map.get(chunk_record.document_id)
|
||||
doc_uri = doc.uri if doc else None
|
||||
doc_title = doc.title if doc else None
|
||||
doc_meta = doc.metadata if doc else "{}"
|
||||
|
||||
md = json.loads(chunk_record.metadata)
|
||||
|
||||
chunk = Chunk(
|
||||
id=chunk_record.id,
|
||||
document_id=chunk_record.document_id,
|
||||
content=chunk_record.content,
|
||||
metadata=md,
|
||||
metadata=json.loads(chunk_record.metadata),
|
||||
order=chunk_record.order,
|
||||
document_uri=doc_uri,
|
||||
document_title=doc_title,
|
||||
document_meta=json.loads(doc_meta),
|
||||
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 "{}"),
|
||||
)
|
||||
|
||||
# Get score from arrow result
|
||||
score = scores[i] if i < len(scores) else 1.0
|
||||
|
||||
chunks_with_scores.append((chunk, score))
|
||||
|
||||
return chunks_with_scores
|
||||
|
|
|
|||
Loading…
Reference in a new issue