diff --git a/geaflow-ai/docs/feature-resident-keyword-index.md b/geaflow-ai/docs/feature-resident-keyword-index.md new file mode 100644 index 000000000..d27beef73 --- /dev/null +++ b/geaflow-ai/docs/feature-resident-keyword-index.md @@ -0,0 +1,462 @@ +# Feature：关键词检索索引常驻化与就地增量维护 + +> 模块：`geaflow-ai`　状态：已实现、已验证 + +--- + +## 1. 摘要 + +| 项 | 内容 | +|---|---| +| 能力 | 关键词检索使用按图常驻的 Lucene 索引；查询路径不做索引构建，写入以增量方式就地维护 | +| 手段 | 索引常驻化 + Lucene 主键 update / delete by term + 文本化结果条目级记忆化 + 版本窗口校验 | +| 读效果 | 10000 顶点，每查询 **0.38 ~ 0.60 ms**（每查询重建方案 101.6 ~ 105.7 ms） | +| 写效果 | 5000 顶点、写查交替，每轮 **1.36 ~ 1.40 ms**（失效重建方案 46.3 ~ 46.6 ms）；全量构建次数 41 → **1** | +| 语义 | 召回结果与每查询重建方案一致，由等价性测试保证 | +| 规模 | 新增 `ResidentSearchIndex`、`VertexVersionWindow`；改动 `geaflow-ai` 内 14 个主干文件 | +| 验证 | `mvn -pl geaflow-ai -am clean install` 全绿，19 个测试通过，Checkstyle 0 违规，RAT Unapproved 0 | + +--- + +## 2. 要解决的问题 + +关键词检索入口是 `SessionOperator.apply()`，会话首轮（子图为空）走全局检索分支。原实现每次查询都从零构建索引： + +``` +scanVertex() 遍历全图顶点 + └─ indexStore.getEntityIndex(v) 每个顶点做一次 verbalize +new GraphSearchStore() 新建 Lucene 内存索引（方法内局部变量） + └─ indexVertex(...) × V 逐个 addDoc +close() 关闭 writer 让文档对 reader 可见 +search(query) 真正的检索 + 方法返回，索引被 GC +``` + +单次冷查询代价 `O(V × (verbalize + 分词 + 倒排写入))`，真正的检索只占其中极小一部分。四个叠加的代价点： + +1. **索引每查询重建后丢弃**。`SearchStore` 用 `ByteBuffersDirectory`（纯内存），且是方法内局部变量。 +2. **`close()` 被当作 flush 手段**，索引生命周期被固定为一次性。 +3. **每查询把整张图重新文本化一遍**。`EntityAttributeIndexStore.getEntityIndex()` 在全图扫描中被调用 V 次，每次构造 `SubGraph`、渲染 prompt、分配大量临时字符串。 +4. **多轮会话白跑一次全局检索**。`apply()` 无条件先调用全局检索，但子图非空的热路径不引用其返回值，结果被直接丢弃。 + +同时 `EmbeddingOperator` 有同类问题：为组装候选集遍历全图顶点，而没有 embedding 的顶点在打分阶段本就会被跳过，扫描是多余的。 + +--- + +## 3. 设计 + +### 3.1 目标与约束 + +| 目标 | 约束 | +|---|---| +| 查询路径上不做索引构建 | 召回结果与每查询重建方案一致 | +| 写入代价与变更量成正比，而非与图规模成正比 | 不修改 `SearchOperator` / `IndexStore` / `GraphAccessor` 既有方法签名 | +| 图变更后不返回脏数据 | 保持 Lucene 8.11.2 / JDK 8 兼容，不触碰 Lucene 9 + JDK 11 这个待决议约束 | + +### 3.2 设计依据：成熟实现怎么做 + +倒排索引与向量索引处理「派生索引 + 可变主数据」都收敛到同一套结构：**索引分段且不可变 → 变更用增量表达（更新 = 标记删除 + 新增，删除只打标记）→ 后台异步合并回收**。没有产品用「整体失效 + 懒重建」。 + +- **Lucene / Elasticsearch**：文档被删或更新时只在段级位图标一位，检索跳过被标记的文档；NRT reader 重开耗时与上次重开以来的变更量成正比，而非与索引总量成正比。ES 在其上用 `refresh_interval`（默认 1 秒）攒批，并靠后台合并控制段数量。 +- **关系库二级索引**：DML 在同一事务内维护索引项，没有「失效」概念；Postgres GIN 的 `fastupdate` 本质是攒批。 +- **向量库**：Milvus 新写入进 growing segment，转 sealed 后异步建索引，索引就绪前退回原始向量暴力检索；HNSW 删不掉节点（会破坏图连通性），统一打墓碑 + 周期性 compaction。 + +**对照 hugegraph-ai**：它的查询路径同样零构建成本 —— 索引在算子构造期由 `FaissVectorIndex.from_name()` 从磁盘一次性加载，算子随 pipeline 被 `Scheduler` 的 `GPipelineManager` 池化复用，写入是独立的离线 flow（`BuildVectorIndexFlow` 等）。值得注意的是它用的是 `faiss.IndexFlatL2`，同样是暴力精确检索而非 ANN，所以优势不在算法层，而在四个工程决策：索引常驻、索引对象小（只索引顶点 ID 文本而非完整 verbalization）、先 Gremlin 精确匹配未命中才走向量、图扩展下推给图数据库。 + +反过来它也有短板：`node_init` 只在 pipeline 首次创建时执行，**没有任何失效机制** —— 离线重建索引后，池中已存在的 pipeline 仍持有旧快照。它能忍受这一点是因为定位是批式知识库（先建索引再问答）；而 `geaflow-ai` 的 GraphMemory 同进程同图可以边写边查，必须有失效机制，这部分没有现成实现可抄。 + +### 3.3 文档集严格等价 + +常驻索引收录的文档集 = **所有 `getEntityIndex()` 返回非空的顶点**，与每查询重建方案完全相同： + +- 边被排除，与全局检索只 `scanVertex()` 的语义一致 +- 全量构建期间用一个**局部** `Set` 去重，等价于原实现以 `Map` 键去重；该集合随构建结束即丢弃，不作为常驻状态保留（见 §3.11） +- Lucene 查询串、`topN`（`Constants.GRAPH_SEARCH_STORE_DEFAULT_TOPN = 30`）、`StandardAnalyzer` 均不变 + +文档集、查询、打分器三者相同，因此召回相同。 + +### 3.4 写入以 Lucene 主键增量维护 + +Lucene 原生支持增量维护，前提是每个文档有一个可精确定位的 term。因此每个文档带一个不分词主键字段 `SearchConstants.KEY`，取值为 `ModelUtils.getGraphEntityKey(entity)`（顶点 `V{id}{label}`，边 `E{src}{label}{dst}`）—— 这个 key 生成逻辑本已存在，`EmbeddingIndexStore` 就用它做 jsonl 行键。 + +| 操作 | 实现 | 代价 | +|---|---|---| +| 全量构建 | `addDocument`（扫描保证每顶点只出现一次，无需去重开销） | O(V)，仅首次 | +| 新增 / 更新 | `updateDocument(new Term(KEY, key), doc)` | O(变更量) | +| 删除 | `deleteDocuments(new Term(KEY, key))`，只在段级位图打标记 | O(1) | +| 可见性 | 批次结束后一次 `refresh()`（NRT reader 重开，不 commit，见 §3.9） | O(变更量) | + +`updateDocument` 对「已存在」与「不存在」处理一致，`deleteDocuments` 对不存在的 key 也是 no-op，因此 **`onEntitiesUpserted` / `onEntitiesRemoved` 都是幂等的，调用方不需要提供精确增量**，重复上报同一实体不会产生重复文档。这消除了「维护正确性依赖调用方给出准确 delta」的隐式契约（但仍要求调用方给出**完整**的变更集合，见 §3.5）。 + +服务层因此拆成三个语义明确的入口，而不是一个布尔开关： + +``` +GraphMemoryServer + ├─ onEntitiesUpserted(entities, window) 新增与更新，就地 upsert + ├─ onEntitiesRemoved(entities, window) 删除，就地 delete + └─ onSchemaChanged() 无法按实体表达的变更，整体失效 +``` + +只有 schema 变更走整体失效 —— 它会改变每一个实体的 verbalize 结果。 + +### 3.5 版本窗口：由写方声明变更范围，而非读方事后猜 + +图可以被绕过服务层直接改写（例如经 `MemoryMutableGraph`），这类变更没有任何钩子会触发。因此由图自身维护版本号： + +``` +MemoryGraph + ├─ version 任何变更 +1（点、边、schema） + └─ vertexVersion 仅点与顶点 schema 变更 +1 + +GraphAccessor + ├─ getGraphVersion() 默认 VERSION_UNSUPPORTED (-1) + └─ getVertexVersion() 默认委托 getGraphVersion() + +VerbalizationFunction + ├─ getSourceVersion() 默认 VERSION_UNSUPPORTED + └─ getSourceVertexVersion() 默认委托 getSourceVersion() + SubgraphSemanticPromptFunction 两者分别透传 accessor 的 graph / vertex 版本 +``` + +**关键点：不能事后读当前版本来当作「我已全部应用」。** 就地维护完把 `builtVersion` 推进到 `getVertexVersion()` 的当前值，会把这期间任何未上报的改图一并当作已应用，于是那些文档**永久**缺失 —— 兜底只在「越界改图之后再也没有钩子写入」时才成立，一次正常写入就把它吞掉了。 + +所以改成由写方声明范围，读方校验： + +```java +// 写方 +VertexVersionWindow window = server.openVertexVersionWindow(); // 记录写前 vertexVersion +... 改图 ... +server.onEntitiesUpserted(entities, window.seal()); // 记录写后 vertexVersion + +// 读方（ResidentSearchIndex.applyWrite） +if (window == null || !window.covers(builtVersion)) { + invalidateLocked(); // 无法证明这批变更是完整的 → 重建，不冒脏数据的风险 + return; +} +... 就地 upsert / delete ... +builtVersion = window.getTo(); +``` + +`covers()` 要求三件事同时成立：窗口已封（`seal()` 调过）、`from == builtVersion`（这批变更正好接在索引已接受的版本之后）、`to == 当前 vertexVersion`（封窗之后没有别人再动过图）。任一不成立就整体重建。 + +覆盖情况： + +| 场景 | 结果 | +|---|---| +| 只走钩子写入 | `covers()` 成立，全程就地维护，`buildCount` 不增长 | +| 越界改图，之后无钩子写入 | 冷查询时 `vertexVersion != builtVersion` → 重建 | +| 越界改图，之后有钩子写入 | `from != builtVersion` → 重建（**修复前会被吞掉**，有回归测试 `testUnreportedMutationIsNotSwallowedByALaterReportedWrite`） | +| 封窗后到钩子执行之间有人改图 | `to != 当前版本` → 重建 | +| 写方自己的改图与封窗之间有人改图 | **未覆盖**，需要写方之间自行串行化，见 §6.4 | + +**版本不可用时自动降级**：返回 `VERSION_UNSUPPORTED` 的 accessor（如 `EmptyGraphAccessor`，及任何未实现版本上报的实现）使常驻索引每次重建、缓存完全不启用 —— 退化为每查询重建行为。 + +`MemoryGraph` 中失败的变更同样 bump 版本：过度失效只是慢，漏失效是正确性缺陷。 + +**`vertexVersion` 与 `version` 分离**：顶点的 verbalize 只读该顶点自身与 schema（`schema.getPrompt(vertex)`），所以常驻索引与顶点文本化缓存都 watch `vertexVersion`，边写入不会使它们失效 —— 这在 consolidate 场景下有实际意义，一次插入会带来约 30 次 `addEdge`。边的 verbalize 会读取两端顶点（`schema.getPrompt(edge, start, end)`），依赖面更广，因此边的缓存条目 watch 全局 `version`。 + +**边 schema 只推进 `version`**：新增一个边 schema 不可能改变任何已存在顶点的 verbalize 结果，所以 `MemoryGraph.registerEdgeSchema()` 走 `bumpEdgeVersion()`。否则 consolidate 首次插入时注册 `consolidate_keyword_edge` 会白白使顶点索引失效一次。顶点 schema 与整体 `setGraphSchema()` 仍推进 `vertexVersion`。 + +schema 变更的写入口收回图内部（`registerVertexSchema` / `registerEdgeSchema`），`bumpVersion` / `bumpEdgeVersion` 因此是 `private` —— 版本推进不再是任何调用方都能触发的公开动作，否则「派生结构是否新鲜」可以被外部随手改写。 + +### 3.5.1 文本化缓存：条目级版本戳，且不加锁 + +`EntityAttributeIndexStore` 的记忆化缓存**每个条目带自己的源版本戳**，命中时逐条比对，不一致就只换这一条。 + +用「整个缓存共享一个版本、不一致即 `clear()`」会让任何一次写入丢掉此前记忆的全部条目（默认上限 20 万条）—— 恰好是写入最频繁的 consolidate 路径受损最重。回归测试 `testEdgeWriteKeepsMemoizedVertexVerbalizations` 固定这一行为：19 次边写入后顶点条目仍然命中，而改写该顶点本身只淘汰它一条。 + +**容器是 `ConcurrentHashMap`，读写路径都不持锁。** 被记忆的是纯函数、缓存值不可变，所以互斥对正确性不是必需的： + +```java +CachedIndex cached = cache.get(entity); +if (cached != null && cached.version == version) { hit; return cached.vectors; } +List computed = computeEntityIndex(entity); // 在 map 之外算 +cache.put(entity, new CachedIndex(computed, version)); // 覆盖旧版本条目，无需先 remove +``` + +两个线程同时未命中同一实体时会各算一遍，但同一版本下算出的东西相同、谁写进去都成立，输的一方只是白做一次；这比让每次查找都进临界区更合适。同理不需要「先查一次、算完再查一次」的 double-check —— 那个二次检查原本是为了配合「整个缓存共享一个版本」的字段，条目级版本戳之后它已经没有对应的竞态要防。 + +需要说明的取舍： + +- 上限变成**近似**、淘汰**不再是 LRU**（按 map 迭代顺序丢一批）。对纯函数的记忆化，淘汰错一条的代价只是重算一次 +- `LinkedHashMap(accessOrder=true)` 的 LRU 语义要求连读操作都独占（`get` 会改动链表），这正是原来那把锁的来源 +- 计数器改 `LongAdder` + +我**没有**测到可靠的吞吐差异：临界区只有一次 map `get`，在 1/4/8/16 线程下监视器都不是瓶颈，两种实现的差距被运行间抖动盖过（8 线程 3.2M 次查找的 wall time 在两侧都落在 230~340 ms）。所以这条改动的理由是「不必要的互斥就不该有，且读路径上不再留共享阻塞点」，不是性能数字。新增 `testConcurrentVerbalizationLookupsAreConsistent` 固定并发下的一致性：8 线程 × 50 轮 × 200 实体，内容全部一致、计数不重不漏、条目数恰好等于实体数。 + +因此服务层钩子里**不需要**再逐实体 `invalidateCache(entity)`：写入本身已经让受影响的条目过期，那个循环是纯开销。`invalidateCache()` 只保留给版本无法描述的变更（替换 verbalization function、schema 变更）。 + +### 3.6 更新与失效时机 + +全部懒执行，没有定时任务，也没有后台重建线程。 + +| 时机 | 行为 | +|---|---| +| 首次冷查询 | `ensureGlobalIndex()` 全量构建一次 | +| 后续冷查询，`vertexVersion` 未变 | 直接复用，零构建开销 | +| 后续冷查询，`vertexVersion` 已变 | 整体重建 | +| `/graph/insertEntity`（新增或更新），索引已建，窗口可信 | **就地 upsert** + 批次末一次 `refresh()`，`builtVersion` 推进到 `window.getTo()`，不重建 | +| `/graph/delEntity`，索引已建，窗口可信 | **就地 delete**，同上，不重建 | +| 窗口不可信（缺失 / 未封 / 与 `builtVersion` 不接续 / 封窗后又被改） | `invalidate()` → 下次冷查询重建 | +| 上述写入，索引尚未构建 | no-op，首次查询构建时一并收录 | +| `/graph/addEntitySchema` | `onSchemaChanged()` → 缓存清空 + 索引失效 → 下次冷查询重建 | +| 绕过服务层直接改图 | 无钩子，但 `vertexVersion` 已变 → 下次冷查询比对失败 → 重建 | +| 仅写边（含注册边 schema） | 只 bump `version`、不 bump `vertexVersion` → **不触发重建** | +| 写入实体的索引内容变为空 | 就地 delete 掉原文档（非文档实体不应留在索引里） | +| accessor 返回 `VERSION_UNSUPPORTED` | 每次冷查询重建，退化为每查询重建行为 | + +文本化缓存的淘汰发生在下一次 `getEntityIndex()` 命中该条目时，不是写入时立即执行。 + +### 3.7 检索与校验同一次加锁完成 + +`ResidentSearchIndex.searchWithIndex()` 在**同一次加锁内**完成「确保索引有效」与「检索」。拆成两次调用会留下窗口：并发写入可以在两者之间使索引失效，让查询落到不存在的索引上。 + +锁是 `ReentrantReadWriteLock`：快路径（索引已建且版本一致）只持读锁，因此并发查询不互相串行；只有需要构建、失效或写入时才升级到写锁。`SearchStore` 的 reader / searcher 字段为此声明成 `volatile`，并保证「每批写入后必定 `refresh()`」，使读路径上的 `ensureSearcher()` 在稳态下是 no-op、不会去改动 store。 + +### 3.8 仅冷路径使用常驻索引 + +热路径（子图非空）的语义是「**在子图扩展集内取 top-30**」。改为查全局索引再与扩展集求交，得到的是「全图 top-30 ∩ 扩展集」，结果不同。因此热路径保留一次性小索引，这是语义要求。其代价受控：扩展集规模受子图大小 × 度数约束，且同样受益于文本化缓存。 + +### 3.9 NRT 刷新替代 `close()`，且不做 `commit()` + +索引长期存活就不能关闭 writer。`SearchStore.refresh()` 直接从 writer 开 reader： + +```java +public void refresh() throws IOException { + if (writeStats) { + if (readStats && nearRealTimeReader) { + DirectoryReader newReader = DirectoryReader.openIfChanged(reader, writer, true); + ... // 换 reader / searcher，关旧 reader + } else { + reader = DirectoryReader.open(writer); // 首次：NRT reader + ... + } + pendingWrite = false; + return; + } + // 无 writer（空索引）时才从 directory 开，IndexNotFoundException 由 GraphSearchStore 吞掉 + ... +} +``` + +**为什么不 commit**：`ByteBuffersDirectory` 是纯内存目录，commit point 换不到任何持久性，只是每批写入多做一次工作。实测去掉 commit 后写查交替从 2.14 ~ 2.42 ms/轮 降到 1.72 ~ 1.86 ms/轮（同机同轮次对比）。writer 关闭时 Lucene 自身会 commit，所以 `close()` 之后目录仍是可读的。 + +`close()` 只负责真正释放，并把 writer / reader 引用清空。`initWriter()` 每次新建 `IndexWriterConfig` —— Lucene 禁止把已交给某个 writer 的 config 再交给下一个，复用会抛 `IllegalStateException`；这个类既然对外宣称长生命周期，就不能留这种「close 后不能再用」的隐式陷阱。 + +### 3.10 参照实现 + +每查询重建的逻辑保留为 `SessionOperator.searchWithGlobalGraphByRebuild()`（package-private），用于等价性测试的对照组，以及未提供常驻索引时的兜底。`SessionOperator` 的两参数构造函数保持可用。 + +### 3.11 常驻状态只有索引本身 + +`ResidentSearchIndex` 不再额外保留 `Set indexedEntities`： + +- 「实体索引内容变空 → 删掉旧文档」原先靠这个集合判断是否曾经收录过，但 delete-by-term 本身幂等，无条件删一次即可 +- 文档数改为直接问 Lucene（`SearchStore.numDocs()`），不再维护一个需要自行镜像 update / delete 语义的计数器（原 `GraphSearchStore.entityNum` 既无读取方，upsert 已存在文档时还会多计、删除时不会减） + +于是常驻状态就是 Lucene 索引本身。一份与图等大的 `HashSet` 在千万点级图上是数百 MB 的稳态堆，去掉它直接改善 §6.3。 + +--- + +## 4. 改动点 + +### 4.1 新增 + +| 文件 | 职责 | +|---|---| +| `operator/ResidentSearchIndex.java` | 按图常驻的关键词索引：`ensureGlobalIndex()` 懒构建 + 版本校验、`searchWithIndex()` 校验与检索一次加锁（读锁快路径）、`onEntitiesUpserted()` / `onEntitiesRemoved()` 带窗口校验的就地增量维护、`invalidate()` 整体失效；暴露 `buildCount` / `upsertCount` / `removeCount` / `indexedEntityNum` 供测试与观测 | +| `graph/VertexVersionWindow.java` | 一批变更所声明的顶点版本区间：`open()` 记写前版本、`seal()` 记写后版本、`covers(acceptedVersion)` 判定这批变更是否可信为完整 | +| `test/operator/ResidentSearchIndexTest.java` | 读写等价性与性能、插入 / 更新 / 删除就地生效、幂等、未通知改图触发重建（含被后续写入吞掉的回归）、边写入不失效且不清缓存、写查交替、缓存计数、无锁缓存并发一致性，13 例 | +| `test/operator/EmbeddingCandidateSetTest.java` | 向量候选集两条收集路径的等价性，1 例 | + +### 4.2 修改 + +**索引与检索** + +| 文件 | 改动 | +|---|---| +| `operator/SearchStore.java` | 新增 `refresh()`（NRT `DirectoryReader.open(writer)` / `openIfChanged(reader, writer, true)`，不 commit）与 `ensureSearcher()`；新增 `updateDoc()` / `deleteDoc()`（by term）、`numDocs()`；`addDoc(kv, exactField)` 支持把主键写成不分词 `StringField`；`reader` 类型 `IndexReader` → `DirectoryReader`；reader / searcher / 状态标记改 `volatile` 以支持并发检索；`initWriter()` 每次新建 `IndexWriterConfig`；移除未使用的 `getConfig()`；`close()` 只做真正释放并清空引用 | +| `operator/GraphSearchStore.java` | 文档带 `SearchConstants.KEY` 主键；新增 `upsertVertex()` / `upsertEdge()` / `removeEntity()` / `refresh()`（吞掉空索引的 `IndexNotFoundException`）/ `getDocNum()`；抽出 `vertexDoc()` / `edgeDoc()` / `writeDoc()` 去重；缓存 schema 的点 / 边 label 集合，不再每次检索用 stream 重算；删除只写不读且计数错误的 `entityNum` | +| `operator/SearchConstants.java` | 新增 `KEY` 字段名 | +| `operator/SessionOperator.java` | 新增三参构造接收 `ResidentSearchIndex`；冷路径改走 `searchWithIndex()`；全局检索调用移入冷分支；原逻辑重命名为 `searchWithGlobalGraphByRebuild()`；热路径 `close()` → `refresh()`，检索后 `closeQuietly()` | +| `operator/EmbeddingOperator.java` | 全局检索调用移入冷分支；抽出 `collectGlobalCandidates()`，优先用 `indexStore.getIndexedEntities()` 并按图解析每个实体（过滤已删除的残留索引项、取当前顶点对象），不可用时退回全图扫描 | +| `index/IndexStore.java` | 新增可选 `default Collection getIndexedEntities()`，返回 `null` 表示无法枚举 | +| `index/EmbeddingIndexStore.java` | 实现 `getIndexedEntities()`，返回 key 集合的**快照**（活视图会在并发写入时抛 `ConcurrentModificationException`） | +| `index/EntityAttributeIndexStore.java` | 记忆化改 `ConcurrentHashMap`，读写路径均不持锁、无 double-check（见 §3.5.1）；**每条目带源版本戳**、逐条比对逐条覆盖；顶点条目 watch `getSourceVertexVersion()`、边条目 watch `getSourceVersion()`；按近似上限批量淘汰（`enforceBound()`）；计数器改 `LongAdder`；`invalidateCache()` / `invalidateCache(entity)`；上限直接读 `Constants`，不再在构造期拷进 `final` 字段；版本不可用时不缓存 | + +**图版本号** + +| 文件 | 改动 | +|---|---| +| `graph/io/MemoryGraph.java` | 新增 `version` / `vertexVersion`（`AtomicLong`）及 `getVersion()` / `getVertexVersion()`；点操作走 `bumped()`，边操作走 `edgeBumped()`；新增 `registerVertexSchema()`（推进 `vertexVersion`）/ `registerEdgeSchema()`（只推进 `version`）；`bumpVersion()` / `bumpEdgeVersion()` 收为 `private` | +| `graph/GraphAccessor.java` | 新增常量 `VERSION_UNSUPPORTED` 与 `default getGraphVersion()` / `getVertexVersion()` | +| `graph/LocalMemoryGraphAccessor.java` | 覆写两个版本方法，委托 `MemoryGraph` | +| `graph/MemoryMutableGraph.java` | `addVertexSchema()` / `addEdgeSchema()` 校验后改为调用 `MemoryGraph.register*Schema()`，不再直接改 `graph.entities` | +| `verbalization/VerbalizationFunction.java` | 新增 `default getSourceVersion()` 与 `default getSourceVertexVersion()` | +| `verbalization/SubgraphSemanticPromptFunction.java` | 分别覆写两者，透传 accessor 的 graph / vertex 版本 | + +**服务层** + +| 文件 | 改动 | +|---|---| +| `GraphMemoryServer.java` | 新增 `IdentityHashMap`（`synchronizedMap` 包装）；`addIndexStore()` 为关键词索引存注册常驻索引；`search()` 注入常驻索引；新增 `openVertexVersionWindow()` 与带窗口的 `onEntitiesUpserted()` / `onEntitiesRemoved()`、`onSchemaChanged()`；去掉钩子里逐实体清缓存的无效循环 | +| `GeaFlowMemoryServer.java` | `/graph/insertEntity` 改图前开窗、改完即封窗并上报 upsert（consolidate 移到上报之后，使窗口不含它的写入）；`/graph/delEntity` 同理走 delete；只有 `/graph/addEntitySchema` 走整体失效 | +| `common/config/Constants.java` | 新增 `ENTITY_ATTRIBUTE_INDEX_CACHE_MAX_SIZE = 200000` | + +--- + +## 5. 测试结果 + +环境：macOS arm64、OpenJDK 21.0.11、Maven 3.9.16、Lucene 8.11.2、`topN = 30`。性能数据为连续 3 轮取值范围。 + +### 5.1 正确性 + +`ResidentSearchIndexTest`（13 例）+ `EmbeddingCandidateSetTest`（1 例），全部通过： + +| 断言 | 说明 | +|---|---| +| 召回等价（常驻 vs 重建） | 10000 顶点、10 组查询，结果集完全一致 | +| 召回等价（有缓存 vs 无缓存） | 文本化缓存不改变召回 | +| `buildCount == 1` | 10 次查询后全图索引只构建 1 次 | +| 插入就地生效 | 新点写入后立即可检索，`upsertCount == 1`，文档数 +1，**不重建** | +| 更新就地生效 | 新内容可检索，**被替换的旧文档不再可检索**，文档数不变，**不重建** | +| 删除就地生效 | 被删文档不再可检索，`removeCount == 1`，文档数 −1，**不重建** | +| upsert 幂等 | 同一实体重复上报 3 次，文档数不变、无重复文档、不重建 | +| 未通知改图触发重建 | 绕过索引直接改图后，版本兜底强制重建并返回新内容（`buildCount == 2`） | +| **未通知改图不被后续写入吞掉** | 越界改图后再走一次正常 upsert，越界写入的文档仍能被检索到，`buildCount == 2`、文档数 22 —— 事后读当前版本的做法会让这条永久缺失 | +| 边写入不失效 | 写边后 `buildCount` 不变 | +| **边写入不清空文本化缓存** | 注册边 schema + 19 次写边后顶点条目仍 hit；改写该顶点本身只淘汰它一条（`cacheSize == 1`） | +| 写查交替全程不重建 | 40 轮「写入 + 查询」，每轮召回与失效重建方案逐轮一致，全程 `buildCount == 1` | +| 缓存计数 | 首次 miss、再次 hit、单条失效后再次 miss | +| **无锁缓存并发一致** | 8 线程 × 50 轮 × 200 实体并发查找：内容全部一致、hit + miss 恰等于调用次数、条目数恰等于实体数 | +| 向量候选集等价 | 枚举索引实体 vs 全图扫描，5 组查询结果逐位一致；数据集含「未索引」与「已索引但无向量」两类顶点 | + +结果比对用集合而非列表：每查询重建方案从 `HashMap` 迭代喂 Lucene，文档顺序不确定，并列打分的顺序不稳定。测试查询的命中数控制在 `topN = 30` 以内，规避截断带来的顺序敏感。 + +### 5.2 读性能：10000 顶点 / 10 次查询 + +| 配置 | 每查询耗时 | +|---|---| +| A 每查询重建 + 无文本化缓存 | 101.6 ~ 105.7 ms | +| B 每查询重建 + 文本化缓存 | 86.4 ~ 89.2 ms | +| C 常驻索引（稳态） | **0.38 ~ 0.60 ms** | + +- C 的一次性构建 97.7 ~ 106.4 ms，仅首次查询承担，稳态相比 A 约 **170 ~ 280 倍** +- B 的缓存命中 90000/100000，仅带来约 18% 改善 —— 本用例内容规模下主要成本是 Lucene 建索引而非文本化，索引常驻是主因，缓存是次要项。注意这个比例只对「图只读」的读基准成立；写入频繁时缓存的价值取决于失效粒度，见 §3.5.1 + +**基准场景前提**（该数字的适用边界）：合成图、单一顶点标签、无边、短文本单属性、未设 `PromptFormatter`；测量期间图只读；仅测冷路径全局检索本身，不含 `apply()` 其余部分、会话处理、结果 verbalize 与 HTTP 开销。 + +### 5.3 写性能：5000 顶点、40 轮「写入 + 查询」交替 + +就地增量维护存在的理由所在。对照组为失效重建方案（每次写入后 `invalidate()`，下次查询重建）： + +| 配置 | 每轮耗时 | 全量构建次数 | +|---|---|---| +| D 写入即失效，查询时重建 | 46.3 ~ 46.6 ms | 41 | +| E 就地增量维护 | **1.36 ~ 1.40 ms** | **1** | + +约 **33 倍**，且构建次数与写入次数解耦。每轮召回逐轮比对一致。 + +E 的每轮 1.4 ms 高于纯读稳态的 0.4 ms，成本落在写入与 `refresh()` 一侧，不在检索一侧。把 800 轮写查交替按 100 轮分桶、写与查分开计时可以看到这一点（去 commit 之前的数据）： + +| 轮次 | 写入 + refresh | 检索 | +|---|---|---| +| 1..100 | 2.045 ms | 0.754 ms | +| 201..300 | 1.386 ms | 0.396 ms | +| 401..500 | 0.928 ms | 0.215 ms | +| 701..800 | 0.998 ms | 0.260 ms | + +检索耗时全程不升反降 —— 段数量增长并没有拖慢检索，Lucene 默认的 `TieredMergePolicy` 已经在后台合并。因此这里的可优化项是 refresh 本身（去掉 commit 已经拿到约 20%，见 §3.9）与跨批次攒批，而不是自己实现段合并策略。参见 §6.5 的更正。 + +### 5.4 规模敏感性 + +| 顶点数 | 每查询重建 | 常驻索引稳态 | +|---|---|---| +| 5000 | 45.4 ~ 50.6 ms | 0.31 ~ 0.33 ms | +| 20000 | 176.6 ~ 179.3 ms | 0.51 ~ 0.55 ms | + +顶点数 4 倍 → 重建路径约 3.7 倍（线性，确认 O(V)）；常驻路径基本持平。延迟特征从「随图规模线性增长」变为「基本不随图规模增长」。 + +### 5.5 既有回归用例 + +| 测试 | 每查询重建 | 常驻索引 | +|---|---|---| +| `MutableGraphTest`（多轮会话 + 频繁改图） | 0.887 s | **0.129 s** | +| `GraphMemoryTest`（LDBC 数据集，严格内容断言） | 0.599 s | 0.592 s | +| `MemoryServerTest`（HTTP 端到端，532 chunk 导入） | 5.717 s | 5.596 s | + +`GraphMemoryTest` 的严格内容断言全部通过，是召回未变化的额外佐证。`MemoryServerTest` 基本不获益，原因见 §6.1。 + +### 5.6 全量验证 + +`mvn -B -pl geaflow-ai -am clean install`：Reactor 12 个模块全部 SUCCESS；`geaflow-ai` **19 个测试通过，0 失败 0 错误**；Checkstyle **0 违规**；Apache RAT **Unapproved 0**。 + +--- + +## 6. 已知限制 + +### 6.1 consolidate 写入路径每次插入重建全部检索状态 + +`KeywordRelationFunction.eval()` 每次调用都 `new EntityAttributeIndexStore()` + `new GraphMemoryServer()`，然后对无关联边的顶点做全局检索。它由 `ConsolidateServer` 在每次 `/graph/insertEntity` 时执行，因此**单次插入代价 O(V)、整体导入 O(V²)**。 + +实测：`MemoryServerTest` 导入 532 个 chunk，触发 532 次全图索引构建（单次从 2 ms 增至 9 ms）。这是 §5.5 中该测试无改善的原因，也是目前**唯一一条仍在按 O(V) 重建索引的路径**。 + +修复需让 consolidate 复用按图的检索状态并以增量方式接收新实体，会改动 `ConsolidateFunction.eval` 的契约，建议单独立项。 + +### 6.2 向量检索无 ANN + +`EmbeddingIndexStore` 仍是 `HashMap` + 全候选集余弦计算。本特性只消除了「为组装候选集而扫全图」的开销，算法复杂度仍是 O(N·d)。引入 ANN 需先解决 Lucene 8.11.2 → 9.8.0 与随之而来的 JDK 11 约束（`geaflow-store-vector` 的 `GraphVectorIndex` 正因此被 `-Pjdk8` CI 构建排除），建议先做可插拔 SPI。 + +### 6.3 常驻索引无内存上界，缓存上限是近似值 + +文本化缓存有条数上限（`Constants.ENTITY_ATTRIBUTE_INDEX_CACHE_MAX_SIZE`，默认 20 万条），Lucene `ByteBuffersDirectory` 内存索引没有。索引本身是新增的稳态堆占用（原来是瞬态的），大图上需评估改用磁盘 `Directory`。 + +缓存上限有两处不足：按条数而非字节计（20 万条在 LDBC 量级的 verbalization 下可能偏大，需要一个字节预算口径），以及为了让读路径不持锁而放弃了精确上限与 LRU 顺序（§3.5.1）。两者都指向同一个解法：换成带权重与淘汰策略的缓存库（Caffeine 之类），而不是继续在手写 map 上加码。 + +与图等大的辅助集合已经去掉（§3.11），所以常驻状态只剩索引本身。 + +### 6.4 并发只做到「查询不互相串行」 + +`ResidentSearchIndex` 用 `ReentrantReadWriteLock`：并发查询走读锁互不阻塞，但构建与写入仍持写锁、期间检索被阻塞（10000 顶点冷建约 100 ms）。 + +这里为什么不像文本化缓存那样退到「一个 `volatile` 原子引用、线程之间互不干扰」？因为两者的对象生命周期不同。缓存值是不可变的、被替换后旧值仍然完全可用，读方拿到哪个版本都成立；而 `invalidate()` 会 `close()` 掉 `GraphSearchStore`，一个无锁读方若正好持有它的引用就会撞上 `AlreadyClosedException`。要在无锁的前提下安全回收，必须知道「还有没有人在读」——也就是引用计数，即 Lucene 的 `SearcherManager` / `ReferenceManager`。所以这一步的正确形态是换成 `SearcherManager`，而不是把读写锁直接摘掉。 + +`GraphMemoryServer.residentIndexes` 已用 `synchronizedMap` 包装，但服务端本身仍是「全局静态 `CACHE`、每请求 new 一个 `MemoryMutableGraph`」，多写方并发时 §3.5 表格最后一行的窗口（写方自己改图与封窗之间被别人插入）无法覆盖 —— 服务化时需要按图的写锁把写入串行化。 + +### 6.5 无跨批次刷新攒批（原「无段合并」结论已更正） + +原先此处认为「段数量增长会拖慢检索，需要自己实现段合并策略」。800 轮写查交替的分桶实测（§5.3）**不支持**这个结论：检索耗时全程不升反降，Lucene 默认的 `TieredMergePolicy` 已经在后台合并段、并按 `deletesPctAllowed` 回收被标记删除的文档，不需要额外的合并策略。 + +真正剩下的是刷新攒批：目前一批写入（一次 HTTP 请求）刷新一次，`commit()` 已去掉（§3.9），但批量导入场景下仍然是每请求一次 reader 重开。按时间或按变更量延迟刷新（对应 Elasticsearch 的 `refresh_interval`）能把这部分摊薄，代价是可见性延迟。 + +### 6.6 版本号仅内存图实现 + +只有 `MemoryGraph` / `LocalMemoryGraphAccessor` 上报版本。未来 `GeaFlowStateGraphAccessor` 若不实现 `getVertexVersion()`，常驻索引会退化为每查询重建 —— 安全但无收益。接引擎时须一并实现版本上报。 + +--- + +## 7. 后续建议 + +按投入产出排序： + +1. **打通 HTTP 层 embedding 通路** —— `GeaFlowMemoryServer.createGraph()` 未注册 `EmbeddingIndexStore`，`execQuery()` 也从不产生 `EmbeddingVector`，导致线上路径的向量检索完全没生效。改动极小，属功能缺陷而非优化。 +2. **修 §6.1 的 consolidate 写入路径** —— 让它复用检索状态，把导入从 O(V²) 降到 O(V)。 +3. **换 `SearcherManager`**（§6.4）—— 让构建期也不阻塞检索，顺带把 reader 生命周期交给引用计数管理。 +4. **刷新攒批**（§6.5）—— 批量导入场景摊薄 reader 重开成本。 +5. **热路径改用扩展集过滤** —— 用 `TermInSetQuery` 把扩展集的 `KEY` 挂成 filter 查常驻索引，热路径也不必每查询建小索引。注意 BM25 语料统计会从小索引变成全局索引，排序与 §3.8 的重建路径不再严格等价，需要先做 benchmark 与召回对比再决定。 +6. **加前置精确匹配** —— query 命中实体 ID / label 时直接定位，跳过全量比对（对应 hugegraph-ai 的 `_exact_match_vids`）。 +7. **缩小索引对象** —— 区分「实体 ID / 名称索引」与「完整文本索引」两级，先在小索引上召回候选再精排（对应 hugegraph-ai 只索引 `graph_vids` 的做法）。 +8. **建评测基线** —— 缺少评测集，后续检索质量优化无法验证。 + +--- + +## 附：参考来源 + +- [Lucene's Handling of Deleted Documents — Elastic](https://www.elastic.co/blog/lucenes-handling-of-deleted-documents) +- [Lucene's near-real-time search is fast! — DZone](https://dzone.com/articles/lucenes-near-real-time-search) +- [Elasticsearch refresh_interval 说明 — pulse.support](https://pulse.support/kb/what-is-elasticsearch-refresh-interval) +- [Elasticsearch merge storms — Netdata](https://www.netdata.cloud/guides/elasticsearch/elasticsearch-merge-storms/) +- [Milvus data processing 文档](https://milvus.io/docs/data_processing.md) +- [HNSW 删除与墓碑机制分析 — tianpan.co](https://tianpan.co/blog/2026-05-09-retrieval-cascade-failure-document-deletion-rag) + +上述来源内容均经改写与摘要，以符合许可要求。 diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java index 5121ae7c9..1dacb8279 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java @@ -123,6 +123,11 @@ public String addSchema(@Param("graphName") String graphName, } else { throw new RuntimeException("Cannot add schema: " + input); } + GraphMemoryServer schemaServer = CACHE.getServerByName(graphName); + if (schemaServer != null) { + // Verbalization is schema driven, so cached prompts and the index must be dropped. + schemaServer.onSchemaChanged(); + } return "addSchema has been called, schemaName: " + schemaName; } @@ -150,9 +155,16 @@ public String addEntity(@Param("graphName") String graphName, if (!(graph instanceof MemoryGraph)) { throw new RuntimeException("Graph cannot modify."); } + GraphMemoryServer insertServer = CACHE.getServerByName(graphName); + if (insertServer == null || insertServer.getGraphAccessors().isEmpty()) { + throw new RuntimeException("Server or graph accessor not available for graph: " + graphName); + } MemoryMutableGraph memoryMutableGraph = new MemoryMutableGraph((MemoryGraph) graph); List graphEntities = SeDeUtil.deserializeEntities(input); + // Opened before the writes and sealed right after them, so the resident index can verify + // that these entities really are every vertex level change it has not seen yet. + VertexVersionWindow window = insertServer.openVertexVersionWindow(); for (GraphEntity entity : graphEntities) { if (entity instanceof GraphVertex) { memoryMutableGraph.addVertex(((GraphVertex) entity).getVertex()); @@ -160,10 +172,8 @@ public String addEntity(@Param("graphName") String graphName, memoryMutableGraph.addEdge(((GraphEdge) entity).getEdge()); } } - GraphMemoryServer insertServer = CACHE.getServerByName(graphName); - if (insertServer == null || insertServer.getGraphAccessors().isEmpty()) { - 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()); CACHE.getConsolidateServer().executeConsolidateTask( insertServer.getGraphAccessors().get(0), memoryMutableGraph); return "Success to add entities, num: " + graphEntities.size(); @@ -182,6 +192,9 @@ public String deleteEntity(@Param("graphName") String graphName, } MemoryMutableGraph memoryMutableGraph = new MemoryMutableGraph((MemoryGraph) graph); List graphEntities = SeDeUtil.deserializeEntities(input); + GraphMemoryServer deleteServer = CACHE.getServerByName(graphName); + VertexVersionWindow window = deleteServer == null + ? null : deleteServer.openVertexVersionWindow(); for (GraphEntity entity : graphEntities) { if (entity instanceof GraphVertex) { memoryMutableGraph.removeVertex(entity.getLabel(), @@ -190,6 +203,10 @@ public String deleteEntity(@Param("graphName") String graphName, memoryMutableGraph.removeEdge(((GraphEdge) entity).getEdge()); } } + if (deleteServer != null) { + // Deletes are applied to the index in place, no rebuild needed. + deleteServer.onEntitiesRemoved(graphEntities, window.seal()); + } return "Success to remove entities, num: " + graphEntities.size(); } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/GraphMemoryServer.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/GraphMemoryServer.java index 57a8d37a2..956345818 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/GraphMemoryServer.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/GraphMemoryServer.java @@ -20,16 +20,21 @@ package org.apache.geaflow.ai; import java.util.ArrayList; +import java.util.Collections; import java.util.HashSet; +import java.util.IdentityHashMap; import java.util.List; +import java.util.Map; import java.util.Set; import java.util.stream.Collectors; import org.apache.geaflow.ai.graph.GraphAccessor; import org.apache.geaflow.ai.graph.GraphEntity; +import org.apache.geaflow.ai.graph.VertexVersionWindow; import org.apache.geaflow.ai.index.EmbeddingIndexStore; import org.apache.geaflow.ai.index.EntityAttributeIndexStore; import org.apache.geaflow.ai.index.IndexStore; import org.apache.geaflow.ai.operator.EmbeddingOperator; +import org.apache.geaflow.ai.operator.ResidentSearchIndex; import org.apache.geaflow.ai.operator.SearchOperator; import org.apache.geaflow.ai.operator.SessionOperator; import org.apache.geaflow.ai.search.VectorSearch; @@ -44,6 +49,13 @@ public class GraphMemoryServer { private final List graphAccessors = new ArrayList<>(); private final List indexStores = new ArrayList<>(); + /** + * Keyword indexes kept alive across queries, one per keyword index store. Without this the + * global keyword index would be rebuilt from a full graph scan on every single query. + */ + private final Map residentIndexes = + Collections.synchronizedMap(new IdentityHashMap<>()); + public void addGraphAccessor(GraphAccessor graph) { if (graph != null) { graphAccessors.add(graph); @@ -57,6 +69,9 @@ public List getGraphAccessors() { public void addIndexStore(IndexStore indexStore) { if (indexStore != null) { indexStores.add(indexStore); + if (indexStore instanceof EntityAttributeIndexStore) { + residentIndexes.put(indexStore, new ResidentSearchIndex()); + } } } @@ -86,7 +101,8 @@ public String search(VectorSearch search) { } for (IndexStore indexStore : indexStores) { if (indexStore instanceof EntityAttributeIndexStore) { - SessionOperator searchOperator = new SessionOperator(graphAccessors.get(0), indexStore); + SessionOperator searchOperator = new SessionOperator(graphAccessors.get(0), + indexStore, residentIndexes.get(indexStore)); applySearch(sessionId, searchOperator, search); } if (indexStore instanceof EmbeddingIndexStore) { @@ -121,6 +137,76 @@ public Context verbalize(String sessionId, VerbalizationFunction verbalizationFu return new Context(stringBuilder.toString()); } + /** + * Captures the vertex version before a batch of graph writes. Pass the sealed window to + * {@link #onEntitiesUpserted} / {@link #onEntitiesRemoved} so the derived structures can tell + * whether the reported entities really are everything that changed. + */ + public VertexVersionWindow openVertexVersionWindow() { + return VertexVersionWindow.open(graphAccessors.isEmpty() ? null : graphAccessors.get(0)); + } + + /** + * Applies written entities to the derived structures in place. Handles both new and rewritten + * entities, so callers do not need to distinguish them. + * + *

