Skip to content

Commit 5d61de3

Browse files
normen662normen662
committed
adding new distance rank comparison and some index helpers
1 parent bf6a02a commit 5d61de3

File tree

10 files changed

+1145
-47
lines changed

10 files changed

+1145
-47
lines changed

‎fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java

Lines changed: 0 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -23,12 +23,6 @@
2323
import javax.annotation.Nonnull;
2424

2525
public interface Metric {
26-
ManhattanMetric MANHATTAN_METRIC = new ManhattanMetric();
27-
EuclideanMetric EUCLIDEAN_METRIC = new EuclideanMetric();
28-
EuclideanSquareMetric EUCLIDEAN_SQUARE_METRIC = new EuclideanSquareMetric();
29-
CosineMetric COSINE_METRIC = new CosineMetric();
30-
DotProductMetric DOT_PRODUCT_METRIC = new DotProductMetric();
31-
3226
double distance(Double[] vector1, Double[] vector2);
3327

3428
default double comparativeDistance(Double[] vector1, Double[] vector2) {
@@ -54,31 +48,6 @@ private static void validate(Double[] vector1, Double[] vector2) {
5448
}
5549
}
5650

57-
@Nonnull
58-
static ManhattanMetric manhattanMetric() {
59-
return Metric.MANHATTAN_METRIC;
60-
}
61-
62-
@Nonnull
63-
static EuclideanMetric euclideanMetric() {
64-
return Metric.EUCLIDEAN_METRIC;
65-
}
66-
67-
@Nonnull
68-
static EuclideanSquareMetric euclideanSquareMetric() {
69-
return Metric.EUCLIDEAN_SQUARE_METRIC;
70-
}
71-
72-
@Nonnull
73-
static CosineMetric cosineMetric() {
74-
return Metric.COSINE_METRIC;
75-
}
76-
77-
@Nonnull
78-
static DotProductMetric dotProductMetric() {
79-
return Metric.DOT_PRODUCT_METRIC;
80-
}
81-
8251
class ManhattanMetric implements Metric {
8352
@Override
8453
public double distance(final Double[] vector1, final Double[] vector2) {

‎fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metrics.java

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
/*
2+
* Metric.java
3+
*
4+
* This source file is part of the FoundationDB open source project
5+
*
6+
* Copyright 2015-2025 Apple Inc. and the FoundationDB project authors
7+
*
8+
* Licensed under the Apache License, Version 2.0 (the "License");
9+
* you may not use this file except in compliance with the License.
10+
* You may obtain a copy of the License at
11+
*
12+
* http://www.apache.org/licenses/LICENSE-2.0
13+
*
14+
* Unless required by applicable law or agreed to in writing, software
15+
* distributed under the License is distributed on an "AS IS" BASIS,
16+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17+
* See the License for the specific language governing permissions and
18+
* limitations under the License.
19+
*/
20+
21+
package com.apple.foundationdb.async.hnsw;
22+
23+
import javax.annotation.Nonnull;
24+
25+
public enum Metrics {
26+
MANHATTAN_METRIC(new Metric.ManhattanMetric()),
27+
EUCLIDEAN_METRIC(new Metric.EuclideanMetric()),
28+
EUCLIDEAN_SQUARE_METRIC(new Metric.EuclideanSquareMetric()),
29+
COSINE_METRIC(new Metric.CosineMetric()),
30+
DOT_PRODUCT_METRIC(new Metric.DotProductMetric());
31+
32+
@Nonnull
33+
private final Metric metric;
34+
35+
Metrics(@Nonnull final Metric metric) {
36+
this.metric = metric;
37+
}
38+
39+
@Nonnull
40+
public Metric getMetric() {
41+
return metric;
42+
}
43+
}

‎fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* RTreeModificationTest.java
2+
* HNSWModificationTest.java
33
*
44
* This source file is part of the FoundationDB open source project
55
*
@@ -172,7 +172,8 @@ public void testBasicInsert() {
172172

173173
final int dimensions = 128;
174174
final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool(),
175-
HNSW.DEFAULT_CONFIG.toBuilder().setMetric(Metric.EUCLIDEAN_METRIC).setM(32).setMMax(32).setMMax0(64).build(),
175+
HNSW.DEFAULT_CONFIG.toBuilder().setMetric(Metrics.EUCLIDEAN_METRIC.getMetric())
176+
.setM(32).setMMax(32).setMMax0(64).build(),
176177
OnWriteListener.NOOP, onReadListener);
177178

178179
for (int i = 0; i < 1000;) {
@@ -243,7 +244,7 @@ private int insertBatch(final HNSW hnsw, final int batchSize,
243244
@Test
244245
@Timeout(value = 150, unit = TimeUnit.MINUTES)
245246
public void testSIFTInsert10k() throws Exception {
246-
final Metric metric = Metric.EUCLIDEAN_METRIC;
247+
final Metric metric = Metrics.EUCLIDEAN_METRIC.getMetric();
247248
final int k = 10;
248249
final AtomicLong nextNodeIdAtomic = new AtomicLong(0L);
249250

@@ -328,7 +329,7 @@ public void testSIFTInsert10k() throws Exception {
328329
@Test
329330
@Timeout(value = 150, unit = TimeUnit.MINUTES)
330331
public void testSIFTInsert10kWithBatchInsert() throws Exception {
331-
final Metric metric = Metric.EUCLIDEAN_METRIC;
332+
final Metric metric = Metrics.EUCLIDEAN_METRIC.getMetric();
332333
final int k = 10;
333334
final AtomicLong nextNodeIdAtomic = new AtomicLong(0L);
334335

@@ -440,7 +441,7 @@ public void testManyRandomVectors() {
440441
final HalfVector randomVector = createRandomVector(random, 768);
441442
final Tuple vectorTuple = StorageAdapter.tupleFromVector(randomVector);
442443
final Vector<Half> roundTripVector = StorageAdapter.vectorFromTuple(vectorTuple);
443-
Vector.comparativeDistance(Metric.EuclideanMetric.EUCLIDEAN_METRIC, randomVector, roundTripVector);
444+
Vector.comparativeDistance(Metrics.EUCLIDEAN_METRIC.getMetric(), randomVector, roundTripVector);
444445
Assertions.assertEquals(randomVector, roundTripVector);
445446
}
446447
}
@@ -453,7 +454,8 @@ public void testSIFTVectors() throws Exception {
453454
final TestOnReadListener onReadListener = new TestOnReadListener();
454455

455456
final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool(),
456-
HNSW.DEFAULT_CONFIG.toBuilder().setMetric(Metric.EUCLIDEAN_METRIC).setM(32).setMMax(32).setMMax0(64).build(),
457+
HNSW.DEFAULT_CONFIG.toBuilder().setMetric(Metrics.EUCLIDEAN_METRIC.getMetric())
458+
.setM(32).setMMax(32).setMMax0(64).build(),
457459
OnWriteListener.NOOP, onReadListener);
458460

459461

@@ -481,7 +483,8 @@ public void testSIFTVectors() throws Exception {
481483
halfs[c] = HNSWHelpers.halfValueOf(Double.parseDouble(value));
482484
}
483485
final HalfVector newVector = new HalfVector(halfs);
484-
final double distance = Vector.comparativeDistance(Metric.EUCLIDEAN_METRIC, referenceVector, newVector);
486+
final double distance = Vector.comparativeDistance(Metrics.EUCLIDEAN_METRIC.getMetric(),
487+
referenceVector, newVector);
485488
count++;
486489
final double delta = distance - mean;
487490
mean += delta / count;
@@ -500,7 +503,7 @@ public void testSIFTVectors() throws Exception {
500503
@ValueSource(ints = {2, 3, 10, 100, 768})
501504
public void testManyVectorsStandardDeviation(final int dimensionality) {
502505
final Random random = new Random();
503-
final Metric metric = Metric.EuclideanMetric.EUCLIDEAN_METRIC;
506+
final Metric metric = Metrics.EUCLIDEAN_METRIC.getMetric();
504507
long count = 0L;
505508
double mean = 0.0d;
506509
double mean2 = 0.0d;

‎fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/IndexOptions.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -223,6 +223,15 @@ public class IndexOptions {
223223
*/
224224
public static final String RTREE_USE_NODE_SLOT_INDEX = "rtreeUseNodeSlotIndex";
225225

226+
public static final String HNSW_METRIC = "hnswMetric";
227+
public static final String HNSW_M = "hnswM";
228+
public static final String HNSW_M_MAX = "hnswMax";
229+
public static final String HNSW_M_MAX_0 = "hnswMax0";
230+
public static final String HNSW_EF_SEARCH = "hnswEfSearch";
231+
public static final String HNSW_EF_CONSTRUCTION = "hnswEfConstruction";
232+
public static final String HNSW_EXTEND_CANDIDATES = "hnswExtendCandidates";
233+
public static final String HNSW_KEEP_PRUNED_CONNECTIONS = "hnswKeepPrunedConnections";
234+
226235
private IndexOptions() {
227236
}
228237
}

‎fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/IndexTypes.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,11 @@ public class IndexTypes {
164164
*/
165165
public static final String MULTIDIMENSIONAL = "multidimensional";
166166

167+
/**
168+
* An index using an HNSW structure.
169+
*/
170+
public static final String VECTOR = "vector";
171+
167172
private IndexTypes() {
168173
}
169174
}

‎fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexHelper.java

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
/*
2+
* VectorIndexHelper.java
3+
*
4+
* This source file is part of the FoundationDB open source project
5+
*
6+
* Copyright 2023 Apple Inc. and the FoundationDB project authors
7+
*
8+
* Licensed under the Apache License, Version 2.0 (the "License");
9+
* you may not use this file except in compliance with the License.
10+
* You may obtain a copy of the License at
11+
*
12+
* http://www.apache.org/licenses/LICENSE-2.0
13+
*
14+
* Unless required by applicable law or agreed to in writing, software
15+
* distributed under the License is distributed on an "AS IS" BASIS,
16+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17+
* See the License for the specific language governing permissions and
18+
* limitations under the License.
19+
*/
20+
21+
package com.apple.foundationdb.record.provider.foundationdb.indexes;
22+
23+
import com.apple.foundationdb.annotation.API;
24+
import com.apple.foundationdb.async.hnsw.HNSW;
25+
import com.apple.foundationdb.async.hnsw.Metrics;
26+
import com.apple.foundationdb.record.metadata.Index;
27+
import com.apple.foundationdb.record.metadata.IndexOptions;
28+
import com.apple.foundationdb.record.provider.common.StoreTimer;
29+
30+
import javax.annotation.Nonnull;
31+
32+
/**
33+
* Helper functions for index maintainers that use a {@link HNSW}.
34+
*/
35+
@API(API.Status.EXPERIMENTAL)
36+
public class VectorIndexHelper {
37+
private VectorIndexHelper() {
38+
}
39+
40+
/**
41+
* Parse standard options into {@link HNSW.Config}.
42+
* @param index the index definition to get options from
43+
* @return parsed config options
44+
*/
45+
public static HNSW.Config getConfig(@Nonnull final Index index) {
46+
final HNSW.ConfigBuilder builder = HNSW.newConfigBuilder();
47+
final String hnswMetricOption = index.getOption(IndexOptions.HNSW_METRIC);
48+
if (hnswMetricOption != null) {
49+
builder.setMetric(Metrics.valueOf(hnswMetricOption).getMetric());
50+
}
51+
final String hnswMOption = index.getOption(IndexOptions.HNSW_M);
52+
if (hnswMOption != null) {
53+
builder.setM(Integer.parseInt(hnswMOption));
54+
}
55+
final String hnswMMaxOption = index.getOption(IndexOptions.HNSW_M_MAX);
56+
if (hnswMMaxOption != null) {
57+
builder.setMMax(Integer.parseInt(hnswMMaxOption));
58+
}
59+
final String hnswMMax0Option = index.getOption(IndexOptions.HNSW_M_MAX_0);
60+
if (hnswMMax0Option != null) {
61+
builder.setMMax0(Integer.parseInt(hnswMMax0Option));
62+
}
63+
final String hnswEfSearchOption = index.getOption(IndexOptions.HNSW_EF_SEARCH);
64+
if (hnswEfSearchOption != null) {
65+
builder.setEfSearch(Integer.parseInt(hnswEfSearchOption));
66+
}
67+
final String hnswEfConstructionOption = index.getOption(IndexOptions.HNSW_EF_CONSTRUCTION);
68+
if (hnswEfConstructionOption != null) {
69+
builder.setEfConstruction(Integer.parseInt(hnswEfConstructionOption));
70+
}
71+
final String hnswExtendCandidatesOption = index.getOption(IndexOptions.HNSW_EXTEND_CANDIDATES);
72+
if (hnswExtendCandidatesOption != null) {
73+
builder.setExtendCandidates(Boolean.parseBoolean(hnswExtendCandidatesOption));
74+
}
75+
final String hnswKeepPrunedConnectionsOption = index.getOption(IndexOptions.HNSW_KEEP_PRUNED_CONNECTIONS);
76+
if (hnswKeepPrunedConnectionsOption != null) {
77+
builder.setKeepPrunedConnections(Boolean.parseBoolean(hnswKeepPrunedConnectionsOption));
78+
}
79+
80+
return builder.build();
81+
}
82+
83+
/**
84+
* Instrumentation events specific to R-tree index maintenance.
85+
*/
86+
public enum Events implements StoreTimer.DetailEvent {
87+
VECTOR_SCAN("scanning the HNSW of a vector index"),
88+
VECTOR_SKIP_SCAN("skip scan the prefix tuples of a vector index scan"),
89+
VECTOR_MODIFICATION("modifying the HNSW of a vector index");
90+
91+
private final String title;
92+
private final String logKey;
93+
94+
Events(String title, String logKey) {
95+
this.title = title;
96+
this.logKey = (logKey != null) ? logKey : StoreTimer.DetailEvent.super.logKey();
97+
}
98+
99+
Events(String title) {
100+
this(title, null);
101+
}
102+
103+
@Override
104+
public String title() {
105+
return title;
106+
}
107+
108+
@Override
109+
@Nonnull
110+
public String logKey() {
111+
return this.logKey;
112+
}
113+
}
114+
}

0 commit comments

Comments
 (0)