feat(ai): make the keyword search index resident and incremental - #825
feat(ai): make the keyword search index resident and incremental#825E2ern1ty wants to merge 4 commits into
Conversation
…lace The global keyword search rebuilt a throw-away Lucene index from a full graph scan on every query, so index construction cost was paid on the query path and discarded afterwards. Verbalization was also redone for the whole graph per query, and multi-round sessions ran a global search whose result was dropped. Follow the standard inverted index maintenance model instead of invalidate-and-rebuild: - Add ResidentSearchIndex, a graph scoped index built at most once and kept alive across queries, held by GraphMemoryServer per keyword index store. - Give each document a non analyzed primary key (ModelUtils.getGraphEntityKey) so writes map to Lucene update/delete by term. Cost is proportional to the change, not to graph size, and upserts are idempotent, so callers do not have to supply an exact delta. - Split the server hook into onEntitiesUpserted / onEntitiesRemoved / onSchemaChanged; only a schema change needs wholesale invalidation. - Track validity against a graph vertex version so mutations made outside the server (for example directly through MemoryMutableGraph) are detected and force a rebuild instead of serving stale results. Accessors that cannot report a version degrade to the previous per-query rebuild behaviour. - Memoize entity verbalization in a bounded, version aware LRU cache. - Replace the close()-as-flush pattern with a near real-time refresh (commit + openIfChanged), and make ensure+search atomic under one lock. - Move the discarded global search out of the multi-round session path. - Let EmbeddingOperator enumerate what the index store holds instead of scanning the graph, resolving each entity so deleted leftovers are filtered. Recall is unchanged: the document set, query string, analyzer and topN all stay the same, and equivalence against the rebuild path is asserted by tests. The rebuild path is kept as SessionOperator.searchWithGlobalGraphByRebuild for that purpose. Measured on 10000 vertices: 98.4~119.5 ms per query -> 0.41~0.46 ms steady state. On 5000 vertices with writes interleaved with queries: 44.5~47.1 ms per round -> 1.60~1.81 ms, and full builds drop from 41 to 1. MutableGraphTest goes from 0.887 s to 0.186 s. See geaflow-ai/docs/feature-resident-keyword-index.md for the design, the change list, full measurements and known limitations. The consolidate write path still rebuilds all retrieval state per insert and is tracked there.
Review of the previous commit found the version guard does not actually guard. It read the current vertex version after applying a write batch and adopted it as "everything is applied", so any mutation made outside the reporting path was silently accepted as already indexed. One later reported write was enough to swallow it, and the missing document never came back. Reproduced: build the index, add a vertex directly through MemoryMutableGraph, then upsert an unrelated vertex through the server; the first vertex becomes permanently unsearchable. Make the writer state the range instead of the reader guessing it. Add VertexVersionWindow, opened before a batch of writes and sealed after them. A batch is applied in place only when the window proves it is complete: sealed, starting exactly at the version the index last accepted, and ending at the version the graph still reports. Anything else rebuilds. The window opened by the writer between its own writes remains a blind spot, which now needs writers to serialize rather than being papered over. Also in this area: - Register an edge schema against the edge version only. It cannot change how an existing vertex is verbalized, and bumping the vertex version made consolidate invalidate the index on its first insert for nothing. - Move schema registration into MemoryGraph so bumpVersion can be private; advancing the version was a public operation any caller could trigger. Concurrency and cost: - Search under a read lock instead of one monitor covering build, write and search, so concurrent queries no longer serialize. - Memoize verbalization in a ConcurrentHashMap with no lock on either path. The memoized function is pure and its value immutable, so exclusion is not needed for correctness; racing threads at worst duplicate work. This drops a double check that only existed to protect a cache wide version field and a remove that the following put already did. Cost: the bound is now approximate and eviction is not LRU. No throughput difference was measurable at 1 to 16 threads, the critical section is a single map get. - Stamp each cache entry with its source version, vertex entries against the vertex version. Sharing one version discarded every memoized entry on any write, worst exactly where writes are frequent: consolidate issues about thirty edge writes per inserted entity. - Drop IndexWriter#commit from refresh and open the reader from the writer. A commit point buys no durability on an in memory directory. Measured 2.14~2.42 -> 1.72~1.86 ms per round on writes interleaved with queries. - Drop the graph sized Set the index kept alongside Lucene. Delete by term is idempotent so it guarded nothing, and the document count can be read from Lucene. Also removes GraphSearchStore.entityNum, which had no reader and miscounted both upserts and deletes. - Snapshot getIndexedEntities instead of returning a live key set view, which threw ConcurrentModificationException under concurrent writes. - Build a fresh IndexWriterConfig per writer; Lucene rejects reuse, so a store that advertises a long life could not be reopened after close. - Cache the schema label sets per search, guard residentIndexes with synchronizedMap, read the cache bound from Constants at use time so it is configurable, and translate the remaining Chinese Javadoc. Corrects the design doc: it blamed the 1.7 ms write round on segment growth slowing search and proposed a merge policy. 800 rounds bucketed say otherwise, search cost falls from 0.75 to 0.21 ms while write plus refresh dominates, so Lucene's default TieredMergePolicy already handles merging. The open item is refresh batching, not a hand written merge policy. Steady state read is 0.38~0.60 ms per query on 10000 vertices, writes interleaved with queries 1.36~1.40 ms per round with one full build. 20 tests pass, three of them new: the swallowed mutation regression, edge writes keeping memoized vertex verbalizations, and consistency of the lock free cache under 8 threads.
| throw new RuntimeException("Server or graph accessor not available for graph: " + graphName); | ||
| } | ||
| // Maintain the resident keyword index in place instead of rebuilding it on next query. | ||
| insertServer.onEntitiesUpserted(graphEntities, window.seal()); |
There was a problem hiding this comment.
graphEntities contains every requested entity, not only the entities that were successfully written. The return values from addVertex / addEdge above are ignored.
For example, VertexGroup.addVertex rejects a duplicate ID. If a duplicate insert carries different text, the graph keeps the old vertex, but this call still upserts the rejected request object into Lucene. The resident index then matches the new text while returning the old graph vertex, and the old text is no longer searchable. Because failed mutations also advance vertexVersion, the version window still passes and this divergence persists.
Please either fail the request on mutation errors or collect only successful writes and resolve the current entity from the graph before updating the index. A regression test with a duplicate ID and different content would catch this.
There was a problem hiding this comment.
Confirmed and fixed. Thanks, this one was reachable and I had missed the case where the graph keeps a different entity under the same key.
Fixed in two layers, because the caller and the index were both trusting the wrong thing:
GeaFlowMemoryServer.addEntitynow checks the result of every mutation and reports only the entities the graph accepted.deleteEntityhad the mirror of this defect, so it got the same treatment: a rejected removal was still reported, which dropped the document of a vertex that is in fact still in the graph.ResidentSearchIndexno longer indexes the reported object at all.reconcile()resolves each reported entity through the accessor and indexes what the graph holds, dropping the document when the graph holds nothing.onEntitiesUpsertedandonEntitiesRemovednow share that one path, so whether an entity ends up indexed or removed is decided by the graph rather than by which hook the caller invoked.
Regression test as suggested, duplicate id with different content: ResidentSearchIndexTest#testEntityRejectedByTheGraphIsNotIndexed. It asserts the rejected text never becomes searchable and the kept text still is, and it also pins the detail that made this subtle, that the failed mutation still moves the version so the window check passes. Reverting reconcile() to use the reported object fails it:
AssertionFailedError: content the graph rejected must not become searchable ==> expected: <true> but was: <false>
Also checked end to end against a running server: a duplicate insert now answers Added entities, num: 0, rejected: 1, and a delete of an absent id answers Removed entities, num: 0, rejected: 1, with the previously stored content still retrievable in both cases.
| private Map<String, String> vertexDoc(GraphVertex graphVertex, List<IVector> indexVectors) { | ||
| Map<String, String> kv = new HashMap<>(); | ||
| Vertex vertex = graphVertex.getVertex(); | ||
| kv.put(SearchConstants.KEY, ModelUtils.getGraphEntityKey(graphVertex)); |
There was a problem hiding this comment.
ModelUtils.getGraphEntityKey is not collision-free, but this PR now uses it as the Lucene primary key for update and delete.
For vertices it concatenates PREFIX_V + id + label without lengths or escaping. For example, (id="a", label="bc") and (id="ab", label="c") both produce Vabc. Both documents can exist after the initial addDocument build, but updating or deleting either one by this term will delete both and then optionally insert only one.
Please introduce an injective key encoding, such as length-prefixed components or another unambiguous structured encoding, and add a multi-label collision test covering both update and delete.
There was a problem hiding this comment.
Confirmed and fixed. Key components are now length prefixed, so the encoding can be parsed back unambiguously and is therefore injective.
Tests in EntityKeyCollisionTest: vertex keys, edge keys, injectivity across every way of splitting one string into id and label, and update and delete each leaving a colliding partner's document intact. Reverting to the old concatenation fails six of them, including both of the cases you asked for:
testUpdatingOneEntityKeepsTheCollidingOneSearchable ==> expected: <[ab]> but was: <[]>
testDeletingOneEntityKeepsTheCollidingOneSearchable ==> expected: <[ab]> but was: <[]>
One consequence worth surfacing, since it is not visible from this file. The same key is the row key of EmbeddingIndexStore's jsonl index file, so changing the encoding made every existing entry fail to match. That is silent: unmatched entries are skipped, the store concludes the entities are not indexed yet and re-requests embeddings for the whole graph. It showed up as GraphMemoryTest failing against the committed LDBC index file.
Rather than require everyone to rebuild their index, loading now accepts the legacy key as a fallback, and where two entities share one legacy key it drops the entry instead of guessing which entity it belongs to. ModelUtils.getLegacyGraphEntityKey exists only for reading such files; everything written from now on uses the new key. The ambiguous branch is covered by EntityKeyCollisionTest#testAmbiguousLegacyKeysAreDroppedNotGuessed.
I verified the compatibility path against a real embedding service rather than a mock: an index file carrying the old keys matched all of its entries by the legacy key, requested no new embeddings, and produced results identical to the same index written with the new keys.
| built.refresh(); | ||
| store = built; | ||
| globalIndexBuilt = true; | ||
| builtVersion = version; |
There was a problem hiding this comment.
The graph version is sampled only before the full scan. A graph write can happen after that read or during scanVertex(), because the resident-index lock is not shared with MemoryMutableGraph.
In that case the newly built index may contain a mixed snapshot or miss the concurrent write, but it is still published with the old version. searchWithIndex immediately searches this store in the same call, before another version check can force a rebuild, so the current query can serve stale results.
Please read the vertex version again after the scan and publish the built store only when the start and end versions match; otherwise close it and retry or invalidate it. A latch-based test that mutates the graph in the middle of a blocked scan would cover this race.
There was a problem hiding this comment.
Confirmed and fixed, along the lines you suggested. The vertex version is now read again after the scan, and the store is published only when the two readings agree; otherwise the attempt is closed and retried, bounded at three tries. Added getDiscardedBuildCount so a discarded attempt is observable rather than invisible.
For the case where the graph is written continuously and no scan can be proven consistent, the last attempt is published with VERSION_UNSUPPORTED instead of a version it cannot vouch for. That makes every later query rebuild rather than reuse it, which is the same safe degradation the class already uses for accessors that cannot report a version, and it is logged at warn level.
Latch based test as suggested: ResidentSearchIndexTest#testBuildInterruptedByAConcurrentWriteIsNotPublishedAsConsistent, which blocks the scan at its first entity, writes a vertex, then releases it.
One note in case it is useful to other reviewers: the assertion has to be on the result of the query that triggered the build. My first attempt asserted on a subsequent query, and that version passed even with the fix reverted, because a later query finds the version inconsistent and rebuilds on its own. Asserting on the triggering query does fail without the fix:
AssertionFailedError: the query that triggered the build must not answer from a mid scan snapshot
==> expected: <[id-mid]> but was: <[]>
| // A stale entry needs no explicit removal, the put below replaces it. | ||
| List<IVector> computed = computeEntityIndex(entity); | ||
| cacheMiss.increment(); | ||
| verbalizationCache.put(entity, new CachedIndex(computed, version)); |
There was a problem hiding this comment.
The cached value is computed from the caller-supplied GraphEntity, but the stamp is the graph's current global version. Those two do not necessarily describe the same snapshot.
GraphVertex.equals ignores attribute values and compares only ID and label. A concrete stale-cache sequence is:
- retain a wrapper for vertex A at version 0;
- update A in the graph, advancing to version 1;
- call
getEntityIndexwith the retained old wrapper; - the old text is computed and cached under version 1;
- a lookup using the current wrapper compares equal and returns the stale text as a cache hit.
Please tie the cache entry to the actual entity content, or resolve the current entity from the accessor before computing and publishing it. The test should assert the returned text after updating a vertex while retaining its old wrapper.
There was a problem hiding this comment.
Confirmed and fixed. Each cache entry now carries the content it was computed from, and a hit requires both the source version and that content to match, so a superseded wrapper can no longer have its result served to a lookup made with the current one.
I did not give the store an accessor, and instead made the contract explicit in its Javadoc: getEntityIndex returns the index of the entity handed to it, and a caller that needs whatever the graph currently holds under that label and id has to resolve it first. ResidentSearchIndex now does exactly that, so the index is authoritative on graph content while this class stays a memoized pure function. That also closes the related case raised on GeaFlowMemoryServer.
Test as suggested, asserting the returned text after updating a vertex while holding on to the old wrapper: ResidentSearchIndexTest#testStaleEntityWrapperDoesNotPoisonTheCache. It first seeds the cache through the stale wrapper, then asserts the current wrapper is not served the superseded content. Reverting the content comparison fails it:
AssertionFailedError: the current entity must not be served the superseded content as a cache hit
==> expected: <true> but was: <false>
Unrelated to your point but in the same method: the cache no longer takes a lock on either path. It memoizes a pure function whose value is immutable, so racing threads at worst duplicate work, and the double checked lookup that used to be there only existed to protect a cache wide version field that the per entry stamps replaced. The trade is that the size bound became approximate and eviction is no longer LRU, which is recorded in the design doc. Concurrency of the lock free version is pinned by ResidentSearchIndexTest#testConcurrentVerbalizationLookupsAreConsistent.
Four defects found in review, each with a regression test that fails without the fix. Index follows the graph, not the reported object. A reported entity is a request the graph may have rejected: VertexGroup.addVertex refuses a duplicate id, and a failed mutation still advances the version, so the version window passed and the rejected object was indexed. The new text then matched while retrieval resolved and returned the vertex the graph had kept, and the text actually in the graph stopped being searchable. ResidentSearchIndex now resolves every reported entity through the accessor and indexes what the graph holds, dropping the document when the graph holds nothing; upsert and remove share that one path, so the graph decides the outcome rather than which hook was called. GeaFlowMemoryServer checks mutation results and reports only what was accepted, on the delete path too, where reporting a rejected removal would have dropped a document that is still supposed to be searchable. Collision free entity key. getGraphEntityKey is the Lucene primary key, but it plainly concatenated id and label, so (id="a", label="bc") and (id="ab", label="c") both yielded Vabc and a write to one destroyed the other's document. Components are now length prefixed, which makes the encoding parseable and therefore injective. Index files written with the old encoding are still read: the legacy key is accepted as a fallback, and where two entities share one it is dropped rather than guessed. Verified against a real embedding service, an existing index file matches all its entries by the legacy key, requests no new embeddings, and returns identical results. Build snapshot is verified before publication. The vertex version was sampled only before the full scan, so a concurrent write produced a snapshot that could not be attributed to that version, and the query that triggered the build answered from a graph state that never existed. The version is now read again after the scan and the store is published only when the two agree, otherwise the attempt is discarded and retried. If no stable snapshot can be obtained the index is published as unverified so later queries rebuild instead of reusing it. Adds getDiscardedBuildCount. Cache entries are bound to content. Entity equality is label and id only, so a caller holding a superseded wrapper shared a cache key with the current one: the stale text was computed, stored under the current version, and served to the next lookup as a hit. Entries now carry the content they were computed from and a hit requires it to match. The contract is documented: the store returns the index of the entity handed to it, and callers needing the graph's current entity must resolve it first, as the resident index now does. Also records in the design doc that the equivalence claim is against the per-query rebuild path and does not mean keyword matching itself is correct: SessionOperator builds the query string with KeywordVector.toString(), so a query issued through the service carries tokens every document also has and matches everything. That line predates this work and is untouched here; both compared paths take the same query string and the tests pass unwrapped ones.
2 participants
What changes were proposed in this pull request?
The global keyword search in
geaflow-airebuilt a throw-away Lucene index from a full graph scan on every query, then discarded it. Index construction cost sat entirely on the query path, verbalization was redone for the whole graph per query, and multi-round sessions ran a global search whose result was thrown away.This makes the index resident per graph and maintains it incrementally, following the standard inverted index model rather than invalidate-and-rebuild.
Retrieval
ResidentSearchIndex, built at most once and kept alive across queries, held byGraphMemoryServerper keyword index store.ModelUtils.getGraphEntityKey), so writes map to Lucene update/delete by term. Cost is proportional to the change, not to graph size, and both are idempotent, so callers need not supply an exact delta.onEntitiesUpserted/onEntitiesRemoved/onSchemaChanged. Only a schema change needs wholesale invalidation.close()-as-flush pattern with a near-real-time reader opened from the writer. Nocommit(): the directory is in memory, so a commit point buys no durability.ConcurrentHashMap, each entry stamped with the source version it was computed from, so a write invalidates only the entries it affects.EmbeddingOperatorenumerate what the index store holds instead of scanning the graph, resolving each entity so deleted leftovers are filtered out.Collision free entity key
The entity key doubles as the Lucene primary key for update and delete, so it must identify exactly one entity. It previously concatenated id and label plainly, which is not injective:
(id="a", label="bc")and(id="ab", label="c")both producedVabc, and writing either one would delete the other's document. Components are now length prefixed, which makes the encoding parseable and therefore collision free.The same key is the row key of
EmbeddingIndexStore's jsonl index file, so this would otherwise have made every existing entry fail to match, silently discarding all stored embeddings and re-requesting them. Loading therefore accepts the legacy key as a fallback, and drops it where two entities share one rather than guessing.ModelUtils.getLegacyGraphEntityKeyexists only for reading such files; new writes use the new key. Verified against a real embedding service: an index file carrying old keys matched all entries by the legacy key, requested no new embeddings, and returned results identical to the same index written with new keys.Not serving stale results
A resident structure over a mutable graph needs to know when it went stale, including when the graph is mutated outside the server (for example directly through
MemoryMutableGraph).MemoryGraphtherefore maintains a version, split into a general one and a vertex-only one so that edge writes do not invalidate a vertex-only index.Reading the current version after applying a batch is not sufficient: it accepts any unreported change as already applied, and the affected documents then go missing permanently. So the writer states the range and the index verifies it —
VertexVersionWindowis opened before a batch of writes and sealed after them, and a batch is applied in place only when it is provably complete (sealed, starting exactly at the version the index last accepted, ending at the version the graph still reports). Anything else falls back to a rebuild. Accessors that cannot report a version degrade to the previous per-query rebuild behaviour.Recall is unchanged relative to the rebuild path. The document set, query string, analyzer and
topNall stay the same, and the equivalence is asserted by tests; the rebuild path is kept asSessionOperator.searchWithGlobalGraphByRebuildfor that purpose. To be precise about what that does and does not claim: it does not mean keyword matching itself is correct.SessionOperatorbuilds the Lucene query string withKeywordVector.toString(), so a query issued through the service carries tokens that every document also carries and therefore matches everything, with the search term only affecting the ranking. That line predates this change and is untouched here; both compared paths receive the same query string. Recorded as a known limitation in the design doc (§6.6).Measured (macOS arm64, 10000 vertices,
topN = 30, ranges over 3 runs):MutableGraphTestLatency also stops tracking graph size: 5000 -> 20000 vertices moves the rebuild path 45 -> 177 ms per query but the resident path only 0.31 -> 0.55 ms.
geaflow-ai/docs/feature-resident-keyword-index.mdcarries the full design, the reasoning against comparable systems (Lucene/ES, Milvus, hugegraph-ai), the complete change list, measurements and limitations.Deliberately out of scope, all recorded in the doc:
KeywordRelationFunction.eval()builds a fresh index store and server per insert, making import O(V^2). It is the only path still rebuilding per O(V). Fixing it changes theConsolidateFunction.evalcontract and deserves its own change.SearcherManager), not just dropping the lock.How was this PR tested?
mvn -B -pl geaflow-ai -am clean install: all Reactor modules SUCCESS, 30 tests pass, Checkstyle 0 violations, Apache RAT Unapproved 0.New tests,
ResidentSearchIndexTest(16),EntityKeyCollisionTest(7) andEmbeddingCandidateSetTest(1). Every defect found in review has a test that fails when its fix is reverted, verified one at a time:Recall equivalence against the per-query rebuild path, 10000 vertices and 10 queries, and separately with the verbalization cache on versus off. Compared as sets, since the rebuild path feeds Lucene from
HashMapiteration and its tie-break order was never deterministic; match counts are kept undertopNto avoid truncation sensitivity.Insert / update / delete take effect in place: new content searchable, superseded documents no longer searchable, document count moves correctly, and
buildCountdoes not change.Upsert is idempotent: replaying the same write three times leaves one document.
A mutation that bypasses the hook forces a rebuild, and is still not lost when a later reported write arrives — the case that motivated
VertexVersionWindow.Edge writes invalidate neither the vertex index nor the memoized vertex verbalizations, while rewriting the vertex itself evicts exactly one entry.
40 rounds of write-then-query: recall matches the invalidate-and-rebuild reference every round, with one build for the whole run.
The lock-free verbalization cache stays consistent under 8 threads x 50 rounds x 200 entities: identical content everywhere, hit + miss exactly equal to the call count, one entry per entity.
An entity the graph rejected is not indexed: a duplicate id carrying different content leaves the graph holding the previous entity, so the index follows the graph rather than the reported request object. Same on the delete path, where reporting a rejected removal would drop a document that is still supposed to be searchable.
Key collisions do not let one entity's update or delete destroy another's document, including across every way of splitting one string into id and label, and an ambiguous legacy key is dropped rather than guessed.
A build whose scan was interrupted by a concurrent write is not published under the version read before it: the query that triggered the build must not answer from a mid-scan snapshot.
A superseded entity wrapper does not poison the memoized verbalization, since entity equality is label and id only.
The lock-free verbalization cache stays consistent under 8 threads x 50 rounds x 200 entities: identical content everywhere, hit + miss exactly equal to the call count, one entry per entity.
Existing
GraphMemoryTest(LDBC data, strict content assertions),MemoryServerTest(HTTP end to end, 532-chunk import) andMutableGraphTestall pass unchanged, which is additional evidence recall did not move.Beyond the unit tests, the HTTP paths were exercised against a running server, and the embedding store and operator against a real embedding service, which is how the index-file compatibility above was checked. Not verified in a production environment.