feat: Add partial reindexing#225
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests.
🚀 New features to boost your workflow:
|
Confidence Score: 4/5Safe to merge; the incremental reindex logic is well-designed and thoroughly tested, with no observed correctness issues in the happy path. The core indexing and search paths are correct, the cache validation is thorough, and the test suite covers the main edge cases. All observations are non-blocking: BM25 IDF uses the full document count rather than the scored-document count (an API fragility, not a current bug), the
Reviews (1): Last reviewed commit: "Optimizations" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
stephantul
left a comment
stephantul
left a comment
There was a problem hiding this comment.
Looks good to me! Lots of small comments, but nothing major.
| for index, _, count in embedding_parts: | ||
| vector_parts[index] = changed_embeddings[offset : offset + count] | ||
| offset += count | ||
| same_vector_layout = len(manifest) == len(previous_manifest) and all( |
There was a problem hiding this comment.
Maybe turn this into a function
There was a problem hiding this comment.
Agree, I turned this into a helper
| found_version = metadata.get("cache_version") | ||
| if found_version != CACHE_FORMAT_VERSION: | ||
| raise ValueError( | ||
| f"Unsupported index format {found_version!r}; expected {CACHE_FORMAT_VERSION}. Rebuild the index." |
There was a problem hiding this comment.
Maybe make it clear to the user how to rebuild the index? Or do it for them automatically? A possible risk with raising here is that an agent runs into this error and keeps searching and never actually runs semble anymore.
There was a problem hiding this comment.
Yeah makes sense, I've updated the error message
| chunks = [] | ||
| for chunk_item in chunk_data: | ||
| chunks.append(Chunk.from_dict(chunk_item)) | ||
| if len(chunks) != len(bm25_index.doc_order) or len(chunks) != semantic_index.vectors.shape[0]: |
There was a problem hiding this comment.
you can chain this: if not (len(chunks) == x == y)
| vectors = SelectableBasicBackend.load(persistence_path.semantic_index).vectors | ||
| bm25_index = BM25.load(persistence_path.bm25_index) | ||
| chunk_count = len(chunks) | ||
| if vectors.shape[0] != chunk_count or len(bm25_index.doc_order) != chunk_count: |
There was a problem hiding this comment.
this is a repetition of a snippet I also commented on, maybe turn it into a function
| return None | ||
|
|
||
| return PreviousIndex(chunks=chunks, vectors=vectors, manifest=manifest, bm25_index=bm25_index) | ||
| except (OSError, orjson.JSONDecodeError, json.JSONDecodeError, KeyError, TypeError, ValueError): |
There was a problem hiding this comment.
maybe turn this into a constant? I'm not sure if that's the right way. Also maybe we expect this to log something useful? This seems to point towards an unexpected path, rather than an expected cache invalidation.
There was a problem hiding this comment.
I've added a debug log with the exception context before falling back, that should be helpful.I'm not sure about a constant though, this is only used once here and I think it's pretty specific for the cache
|
|
||
| mtime: float | ||
| start: int | ||
| count: int |
There was a problem hiding this comment.
should this not be length? As in start + length == end?
There was a problem hiding this comment.
Hmm it represents the number of chunks so I think count is a bit more intuitive
| for indexed_path, entry in manifest.items(): | ||
| if ( | ||
| entry.start != next_start | ||
| or entry.count < 0 |
There was a problem hiding this comment.
when can entry.count be < 0? Should this not fail when creating the manifest? It seems to me that a manifest item with a negative count can't be valid anyway.
There was a problem hiding this comment.
Yeah you're right this is overly defensive, we control the cache so this won't ever happen in practice, I'll get rid of it
| mtime = file_path.stat().st_mtime | ||
| previous_entry = previous_manifest.get(indexed_path) | ||
|
|
||
| if previous is not None and previous_entry is not None and previous_entry.mtime == mtime: |
There was a problem hiding this comment.
The previous_entry.mtime is a float comparison. I wonder if this ever doesn't work?
There was a problem hiding this comment.
Yess good catch, will update
| previous_entry = previous_manifest.get(indexed_path) | ||
|
|
||
| if previous is not None and previous_entry is not None and previous_entry.mtime == mtime: | ||
| file_chunks = previous.chunks[previous_entry.start : previous_entry.start + previous_entry.count] |
There was a problem hiding this comment.
maybe create an end property to be used. e.g., end is defined as start + count.
| _reindex_file(bm25_index, indexed_path, file_chunks, previous_entry) | ||
|
|
||
| embedding_parts.append((len(vector_parts), len(chunks), len(file_chunks))) | ||
| vector_parts.append(np.empty((0, model.dim), dtype=np.float32)) |
There was a problem hiding this comment.
I don't understand this yet, why do we need an empty thing here?
There was a problem hiding this comment.
To make batching work across different files (you need to preserve the fileorder), but tbh, this is kind of overtuned, I'll just embed per file. The chances of having to re-embed multiple files is small and even if you have to do it it's superfast
2 participants
This PR adds partial reindexing support. BM25s is dropped in favor of a custom BM25 implementation that supports adding/removing documents without full rebuilds. Results on our benchmarks:
The scores are basically the same. Queries are a bit slower but still so fast that I don't think it matters. Now the big improvement, tested on Zig (one of the larger repos in our benchmark):