catlog22/Claude-Code-Workflow
1
0
Fork 0
mirror of https://github.com/catlog22/Claude-Code-Workflow.git synced 2026-02-13 02:41:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: Optimize FTS search with batch symbol fetching and improved excerpt generation

Browse Source
This commit is contained in:
catlog22
2025-12-19 15:35:31 +08:00
parent 2f0cce0089
commit 3428642d04
1 changed files with 252 additions and 294 deletions
Download Patch File Download Diff File Expand all files Collapse all files

546
codex-lens/src/codexlens/storage/dir_index.py
View File

@@ -1221,6 +1221,237 @@ class DirIndexStore:
return block_content, start_line, end_line return block_content, start_line, end_line
def _batch_fetch_symbols(
self, conn: sqlite3.Connection, file_ids: List[int]
) -> Dict[int, List[Tuple[int, int, str, str]]]:
"""Batch fetch all symbols for multiple files in a single query.
Args:
conn: Database connection
file_ids: List of file IDs to fetch symbols for
Returns:
Dictionary mapping file_id to list of (start_line, end_line, name, kind) tuples
"""
if not file_ids:
return {}
# Build placeholder string for IN clause
placeholders = ','.join('?' for _ in file_ids)
rows = conn.execute(
f"""
SELECT file_id, start_line, end_line, name, kind
FROM symbols
WHERE file_id IN ({placeholders})
ORDER BY file_id, (end_line - start_line) ASC
""",
file_ids,
).fetchall()
# Organize symbols by file_id
symbols_by_file: Dict[int, List[Tuple[int, int, str, str]]] = {fid: [] for fid in file_ids}
for row in rows:
symbols_by_file[row["file_id"]].append(
(row["start_line"], row["end_line"], row["name"], row["kind"])
)
return symbols_by_file
def _find_containing_symbol_from_cache(
self, symbols: List[Tuple[int, int, str, str]], line_num: int
) -> Optional[Tuple[int, int, str, str]]:
"""Find the smallest symbol containing the given line number from cached symbols.
Args:
symbols: List of (start_line, end_line, name, kind) tuples, sorted by size
line_num: 1-based line number
Returns:
Tuple of (start_line, end_line, symbol_name, symbol_kind) or None
"""
for start_line, end_line, name, kind in symbols:
if start_line <= line_num <= end_line:
return (start_line, end_line, name, kind)
return None
def _generate_centered_excerpt(
self, content: str, match_line: int, start_line: int, end_line: int, max_chars: int = 200
) -> str:
"""Generate excerpt centered around the match line.
Args:
content: Full file content
match_line: 1-based line where match occurred
start_line: 1-based start line of the code block
end_line: 1-based end line of the code block
max_chars: Maximum characters for excerpt
Returns:
Excerpt string centered around the match
"""
lines = content.split('\n')
total_lines = len(lines)
# Ensure match_line is within bounds
match_line = max(1, min(match_line, total_lines))
# Calculate context window (2 lines before, 2 lines after the match)
ctx_start = max(start_line, match_line - 2)
ctx_end = min(end_line, match_line + 2)
# Extract and join lines
excerpt_lines = lines[ctx_start - 1:ctx_end]
excerpt = '\n'.join(excerpt_lines)
# Truncate if too long
if len(excerpt) > max_chars:
excerpt = excerpt[:max_chars] + "..."
return excerpt
def _search_internal(
self,
query: str,
fts_table: str,
limit: int = 20,
return_full_content: bool = False,
context_lines: int = 10,
) -> List[SearchResult]:
"""Internal unified search implementation for all FTS modes.
Optimizations:
- Fast path: Direct FTS query with snippet() for location-only results
- Full content path: Batch fetch symbols to eliminate N+1 queries
- Centered excerpt generation for better context
Args:
query: FTS5 query string
fts_table: FTS table name ('files_fts_exact' or 'files_fts_fuzzy')
limit: Maximum results to return
return_full_content: If True, include full code block in content field
context_lines: Lines of context when no symbol contains the match
Returns:
List of SearchResult objects
"""
with self._lock:
conn = self._get_connection()
# Fast path: location-only results (no content processing)
if not return_full_content:
try:
rows = conn.execute(
f"""
SELECT rowid, full_path, bm25({fts_table}) AS rank,
snippet({fts_table}, 2, '', '', '...', 30) AS excerpt
FROM {fts_table}
WHERE {fts_table} MATCH ?
ORDER BY rank
LIMIT ?
""",
(query, limit),
).fetchall()
except sqlite3.DatabaseError as exc:
raise StorageError(f"FTS search failed: {exc}") from exc
results: List[SearchResult] = []
for row in rows:
rank = float(row["rank"]) if row["rank"] is not None else 0.0
score = abs(rank) if rank < 0 else 0.0
results.append(
SearchResult(
path=row["full_path"],
score=score,
excerpt=row["excerpt"],
)
)
return results
# Full content path with batch optimization
# Step 1: Get file_ids and ranks (lightweight query)
try:
id_rows = conn.execute(
f"""
SELECT rowid AS file_id, bm25({fts_table}) AS rank
FROM {fts_table}
WHERE {fts_table} MATCH ?
ORDER BY rank
LIMIT ?
""",
(query, limit),
).fetchall()
except sqlite3.DatabaseError as exc:
raise StorageError(f"FTS search failed: {exc}") from exc
if not id_rows:
return []
file_ids = [row["file_id"] for row in id_rows]
ranks_by_id = {row["file_id"]: row["rank"] for row in id_rows}
# Step 2: Batch fetch all symbols for matched files (eliminates N+1)
symbols_by_file = self._batch_fetch_symbols(conn, file_ids)
# Step 3: Process each file on-demand (reduces memory)
results: List[SearchResult] = []
for file_id in file_ids:
# Fetch file content on-demand
file_row = conn.execute(
"SELECT full_path, content FROM files WHERE id = ?",
(file_id,),
).fetchone()
if not file_row:
continue
file_path = file_row["full_path"]
content = file_row["content"] or ""
rank = ranks_by_id.get(file_id, 0.0)
score = abs(rank) if rank < 0 else 0.0
# Find matching lines
match_lines = self._find_match_lines(content, query)
first_match_line = match_lines[0] if match_lines else 1
# Find symbol from cached symbols (no extra SQL query)
file_symbols = symbols_by_file.get(file_id, [])
symbol_info = self._find_containing_symbol_from_cache(file_symbols, first_match_line)
if symbol_info:
start_line, end_line, symbol_name, symbol_kind = symbol_info
else:
# No symbol found, use context around match
lines = content.split('\n')
total_lines = len(lines)
start_line = max(1, first_match_line - context_lines)
end_line = min(total_lines, first_match_line + context_lines)
symbol_name = None
symbol_kind = None
# Extract code block
block_content, start_line, end_line = self._extract_code_block(
content, start_line, end_line
)
# Generate centered excerpt (improved quality)
excerpt = self._generate_centered_excerpt(
content, first_match_line, start_line, end_line
)
results.append(
SearchResult(
path=file_path,
score=score,
excerpt=excerpt,
content=block_content,
start_line=start_line,
end_line=end_line,
symbol_name=symbol_name,
symbol_kind=symbol_kind,
)
)
return results
def search_fts( def search_fts(
self, self,
query: str, query: str,
@@ -1255,107 +1486,14 @@ class DirIndexStore:
Raises: Raises:
StorageError: If FTS search fails StorageError: If FTS search fails
""" """
# Only enhance query if explicitly requested (not default behavior)
# Best practice: Let users control wildcards manually
final_query = self._enhance_fts_query(query) if enhance_query else query final_query = self._enhance_fts_query(query) if enhance_query else query
return self._search_internal(
with self._lock: query=final_query,
conn = self._get_connection() fts_table='files_fts_exact',
limit=limit,
# Fast path: location-only results (no content processing) return_full_content=return_full_content,
if not return_full_content: context_lines=context_lines,
try: )
rows = conn.execute(
"""
SELECT rowid, full_path, bm25(files_fts_exact) AS rank,
snippet(files_fts_exact, 2, '', '', '...', 30) AS excerpt
FROM files_fts_exact
WHERE files_fts_exact MATCH ?
ORDER BY rank
LIMIT ?
""",
(final_query, limit),
).fetchall()
except sqlite3.DatabaseError as exc:
raise StorageError(f"FTS search failed: {exc}") from exc
results: List[SearchResult] = []
for row in rows:
rank = float(row["rank"]) if row["rank"] is not None else 0.0
score = abs(rank) if rank < 0 else 0.0
results.append(
SearchResult(
path=row["full_path"],
score=score,
excerpt=row["excerpt"],
)
)
return results
# Full content path: fetch content and find containing symbols
try:
rows = conn.execute(
"""
SELECT f.id AS file_id, f.full_path, f.content,
bm25(files_fts_exact) AS rank
FROM files_fts_exact
JOIN files f ON files_fts_exact.rowid = f.id
WHERE files_fts_exact MATCH ?
ORDER BY rank
LIMIT ?
""",
(final_query, limit),
).fetchall()
except sqlite3.DatabaseError as exc:
raise StorageError(f"FTS search failed: {exc}") from exc
results: List[SearchResult] = []
for row in rows:
file_id = row["file_id"]
file_path = row["full_path"]
content = row["content"] or ""
rank = float(row["rank"]) if row["rank"] is not None else 0.0
score = abs(rank) if rank < 0 else 0.0
# Find matching lines
match_lines = self._find_match_lines(content, final_query)
first_match_line = match_lines[0] if match_lines else 1
# Find symbol containing the first match
symbol_info = self._find_containing_symbol(conn, file_id, first_match_line)
if symbol_info:
start_line, end_line, symbol_name, symbol_kind = symbol_info
else:
# No symbol found, use context around match
lines = content.split('\n')
total_lines = len(lines)
start_line = max(1, first_match_line - context_lines)
end_line = min(total_lines, first_match_line + context_lines)
symbol_name = None
symbol_kind = None
# Extract code block
block_content, start_line, end_line = self._extract_code_block(
content, start_line, end_line
)
# Generate excerpt (first 200 chars)
excerpt = block_content[:200] + "..." if len(block_content) > 200 else block_content
results.append(
SearchResult(
path=file_path,
score=score,
excerpt=excerpt,
content=block_content,
start_line=start_line,
end_line=end_line,
symbol_name=symbol_name,
symbol_kind=symbol_kind,
)
)
return results
def search_fts_exact( def search_fts_exact(
self, self,
@@ -1379,103 +1517,13 @@ class DirIndexStore:
Raises: Raises:
StorageError: If FTS search fails StorageError: If FTS search fails
""" """
with self._lock: return self._search_internal(
conn = self._get_connection() query=query,
fts_table='files_fts_exact',
# Fast path: location-only results (no content processing) limit=limit,
if not return_full_content: return_full_content=return_full_content,
try: context_lines=context_lines,
rows = conn.execute( )
"""
SELECT rowid, full_path, bm25(files_fts_exact) AS rank,
snippet(files_fts_exact, 2, '', '', '...', 30) AS excerpt
FROM files_fts_exact
WHERE files_fts_exact MATCH ?
ORDER BY rank
LIMIT ?
""",
(query, limit),
).fetchall()
except sqlite3.DatabaseError as exc:
raise StorageError(f"FTS exact search failed: {exc}") from exc
results: List[SearchResult] = []
for row in rows:
rank = float(row["rank"]) if row["rank"] is not None else 0.0
score = abs(rank) if rank < 0 else 0.0
results.append(
SearchResult(
path=row["full_path"],
score=score,
excerpt=row["excerpt"],
)
)
return results
# Full content path: fetch content and find containing symbols
try:
rows = conn.execute(
"""
SELECT f.id AS file_id, f.full_path, f.content,
bm25(files_fts_exact) AS rank
FROM files_fts_exact
JOIN files f ON files_fts_exact.rowid = f.id
WHERE files_fts_exact MATCH ?
ORDER BY rank
LIMIT ?
""",
(query, limit),
).fetchall()
except sqlite3.DatabaseError as exc:
raise StorageError(f"FTS exact search failed: {exc}") from exc
results: List[SearchResult] = []
for row in rows:
file_id = row["file_id"]
file_path = row["full_path"]
content = row["content"] or ""
rank = float(row["rank"]) if row["rank"] is not None else 0.0
score = abs(rank) if rank < 0 else 0.0
# Find matching lines
match_lines = self._find_match_lines(content, query)
first_match_line = match_lines[0] if match_lines else 1
# Find symbol containing the first match
symbol_info = self._find_containing_symbol(conn, file_id, first_match_line)
if symbol_info:
start_line, end_line, symbol_name, symbol_kind = symbol_info
else:
# No symbol found, use context around match
lines = content.split('\n')
total_lines = len(lines)
start_line = max(1, first_match_line - context_lines)
end_line = min(total_lines, first_match_line + context_lines)
symbol_name = None
symbol_kind = None
# Extract code block
block_content, start_line, end_line = self._extract_code_block(
content, start_line, end_line
)
# Generate excerpt (first 200 chars)
excerpt = block_content[:200] + "..." if len(block_content) > 200 else block_content
results.append(
SearchResult(
path=file_path,
score=score,
excerpt=excerpt,
content=block_content,
start_line=start_line,
end_line=end_line,
symbol_name=symbol_name,
symbol_kind=symbol_kind,
)
)
return results
def search_fts_fuzzy( def search_fts_fuzzy(
self, self,
@@ -1499,103 +1547,13 @@ class DirIndexStore:
Raises: Raises:
StorageError: If FTS search fails StorageError: If FTS search fails
""" """
with self._lock: return self._search_internal(
conn = self._get_connection() query=query,
fts_table='files_fts_fuzzy',
# Fast path: location-only results (no content processing) limit=limit,
if not return_full_content: return_full_content=return_full_content,
try: context_lines=context_lines,
rows = conn.execute( )
"""
SELECT rowid, full_path, bm25(files_fts_fuzzy) AS rank,
snippet(files_fts_fuzzy, 2, '', '', '...', 30) AS excerpt
FROM files_fts_fuzzy
WHERE files_fts_fuzzy MATCH ?
ORDER BY rank
LIMIT ?
""",
(query, limit),
).fetchall()
except sqlite3.DatabaseError as exc:
raise StorageError(f"FTS fuzzy search failed: {exc}") from exc
results: List[SearchResult] = []
for row in rows:
rank = float(row["rank"]) if row["rank"] is not None else 0.0
score = abs(rank) if rank < 0 else 0.0
results.append(
SearchResult(
path=row["full_path"],
score=score,
excerpt=row["excerpt"],
)
)
return results
# Full content path: fetch content and find containing symbols
try:
rows = conn.execute(
"""
SELECT f.id AS file_id, f.full_path, f.content,
bm25(files_fts_fuzzy) AS rank
FROM files_fts_fuzzy
JOIN files f ON files_fts_fuzzy.rowid = f.id
WHERE files_fts_fuzzy MATCH ?
ORDER BY rank
LIMIT ?
""",
(query, limit),
).fetchall()
except sqlite3.DatabaseError as exc:
raise StorageError(f"FTS fuzzy search failed: {exc}") from exc
results: List[SearchResult] = []
for row in rows:
file_id = row["file_id"]
file_path = row["full_path"]
content = row["content"] or ""
rank = float(row["rank"]) if row["rank"] is not None else 0.0
score = abs(rank) if rank < 0 else 0.0
# Find matching lines
match_lines = self._find_match_lines(content, query)
first_match_line = match_lines[0] if match_lines else 1
# Find symbol containing the first match
symbol_info = self._find_containing_symbol(conn, file_id, first_match_line)
if symbol_info:
start_line, end_line, symbol_name, symbol_kind = symbol_info
else:
# No symbol found, use context around match
lines = content.split('\n')
total_lines = len(lines)
start_line = max(1, first_match_line - context_lines)
end_line = min(total_lines, first_match_line + context_lines)
symbol_name = None
symbol_kind = None
# Extract code block
block_content, start_line, end_line = self._extract_code_block(
content, start_line, end_line
)
# Generate excerpt (first 200 chars)
excerpt = block_content[:200] + "..." if len(block_content) > 200 else block_content
results.append(
SearchResult(
path=file_path,
score=score,
excerpt=excerpt,
content=block_content,
start_line=start_line,
end_line=end_line,
symbol_name=symbol_name,
symbol_kind=symbol_kind,
)
)
return results
def search_files_only(self, query: str, limit: int = 20) -> List[str]: def search_files_only(self, query: str, limit: int = 20) -> List[str]:
"""Fast FTS search returning only file paths (no snippet generation). """Fast FTS search returning only file paths (no snippet generation).