Skip to content

Commit 08d57a6

Browse files
hahahahbennyhahahahbenny
committed
Add unit tests for vector index and integrate vector index operations into the original server pipeline
1 parent e2c29b3 commit 08d57a6

20 files changed

Lines changed: 1036 additions & 97 deletions

File tree

‎hugegraph-server/hugegraph-api/src/main/java/org/apache/hugegraph/auth/HugeGraphAuthProxy.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@
8787
import org.apache.hugegraph.util.E;
8888
import org.apache.hugegraph.util.Log;
8989
import org.apache.hugegraph.util.RateLimiter;
90+
import org.apache.hugegraph.vector.VectorIndexManager;
9091
import org.apache.tinkerpop.gremlin.process.computer.GraphComputer;
9192
import org.apache.tinkerpop.gremlin.process.traversal.Bytecode;
9293
import org.apache.tinkerpop.gremlin.process.traversal.Bytecode.Instruction;
@@ -834,6 +835,11 @@ public RaftGroupManager raftGroupManager() {
834835
return this.hugegraph.raftGroupManager();
835836
}
836837

838+
@Override
839+
public VectorIndexManager<Id> vectorIndexManager() {
840+
return this.hugegraph.vectorIndexManager();
841+
}
842+
837843
@Override
838844
public void registerRpcServices(RpcServiceConfig4Server serverConfig,
839845
RpcServiceConfig4Client clientConfig) {

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/HugeGraph.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@
5656
import org.apache.hugegraph.type.HugeType;
5757
import org.apache.hugegraph.type.define.GraphMode;
5858
import org.apache.hugegraph.type.define.GraphReadMode;
59+
import org.apache.hugegraph.vector.VectorIndexManager;
5960
import org.apache.tinkerpop.gremlin.process.traversal.TraversalStrategies;
6061
import org.apache.tinkerpop.gremlin.structure.Edge;
6162
import org.apache.tinkerpop.gremlin.structure.Graph;
@@ -278,6 +279,8 @@ public interface HugeGraph extends Graph {
278279

279280
RaftGroupManager raftGroupManager();
280281

282+
VectorIndexManager<Id> vectorIndexManager();
283+
281284
void proxy(HugeGraph graph);
282285

283286
boolean sameAs(HugeGraph graph);

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/StandardHugeGraph.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,7 @@
108108
import org.apache.hugegraph.util.LockUtil;
109109
import org.apache.hugegraph.util.Log;
110110
import org.apache.hugegraph.variables.HugeVariables;
111+
import org.apache.hugegraph.vector.VectorIndexManager;
111112
import org.apache.tinkerpop.gremlin.process.computer.GraphComputer;
112113
import org.apache.tinkerpop.gremlin.structure.Edge;
113114
import org.apache.tinkerpop.gremlin.structure.Graph;
@@ -190,6 +191,7 @@ public class StandardHugeGraph implements HugeGraph {
190191
private Date createTime;
191192
private Date updateTime;
192193
private KvStore kvStore;
194+
private VectorIndexManager<Id> vectorIndexManager;
193195

194196
public StandardHugeGraph(HugeConfig config) {
195197
this.params = new StandardHugeGraphParams();
@@ -1243,6 +1245,11 @@ public RaftGroupManager raftGroupManager() {
12431245
return provider.raftNodeManager();
12441246
}
12451247

1248+
@Override
1249+
public VectorIndexManager<Id> vectorIndexManager() {
1250+
return this.vectorIndexManager;
1251+
}
1252+
12461253
@Override
12471254
public HugeConfig configuration() {
12481255
return this.configuration;

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/serializer/BinarySerializer.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -392,7 +392,7 @@ protected void parseIndexName(HugeGraph graph, ConditionQuery query,
392392
HugeIndex index, Object fieldValues) {
393393
boolean isVectorIndex = index.type() != HugeType.VECTOR_INDEX_MAP;
394394
for (BackendColumn col : entry.columns()) {
395-
if(isVectorIndex && isVectorDleted(col.value) ||
395+
if((isVectorIndex && isVectorDleted(col.value)) ||
396396
indexFieldValuesUnmatched(col.value, fieldValues)){
397397
// Skip if field-values is not matched (just the same hash)
398398
continue;

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/tx/AbstractTransaction.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717

1818
package org.apache.hugegraph.backend.tx;
1919

20+
import java.util.HashMap;
21+
import java.util.Map;
2022
import java.util.Set;
2123

2224
import org.apache.commons.lang3.StringUtils;
@@ -67,6 +69,9 @@ public abstract class AbstractTransaction implements Transaction {
6769

6870
protected final AbstractSerializer serializer;
6971

72+
protected final ThreadLocal<Map<Id, Boolean>> vectorIndexChanges =
73+
ThreadLocal.withInitial(HashMap::new);
74+
7075
public AbstractTransaction(HugeGraphParams graph, BackendStore store) {
7176
E.checkNotNull(graph, "graph");
7277
E.checkNotNull(store, "store");
@@ -234,12 +239,20 @@ public void commit() throws BackendException {
234239
this.committing = true;
235240
try {
236241
this.commit2Backend();
242+
signalChangedIndexes();
237243
} finally {
238244
this.committing = false;
239245
this.reset();
240246
}
241247
}
242248

249+
private void signalChangedIndexes() {
250+
Map<Id, Boolean> changes = vectorIndexChanges.get();
251+
if (changes.isEmpty()) return;
252+
253+
changes.keySet().forEach(id -> graph().vectorIndexManager().signal(id));
254+
}
255+
243256
@Override
244257
public void commitIfGtSize(int size) throws BackendException {
245258
if (this.mutationSize() >= size) {
@@ -296,6 +309,9 @@ protected void reset() {
296309
if (this.mutation == null || !this.mutation.isEmpty()) {
297310
this.mutation = new BackendMutation();
298311
}
312+
if (!this.vectorIndexChanges.get().isEmpty()) {
313+
vectorIndexChanges.get().clear();
314+
}
299315
}
300316

301317
protected BackendMutation mutation() {

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/tx/GraphIndexTransaction.java‎

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -81,12 +81,10 @@
8181
import org.apache.hugegraph.task.EphemeralJobQueue;
8282
import org.apache.hugegraph.type.HugeType;
8383
import org.apache.hugegraph.type.define.Action;
84-
import org.apache.hugegraph.type.define.DataType;
8584
import org.apache.hugegraph.type.define.HugeKeys;
8685
import org.apache.hugegraph.type.define.IndexType;
8786
import org.apache.hugegraph.util.CollectionUtil;
8887
import org.apache.hugegraph.util.E;
89-
import org.apache.hugegraph.util.HashUtil;
9088
import org.apache.hugegraph.util.InsertionOrderUtil;
9189
import org.apache.hugegraph.util.LockUtil;
9290
import org.apache.hugegraph.util.LongEncoding;
@@ -324,11 +322,15 @@ protected void updateIndex(Id ilId, HugeElement element, boolean removed) {
324322
* The column is garbage-collected once the data has been flushed to disk.
325323
* generate new vector id from the context
326324
*/
327-
byte[] vectorId = HashUtil.hash(elementId.asBytes());
328325

329-
this.updateVectorIndex(indexLabel, HugeIndex.bytes2number(vectorId, DataType.INT.clazz()),
330-
elementId, expiredTime, removed);
331-
break;
326+
try(HugeGraph graph = this.graph()){
327+
int vectorId = graph.vectorIndexManager().getNextVectorId(indexLabel.id());
328+
this.updateVectorIndex(indexLabel, vectorId, elementId, expiredTime, removed);
329+
break;
330+
} catch (Exception e) {
331+
throw new RuntimeException(e);
332+
}
333+
332334
default:
333335
throw new AssertionError(String.format(
334336
"Unknown index type '%s'", indexLabel.indexType()));
@@ -340,11 +342,18 @@ private void updateVectorIndex(IndexLabel indexLabel, Object vectorId, Id elemen
340342

341343
HugeVectorIndexMap indexMap = new HugeVectorIndexMap(this.graph(), indexLabel, removed);
342344
indexMap.fieldValues(vectorId);
343-
indexMap.elementIds(elementId, expiredTime);
344345

345-
this.doAppend(this.serializer.writeIndex(indexMap));
346-
// writeIndex
347-
this.doAppend(this.serializer.writeVectorSequence(indexMap));
346+
try(HugeGraph graph = this.graph()){
347+
vectorIndexChanges.get().put(indexLabel.id(), true);
348+
long sequence = graph.vectorIndexManager().getNextSequence(indexLabel.id());
349+
indexMap.sequence(sequence);
350+
indexMap.elementIds(elementId, expiredTime);
351+
this.doAppend(this.serializer.writeIndex(indexMap));
352+
// writeIndex
353+
this.doAppend(this.serializer.writeVectorSequence(indexMap));
354+
} catch (Exception e) {
355+
throw new RuntimeException(e);
356+
}
348357
}
349358

350359
private void updateIndex(IndexLabel indexLabel, Object propValue,

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/schema/IndexLabel.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -275,6 +275,8 @@ public interface Builder extends SchemaBuilder<IndexLabel> {
275275

276276
Builder unique();
277277

278+
Builder vector();
279+
278280
Builder on(HugeType baseType, String baseValue);
279281

280282
Builder indexType(IndexType indexType);

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/vector/ServerVectorRuntime.java‎

Lines changed: 34 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,6 @@
3636
import io.github.jbellis.jvector.disk.SimpleMappedReader;
3737
import io.github.jbellis.jvector.graph.GraphIndexBuilder;
3838
import io.github.jbellis.jvector.graph.GraphSearcher;
39-
import io.github.jbellis.jvector.graph.MapRandomAccessVectorValues;
4039
import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
4140
import io.github.jbellis.jvector.graph.SearchResult;
4241
import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
@@ -71,8 +70,8 @@ public void update(Id vectorIndexLableId, Iterator<VectorRecord> records) {
7170
}
7271

7372
@Override
74-
public Iterator<Integer> search(Id indexlabelId, float[] queryVector, int topK) {
75-
IndexContext<Id> context = obtainContext(indexlabelId);
73+
public Iterator<Integer> search(Id indexLabelId, float[] queryVector, int topK) {
74+
IndexContext<Id> context = obtainContext(indexLabelId);
7675
VectorTypeSupport vectorTypeSupport =
7776
VectorizationProvider.getInstance().getVectorTypeSupport();
7877
VectorFloat<?> vector = vectorTypeSupport.createFloatVector(queryVector);
@@ -89,6 +88,12 @@ public Iterator<Integer> search(Id indexlabelId, float[] queryVector, int topK)
8988
return Arrays.stream(sr.getNodes()).mapToInt(c -> c.node).iterator();
9089
}
9190

91+
@Override
92+
public boolean isUpdateMetaData(Id indexLabelId) {
93+
IndexContext<Id> context = obtainContext(indexLabelId);
94+
return context.metaData.isUpdateFromLog();
95+
}
96+
9297
private void handleDelete(VectorRecord record, IndexContext<Id> context) {
9398
// now just mark the record deleted
9499
// TODO： add the counter to save the deleted count
@@ -101,52 +106,56 @@ private void handleDelete(VectorRecord record, IndexContext<Id> context) {
101106
private void handleBuilding(VectorRecord record, IndexContext<Id> context) {
102107

103108
if (context.vectors.getVector(record.getVectorId()) != null) {
104-
context.builder.markNodeDeleted(record.getVectorId());
105-
} else {
106-
VectorTypeSupport vectorTypeSupport =
107-
VectorizationProvider.getInstance().getVectorTypeSupport();
108-
VectorFloat<?> vector = vectorTypeSupport.createFloatVector(record.getVectorData());
109-
context.builder.addGraphNode(record.getVectorId(), vector);
109+
return;
110110
}
111+
112+
VectorTypeSupport vectorTypeSupport =
113+
VectorizationProvider.getInstance().getVectorTypeSupport();
114+
VectorFloat<?> vector = vectorTypeSupport.createFloatVector(record.getVectorData());
115+
116+
context.builder.addGraphNode(record.getVectorId(), vector);
117+
111118
}
112119

113120
@Override
114-
protected IndexContext<Id> createNewContext(Id indexlabelId) {
121+
public IndexContext<Id> createNewContext(Id indexLabelId) {
115122

116-
if (!checkPathValid(indexlabelId)) {
117-
return getNewContext(indexlabelId, null);
123+
if (!checkPathValid(indexLabelId)) {
124+
return getNewContext(indexLabelId, null);
118125
}
119126
// construct the dataPath to read the index and sequence
120127
Path currentPathDir = null;
121128
try {
122-
currentPathDir = getOnDiskIndexDirPath(indexlabelId);
129+
currentPathDir = getOnDiskIndexDirPath(indexLabelId);
123130
} catch (IOException e) {
124-
throw new RuntimeException("Failed to resolve index dir for " + indexlabelId.asString(), e);
131+
throw new RuntimeException("Failed to resolve index dir for " + indexLabelId.asString(), e);
125132
}
126133
// get the index and json
127134
try (ReaderSupplier rs = new SimpleMappedReader.Supplier(currentPathDir.resolve(INDEX_FILE_NAME));
128135
OnDiskGraphIndex index = OnDiskGraphIndex.load(rs)) {
129136
RandomAccessVectorValues ravv = index.getView();
130-
IndexContext<Id> context = getNewContext(indexlabelId, ravv);
131-
vectorMap.put(indexlabelId, context);
137+
IndexContext<Id> context = getNewContext(indexLabelId, ravv);
132138
Path currentJsonPath = currentPathDir.resolve(META_FILE_NAME);
133139
String jsonMetaData = Files.readString(currentJsonPath);
140+
// get the metaData from disk but not update from rocksdb
134141
IndexContext.IndexContextMetaData metaData = JsonUtilCommon.fromJson(jsonMetaData,
135142
IndexContext.IndexContextMetaData.class);
143+
metaData.setUpdateFromLog(false);
136144
context.setMetaData(metaData);
145+
vectorMap.put(indexLabelId, context);
137146
return context;
138147
} catch (FileNotFoundException e) {
139148
System.err.println("Index file not found: " + currentPathDir);
140-
IndexContext<Id> empty = getNewContext(indexlabelId, null);
141-
vectorMap.put(indexlabelId, empty);
149+
IndexContext<Id> empty = getNewContext(indexLabelId, null);
150+
vectorMap.put(indexLabelId, empty);
142151
return empty;
143152
} catch (IOException e) {
144153
throw new RuntimeException("Read index failed: " + currentPathDir, e);
145154
}
146155
}
147156

148157
@Override
149-
protected String idToString(Id id) {
158+
public String idToString(Id id) {
150159
return id.asString();
151160
}
152161

@@ -161,24 +170,25 @@ IndexContext<Id> getNewContext(Id vectorIndexLableId, RandomAccessVectorValues r
161170
int dimension = (int) userData.get("dimension");
162171
VectorSimilarityFunction similarityFunction =
163172
getSimilarityFunction((String) userData.get("similarityFunction"));
164-
173+
/* use the magic number first
174+
* TODO: use the config to set */
165175
int M = userData.containsKey("M") ? (int) userData.get("M") : 16;
166176
int beamWidth = userData.containsKey("beamWidth") ? (int) userData.get("beamWidthM") : 100;
167177
float neighborOverflow = userData.containsKey("neighborOverflow") ?
168178
(float) userData.get("neighborOverflow") : (float) 1.2;
169179
float alpha = userData.containsKey("alpha") ? (float) userData.get("alpha") : (float) 1.2;
170180

171181
RandomAccessVectorValues vectorValueMap = ravv != null ? ravv:
172-
new MapRandomAccessVectorValues(new HashMap<>(), dimension);
182+
new UpdatableRandomAccessVectorValues(new HashMap<>(), dimension);
173183

174184
GraphIndexBuilder builder = new GraphIndexBuilder(vectorValueMap, similarityFunction,
175185
M, beamWidth, neighborOverflow, alpha);
176-
177-
return new IndexContext<Id>(vectorIndexLableId, vectorValueMap, builder,0,
186+
// need to test the vectorValueMap is updatable
187+
return new IndexContext<Id>(vectorIndexLableId, (UpdatableRandomAccessVectorValues) vectorValueMap, builder,0,
178188
dimension, similarityFunction);
179189
}
180190

181-
VectorSimilarityFunction getSimilarityFunction(String similarityFunction) {
191+
public VectorSimilarityFunction getSimilarityFunction(String similarityFunction) {
182192
similarityFunction = similarityFunction.toUpperCase();
183193
switch (similarityFunction) {
184194
case "EUCLIDEAN":

0 commit comments

Comments
 (0)