Repository navigation
Conversation
Persist embedding search entries in per-note shard files so that updating a paragraph rewrites only the affected note instead of the entire index. Add note-level dirty tracking and locking, atomic shard writes, targeted recovery for missing or corrupt shards, orphan cleanup, and migration from the legacy single-file index. Add tests covering shard isolation, migration, recovery, deletion races, and partial-load failures.
tbonelee
left a comment
tbonelee
left a comment
There was a problem hiding this comment.
It would be good to keep comments only where they explain a "why" the code can't express. For example, the reason query() doesn't take a lock (EmbeddingSearch.java L577-581) is worth keeping. The following look like they could be removed or replaced with code:
- Ticket numbers and change history:
EmbeddingSearch.javaL146, L962, L1169, andEmbeddingSearchShardingTest.javaL173, L363-364, L701-702. The history already lives in the commit message and JIRA. - Javadoc that argues design decisions at length:
EmbeddingSearch.javaL785-791, L961-971, L1176-1185. The last one also doesn't match the current behavior, whereloadAllShardsrebuilds notes with missing shards. shouldBootstrapIndex(L842-854): ten lines of javadoc for a one-line body,zConf.isIndexRebuild(), which could be inlined at the call site.
| List<Map.Entry<String, IndexEntry>> entries = index.entrySet().stream() | ||
| .filter(e -> e.getKey().equals(noteId) || e.getKey().startsWith(noteId + "/")) | ||
| .collect(Collectors.toList()); |
There was a problem hiding this comment.
saveNoteShard, commitNoteEntries (L1160) and writeShardToDir (L1262) each scan the whole index with startsWith(noteId + "/") for every note, so handling many notes at once costs (number of notes × total entries).
Measured locally without the model (skipModel), with 2,000 notes × 25 paragraphs (50K entries): a full flush went from 270ms to 5.1s, and loading shards on startup from 133ms to 1.3s. With 4,000 notes they were 29s and 5.2s. Load was measured with a BufferedInputStream to leave out file read cost. Full bootstrap, zeppelin.search.index.rebuild and migration all go through this path.
If the index were kept per note, e.g. Map<noteId, Map<docId, IndexEntry>>, all three could look only at that note's map, and the prefix scans in deleteNoteIndex and updateNoteIndex would go away as well. What do you think?
| * left in place (so the next restart retries from scratch; shard writes are keyed by | ||
| * noteId and idempotent, so repeating this is safe) and the staging directory is cleared. | ||
| */ | ||
| private void migrateLegacyIndexIfPresent() { |
There was a problem hiding this comment.
#5218 (ZEPPELIN-6411) isn't in any release yet, and semantic search is opt-in, requiring its own config and a model install. So only users who turned it on with a master snapshot have an embedding_index.bin. Since the index can be rebuilt from the notebooks, would it be enough to delete the legacy file and run a full bootstrap? That would let us drop loadLegacyIndex, writeShardToDir, the staging directory handling and the related tests.
Even if we keep the migration, I don't think the staging directory is needed. If migration fails partway and only some notes' shards are left, that is the same as a few shards having gone missing, which loadAllShards (L1084-1090) already handles by reindexing only the notes without a shard. I changed the migration to write directly into notes/ and made it fail on one note, and both the migrated note and the rebuilt note were searchable.
This branch has not been deployed
2 participants
What is this PR for?
EmbeddingSearchcurrently keeps all entries in a singleembedding_index.binfile. Although paragraph embeddings are updated incrementally in memory, every flush serializes and rewrites the entire index.This PR shards the persisted index by note. Updating a paragraph now rewrites only the shard belonging to the affected note, while semantic search continues to operate across the complete in-memory index.
What type of PR is it?
Bug Fix
What changes were made?
Per-note persistence
notes/Shard lifecycle and recovery
Legacy index migration
embedding_index.binusing an all-or-nothing staging directoryTests and documentation
Todos
What is the Jira issue?
How should this be tested?
./mvnw test -pl zeppelin-server --am \ -Dtest=EmbeddingSearchShardingTest \ -Dsurefire.failIfNoSpecifiedTests=falseEmbeddingSearchShardingTestcontains 20 tests covering:All 20 tests pass.
Screenshots (if appropriate)
Not applicable.
Questions:
docs/embedding-search.mdhas been updated to describe per-note persistence.