Memoized verbalizations need no explicit invalidation here: every entry carries the source + * version it was computed from, so the write itself makes the affected entries stale. + * + * @param window version range the batch covers, obtained from + * {@link #openVertexVersionWindow()} and sealed after the writes + */ + public void onEntitiesUpserted(List entities, VertexVersionWindow window) { + if (entities == null || entities.isEmpty() || graphAccessors.isEmpty()) { + return; + } + for (IndexStore indexStore : indexStores) { + if (!(indexStore instanceof EntityAttributeIndexStore)) { + continue; + } + ResidentSearchIndex residentIndex = residentIndexes.get(indexStore); + if (residentIndex != null) { + residentIndex.onEntitiesUpserted(graphAccessors.get(0), entities, indexStore, + window); + } + } + } + + /** + * Applies removed entities to the derived structures in place. + */ + public void onEntitiesRemoved(List entities, VertexVersionWindow window) { + if (entities == null || entities.isEmpty() || graphAccessors.isEmpty()) { + return; + } + for (IndexStore indexStore : indexStores) { + if (!(indexStore instanceof EntityAttributeIndexStore)) { + continue; + } + ResidentSearchIndex residentIndex = residentIndexes.get(indexStore); + if (residentIndex != null) { + residentIndex.onEntitiesRemoved(graphAccessors.get(0), entities, window); + } + } + } + + /** + * Drops the derived structures wholesale. Used for changes that cannot be expressed per entity, + * such as a schema change altering how every entity is verbalized. + */ + public void onSchemaChanged() { + for (IndexStore indexStore : indexStores) { + if (!(indexStore instanceof EntityAttributeIndexStore)) { + continue; + } + ((EntityAttributeIndexStore) indexStore).invalidateCache(); + ResidentSearchIndex residentIndex = residentIndexes.get(indexStore); + if (residentIndex != null) { + residentIndex.invalidate(); + } + } + } + public List getSessionEntities(String sessionId) { List subGraphList = sessionManagement.getSubGraph(sessionId); Set entitySet = new HashSet<>(); diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/common/config/Constants.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/common/config/Constants.java index ef2646035..439f517d9 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/common/config/Constants.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/common/config/Constants.java @@ -47,6 +47,9 @@ public class Constants { public static int EMBEDDING_OPERATE_DEFAULT_TOPN = 50; public static int GRAPH_SEARCH_STORE_DEFAULT_TOPN = 30; + // Max number of memoized entity verbalizations kept by EntityAttributeIndexStore. + public static int ENTITY_ATTRIBUTE_INDEX_CACHE_MAX_SIZE = 200000; + public static String CONSOLIDATE_KEYWORD_RELATION_LABEL = "consolidate_keyword_edge"; public static String PREFIX_COMMON_KEYWORDS = "common_keywords"; } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/GraphAccessor.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/GraphAccessor.java index 82dd1cd57..369b60bf6 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/GraphAccessor.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/GraphAccessor.java @@ -25,6 +25,34 @@ public interface GraphAccessor { + /** + * Returned by {@link #getGraphVersion()} when the accessor cannot report content changes. + * Callers must then treat every read as potentially different and skip caching. + */ + long VERSION_UNSUPPORTED = -1L; + + /** + * A monotonically increasing counter bumped on every content or schema change of the underlying + * graph. Derived structures (verbalization caches, keyword indexes) compare it to decide whether + * they are still valid, so that direct mutations of the graph cannot silently go unnoticed. + * + * @return current graph version, or {@link #VERSION_UNSUPPORTED} if change tracking is not + * available for this accessor + */ + default long getGraphVersion() { + return VERSION_UNSUPPORTED; + } + + /** + * Like {@link #getGraphVersion()} but only advanced by vertex and schema changes. Structures + * derived from vertices alone can watch this and survive edge writes. + * + * @return current vertex version, defaults to {@link #getGraphVersion()} + */ + default long getVertexVersion() { + return getGraphVersion(); + } + GraphSchema getGraphSchema(); GraphVertex getVertex(String label, String id); diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/LocalMemoryGraphAccessor.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/LocalMemoryGraphAccessor.java index 98614bc4a..6a0a8d6ed 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/LocalMemoryGraphAccessor.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/LocalMemoryGraphAccessor.java @@ -46,6 +46,16 @@ public LocalMemoryGraphAccessor(MemoryGraph memoryGraph) { this.graph = memoryGraph; } + @Override + public long getGraphVersion() { + return graph.getVersion(); + } + + @Override + public long getVertexVersion() { + return graph.getVertexVersion(); + } + @Override public GraphSchema getGraphSchema() { return graph.getGraphSchema(); diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/MemoryMutableGraph.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/MemoryMutableGraph.java index 37b9a6d93..7a532c0de 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/MemoryMutableGraph.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/MemoryMutableGraph.java @@ -19,7 +19,6 @@ package org.apache.geaflow.ai.graph; -import java.util.ArrayList; import org.apache.commons.lang3.StringUtils; import org.apache.geaflow.ai.common.ErrorCode; import org.apache.geaflow.ai.graph.io.*; @@ -80,8 +79,7 @@ public int addVertexSchema(VertexSchema vertexSchema) { if (this.graph.entities.get(vertexSchema.getLabel()) != null) { return ErrorCode.GRAPH_ADD_VERTEX_SCHEMA_FAILED; } - this.graph.getGraphSchema().addVertex(vertexSchema); - this.graph.entities.put(vertexSchema.getLabel(), new VertexGroup(vertexSchema, new ArrayList<>())); + this.graph.registerVertexSchema(vertexSchema); return ErrorCode.SUCCESS; } @@ -103,8 +101,7 @@ public int addEdgeSchema(EdgeSchema edgeSchema) { if (this.graph.entities.get(edgeSchema.getLabel()) != null) { return ErrorCode.GRAPH_ADD_EDGE_SCHEMA_FAILED; } - this.graph.getGraphSchema().addEdge(edgeSchema); - this.graph.entities.put(edgeSchema.getLabel(), new EdgeGroup(edgeSchema, new ArrayList<>())); + this.graph.registerEdgeSchema(edgeSchema); return ErrorCode.SUCCESS; } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/VertexVersionWindow.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/VertexVersionWindow.java new file mode 100644 index 000000000..7a92a013c --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/VertexVersionWindow.java @@ -0,0 +1,108 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.ai.graph; + +/** + * The vertex version range that a reported batch of entity changes claims to cover. + * + *

A derived structure such as a resident keyword index can only apply a batch of changes in + * place if it can prove the batch is the complete set of vertex level changes since the + * structure was last known to be correct. Reading the current version after the fact is not a + * proof: any change made outside the reporting path would be silently accepted as already applied, + * and the structure would keep serving stale results forever. + * + *

So the writer states the range instead of the reader guessing it: + * + * 

+ * VertexVersionWindow window = VertexVersionWindow.open(accessor);
+ * ... mutate the graph ...
+ * index.onEntitiesUpserted(entities, window.seal());
+ *
+ * + *

The consumer accepts the batch only when {@link #getFrom()} matches the version it last + * accepted and {@link #getTo()} still matches the graph, otherwise it falls back to a + * full rebuild. Changes that slip in between the writer's own mutations and {@link #seal()} cannot + * be detected this way; writers that mutate a graph concurrently must serialize themselves. + */ +public final class VertexVersionWindow { + + /** Marks a window that has not been sealed yet, and can therefore not be trusted. */ + private static final long UNSEALED = Long.MIN_VALUE; + + private final GraphAccessor accessor; + private final long from; + private final long to; + + private VertexVersionWindow(GraphAccessor accessor, long from, long to) { + this.accessor = accessor; + this.from = from; + this.to = to; + } + + /** + * Captures the vertex version before a batch of writes. + * + * @param accessor graph the writes will be applied to, may be {@code null} + */ + public static VertexVersionWindow open(GraphAccessor accessor) { + long from = accessor == null ? GraphAccessor.VERSION_UNSUPPORTED : accessor.getVertexVersion(); + return new VertexVersionWindow(accessor, from, UNSEALED); + } + + /** + * Captures the vertex version after the batch of writes. Call this as close to the last write + * as possible: everything between the write and this call is a blind spot. + */ + public VertexVersionWindow seal() { + long to = accessor == null ? GraphAccessor.VERSION_UNSUPPORTED : accessor.getVertexVersion(); + return new VertexVersionWindow(accessor, from, to); + } + + public boolean isSealed() { + return to != UNSEALED; + } + + public long getFrom() { + return from; + } + + public long getTo() { + return to; + } + + /** + * Whether this window can be trusted to describe every vertex level change between + * {@code acceptedVersion} and now. + * + * @param acceptedVersion version the consumer last accepted as fully applied + */ + public boolean covers(long acceptedVersion) { + if (!isSealed() || from != acceptedVersion) { + return false; + } + // Anything that moved the version after the window was sealed is not described by it. + return accessor == null || accessor.getVertexVersion() == to; + } + + @Override + public String toString() { + return "VertexVersionWindow{from=" + from + ", to=" + (isSealed() ? to : "unsealed") + '}'; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/io/MemoryGraph.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/io/MemoryGraph.java index c1ae84b1e..421ed845b 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/io/MemoryGraph.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/io/MemoryGraph.java @@ -20,6 +20,7 @@ package org.apache.geaflow.ai.graph.io; import java.util.*; +import java.util.concurrent.atomic.AtomicLong; import org.apache.geaflow.ai.common.ErrorCode; import org.apache.geaflow.ai.graph.Graph; @@ -28,11 +29,63 @@ public class MemoryGraph implements Graph { public GraphSchema graphSchema; public Map entities; + /** + * Bumped on every content or schema change so that derived structures such as verbalization + * caches and keyword indexes can detect that they went stale, even when the graph is mutated + * directly instead of through a server API. + */ + private final AtomicLong version = new AtomicLong(); + + /** + * Bumped only by vertex and schema changes. Structures that depend on vertices alone, such as + * the global keyword index, can watch this instead of {@link #version} and stay valid across + * edge writes. + */ + private final AtomicLong vertexVersion = new AtomicLong(); + public MemoryGraph(GraphSchema graphSchema, Map entities) { this.graphSchema = graphSchema; this.entities = entities; } + public long getVersion() { + return version.get(); + } + + public long getVertexVersion() { + return vertexVersion.get(); + } + + /** + * Registers a vertex schema and its empty entity group. Callers are expected to have validated + * the label first. Advances the vertex version, since vertex verbalization is schema driven. + */ + public void registerVertexSchema(VertexSchema vertexSchema) { + graphSchema.addVertex(vertexSchema); + entities.put(vertexSchema.getLabel(), new VertexGroup(vertexSchema, new ArrayList<>())); + bumpVersion(); + } + + /** + * Registers an edge schema and its empty entity group. Only the general version is advanced: an + * edge schema cannot change how an existing vertex is verbalized, so vertex only derived + * structures stay valid. + */ + public void registerEdgeSchema(EdgeSchema edgeSchema) { + graphSchema.addEdge(edgeSchema); + entities.put(edgeSchema.getLabel(), new EdgeGroup(edgeSchema, new ArrayList<>())); + bumpEdgeVersion(); + } + + private void bumpVersion() { + version.incrementAndGet(); + vertexVersion.incrementAndGet(); + } + + private void bumpEdgeVersion() { + version.incrementAndGet(); + } + @Override public GraphSchema getGraphSchema() { return graphSchema; @@ -40,6 +93,7 @@ public GraphSchema getGraphSchema() { public void setGraphSchema(GraphSchema graphSchema) { this.graphSchema = graphSchema; + bumpVersion(); } private EntityGroup getEntity(String entityName) { @@ -75,7 +129,7 @@ public int removeVertex(String label, String id) { return ErrorCode.GRAPH_ENTITY_GROUP_NOT_MATCH; } VertexGroup vertexGroup = (VertexGroup) vg; - return vertexGroup.removeVertex(id); + return bumped(vertexGroup.removeVertex(id)); } @Override @@ -89,7 +143,7 @@ public int updateVertex(Vertex newVertex) { return ErrorCode.GRAPH_ENTITY_GROUP_NOT_MATCH; } VertexGroup vertexGroup = (VertexGroup) vg; - return vertexGroup.updateVertex(newVertex); + return bumped(vertexGroup.updateVertex(newVertex)); } @Override @@ -103,7 +157,7 @@ public int addVertex(Vertex newVertex) { return ErrorCode.GRAPH_ENTITY_GROUP_NOT_MATCH; } VertexGroup vertexGroup = (VertexGroup) vg; - return vertexGroup.addVertex(newVertex); + return bumped(vertexGroup.addVertex(newVertex)); } @Override @@ -126,7 +180,7 @@ public int removeEdge(Edge edge) { return ErrorCode.GRAPH_ENTITY_GROUP_NOT_MATCH; } EdgeGroup edgeGroup = (EdgeGroup) vg; - return edgeGroup.removeEdge(edge); + return edgeBumped(edgeGroup.removeEdge(edge)); } @Override @@ -140,7 +194,7 @@ public int addEdge(Edge newEdge) { return ErrorCode.GRAPH_ENTITY_GROUP_NOT_MATCH; } EdgeGroup edgeGroup = (EdgeGroup) vg; - return edgeGroup.addEdge(newEdge); + return edgeBumped(edgeGroup.addEdge(newEdge)); } @Override @@ -166,6 +220,24 @@ public Iterator scanVertex() { return new CompositeIterator<>(iterators); } + /** + * Marks the graph as changed and passes the mutation result through. The version is bumped even + * for failed mutations: over invalidation is cheap, a missed invalidation is a correctness bug. + */ + private int bumped(int mutationResult) { + bumpVersion(); + return mutationResult; + } + + /** + * Same as {@link #bumped} but only advances the general version, leaving vertex only derived + * structures valid. + */ + private int edgeBumped(int mutationResult) { + bumpEdgeVersion(); + return mutationResult; + } + static class CompositeIterator implements Iterator { private final List> iterators; diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/EmbeddingIndexStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/EmbeddingIndexStore.java index 8fae2e1f0..30835751b 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/EmbeddingIndexStore.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/EmbeddingIndexStore.java @@ -241,6 +241,22 @@ private void flushBatchIndex(List newItemStrings, boolean force) { } } + /** + * The store knows exactly which entities it holds embeddings for, so retrieval does not need + * to scan the whole graph to assemble the candidate set. + * + *

Returns a snapshot rather than a view of the live key set: callers iterate it while other + * requests may still be writing to the store, and a view would fail with + * {@link java.util.ConcurrentModificationException}. + */ + @Override + public Collection getIndexedEntities() { + if (indexStoreMap == null) { + return null; + } + return new ArrayList<>(indexStoreMap.keySet()); + } + @Override public List getEntityIndex(GraphEntity entity) { if (entity != null && indexStoreMap.get(entity) != null) { diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/EntityAttributeIndexStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/EntityAttributeIndexStore.java index aa9823876..f290a8cef 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/EntityAttributeIndexStore.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/EntityAttributeIndexStore.java @@ -20,7 +20,13 @@ package org.apache.geaflow.ai.index; import java.util.ArrayList; +import java.util.Collections; +import java.util.Iterator; import java.util.List; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.LongAdder; +import org.apache.geaflow.ai.common.config.Constants; +import org.apache.geaflow.ai.graph.GraphAccessor; import org.apache.geaflow.ai.graph.GraphEdge; import org.apache.geaflow.ai.graph.GraphEntity; import org.apache.geaflow.ai.graph.GraphVertex; @@ -29,34 +35,160 @@ import org.apache.geaflow.ai.subgraph.SubGraph; import org.apache.geaflow.ai.verbalization.VerbalizationFunction; +/** + * Derives keyword vectors from an entity by verbalizing it. + * + *

Verbalization is deterministic for a given entity, but it is not free: it builds a + * {@link SubGraph}, renders a prompt and allocates intermediate strings. Since retrieval calls + * {@link #getEntityIndex} once per candidate entity on every query, the results are memoized. + * + *

Each entry carries the source version it was computed from, so a write invalidates only the + * entries it actually affects. Comparing a single version for the whole cache would be simpler but + * would throw away everything memoized so far on every write, which is exactly what the write heavy + * paths do (consolidate issues roughly thirty edge writes per inserted entity). + * + *

No locking. The memoized function is pure and its result is immutable, so the cache needs + * no mutual exclusion to be correct: a lookup is one {@link ConcurrentHashMap#get}, and a miss + * computes outside the map and publishes with {@link ConcurrentHashMap#put}. Two threads racing on + * the same entity may both compute it, but for the same version they compute the same thing, so the + * loser only wasted work and either published value is equally valid. A stale entry needs no + * explicit removal either, the put replaces it. + * + *

The price is that the size bound is approximate and eviction is not LRU: entries are dropped in + * map iteration order once the bound is exceeded. For a memoization of a pure function a wrong + * eviction costs one recompute, which is an acceptable price for not having to hold a monitor across + * every lookup. A real eviction policy would need a proper cache library; see the module docs. + */ public class EntityAttributeIndexStore implements IndexStore { + /** Fraction of the bound dropped per eviction pass, so eviction does not run on every put. */ + private static final int EVICTION_BATCH_DIVISOR = 16; + private VerbalizationFunction verbFunc; + private final ConcurrentHashMap verbalizationCache = + new ConcurrentHashMap<>(); + + private final LongAdder cacheHit = new LongAdder(); + private final LongAdder cacheMiss = new LongAdder(); + public void initStore(VerbalizationFunction func) { if (func != null) { this.verbFunc = func; } + invalidateCache(); } @Override public List getEntityIndex(GraphEntity entity) { + if (entity == null) { + return Collections.emptyList(); + } + long version = sourceVersionOf(entity); + if (version == GraphAccessor.VERSION_UNSUPPORTED) { + // The source cannot tell us when it changes, so memoizing would risk stale results. + return computeEntityIndex(entity); + } + CachedIndex cached = verbalizationCache.get(entity); + if (cached != null && cached.version == version) { + cacheHit.increment(); + return cached.vectors; + } + // Computed outside the map: verbalization is the expensive part and must not block others. + // A stale entry needs no explicit removal, the put below replaces it. + List computed = computeEntityIndex(entity); + cacheMiss.increment(); + verbalizationCache.put(entity, new CachedIndex(computed, version)); + enforceBound(); + return computed; + } + + /** + * Keeps the cache near its configured bound. Approximate on purpose: {@code size()} on a + * concurrent map is an estimate, and several threads may evict at once. Overshooting slightly is + * acceptable, holding a lock to be exact is not. + */ + private void enforceBound() { + int max = Constants.ENTITY_ATTRIBUTE_INDEX_CACHE_MAX_SIZE; + int size = verbalizationCache.size(); + if (size <= max) { + return; + } + int toDrop = size - max + Math.max(1, max / EVICTION_BATCH_DIVISOR); + Iterator it = verbalizationCache.keySet().iterator(); + while (toDrop-- > 0 && it.hasNext()) { + it.next(); + it.remove(); + } + } + + /** + * The version an entity's verbalization depends on. + * + *

A vertex is verbalized from itself and the schema, so it only has to watch the vertex + * version and survives edge writes. An edge is verbalized together with both of its endpoints, + * so it has to watch every change. + */ + private long sourceVersionOf(GraphEntity entity) { + return entity instanceof GraphVertex + ? verbFunc.getSourceVertexVersion() + : verbFunc.getSourceVersion(); + } + + private List computeEntityIndex(GraphEntity entity) { + String verbalization; if (entity instanceof GraphVertex) { - String verbalization = verbFunc.verbalize(new SubGraph().addVertex((GraphVertex) entity)); - List sentences = new ArrayList<>(); - sentences.add(verbalization); - KeywordVector keywordVector = new KeywordVector(sentences.toArray(new String[0])); - List results = new ArrayList<>(); - results.add(keywordVector); - return results; + verbalization = verbFunc.verbalize(new SubGraph().addVertex((GraphVertex) entity)); } else { - String verbalization = verbFunc.verbalize(new SubGraph().addEdge((GraphEdge) entity)); - List sentences = new ArrayList<>(); - sentences.add(verbalization); - KeywordVector keywordVector = new KeywordVector(sentences.toArray(new String[0])); - List results = new ArrayList<>(); - results.add(keywordVector); - return results; + verbalization = verbFunc.verbalize(new SubGraph().addEdge((GraphEdge) entity)); + } + KeywordVector keywordVector = new KeywordVector(verbalization); + List results = new ArrayList<>(1); + results.add(keywordVector); + return Collections.unmodifiableList(results); + } + + /** + * Drops all memoized verbalizations. Version stamps already keep stale entries from being + * served, so this is only needed for changes the version does not describe, such as replacing + * the verbalization function itself. + */ + public void invalidateCache() { + verbalizationCache.clear(); + } + + public void invalidateCache(GraphEntity entity) { + if (entity == null) { + return; + } + verbalizationCache.remove(entity); + } + + public long getCacheHit() { + return cacheHit.sum(); + } + + public long getCacheMiss() { + return cacheMiss.sum(); + } + + public int getCacheSize() { + return verbalizationCache.size(); + } + + /** + * A memoized verbalization together with the source version it was computed from. Stamping each + * entry, rather than the cache as a whole, is what lets a single write invalidate one entry + * instead of discarding everything memoized so far. + */ + private static final class CachedIndex { + + private final List vectors; + private final long version; + + private CachedIndex(List vectors, long version) { + this.vectors = vectors; + this.version = version; } } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/IndexStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/IndexStore.java index 53d1d8e51..c015da519 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/IndexStore.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/IndexStore.java @@ -19,6 +19,7 @@ package org.apache.geaflow.ai.index; +import java.util.Collection; import java.util.List; import org.apache.geaflow.ai.graph.GraphEntity; import org.apache.geaflow.ai.index.vector.IVector; @@ -26,4 +27,17 @@ public interface IndexStore { List getEntityIndex(GraphEntity entity); + + /** + * Returns the entities this store actually holds an index for, or {@code null} when the store + * cannot enumerate them (e.g. it derives the index on demand for any entity). + * + *

When available, callers can iterate this collection instead of scanning the whole graph: + * entities absent from the store contribute nothing to recall anyway. + * + * @return indexed entities, or {@code null} if unknown + */ + default Collection getIndexedEntities() { + return null; + } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/EmbeddingOperator.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/EmbeddingOperator.java index 849f53ee4..76b6aeff8 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/EmbeddingOperator.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/EmbeddingOperator.java @@ -55,8 +55,8 @@ public List apply(List subGraphList, VectorSearch search) { } return new ArrayList<>(subGraphList); } - List globalResults = searchWithGlobalGraph(queryEmbeddingVectors); if (subGraphList == null || subGraphList.isEmpty()) { + List globalResults = searchWithGlobalGraph(queryEmbeddingVectors); List startVertices = new ArrayList<>(); for (GraphEntity resEntity : globalResults) { if (resEntity instanceof GraphVertex) { @@ -111,7 +111,37 @@ private List getSubgraphExpand(SubGraph subGraph) { } private List searchWithGlobalGraph(List queryEmbeddingVectors) { + return searchEmbeddings(queryEmbeddingVectors, collectGlobalCandidates()); + } + + /** + * Collects the global candidate set. + * + *

Only vertices that actually carry an embedding can be recalled, because + * {@link #searchEmbeddings} skips entities without vectors. So when the index store can + * enumerate what it holds, iterate that instead of scanning the whole graph: the candidate set + * is identical, but the graph scan and the per-vertex wrapper allocation disappear. + */ + private Map> collectGlobalCandidates() { Map> entityIndexMap = new HashMap<>(); + Collection indexedEntities = indexStore.getIndexedEntities(); + if (indexedEntities != null) { + for (GraphEntity entity : indexedEntities) { + if (!(entity instanceof GraphVertex)) { + continue; + } + // The index store is not notified about deletes, so it can still hold entries for + // vertices the graph no longer has. Resolving each one keeps the candidate set + // identical to the graph scan and yields the current vertex object. + GraphVertex current = graphAccessor.getVertex(entity.getLabel(), + ((GraphVertex) entity).getVertex().getId()); + if (current == null) { + continue; + } + entityIndexMap.put(current, indexStore.getEntityIndex(entity)); + } + return entityIndexMap; + } Iterator vertexIterator = graphAccessor.scanVertex(); while (vertexIterator.hasNext()) { GraphVertex vertex = vertexIterator.next(); @@ -119,8 +149,7 @@ private List searchWithGlobalGraph(List queryEmbeddingVect List vertexIndex = indexStore.getEntityIndex(vertex); entityIndexMap.put(vertex, vertexIndex); } - //recall compute - return searchEmbeddings(queryEmbeddingVectors, entityIndexMap); + return entityIndexMap; } private List searchEmbeddings(List queryEmbeddingVectors, diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/GraphSearchStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/GraphSearchStore.java index 20beb48ad..5acb966ed 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/GraphSearchStore.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/GraphSearchStore.java @@ -21,71 +21,140 @@ import java.util.*; import java.util.stream.Collectors; +import org.apache.geaflow.ai.common.model.ModelUtils; import org.apache.geaflow.ai.graph.GraphAccessor; import org.apache.geaflow.ai.graph.GraphEdge; import org.apache.geaflow.ai.graph.GraphEntity; import org.apache.geaflow.ai.graph.GraphVertex; import org.apache.geaflow.ai.graph.io.Edge; import org.apache.geaflow.ai.graph.io.EdgeSchema; +import org.apache.geaflow.ai.graph.io.GraphSchema; import org.apache.geaflow.ai.graph.io.Vertex; import org.apache.geaflow.ai.graph.io.VertexSchema; import org.apache.geaflow.ai.index.vector.IVector; import org.apache.lucene.analysis.Analyzer; import org.apache.lucene.document.Document; import org.apache.lucene.index.IndexNotFoundException; -import org.apache.lucene.index.IndexWriterConfig; import org.apache.lucene.search.ScoreDoc; import org.apache.lucene.search.TopDocs; import org.apache.lucene.store.Directory; public class GraphSearchStore { - private SearchStore store; - private long entityNum = 0L; + private final SearchStore store; + + /** Label sets of the schema last searched against, cached to avoid rebuilding them per query. */ + private volatile GraphSchema cachedSchema; + private volatile Set cachedVertexLabels; + private volatile Set cachedEdgeLabels; public GraphSearchStore() { this.store = new SearchStore(); } - public boolean indexVertex(GraphVertex graphVertex, List indexVectors) { - Map kv = new HashMap<>(); - Vertex vertex = graphVertex.getVertex(); - kv.put(SearchConstants.ID, vertex.getId()); - kv.put(SearchConstants.LABEL, vertex.getLabel()); - List contents = new ArrayList<>(indexVectors.size()); - for (IVector v : indexVectors) { - contents.add(v.toString()); + /** + * Number of live documents in the index. Cheaper and more trustworthy than tracking a counter + * alongside Lucene, which would have to mirror update and delete semantics. + */ + public int getDocNum() { + try { + return store.numDocs(); + } catch (IndexNotFoundException notFoundException) { + return 0; + } catch (Throwable e) { + throw new RuntimeException("Cannot read search store", e); } - String content = String.join(SearchConstants.DELIMITER, contents); - kv.put(SearchConstants.CONTENT, content); + } + /** + * Makes previously indexed entities searchable without discarding the index. + */ + public void refresh() { try { - store.addDoc(kv); + store.refresh(); + } catch (IndexNotFoundException notFoundException) { + // Nothing has been indexed yet, there is nothing to make visible. } catch (Throwable e) { - throw new RuntimeException("Cannot index vertex to search store", e); + throw new RuntimeException("Cannot refresh search store", e); } - addItem(); - return true; + } + + public boolean indexVertex(GraphVertex graphVertex, List indexVectors) { + return writeDoc(vertexDoc(graphVertex, indexVectors), false); + } + + /** + * Adds or replaces a vertex document in place, keyed by + * {@link ModelUtils#getGraphEntityKey}. Idempotent: calling it twice for the same vertex leaves + * a single document. + */ + public boolean upsertVertex(GraphVertex graphVertex, List indexVectors) { + return writeDoc(vertexDoc(graphVertex, indexVectors), true); } public boolean indexEdge(GraphEdge graphEdge, List indexVectors) { + return writeDoc(edgeDoc(graphEdge, indexVectors), false); + } + + public boolean upsertEdge(GraphEdge graphEdge, List indexVectors) { + return writeDoc(edgeDoc(graphEdge, indexVectors), true); + } + + /** + * Marks the entity's document as deleted. Lucene flips a bit in a per segment bitset, so the + * cost is independent of index size and no rebuild is needed. + */ + public boolean removeEntity(GraphEntity entity) { + if (entity == null) { + return false; + } + try { + store.deleteDoc(SearchConstants.KEY, ModelUtils.getGraphEntityKey(entity)); + } catch (Throwable e) { + throw new RuntimeException("Cannot remove entity from search store", e); + } + return true; + } + + private Map vertexDoc(GraphVertex graphVertex, List indexVectors) { + Map kv = new HashMap<>(); + Vertex vertex = graphVertex.getVertex(); + kv.put(SearchConstants.KEY, ModelUtils.getGraphEntityKey(graphVertex)); + kv.put(SearchConstants.ID, vertex.getId()); + kv.put(SearchConstants.LABEL, vertex.getLabel()); + kv.put(SearchConstants.CONTENT, joinVectors(indexVectors)); + return kv; + } + + private Map edgeDoc(GraphEdge graphEdge, List indexVectors) { Map kv = new HashMap<>(); Edge edge = graphEdge.getEdge(); + kv.put(SearchConstants.KEY, ModelUtils.getGraphEntityKey(graphEdge)); kv.put(SearchConstants.SRC, edge.getSrcId()); kv.put(SearchConstants.DST, edge.getDstId()); kv.put(SearchConstants.LABEL, edge.getLabel()); + kv.put(SearchConstants.CONTENT, joinVectors(indexVectors)); + return kv; + } + + private String joinVectors(List indexVectors) { List contents = new ArrayList<>(indexVectors.size()); for (IVector v : indexVectors) { contents.add(v.toString()); } - String content = String.join(SearchConstants.DELIMITER, contents); - kv.put(SearchConstants.CONTENT, content); + return String.join(SearchConstants.DELIMITER, contents); + } + + private boolean writeDoc(Map kv, boolean upsert) { try { - store.addDoc(kv); + if (upsert) { + store.updateDoc(SearchConstants.KEY, kv.get(SearchConstants.KEY), kv); + } else { + store.addDoc(kv, SearchConstants.KEY); + } } catch (Throwable e) { - throw new RuntimeException("Cannot index vertex to search store", e); + throw new RuntimeException("Cannot index entity to search store", e); } - addItem(); return true; } @@ -94,10 +163,20 @@ public List search(String key1, GraphAccessor graphAccessor) { String query = SearchUtils.formatQuery(key1); TopDocs docs = store.searchDoc(SearchConstants.CONTENT, query); ScoreDoc[] scoreDocArray = docs.scoreDocs; - Set vertexLabels = graphAccessor.getGraphSchema().getVertexSchemaList() - .stream().map(VertexSchema::getLabel).collect(Collectors.toSet()); - Set edgeLabels = graphAccessor.getGraphSchema().getEdgeSchemaList() - .stream().map(EdgeSchema::getLabel).collect(Collectors.toSet()); + GraphSchema schema = graphAccessor.getGraphSchema(); + // Schemas are only ever appended to, so the object plus the two list sizes identify the + // cached label sets. Getting this wrong would silently drop hits, so keep it strict. + if (schema != cachedSchema + || cachedVertexLabels.size() != schema.getVertexSchemaList().size() + || cachedEdgeLabels.size() != schema.getEdgeSchemaList().size()) { + cachedVertexLabels = schema.getVertexSchemaList().stream() + .map(VertexSchema::getLabel).collect(Collectors.toSet()); + cachedEdgeLabels = schema.getEdgeSchemaList().stream() + .map(EdgeSchema::getLabel).collect(Collectors.toSet()); + cachedSchema = schema; + } + Set vertexLabels = cachedVertexLabels; + Set edgeLabels = cachedEdgeLabels; List result = new ArrayList<>(); for (ScoreDoc scoreDoc : scoreDocArray) { int docId = scoreDoc.doc; @@ -126,10 +205,6 @@ public List search(String key1, GraphAccessor graphAccessor) { } } - private void addItem() { - entityNum++; - } - public void close() { try { store.close(); @@ -145,10 +220,4 @@ public Directory getDirectory() { public Analyzer getAnalyzer() { return store.getAnalyzer(); } - - public IndexWriterConfig getConfig() { - return store.getConfig(); - } - - } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/ResidentSearchIndex.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/ResidentSearchIndex.java new file mode 100644 index 000000000..9d483c744 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/ResidentSearchIndex.java @@ -0,0 +1,304 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.ai.operator; + +import java.util.Collections; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Set; +import java.util.concurrent.locks.ReentrantReadWriteLock; +import org.apache.geaflow.ai.graph.GraphAccessor; +import org.apache.geaflow.ai.graph.GraphEntity; +import org.apache.geaflow.ai.graph.GraphVertex; +import org.apache.geaflow.ai.graph.VertexVersionWindow; +import org.apache.geaflow.ai.index.IndexStore; +import org.apache.geaflow.ai.index.vector.IVector; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * A graph scoped keyword index that survives across queries and is maintained incrementally. + * + *

Without it, every global search would build a throw-away {@link GraphSearchStore} from a full + * graph scan, paying index construction cost on the query path and discarding the result. + * + *

Maintenance model. Follows the standard inverted index approach rather than + * invalidate-and-rebuild: the index is built once, then writes are applied in place. An upsert maps + * to Lucene's update-by-term (tombstone plus insert) and a remove maps to delete-by-term (a bit in + * a per segment bitset). Cost is proportional to the change, not to graph size. Because both are + * keyed by {@code ModelUtils.getGraphEntityKey}, they are idempotent, so callers do not need to + * supply an exact delta. + * + *

Document set equivalence. The index contains exactly what a per-query global index + * would contain: every vertex whose {@link IndexStore} entry is non-empty. Edges are excluded, + * matching {@code searchWithGlobalGraph}, so recall is unchanged. + * + *

Version guard. Validity is tracked against {@link GraphAccessor#getVertexVersion()}. + * A write batch is applied in place only when its {@link VertexVersionWindow} proves it describes + * every vertex level change since the version this index last accepted. Anything else, including + * mutations made outside the reporting path (for example directly through + * {@code MemoryMutableGraph}), forces a rebuild rather than serving stale results. Edge writes do + * not invalidate anything, since the document set depends on vertices only. + * + *

Concurrency. Searches take the read lock and run concurrently; building, invalidating + * and applying writes take the write lock. + */ +public class ResidentSearchIndex { + + private static final Logger LOGGER = LoggerFactory.getLogger(ResidentSearchIndex.class); + + private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); + + private volatile GraphSearchStore store; + private volatile boolean globalIndexBuilt = false; + private volatile long builtVersion = GraphAccessor.VERSION_UNSUPPORTED; + + private volatile long buildCount = 0L; + private volatile long upsertCount = 0L; + private volatile long removeCount = 0L; + + /** + * Builds the full graph keyword index if it is absent or has gone stale. + */ + public void ensureGlobalIndex(GraphAccessor graphAccessor, IndexStore indexStore) { + lock.writeLock().lock(); + try { + ensureGlobalIndexLocked(graphAccessor, indexStore); + } finally { + lock.writeLock().unlock(); + } + } + + /** + * Ensures the index is valid and searches it. + * + *

The fast path holds only the read lock, so concurrent queries do not serialize. Validation + * and search happen under the same lock acquisition: with two separate calls a concurrent write + * could invalidate the index in between, leaving the query to fail on a missing index. + */ + public List searchWithIndex(GraphAccessor graphAccessor, IndexStore indexStore, + String query) { + lock.readLock().lock(); + try { + if (isUsableLocked(graphAccessor)) { + return store.search(query, graphAccessor); + } + } finally { + lock.readLock().unlock(); + } + lock.writeLock().lock(); + try { + ensureGlobalIndexLocked(graphAccessor, indexStore); + return store.search(query, graphAccessor); + } finally { + lock.writeLock().unlock(); + } + } + + /** + * Applies written entities to the index in place, without rebuilding it. + * + *

Safe for both new and rewritten entities. No-op before the first build: the entities will + * be picked up by it. + * + * @param window version range the batch claims to cover, see {@link VertexVersionWindow} + */ + public void onEntitiesUpserted(GraphAccessor graphAccessor, List entities, + IndexStore indexStore, VertexVersionWindow window) { + applyWrite(graphAccessor, entities, indexStore, false, window); + } + + /** + * Applies removed entities to the index in place, without rebuilding it. + */ + public void onEntitiesRemoved(GraphAccessor graphAccessor, List entities, + VertexVersionWindow window) { + applyWrite(graphAccessor, entities, null, true, window); + } + + private void applyWrite(GraphAccessor graphAccessor, List entities, + IndexStore indexStore, boolean removed, VertexVersionWindow window) { + if (entities == null || entities.isEmpty()) { + return; + } + lock.writeLock().lock(); + try { + if (!globalIndexBuilt) { + return; + } + if (window == null || !window.covers(builtVersion)) { + // The batch cannot be proven to describe everything that changed, so applying it + // would leave the index quietly missing whatever else happened. Rebuild instead. + LOGGER.info("Resident keyword index cannot accept a write batch, window: {}, " + + "accepted version: {}; rebuilding on next query", window, builtVersion); + invalidateLocked(); + return; + } + boolean changed = false; + for (GraphEntity entity : entities) { + // Only vertices are part of this index; edges merely advance the accepted version. + if (!(entity instanceof GraphVertex)) { + continue; + } + if (removed) { + store.removeEntity(entity); + removeCount++; + changed = true; + continue; + } + List vectors = indexStore.getEntityIndex(entity); + if (vectors == null || vectors.isEmpty()) { + // An entity without index content is not a document; drop any previous one. + // Delete by term is idempotent, so there is no need to track what was indexed. + store.removeEntity(entity); + changed = true; + continue; + } + store.upsertVertex((GraphVertex) entity, vectors); + upsertCount++; + changed = true; + } + if (changed) { + // One refresh per batch rather than per entity: each refresh opens a new segment. + store.refresh(); + } + builtVersion = window.getTo(); + } finally { + lock.writeLock().unlock(); + } + } + + /** + * Drops the index so that the next query rebuilds it. Needed for changes that are not expressed + * per entity, such as a schema change altering how every entity is verbalized. + */ + public void invalidate() { + lock.writeLock().lock(); + try { + invalidateLocked(); + } finally { + lock.writeLock().unlock(); + } + } + + public List search(String query, GraphAccessor graphAccessor) { + lock.readLock().lock(); + try { + if (store == null) { + return Collections.emptyList(); + } + return store.search(query, graphAccessor); + } finally { + lock.readLock().unlock(); + } + } + + public boolean isGlobalIndexBuilt() { + return globalIndexBuilt; + } + + /** + * Number of full graph builds. Stays at 1 for a workload whose writes all go through + * {@link #onEntitiesUpserted} / {@link #onEntitiesRemoved}; used by tests to prove the index is + * neither rebuilt per query nor per write. + */ + public long getBuildCount() { + return buildCount; + } + + public long getUpsertCount() { + return upsertCount; + } + + public long getRemoveCount() { + return removeCount; + } + + /** + * Number of documents currently in the index, read from Lucene rather than tracked separately. + */ + public int getIndexedEntityNum() { + // Write lock rather than read: reading the document count may have to open a reader, which + // mutates the store. + lock.writeLock().lock(); + try { + return store == null ? 0 : store.getDocNum(); + } finally { + lock.writeLock().unlock(); + } + } + + private boolean isUsableLocked(GraphAccessor graphAccessor) { + if (!globalIndexBuilt || store == null) { + return false; + } + long version = graphAccessor.getVertexVersion(); + return version != GraphAccessor.VERSION_UNSUPPORTED && version == builtVersion; + } + + private void ensureGlobalIndexLocked(GraphAccessor graphAccessor, IndexStore indexStore) { + long version = graphAccessor.getVertexVersion(); + if (globalIndexBuilt) { + if (version != GraphAccessor.VERSION_UNSUPPORTED && version == builtVersion) { + return; + } + // Either the graph changed outside the reporting path, or it cannot report changes at + // all. Both force a rebuild, which degrades to per-query rebuild rather than stale + // results. + invalidateLocked(); + } + final long start = System.currentTimeMillis(); + GraphSearchStore built = new GraphSearchStore(); + // Deduplication is only needed while scanning; unlike the index itself this set is not + // retained, so a resident index costs no per vertex heap of its own. + Set seen = new HashSet<>(); + for (Iterator it = graphAccessor.scanVertex(); it.hasNext(); ) { + GraphVertex vertex = it.next(); + List vectors = indexStore.getEntityIndex(vertex); + if (vectors == null || vectors.isEmpty() || !seen.add(vertex)) { + continue; + } + // Plain add during the build: the scan yields each vertex once, so no term lookup for + // duplicate removal is needed and build cost stays as low as possible. + built.indexVertex(vertex, vectors); + } + built.refresh(); + store = built; + globalIndexBuilt = true; + builtVersion = version; + buildCount++; + LOGGER.info("Built resident keyword index, entities: {}, vertexVersion: {}, cost: {} ms", + seen.size(), version, System.currentTimeMillis() - start); + } + + private void invalidateLocked() { + if (store != null) { + try { + store.close(); + } catch (Throwable e) { + LOGGER.warn("Ignore error on closing resident keyword index", e); + } + } + store = null; + globalIndexBuilt = false; + builtVersion = GraphAccessor.VERSION_UNSUPPORTED; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchConstants.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchConstants.java index 6a5ae0770..adc440f4e 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchConstants.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchConstants.java @@ -21,6 +21,12 @@ public class SearchConstants { + /** + * Non analyzed unique document key, holding {@code ModelUtils.getGraphEntityKey(entity)}. + * Having an exact term per entity is what lets Lucene express updates and deletes in place + * instead of forcing the whole index to be rebuilt. + */ + public static String KEY = "key"; public static String LABEL = "label"; public static String ID = "id"; public static String SRC = "src"; diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchStore.java index 81183ae47..3fee68c3d 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchStore.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchStore.java @@ -26,11 +26,12 @@ import org.apache.lucene.analysis.standard.StandardAnalyzer; import org.apache.lucene.document.Document; import org.apache.lucene.document.Field; +import org.apache.lucene.document.StringField; import org.apache.lucene.document.TextField; import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.Term; import org.apache.lucene.queryparser.classic.ParseException; import org.apache.lucene.queryparser.classic.QueryParser; import org.apache.lucene.search.IndexSearcher; @@ -38,56 +39,176 @@ import org.apache.lucene.store.ByteBuffersDirectory; import org.apache.lucene.store.Directory; +/** + * A lightweight in-memory Lucene index wrapper. + * + *

The store is designed to be long lived: writes are made visible to readers via + * {@link #refresh()}, a near-real-time reader reopen, instead of closing the writer. + * {@link #close()} is reserved for releasing the store for good. + * + *

Reader state is volatile so that searches may run concurrently with each other. Writes and + * {@link #refresh()} are not thread safe and must be serialized by the caller. + */ public class SearchStore { private final Directory directory = new ByteBuffersDirectory(); private final Analyzer analyzer = new StandardAnalyzer(); - private final IndexWriterConfig config = new IndexWriterConfig(analyzer); private IndexWriter writer; - private boolean writeStats = false; - private IndexReader reader; - private IndexSearcher searcher; - private boolean readStats = false; + private volatile boolean writeStats = false; + private volatile DirectoryReader reader; + private volatile IndexSearcher searcher; + private volatile boolean readStats = false; + private volatile boolean pendingWrite = false; + /** + * Whether {@link #reader} was opened from the writer, and can therefore be reopened from it. + */ + private volatile boolean nearRealTimeReader = false; public SearchStore() { } public void addDoc(Map kv) throws IOException { + addDoc(kv, null); + } + + /** + * Adds a document. The field named {@code exactField}, if present in {@code kv}, is indexed as a + * non analyzed {@link StringField} so that it can be used as a term for + * {@link #updateDoc} and {@link #deleteDoc}. + */ + public void addDoc(Map kv, String exactField) throws IOException { + initWriter(); + writer.addDocument(buildDoc(kv, exactField)); + pendingWrite = true; + } + + /** + * Replaces the document identified by {@code keyField = keyValue}, or adds it when absent. + * + *

Lucene implements this as a tombstone plus an insert within one call, so the cost is + * proportional to the change, not to the index size. Repeated calls with the same key are + * idempotent. + */ + public void updateDoc(String keyField, String keyValue, Map kv) throws IOException { initWriter(); + writer.updateDocument(new Term(keyField, keyValue), buildDoc(kv, keyField)); + pendingWrite = true; + } + + /** + * Marks the document identified by {@code keyField = keyValue} as deleted. Lucene only flips a + * bit in a per segment bitset; space is reclaimed later by segment merging. + */ + public void deleteDoc(String keyField, String keyValue) throws IOException { + initWriter(); + writer.deleteDocuments(new Term(keyField, keyValue)); + pendingWrite = true; + } + + private Document buildDoc(Map kv, String exactField) { Document doc = new Document(); for (Map.Entry entry : kv.entrySet()) { - doc.add(new TextField(entry.getKey(), entry.getValue(), Field.Store.YES)); + if (exactField != null && exactField.equals(entry.getKey())) { + doc.add(new StringField(entry.getKey(), entry.getValue(), Field.Store.YES)); + } else { + doc.add(new TextField(entry.getKey(), entry.getValue(), Field.Store.YES)); + } } - writer.addDocument(doc); + return doc; } - public TopDocs searchDoc(String field, String content) throws ParseException, IOException { + /** + * Makes pending writes searchable. Safe to call repeatedly; it is a no-op when nothing changed. + * + *

This replaces the previous pattern of calling {@link #close()} before searching, which + * forced the index to be discarded and rebuilt for every query. + * + *

When a writer exists the reader is opened from it (near real time) rather than from the + * directory. That deliberately avoids {@code IndexWriter#commit}: the directory is in memory, + * so a commit point buys no durability and only costs work on every write batch. + */ + public void refresh() throws IOException { + if (writeStats) { + if (readStats && nearRealTimeReader) { + DirectoryReader newReader = DirectoryReader.openIfChanged(reader, writer, true); + if (newReader != null) { + DirectoryReader old = reader; + reader = newReader; + searcher = new IndexSearcher(newReader); + old.close(); + } + } else { + final DirectoryReader old = readStats ? reader : null; + DirectoryReader newReader = DirectoryReader.open(writer); + reader = newReader; + searcher = new IndexSearcher(newReader); + readStats = true; + nearRealTimeReader = true; + if (old != null) { + old.close(); + } + } + pendingWrite = false; + return; + } if (!readStats) { reader = DirectoryReader.open(directory); searcher = new IndexSearcher(reader); readStats = true; + nearRealTimeReader = false; + return; + } + DirectoryReader newReader = DirectoryReader.openIfChanged(reader); + if (newReader != null) { + DirectoryReader old = reader; + reader = newReader; + searcher = new IndexSearcher(newReader); + old.close(); } + } + + /** + * Number of live documents currently visible to readers. Reflects the state as of the last + * {@link #refresh()}. + */ + public int numDocs() throws IOException { + ensureSearcher(); + return reader.numDocs(); + } + + public TopDocs searchDoc(String field, String content) throws ParseException, IOException { + ensureSearcher(); + IndexSearcher current = searcher; QueryParser parser = new QueryParser(field, analyzer); - return searcher.search(parser.parse(content), Constants.GRAPH_SEARCH_STORE_DEFAULT_TOPN); + return current.search(parser.parse(content), Constants.GRAPH_SEARCH_STORE_DEFAULT_TOPN); } public Document getDoc(int docId) { try { - if (!readStats) { - reader = DirectoryReader.open(directory); - searcher = new IndexSearcher(reader); - readStats = true; - } + ensureSearcher(); return searcher.doc(docId); } catch (Throwable e) { return null; } + } + /** + * Opens a reader if one is missing or stale. A no-op once a batch of writes has been followed by + * {@link #refresh()}, which is what keeps concurrent searches from mutating the store. + */ + private void ensureSearcher() throws IOException { + if (!readStats || pendingWrite) { + refresh(); + } } + /** + * Opens the writer on first use. A fresh {@link IndexWriterConfig} is built every time, because + * Lucene rejects reusing a config that has already been handed to a writer. + */ public void initWriter() throws IOException { if (!writeStats) { - writer = new IndexWriter(directory, config); + writer = new IndexWriter(directory, new IndexWriterConfig(analyzer)); writeStats = true; } } @@ -95,15 +216,19 @@ public void initWriter() throws IOException { public void close() throws IOException { if (writeStats) { writer.close(); + writer = null; writeStats = false; + pendingWrite = false; } if (readStats) { reader.close(); + reader = null; readStats = false; + nearRealTimeReader = false; + searcher = null; } } - public Directory getDirectory() { return directory; } @@ -111,8 +236,4 @@ public Directory getDirectory() { public Analyzer getAnalyzer() { return analyzer; } - - public IndexWriterConfig getConfig() { - return config; - } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SessionOperator.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SessionOperator.java index 4800c9341..55374a8d1 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SessionOperator.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SessionOperator.java @@ -36,9 +36,21 @@ public class SessionOperator implements SearchOperator { private final GraphAccessor graphAccessor; private final IndexStore indexStore; + /** + * Graph scoped keyword index reused across queries. When {@code null} the operator falls back + * to building a throw-away index per query, which is only kept for backward compatibility and + * for equivalence testing. + */ + private final ResidentSearchIndex residentIndex; + public SessionOperator(GraphAccessor accessor, IndexStore store) { + this(accessor, store, null); + } + + public SessionOperator(GraphAccessor accessor, IndexStore store, ResidentSearchIndex residentIndex) { this.graphAccessor = Objects.requireNonNull(accessor); this.indexStore = Objects.requireNonNull(store); + this.residentIndex = residentIndex; } @Override @@ -55,8 +67,8 @@ public List apply(List subGraphList, VectorSearch search) { contents.add(v.toString()); } String query = String.join(SearchConstants.DELIMITER, contents); - List globalResults = searchWithGlobalGraph(query); if (subGraphList == null || subGraphList.isEmpty()) { + List globalResults = searchWithGlobalGraph(query); List startVertices = new ArrayList<>(); for (GraphEntity resEntity : globalResults) { if (resEntity instanceof GraphVertex) { @@ -79,10 +91,11 @@ public List apply(List subGraphList, VectorSearch search) { extendEntityIndexMap.put(extendEntity, entityIndex); } } - //recall compute + //recall compute, the candidate set is bounded by the subgraph expansion GraphSearchStore searchStore = initSearchStore(extendEntityIndexMap); - searchStore.close(); + searchStore.refresh(); List matchEntities = searchStore.search(query, graphAccessor); + closeQuietly(searchStore); Set matchEntitiesSet = new HashSet<>(matchEntities); //Apply to subgraph @@ -113,6 +126,17 @@ private List getSubgraphExpand(SubGraph subGraph) { } private List searchWithGlobalGraph(String query) { + if (residentIndex != null) { + return residentIndex.searchWithIndex(graphAccessor, indexStore, query); + } + return searchWithGlobalGraphByRebuild(query); + } + + /** + * Legacy behaviour: scan the whole graph, build a throw-away index, search it, drop it. + * Retained as the reference implementation for equivalence tests. + */ + List searchWithGlobalGraphByRebuild(String query) { Map> entityIndexMap = new HashMap<>(); Iterator vertexIterator = graphAccessor.scanVertex(); while (vertexIterator.hasNext()) { @@ -123,8 +147,10 @@ private List searchWithGlobalGraph(String query) { } //recall compute GraphSearchStore searchStore = initSearchStore(entityIndexMap); - searchStore.close(); - return searchStore.search(query, graphAccessor); + searchStore.refresh(); + List result = searchStore.search(query, graphAccessor); + closeQuietly(searchStore); + return result; } private GraphSearchStore initSearchStore(Map> entityIndexMap) { @@ -140,4 +166,12 @@ private GraphSearchStore initSearchStore(Map> entityI } return searchStore; } + + private void closeQuietly(GraphSearchStore searchStore) { + try { + searchStore.close(); + } catch (Throwable ignored) { + // A throw-away store leaking is not worth failing the query for. + } + } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/verbalization/SubgraphSemanticPromptFunction.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/verbalization/SubgraphSemanticPromptFunction.java index 0b3407c06..f14c0b94b 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/verbalization/SubgraphSemanticPromptFunction.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/verbalization/SubgraphSemanticPromptFunction.java @@ -38,6 +38,16 @@ public SubgraphSemanticPromptFunction(GraphAccessor accessor) { this.graphAccessor = Objects.requireNonNull(accessor); } + @Override + public long getSourceVersion() { + return graphAccessor.getGraphVersion(); + } + + @Override + public long getSourceVertexVersion() { + return graphAccessor.getVertexVersion(); + } + @Override public String verbalize(SubGraph subGraph) { if (subGraph == null || subGraph.getGraphEntityList().isEmpty()) { diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/verbalization/VerbalizationFunction.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/verbalization/VerbalizationFunction.java index 6e90bc17a..b1cef0fb6 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/verbalization/VerbalizationFunction.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/verbalization/VerbalizationFunction.java @@ -20,11 +20,34 @@ package org.apache.geaflow.ai.verbalization; import java.util.List; +import org.apache.geaflow.ai.graph.GraphAccessor; import org.apache.geaflow.ai.graph.GraphEntity; import org.apache.geaflow.ai.subgraph.SubGraph; public interface VerbalizationFunction { + /** + * Version of the data this function renders, see {@link GraphAccessor#getGraphVersion()}. + * Consumers memoizing verbalizations use it to drop stale entries. + * + * @return current source version, or {@link GraphAccessor#VERSION_UNSUPPORTED} when the + * function cannot tell whether its source changed, in which case results must not be cached + */ + default long getSourceVersion() { + return GraphAccessor.VERSION_UNSUPPORTED; + } + + /** + * Like {@link #getSourceVersion()} but only advanced by changes that can affect how a vertex is + * rendered, see {@link GraphAccessor#getVertexVersion()}. Consumers memoizing vertex + * verbalizations watch this one so that edge writes do not invalidate them. + * + * @return current vertex source version, defaults to {@link #getSourceVersion()} + */ + default long getSourceVertexVersion() { + return getSourceVersion(); + } + String verbalize(SubGraph subGraph); List verbalize(GraphEntity entity); diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/operator/EmbeddingCandidateSetTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/operator/EmbeddingCandidateSetTest.java new file mode 100644 index 000000000..7467e7229 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/operator/EmbeddingCandidateSetTest.java @@ -0,0 +1,153 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.ai.operator; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.apache.geaflow.ai.graph.GraphEntity; +import org.apache.geaflow.ai.graph.GraphVertex; +import org.apache.geaflow.ai.graph.LocalMemoryGraphAccessor; +import org.apache.geaflow.ai.index.IndexStore; +import org.apache.geaflow.ai.index.vector.EmbeddingVector; +import org.apache.geaflow.ai.index.vector.IVector; +import org.apache.geaflow.ai.graph.io.EntityGroup; +import org.apache.geaflow.ai.graph.io.GraphSchema; +import org.apache.geaflow.ai.graph.io.MemoryGraph; +import org.apache.geaflow.ai.graph.io.Vertex; +import org.apache.geaflow.ai.graph.io.VertexGroup; +import org.apache.geaflow.ai.graph.io.VertexSchema; +import org.apache.geaflow.ai.search.VectorSearch; +import org.apache.geaflow.ai.subgraph.SubGraph; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** + * {@link EmbeddingOperator} may collect its global candidate set either by scanning the graph or by + * asking the index store what it holds. Both must recall exactly the same entities. + */ +public class EmbeddingCandidateSetTest { + + private static final String LABEL = "doc"; + private static final int VERTEX_NUM = 500; + private static final int DIM = 8; + + /** Index store backed by a fixed map, optionally able to enumerate its own content. */ + private static class MapIndexStore implements IndexStore { + + private final Map> data; + private final boolean enumerable; + + MapIndexStore(Map> data, boolean enumerable) { + this.data = data; + this.enumerable = enumerable; + } + + @Override + public List getEntityIndex(GraphEntity entity) { + List vectors = data.get(entity); + return vectors == null ? Collections.emptyList() : vectors; + } + + @Override + public Collection getIndexedEntities() { + return enumerable ? data.keySet() : null; + } + } + + private LocalMemoryGraphAccessor buildGraph() { + GraphSchema schema = new GraphSchema(); + VertexSchema vertexSchema = new VertexSchema(LABEL, "id", Collections.singletonList("text")); + schema.setName("embedding_graph"); + schema.addVertex(vertexSchema); + List vertices = new ArrayList<>(VERTEX_NUM); + for (int i = 0; i < VERTEX_NUM; i++) { + vertices.add(new Vertex(LABEL, "id" + i, Collections.singletonList("text" + i))); + } + Map entities = new HashMap<>(); + entities.put(LABEL, new VertexGroup(vertexSchema, vertices)); + return new LocalMemoryGraphAccessor(new MemoryGraph(schema, entities)); + } + + /** + * Builds embeddings for a deterministic subset of vertices, so the test also covers vertices + * that carry no embedding at all and vertices absent from the store entirely. + */ + private Map> buildIndexData(LocalMemoryGraphAccessor accessor) { + Map> data = new LinkedHashMap<>(); + for (int i = 0; i < VERTEX_NUM; i++) { + if (i % 3 == 2) { + // Not indexed at all. + continue; + } + GraphVertex vertex = accessor.getVertex(LABEL, "id" + i); + Assertions.assertNotNull(vertex); + if (i % 3 == 1) { + // Present but with no vector, must be skipped by recall. + data.put(vertex, Collections.emptyList()); + continue; + } + double[] vec = new double[DIM]; + for (int d = 0; d < DIM; d++) { + vec[d] = Math.sin(i + d) + 1.5; + } + data.put(vertex, Collections.singletonList((IVector) new EmbeddingVector(vec))); + } + return data; + } + + private static List recall(LocalMemoryGraphAccessor accessor, IndexStore store, + double[] query) { + VectorSearch search = new VectorSearch(null, "session"); + search.addVector(new EmbeddingVector(query)); + List result = new EmbeddingOperator(accessor, store).apply(null, search); + List ids = new ArrayList<>(); + for (SubGraph subGraph : result) { + for (GraphEntity entity : subGraph.getGraphEntityList()) { + ids.add(((GraphVertex) entity).getVertex().getId()); + } + } + return ids; + } + + @Test + public void testEnumeratedCandidatesMatchGraphScan() { + LocalMemoryGraphAccessor accessor = buildGraph(); + Map> data = buildIndexData(accessor); + IndexStore scanStore = new MapIndexStore(data, false); + IndexStore enumerableStore = new MapIndexStore(data, true); + + for (int q = 0; q < 5; q++) { + double[] query = new double[DIM]; + for (int d = 0; d < DIM; d++) { + query[d] = Math.cos(q + d) + 1.5; + } + List byScan = recall(accessor, scanStore, query); + List byEnumeration = recall(accessor, enumerableStore, query); + Assertions.assertFalse(byScan.isEmpty(), "query " + q + " recalled nothing"); + Assertions.assertEquals(byScan, byEnumeration, + "enumerating the index store must recall the same entities in the same order"); + } + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/operator/ResidentSearchIndexTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/operator/ResidentSearchIndexTest.java new file mode 100644 index 000000000..06e84fc13 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/operator/ResidentSearchIndexTest.java @@ -0,0 +1,596 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.ai.operator; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import org.apache.geaflow.ai.graph.GraphEntity; +import org.apache.geaflow.ai.graph.GraphVertex; +import org.apache.geaflow.ai.graph.LocalMemoryGraphAccessor; +import org.apache.geaflow.ai.graph.VertexVersionWindow; +import org.apache.geaflow.ai.graph.io.Edge; +import org.apache.geaflow.ai.graph.io.EdgeSchema; +import org.apache.geaflow.ai.graph.io.EntityGroup; +import org.apache.geaflow.ai.graph.io.GraphSchema; +import org.apache.geaflow.ai.graph.io.MemoryGraph; +import org.apache.geaflow.ai.graph.io.Vertex; +import org.apache.geaflow.ai.graph.io.VertexGroup; +import org.apache.geaflow.ai.graph.io.VertexSchema; +import org.apache.geaflow.ai.index.EntityAttributeIndexStore; +import org.apache.geaflow.ai.verbalization.SubgraphSemanticPromptFunction; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Verifies that reusing a resident keyword index is behaviourally equivalent to rebuilding a + * throw-away index per query, and measures the difference. + * + *

Result comparison is done on sets rather than lists on purpose: the legacy path feeds Lucene + * from a {@link HashMap} iteration, so its document order, and therefore its tie break order, was + * never deterministic to begin with. For queries whose match count stays under the Lucene top N + * limit the returned sets must be identical. + */ +public class ResidentSearchIndexTest { + + private static final Logger LOGGER = LoggerFactory.getLogger(ResidentSearchIndexTest.class); + + private static final String LABEL = "doc"; + private static final String EDGE_LABEL = "rel"; + private static final int VERTEX_NUM = 10000; + /** 10000 / 1000 = 10 matches per group token, safely below the Lucene top N of 30. */ + private static final int GROUP_NUM = 1000; + private static final int QUERY_NUM = 10; + + private LocalMemoryGraphAccessor buildGraph(int vertexNum) { + GraphSchema schema = new GraphSchema(); + VertexSchema vertexSchema = new VertexSchema(LABEL, "id", Collections.singletonList("text")); + schema.setName("perf_graph"); + schema.addVertex(vertexSchema); + + List vertices = new ArrayList<>(vertexNum); + for (int i = 0; i < vertexNum; i++) { + vertices.add(newVertex(i)); + } + Map entities = new HashMap<>(); + entities.put(LABEL, new VertexGroup(vertexSchema, vertices)); + return new LocalMemoryGraphAccessor(new MemoryGraph(schema, entities)); + } + + private Vertex newVertex(int i) { + String text = "uniq" + i + " grp" + (i % GROUP_NUM) + " topic" + (i % 17) + + " some filler content for verbalization cost"; + return new Vertex(LABEL, "id" + i, Collections.singletonList(text)); + } + + private EntityAttributeIndexStore newIndexStore(LocalMemoryGraphAccessor accessor) { + EntityAttributeIndexStore store = new EntityAttributeIndexStore(); + store.initStore(new SubgraphSemanticPromptFunction(accessor)); + return store; + } + + private List buildQueries() { + List queries = new ArrayList<>(QUERY_NUM); + for (int i = 0; i < QUERY_NUM; i++) { + queries.add("grp" + (i * 37 % GROUP_NUM)); + } + return queries; + } + + private static String ms(long nanos) { + return String.format("%.2f", nanos / 1_000_000.0); + } + + private static Set idsOf(List entities) { + Set ids = new HashSet<>(); + for (GraphEntity entity : entities) { + Assertions.assertTrue(entity instanceof GraphVertex); + ids.add(((GraphVertex) entity).getVertex().getId()); + } + return ids; + } + + @Test + public void testResidentIndexIsEquivalentAndFaster() { + LocalMemoryGraphAccessor accessor = buildGraph(VERTEX_NUM); + List queries = buildQueries(); + + // Warm up the JIT and the Lucene classes, otherwise the first measured configuration pays + // for class loading and interpretation and the comparison is meaningless. + EntityAttributeIndexStore warmupStore = newIndexStore(accessor); + SessionOperator warmupOperator = new SessionOperator(accessor, warmupStore); + for (int i = 0; i < 3; i++) { + warmupStore.invalidateCache(); + warmupOperator.searchWithGlobalGraphByRebuild(queries.get(0)); + } + + // Baseline: rebuild the index per query, and drop the verbalization cache each time so the + // measurement reflects the original behaviour. + EntityAttributeIndexStore coldStore = newIndexStore(accessor); + SessionOperator coldOperator = new SessionOperator(accessor, coldStore); + Map> baselineResults = new HashMap<>(); + long baselineCost = 0L; + for (String query : queries) { + coldStore.invalidateCache(); + long start = System.nanoTime(); + List result = coldOperator.searchWithGlobalGraphByRebuild(query); + baselineCost += System.nanoTime() - start; + baselineResults.put(query, idsOf(result)); + } + + // Rebuild per query, but keep the verbalization cache. + EntityAttributeIndexStore cachedStore = newIndexStore(accessor); + SessionOperator cachedOperator = new SessionOperator(accessor, cachedStore); + long cachedCost = 0L; + for (String query : queries) { + long start = System.nanoTime(); + List result = cachedOperator.searchWithGlobalGraphByRebuild(query); + cachedCost += System.nanoTime() - start; + Assertions.assertEquals(baselineResults.get(query), idsOf(result), + "verbalization cache must not change recall for query " + query); + } + + // Resident index reused across queries. + EntityAttributeIndexStore residentStore = newIndexStore(accessor); + ResidentSearchIndex residentIndex = new ResidentSearchIndex(); + long residentCost = 0L; + long residentFirstCost = 0L; + long residentSteadyCost = 0L; + for (int i = 0; i < queries.size(); i++) { + String query = queries.get(i); + long start = System.nanoTime(); + residentIndex.ensureGlobalIndex(accessor, residentStore); + List result = residentIndex.search(query, accessor); + long cost = System.nanoTime() - start; + residentCost += cost; + if (i == 0) { + residentFirstCost = cost; + } else { + residentSteadyCost += cost; + } + Assertions.assertEquals(baselineResults.get(query), idsOf(result), + "resident index must not change recall for query " + query); + } + + // The whole point: the full graph index is built once, not once per query. + Assertions.assertEquals(1L, residentIndex.getBuildCount()); + Assertions.assertEquals(VERTEX_NUM, residentIndex.getIndexedEntityNum()); + + // Every query must have actually matched something, otherwise the comparison is vacuous. + for (String query : queries) { + Assertions.assertFalse(baselineResults.get(query).isEmpty(), + "query matched nothing: " + query); + } + + LOGGER.info("=== retrieval cost, vertices={}, queries={} ===", VERTEX_NUM, queries.size()); + LOGGER.info("[A] rebuild per query, no verbalization cache : total {} ms, avg {} ms/query", + ms(baselineCost), ms(baselineCost / queries.size())); + LOGGER.info("[B] rebuild per query, verbalization cached : total {} ms, avg {} ms/query", + ms(cachedCost), ms(cachedCost / queries.size())); + LOGGER.info("[C] resident index : total {} ms, " + + "first query (includes one time build) {} ms, steady state avg {} ms/query", + ms(residentCost), ms(residentFirstCost), ms(residentSteadyCost / (queries.size() - 1))); + LOGGER.info("verbalization cache in [B]: hit={}, miss={}, size={}", + cachedStore.getCacheHit(), cachedStore.getCacheMiss(), cachedStore.getCacheSize()); + + Assertions.assertTrue(residentCost < baselineCost, + "resident index should be cheaper than rebuilding per query, baseline=" + baselineCost + + "ns resident=" + residentCost + "ns"); + } + + /** + * Shows the shape of the problem: the rebuild path grows with the graph, the resident path does + * not. Only logged, not asserted, so the test stays stable on shared CI machines. + */ + @Test + public void testSteadyStateCostDoesNotGrowWithGraphSize() { + for (int vertexNum : new int[] {5000, 20000}) { + LocalMemoryGraphAccessor accessor = buildGraph(vertexNum); + List queries = buildQueries(); + + EntityAttributeIndexStore coldStore = newIndexStore(accessor); + SessionOperator coldOperator = new SessionOperator(accessor, coldStore); + long rebuildCost = 0L; + for (String query : queries) { + coldStore.invalidateCache(); + long start = System.nanoTime(); + coldOperator.searchWithGlobalGraphByRebuild(query); + rebuildCost += System.nanoTime() - start; + } + + EntityAttributeIndexStore residentStore = newIndexStore(accessor); + ResidentSearchIndex residentIndex = new ResidentSearchIndex(); + residentIndex.ensureGlobalIndex(accessor, residentStore); + long steadyCost = 0L; + for (String query : queries) { + long start = System.nanoTime(); + residentIndex.search(query, accessor); + steadyCost += System.nanoTime() - start; + } + Assertions.assertEquals(1L, residentIndex.getBuildCount()); + + LOGGER.info("vertices={} : rebuild avg {} ms/query, resident steady avg {} ms/query", + vertexNum, ms(rebuildCost / queries.size()), ms(steadyCost / queries.size())); + } + } + + @Test + public void testVerbalizationCacheIsUsed() { + LocalMemoryGraphAccessor accessor = buildGraph(100); + EntityAttributeIndexStore store = newIndexStore(accessor); + GraphVertex vertex = accessor.getVertex(LABEL, "id7"); + Assertions.assertNotNull(vertex); + + Assertions.assertEquals(store.getEntityIndex(vertex).toString(), + store.getEntityIndex(vertex).toString()); + Assertions.assertEquals(1L, store.getCacheMiss()); + Assertions.assertEquals(1L, store.getCacheHit()); + + store.invalidateCache(vertex); + store.getEntityIndex(vertex); + Assertions.assertEquals(2L, store.getCacheMiss()); + } + + /** + * The verbalization cache takes no lock, so concurrent lookups must still be consistent: every + * caller sees the same content, every call is counted exactly once, and nothing is lost or + * duplicated in the map. + */ + @Test + public void testConcurrentVerbalizationLookupsAreConsistent() throws Exception { + int entityNum = 200; + int threads = 8; + int roundsPerThread = 50; + LocalMemoryGraphAccessor accessor = buildGraph(entityNum); + EntityAttributeIndexStore store = newIndexStore(accessor); + + List vertices = new ArrayList<>(entityNum); + Map expected = new HashMap<>(); + for (int i = 0; i < entityNum; i++) { + GraphVertex vertex = accessor.getVertex(LABEL, "id" + i); + vertices.add(vertex); + expected.put("id" + i, store.getEntityIndex(vertex).toString()); + } + long baselineCalls = store.getCacheHit() + store.getCacheMiss(); + + ExecutorService pool = Executors.newFixedThreadPool(threads); + CountDownLatch start = new CountDownLatch(1); + List> futures = new ArrayList<>(threads); + for (int t = 0; t < threads; t++) { + futures.add(pool.submit(() -> { + start.await(); + for (int r = 0; r < roundsPerThread; r++) { + for (GraphVertex vertex : vertices) { + String id = vertex.getVertex().getId(); + Assertions.assertEquals(expected.get(id), + store.getEntityIndex(vertex).toString(), + "concurrent lookup returned different content for " + id); + } + } + return null; + })); + } + start.countDown(); + try { + for (Future future : futures) { + future.get(60, TimeUnit.SECONDS); + } + } finally { + pool.shutdownNow(); + } + + long calls = store.getCacheHit() + store.getCacheMiss() - baselineCalls; + Assertions.assertEquals((long) threads * roundsPerThread * entityNum, calls, + "every lookup must be counted exactly once"); + Assertions.assertEquals(entityNum, store.getCacheSize(), + "the graph did not change, so there must be exactly one entry per entity"); + } + + @Test + public void testInsertIsSearchableWithoutRebuild() { + LocalMemoryGraphAccessor accessor = buildGraph(200); + EntityAttributeIndexStore store = newIndexStore(accessor); + ResidentSearchIndex residentIndex = new ResidentSearchIndex(); + residentIndex.ensureGlobalIndex(accessor, store); + Assertions.assertEquals(1L, residentIndex.getBuildCount()); + Assertions.assertTrue(residentIndex.search("zebrafish", accessor).isEmpty()); + + Vertex fresh = new Vertex(LABEL, "id-fresh", + Collections.singletonList("zebrafish appears only here")); + VertexVersionWindow window = write(accessor, () -> accessor.getMutableGraph().addVertex(fresh)); + residentIndex.onEntitiesUpserted(accessor, entities(fresh), store, window); + + Assertions.assertEquals(Collections.singleton("id-fresh"), + idsOf(residentIndex.search("zebrafish", accessor))); + Assertions.assertEquals(1L, residentIndex.getUpsertCount()); + Assertions.assertEquals(201, residentIndex.getIndexedEntityNum()); + assertNoRebuild(residentIndex, accessor, store); + } + + @Test + public void testUpdateIsAppliedInPlaceWithoutRebuild() { + LocalMemoryGraphAccessor accessor = buildGraph(50); + EntityAttributeIndexStore store = newIndexStore(accessor); + ResidentSearchIndex residentIndex = new ResidentSearchIndex(); + residentIndex.ensureGlobalIndex(accessor, store); + Assertions.assertEquals(Collections.singleton("id7"), + idsOf(residentIndex.search("uniq7", accessor))); + + Vertex updated = new Vertex(LABEL, "id7", Collections.singletonList("narwhal now")); + VertexVersionWindow window = write(accessor, + () -> accessor.getMutableGraph().updateVertex(updated)); + residentIndex.onEntitiesUpserted(accessor, entities(updated), store, window); + + // New content is visible, the superseded document is gone, and the doc count is unchanged. + Assertions.assertEquals(Collections.singleton("id7"), + idsOf(residentIndex.search("narwhal", accessor))); + Assertions.assertTrue(residentIndex.search("uniq7", accessor).isEmpty(), + "the replaced document must no longer be searchable"); + Assertions.assertEquals(50, residentIndex.getIndexedEntityNum()); + assertNoRebuild(residentIndex, accessor, store); + } + + @Test + public void testDeleteIsAppliedInPlaceWithoutRebuild() { + LocalMemoryGraphAccessor accessor = buildGraph(50); + EntityAttributeIndexStore store = newIndexStore(accessor); + ResidentSearchIndex residentIndex = new ResidentSearchIndex(); + residentIndex.ensureGlobalIndex(accessor, store); + Assertions.assertEquals(Collections.singleton("id9"), + idsOf(residentIndex.search("uniq9", accessor))); + + Vertex removed = accessor.getVertex(LABEL, "id9").getVertex(); + VertexVersionWindow window = write(accessor, + () -> accessor.getMutableGraph().removeVertex(LABEL, "id9")); + residentIndex.onEntitiesRemoved(accessor, entities(removed), window); + + Assertions.assertTrue(residentIndex.search("uniq9", accessor).isEmpty(), + "the deleted document must no longer be searchable"); + Assertions.assertEquals(1L, residentIndex.getRemoveCount()); + Assertions.assertEquals(49, residentIndex.getIndexedEntityNum()); + assertNoRebuild(residentIndex, accessor, store); + } + + @Test + public void testUpsertIsIdempotent() { + LocalMemoryGraphAccessor accessor = buildGraph(20); + EntityAttributeIndexStore store = newIndexStore(accessor); + ResidentSearchIndex residentIndex = new ResidentSearchIndex(); + residentIndex.ensureGlobalIndex(accessor, store); + + Vertex existing = accessor.getVertex(LABEL, "id3").getVertex(); + for (int i = 0; i < 3; i++) { + // Replaying a batch that changed nothing: the window is empty but still valid. + residentIndex.onEntitiesUpserted(accessor, entities(existing), store, + write(accessor, () -> { })); + } + + // Replaying the same write must not duplicate the document nor trigger a rebuild. + Assertions.assertEquals(Collections.singleton("id3"), + idsOf(residentIndex.search("uniq3", accessor))); + Assertions.assertEquals(20, residentIndex.getIndexedEntityNum()); + assertNoRebuild(residentIndex, accessor, store); + } + + @Test + public void testMutationOutsideTheIndexForcesRebuild() { + LocalMemoryGraphAccessor accessor = buildGraph(20); + EntityAttributeIndexStore store = newIndexStore(accessor); + ResidentSearchIndex residentIndex = new ResidentSearchIndex(); + residentIndex.ensureGlobalIndex(accessor, store); + Assertions.assertEquals(1L, residentIndex.getBuildCount()); + + // Graph changed without notifying the index; only the version guard can catch this. + accessor.getMutableGraph().addVertex(new Vertex(LABEL, "id-hidden", + Collections.singletonList("okapi appears only here"))); + + Assertions.assertEquals(Collections.singleton("id-hidden"), + idsOf(residentIndex.searchWithIndex(accessor, store, "okapi"))); + Assertions.assertEquals(2L, residentIndex.getBuildCount(), + "the version guard must force a rebuild for unnotified mutations"); + } + + /** + * The version guard must survive a later reported write. Reading the current version after + * applying a batch would accept the unreported change as already applied, and the missing + * document would never come back. + */ + @Test + public void testUnreportedMutationIsNotSwallowedByALaterReportedWrite() { + LocalMemoryGraphAccessor accessor = buildGraph(20); + EntityAttributeIndexStore store = newIndexStore(accessor); + ResidentSearchIndex residentIndex = new ResidentSearchIndex(); + residentIndex.ensureGlobalIndex(accessor, store); + Assertions.assertEquals(1L, residentIndex.getBuildCount()); + + // Graph changed without telling the index. + accessor.getMutableGraph().addVertex(new Vertex(LABEL, "id-hidden", + Collections.singletonList("okapi appears only here"))); + + // An unrelated write that does get reported, covering only its own version range. + Vertex other = new Vertex(LABEL, "id-other", Collections.singletonList("lemur here")); + VertexVersionWindow window = write(accessor, + () -> accessor.getMutableGraph().addVertex(other)); + residentIndex.onEntitiesUpserted(accessor, entities(other), store, window); + + Assertions.assertEquals(Collections.singleton("id-hidden"), + idsOf(residentIndex.searchWithIndex(accessor, store, "okapi")), + "the unreported vertex must still be found"); + Assertions.assertEquals(Collections.singleton("id-other"), + idsOf(residentIndex.searchWithIndex(accessor, store, "lemur"))); + Assertions.assertEquals(2L, residentIndex.getBuildCount(), + "a batch that cannot be proven complete must fall back to a rebuild"); + Assertions.assertEquals(22, residentIndex.getIndexedEntityNum()); + } + + /** + * A vertex verbalization depends on the vertex and the schema, not on edges, so writing edges + * must not cost the whole memoized set. Sharing one version for the entire cache would wipe it. + */ + @Test + public void testEdgeWriteKeepsMemoizedVertexVerbalizations() { + LocalMemoryGraphAccessor accessor = buildGraph(20); + EntityAttributeIndexStore store = newIndexStore(accessor); + GraphVertex vertex = accessor.getVertex(LABEL, "id3"); + store.getEntityIndex(vertex); + store.getEntityIndex(vertex); + Assertions.assertEquals(1L, store.getCacheMiss()); + Assertions.assertEquals(1L, store.getCacheHit()); + + accessor.getMutableGraph().addEdgeSchema( + new EdgeSchema(EDGE_LABEL, "srcId", "dstId", Collections.singletonList("rel"))); + for (int i = 0; i < 19; i++) { + accessor.getMutableGraph().addEdge( + new Edge(EDGE_LABEL, "id" + i, "id" + (i + 1), Collections.singletonList("linked"))); + } + + store.getEntityIndex(vertex); + Assertions.assertEquals(1L, store.getCacheMiss(), "edge writes must not evict vertex entries"); + Assertions.assertEquals(2L, store.getCacheHit()); + + // A write to the vertex itself does evict it, and only it. + accessor.getMutableGraph().updateVertex( + new Vertex(LABEL, "id3", Collections.singletonList("rewritten"))); + store.getEntityIndex(vertex); + Assertions.assertEquals(2L, store.getCacheMiss()); + Assertions.assertEquals(1, store.getCacheSize()); + } + + @Test + public void testEdgeWriteDoesNotInvalidateVertexIndex() { + LocalMemoryGraphAccessor accessor = buildGraph(20); + EntityAttributeIndexStore store = newIndexStore(accessor); + ResidentSearchIndex residentIndex = new ResidentSearchIndex(); + residentIndex.ensureGlobalIndex(accessor, store); + + accessor.getMutableGraph().addEdgeSchema( + new EdgeSchema(EDGE_LABEL, "srcId", "dstId", Collections.singletonList("rel"))); + residentIndex.ensureGlobalIndex(accessor, store); + // The schema change is a vertex level change, so a rebuild is expected here. + long afterSchema = residentIndex.getBuildCount(); + + accessor.getMutableGraph().addEdge( + new Edge(EDGE_LABEL, "id1", "id2", Collections.singletonList("linked"))); + residentIndex.ensureGlobalIndex(accessor, store); + Assertions.assertEquals(afterSchema, residentIndex.getBuildCount(), + "an edge write must not invalidate a vertex only index"); + } + + /** + * The write heavy scenario in place maintenance exists for: writes interleaved with queries. + * With invalidate-on-write every query after a write pays a full rebuild; with in place + * maintenance the index is built once and never again. + */ + @Test + public void testInterleavedWritesAndQueriesNeverRebuild() { + // Kept smaller than the read benchmark: the reference configuration rebuilds the whole + // index on every round, so its cost is rounds x O(V). + int vertexNum = 5000; + int rounds = 40; + + LocalMemoryGraphAccessor inPlaceAccessor = buildGraph(vertexNum); + EntityAttributeIndexStore inPlaceStore = newIndexStore(inPlaceAccessor); + ResidentSearchIndex inPlaceIndex = new ResidentSearchIndex(); + inPlaceIndex.ensureGlobalIndex(inPlaceAccessor, inPlaceStore); + + LocalMemoryGraphAccessor invalidateAccessor = buildGraph(vertexNum); + EntityAttributeIndexStore invalidateStore = newIndexStore(invalidateAccessor); + ResidentSearchIndex invalidateIndex = new ResidentSearchIndex(); + invalidateIndex.ensureGlobalIndex(invalidateAccessor, invalidateStore); + + long inPlaceCost = 0L; + long invalidateCost = 0L; + for (int i = 0; i < rounds; i++) { + Vertex fresh = new Vertex(LABEL, "id-new" + i, + Collections.singletonList("grp" + (i % GROUP_NUM) + " freshly written " + i)); + String query = "grp" + (i % GROUP_NUM); + + VertexVersionWindow window = write(inPlaceAccessor, + () -> inPlaceAccessor.getMutableGraph().addVertex(fresh)); + long start = System.nanoTime(); + inPlaceIndex.onEntitiesUpserted(inPlaceAccessor, entities(fresh), inPlaceStore, window); + List inPlaceHit = inPlaceIndex.searchWithIndex(inPlaceAccessor, + inPlaceStore, query); + inPlaceCost += System.nanoTime() - start; + + // Reference behaviour: drop the index on every write and rebuild lazily. + invalidateAccessor.getMutableGraph().addVertex(fresh); + start = System.nanoTime(); + invalidateIndex.invalidate(); + List invalidateHit = invalidateIndex.searchWithIndex(invalidateAccessor, + invalidateStore, query); + invalidateCost += System.nanoTime() - start; + + Assertions.assertEquals(idsOf(invalidateHit), idsOf(inPlaceHit), + "in place maintenance must recall the same entities as rebuilding, round " + i); + Assertions.assertTrue(idsOf(inPlaceHit).contains("id-new" + i), + "the just written vertex must be visible, round " + i); + } + + Assertions.assertEquals(1L, inPlaceIndex.getBuildCount(), + "in place maintenance must never rebuild"); + Assertions.assertEquals(1L + rounds, invalidateIndex.getBuildCount()); + Assertions.assertEquals(rounds, inPlaceIndex.getUpsertCount()); + Assertions.assertEquals(vertexNum + rounds, inPlaceIndex.getIndexedEntityNum()); + + LOGGER.info("=== interleaved write + query, vertices={}, rounds={} ===", vertexNum, rounds); + LOGGER.info("[D] invalidate on write : total {} ms, avg {} ms/round, builds {}", + ms(invalidateCost), ms(invalidateCost / rounds), invalidateIndex.getBuildCount()); + LOGGER.info("[E] in place maintenance: total {} ms, avg {} ms/round, builds {}", + ms(inPlaceCost), ms(inPlaceCost / rounds), inPlaceIndex.getBuildCount()); + } + + /** + * Runs a batch of graph writes inside a version window, the way a caller reporting the batch to + * the index is expected to. + */ + private static VertexVersionWindow write(LocalMemoryGraphAccessor accessor, Runnable writes) { + VertexVersionWindow window = VertexVersionWindow.open(accessor); + writes.run(); + return window.seal(); + } + + private static List entities(Vertex... vertices) { + List list = new ArrayList<>(vertices.length); + for (Vertex vertex : vertices) { + list.add(new GraphVertex(vertex)); + } + return list; + } + + private static void assertNoRebuild(ResidentSearchIndex residentIndex, + LocalMemoryGraphAccessor accessor, + EntityAttributeIndexStore store) { + residentIndex.ensureGlobalIndex(accessor, store); + Assertions.assertEquals(1L, residentIndex.getBuildCount(), + "in place maintenance must keep the accepted version in step, no rebuild expected"); + } +}