Skip to content

refactor(server): optimize rockdb batch query perf #2982

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
imbajin merged 22 commits into from
Jun 5, 2026
Merged

refactor(server): optimize rockdb batch query perf #2982

Show file tree
Hide file tree
Changes from 17 commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
5ba6d4a
optimize: Optimize RocksDB batch query performance
lokidundun Mar 29, 2026
f5405a0
Refactor getByIds to queryByIds in RocksDBTable
lokidundun Mar 31, 2026
939ace0
Modify queryByIds to use super method temporarily
lokidundun Mar 31, 2026
502c7df
Refactor queryByIds to getByIds with HashSet
lokidundun Apr 2, 2026
352f66b
Update RocksDBTables.java
lokidundun Apr 2, 2026
0d9052d
fix: fix ci
lokidundun Apr 6, 2026
45298f9
optimize: optimize the batch query
lokidundun Apr 7, 2026
223fb28
optimize: optimize rockDb query
lokidundun Apr 16, 2026
ce4e2cb
fix: fix ci test
lokidundun Apr 16, 2026
6fe08a7
fix: preserve input id multiplicity in RocksDB multi-get path
Copilot Apr 17, 2026
0391069
fix: fallback getByIds when rocksdb session has pending changes
Copilot Apr 17, 2026
ce802b9
revert: restore rocksdb query files to ce4e2cb state
Copilot Apr 17, 2026
c06225c
fix: fix wrong usage
lokidundun Apr 17, 2026
a99491b
Merge branch 'improvequery' of https://github.com/lokidundun/incubato…
lokidundun Apr 17, 2026
10cb0a4
test: add test for new multi-get path
lokidundun Apr 17, 2026
d33123e
add test for session.hasChanges()
lokidundun Apr 20, 2026
c1b3599
optimize: extract the common logic
lokidundun Apr 26, 2026
0432f58
fix: refine queryByIdsWithGet hasChanges guard and align tests
lokidundun May 29, 2026
cb98e63
fix query test exception
lokidundun Jun 1, 2026
73588c4
test: verify non-consecutive duplicate ids are preserved in edge queries
lokidundun Jun 4, 2026
f2d2e29
fix ci
lokidundun Jun 4, 2026
b11ec06
add detail test
lokidundun Jun 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 23 additions & 4 deletions ...egraph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBTable.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Set;

Expand Down Expand Up @@ -215,6 +216,22 @@ protected BackendColumnIterator queryByIds(RocksDBSessions.Session session,
));
}

protected BackendColumnIterator queryByIdsWithGet(RocksDBSessions.Session session,
Collection<Id> ids) {
if (ids.size() == 1) {
return this.queryById(session, ids.iterator().next());
}

if (!session.hasChanges()) {
return this.getByIds(session, ids);
}

// NOTE: this will lead to lazy create rocksdb iterator
return BackendColumnIterator.wrap(new FlatMapperIterator<>(
ids.iterator(), id -> this.queryById(session, id)
));
}
Comment on lines +219 to +227

protected BackendColumnIterator getById(RocksDBSessions.Session session, Id id) {
byte[] value = session.get(this.table(), id.asBytes());
if (value == null) {
Expand All @@ -224,13 +241,15 @@ protected BackendColumnIterator getById(RocksDBSessions.Session session, Id id)
return BackendColumnIterator.iterator(col);
}

protected BackendColumnIterator getByIds(RocksDBSessions.Session session, Set<Id> ids) {
protected BackendColumnIterator getByIds(RocksDBSessions.Session session,
Collection<Id> ids) {

@imbajin imbajin Apr 18, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

‼️ Set<Id> → Collection<Id> 语义变化需注意

原方法接收 Set<Id>（天然去重），改为 Collection<Id> 后，传入 List 时若含重复 ID，RocksDB multiGet 会对同一 key 重复查询并返回重复结果。

测试 testVertexQueryByIdsWithDuplicateIds 验证了这个行为（id1 返回 2 次），但这与原 Set 语义不一致。需要确认上层 IdQuery 的 ids 是否可能含重复——如果含重复，行为变更可能导致上层重复处理数据。

建议：在 getByIds 入口去重以保持原语义，或者明确文档说明 Collection 含重复返回的行为变更。

@lokidundun lokidundun Apr 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added LinkedHashSet deduplication at the getByIds() entry point while preserving input order; skips if Set is already passed. Test updated to verify dedup behavior. The
backend deduplication does not affect final semantics since GraphTransaction reassembles results by the original input order.

if (ids.size() == 1) {
return this.getById(session, ids.iterator().next());
}

List<byte[]> keys = new ArrayList<>(ids.size());
for (Id id : ids) {
Collection<Id> uniqueIds = ids instanceof Set ? ids : new LinkedHashSet<>(ids);

@imbajin imbajin Jun 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

‼️ 这里的全局去重会让 Edge 的公开查询路径吞掉非相邻重复 id，属于用户可见的结果集回归。

完整路径是：

graph.edges(e1, e2, e1)
        |
        v
GraphTransaction.queryEdgesByIds()
        |
        |  IdQuery.query() 只折叠相邻重复 id
        |  [e1, e2, e1] => query.idsSize()==3, ids.size()==3
        v
edges.isEmpty() && query.idsSize() == ids.size()
        |
        |  命中 fast path，直接返回 backend iterator
        v
RocksDBTables.Edge.queryByIds()
        |
        v
RocksDBTable.getByIds()
        |
        |  new LinkedHashSet<>(ids) 全局去重
        v
[e1, e2, e1] => [e1, e2]

也就是说，GraphTransaction 只有在检测到重复 id 时才会回到按原始 ids 重建结果的慢路径；但当前 fast path 只能识别相邻重复，识别不了 e1, e2, e1 这种非相邻重复。旧的 scan/flat-map 路径会按输入 id 展开 3 次查询，新路径在这里被压成 2 个 key，最终返回结果会少一条。

建议不要在 Edge 的这条 fast path 上做 table-level 全局去重，或者在 transaction 层改成能检测任意重复 id 后再决定是否直接返回 backend iterator。

@lokidundun lokidundun Jun 4, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

List<byte[]> keys = new ArrayList<>(uniqueIds.size());
for (Id id : uniqueIds) {
keys.add(id.asBytes());
}
return session.get(this.table(), keys);
Expand Down Expand Up @@ -309,7 +328,7 @@ protected static BackendEntryIterator newEntryIterator(BackendColumnIterator col
}

protected static BackendEntryIterator newEntryIteratorOlap(
BackendColumnIterator cols, Query query, boolean isOlap) {
BackendColumnIterator cols, Query query, boolean isOlap) {
return new BinaryEntryIterator<>(cols, query, (entry, col) -> {
if (entry == null || !entry.belongToMe(col)) {
HugeType type = query.resultType();
Expand Down
9 changes: 7 additions & 2 deletions ...graph-rocksdb/src/main/java/org/apache/hugegraph/backend/store/rocksdb/RocksDBTables.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,8 +182,7 @@ protected BackendColumnIterator queryById(RocksDBSessions.Session session, Id id
@Override
protected BackendColumnIterator queryByIds(RocksDBSessions.Session session,
Collection<Id> ids) {

@imbajin imbajin Apr 18, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 PR 描述与实际改动不完全匹配

PR 描述提到 "improves the performance of Gremlin queries... when using RPC-based backends such as HBase and HStore"，但代码改动全在 hugegraph-rocksdb 模块。建议更新 PR 描述，准确反映本次优化的范围是 RocksDB 后端。

// TODO: use getByIds() after batch version multi-get is ready
return super.queryByIds(session, ids);
return this.queryByIdsWithGet(session, ids);
}
Comment on lines 182 to 186

@lokidundun lokidundun May 30, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Catch! This PR is RocksDB-only. I’ve updated the description to change closes 2674 to related 2674

}

Expand All @@ -208,6 +207,12 @@ public static Edge in(String database) {
protected BackendColumnIterator queryById(RocksDBSessions.Session session, Id id) {
return this.getById(session, id);
}

@Override
protected BackendColumnIterator queryByIds(RocksDBSessions.Session session,
Collection<Id> ids) {
return this.queryByIdsWithGet(session, ids);
}
}

@imbajin imbajin Apr 18, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Vertex 和 Edge 的 queryByIds 逻辑完全重复

Vertex.queryByIds() 和 Edge.queryByIds() 代码完全一致（5 行相同逻辑）。可以考虑：

  1. 提取到父类 RocksDBTable.queryByIds() 中统一做 hasChanges() 判断
  2. 或者提供一个公共的 helper 方法

这样可以减少重复，也降低后续维护时两处不一致的风险。


public static class IndexTable extends RocksDBTable {
Expand Down
2 changes: 2 additions & 0 deletions hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/UnitTestSuite.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
import org.apache.hugegraph.unit.rocksdb.RocksDBCountersTest;
import org.apache.hugegraph.unit.rocksdb.RocksDBSessionTest;
import org.apache.hugegraph.unit.rocksdb.RocksDBSessionsTest;
import org.apache.hugegraph.unit.rocksdb.RocksDBTableQueryByIdsTest;
import org.apache.hugegraph.unit.serializer.BinaryBackendEntryTest;
import org.apache.hugegraph.unit.serializer.BinaryScatterSerializerTest;
import org.apache.hugegraph.unit.serializer.BinarySerializerTest;
Expand Down Expand Up @@ -141,6 +142,7 @@
RocksDBSessionsTest.class,
RocksDBSessionTest.class,
RocksDBCountersTest.class,
RocksDBTableQueryByIdsTest.class,

/* utils */
VersionTest.class,
Expand Down
Loading
Loading