Skip to content

Commit 8858ab8

Browse files
normen662normen662
committed
added some testing
1 parent f38fe62 commit 8858ab8

File tree

5 files changed

+156
-31
lines changed

5 files changed

+156
-31
lines changed

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

Lines changed: 40 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,15 @@
2020

2121
package com.apple.foundationdb.async.hnsw;
2222

23+
import com.apple.foundationdb.KeyValue;
24+
import com.apple.foundationdb.Range;
2325
import com.apple.foundationdb.ReadTransaction;
26+
import com.apple.foundationdb.StreamingMode;
2427
import com.apple.foundationdb.Transaction;
28+
import com.apple.foundationdb.async.AsyncIterable;
29+
import com.apple.foundationdb.async.AsyncUtil;
2530
import com.apple.foundationdb.subspace.Subspace;
31+
import com.apple.foundationdb.tuple.ByteArrayUtil;
2632
import com.apple.foundationdb.tuple.Tuple;
2733
import com.christianheina.langx.half4j.Half;
2834
import com.google.common.base.Verify;
@@ -31,6 +37,7 @@
3137
import org.slf4j.LoggerFactory;
3238

3339
import javax.annotation.Nonnull;
40+
import javax.annotation.Nullable;
3441
import java.util.List;
3542
import java.util.concurrent.CompletableFuture;
3643

@@ -65,23 +72,28 @@ public StorageAdapter<NodeReferenceWithVector> asInliningStorageAdapter() {
6572
protected CompletableFuture<Node<NodeReference>> fetchNodeInternal(@Nonnull final ReadTransaction readTransaction,
6673
final int layer,
6774
@Nonnull final Tuple primaryKey) {
68-
final byte[] key = getDataSubspace().pack(Tuple.from(layer, primaryKey));
75+
final byte[] keyBytes = getDataSubspace().pack(Tuple.from(layer, primaryKey));
6976

70-
return readTransaction.get(key)
77+
return readTransaction.get(keyBytes)
7178
.thenApply(valueBytes -> {
7279
if (valueBytes == null) {
7380
throw new IllegalStateException("cannot fetch node");
7481
}
75-
76-
final Tuple nodeTuple = Tuple.fromBytes(valueBytes);
77-
final Node<NodeReference> node = nodeFromTuples(primaryKey, nodeTuple);
78-
final OnReadListener onReadListener = getOnReadListener();
79-
onReadListener.onNodeRead(node);
80-
onReadListener.onKeyValueRead(key, valueBytes);
81-
return node;
82+
return nodeFromRaw(primaryKey, keyBytes, valueBytes);
8283
});
8384
}
8485

86+
@Nonnull
87+
private Node<NodeReference> nodeFromRaw(final @Nonnull Tuple primaryKey, @Nonnull final byte[] keyBytes,
88+
@Nonnull final byte[] valueBytes) {
89+
final Tuple nodeTuple = Tuple.fromBytes(valueBytes);
90+
final Node<NodeReference> node = nodeFromTuples(primaryKey, nodeTuple);
91+
final OnReadListener onReadListener = getOnReadListener();
92+
onReadListener.onNodeRead(node);
93+
onReadListener.onKeyValueRead(keyBytes, valueBytes);
94+
return node;
95+
}
96+
8597
@Nonnull
8698
private Node<NodeReference> nodeFromTuples(@Nonnull final Tuple primaryKey,
8799
@Nonnull final Tuple valueTuple) {
@@ -111,7 +123,6 @@ private Node<NodeReference> compactNodeFromTuples(@Nonnull final Tuple primaryKe
111123
return getNodeFactory().create(primaryKey, vector, nodeReferences);
112124
}
113125

114-
115126
@Override
116127
public void writeNodeInternal(@Nonnull final Transaction transaction, @Nonnull final Node<NodeReference> node,
117128
final int layer, @Nonnull final NeighborsChangeSet<NodeReference> neighborsChangeSet) {
@@ -139,4 +150,23 @@ public void writeNodeInternal(@Nonnull final Transaction transaction, @Nonnull f
139150
transaction.set(key, nodeTuple.pack());
140151
getOnWriteListener().onNodeWritten(layer, node);
141152
}
153+
154+
public Iterable<Node<NodeReference>> scanLayer(@Nonnull final ReadTransaction readTransaction, int layer,
155+
@Nullable final Tuple lastPrimaryKey, int maxNumRead) {
156+
final byte[] layerPrefix = getDataSubspace().pack(Tuple.from(layer));
157+
final Range range =
158+
lastPrimaryKey == null
159+
? Range.startsWith(layerPrefix)
160+
: new Range(ByteArrayUtil.strinc(getDataSubspace().pack(Tuple.from(layer, lastPrimaryKey))),
161+
ByteArrayUtil.strinc(layerPrefix));
162+
final AsyncIterable<KeyValue> itemsIterable =
163+
readTransaction.getRange(range, maxNumRead, false, StreamingMode.ITERATOR);
164+
165+
return AsyncUtil.mapIterable(itemsIterable, keyValue -> {
166+
final byte[] key = keyValue.getKey();
167+
final byte[] value = keyValue.getValue();
168+
final Tuple primaryKey = getDataSubspace().unpack(key);
169+
return nodeFromRaw(primaryKey, key, value);
170+
});
171+
}
142172
}

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
package com.apple.foundationdb.async.hnsw;
2222

23+
import com.apple.foundationdb.Database;
2324
import com.apple.foundationdb.ReadTransaction;
2425
import com.apple.foundationdb.Transaction;
2526
import com.apple.foundationdb.annotation.API;
@@ -34,6 +35,7 @@
3435
import com.google.common.collect.Lists;
3536
import com.google.common.collect.Maps;
3637
import com.google.common.collect.Sets;
38+
import com.google.common.collect.Streams;
3739
import com.google.errorprone.annotations.CanIgnoreReturnValue;
3840
import org.slf4j.Logger;
3941
import org.slf4j.LoggerFactory;
@@ -1054,6 +1056,27 @@ private <N extends NodeReference> void writeLonelyNodeOnLayer(@Nonnull final Sto
10541056
debug(l -> l.debug("written lonely node at key={} on layer={}", primaryKey, layer));
10551057
}
10561058

1059+
public void scanLayer(@Nonnull final Database db,
1060+
final int layer,
1061+
final int batchSize,
1062+
@Nonnull final Consumer<Node<? extends NodeReference>> nodeConsumer) {
1063+
final StorageAdapter<? extends NodeReference> storageAdapter = getStorageAdapterForLayer(layer);
1064+
final AtomicReference<Tuple> lastPrimaryKeyAtomic = new AtomicReference<>();
1065+
Tuple newPrimaryKey;
1066+
do {
1067+
final Tuple lastPrimaryKey = lastPrimaryKeyAtomic.get();
1068+
lastPrimaryKeyAtomic.set(null);
1069+
newPrimaryKey = db.run(tr -> {
1070+
Streams.stream(storageAdapter.scanLayer(tr, layer, lastPrimaryKey, batchSize))
1071+
.forEach(node -> {
1072+
nodeConsumer.accept(node);
1073+
lastPrimaryKeyAtomic.set(node.getPrimaryKey());
1074+
});
1075+
return lastPrimaryKeyAtomic.get();
1076+
}, executor);
1077+
} while (newPrimaryKey != null);
1078+
}
1079+
10571080
@Nonnull
10581081
private StorageAdapter<? extends NodeReference> getStorageAdapterForLayer(final int layer) {
10591082
return layer > 0

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

Lines changed: 72 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,17 @@
2525
import com.apple.foundationdb.ReadTransaction;
2626
import com.apple.foundationdb.StreamingMode;
2727
import com.apple.foundationdb.Transaction;
28+
import com.apple.foundationdb.async.AsyncIterable;
2829
import com.apple.foundationdb.async.AsyncUtil;
2930
import com.apple.foundationdb.subspace.Subspace;
31+
import com.apple.foundationdb.tuple.ByteArrayUtil;
3032
import com.apple.foundationdb.tuple.Tuple;
3133
import com.christianheina.langx.half4j.Half;
3234
import com.google.common.collect.ImmutableList;
3335

3436
import javax.annotation.Nonnull;
37+
import javax.annotation.Nullable;
38+
import java.util.List;
3539
import java.util.concurrent.CompletableFuture;
3640

3741
/**
@@ -66,27 +70,35 @@ protected CompletableFuture<Node<NodeReferenceWithVector>> fetchNodeInternal(@No
6670

6771
return AsyncUtil.collect(readTransaction.getRange(Range.startsWith(rangeKey),
6872
ReadTransaction.ROW_LIMIT_UNLIMITED, false, StreamingMode.WANT_ALL), readTransaction.getExecutor())
69-
.thenApply(keyValues -> {
70-
final OnReadListener onReadListener = getOnReadListener();
71-
72-
final ImmutableList.Builder<NodeReferenceWithVector> nodeReferencesWithVectorBuilder = ImmutableList.builder();
73-
for (final KeyValue keyValue : keyValues) {
74-
final byte[] key = keyValue.getKey();
75-
final byte[] value = keyValue.getValue();
76-
onReadListener.onKeyValueRead(key, value);
77-
final Tuple neighborKeyTuple = getDataSubspace().unpack(key);
78-
final Tuple neighborValueTuple = Tuple.fromBytes(value);
79-
80-
final Tuple neighborPrimaryKey = neighborKeyTuple.getNestedTuple(2); // neighbor primary key
81-
final Vector<Half> neighborVector = StorageAdapter.vectorFromTuple(neighborValueTuple); // the entire value is the vector
82-
nodeReferencesWithVectorBuilder.add(new NodeReferenceWithVector(neighborPrimaryKey, neighborVector));
83-
}
84-
85-
final Node<NodeReferenceWithVector> node =
86-
getNodeFactory().create(primaryKey, null, nodeReferencesWithVectorBuilder.build());
87-
onReadListener.onNodeRead(node);
88-
return node;
89-
});
73+
.thenApply(keyValues -> nodeFromRaw(primaryKey, keyValues));
74+
}
75+
76+
@Nonnull
77+
private Node<NodeReferenceWithVector> nodeFromRaw(final @Nonnull Tuple primaryKey, final List<KeyValue> keyValues) {
78+
final OnReadListener onReadListener = getOnReadListener();
79+
80+
final ImmutableList.Builder<NodeReferenceWithVector> nodeReferencesWithVectorBuilder = ImmutableList.builder();
81+
for (final KeyValue keyValue : keyValues) {
82+
nodeReferencesWithVectorBuilder.add(neighborFromRaw(keyValue.getKey(), keyValue.getValue()));
83+
}
84+
85+
final Node<NodeReferenceWithVector> node =
86+
getNodeFactory().create(primaryKey, null, nodeReferencesWithVectorBuilder.build());
87+
onReadListener.onNodeRead(node);
88+
return node;
89+
}
90+
91+
@Nonnull
92+
private NodeReferenceWithVector neighborFromRaw(final @Nonnull byte[] key, final byte[] value) {
93+
final OnReadListener onReadListener = getOnReadListener();
94+
95+
onReadListener.onKeyValueRead(key, value);
96+
final Tuple neighborKeyTuple = getDataSubspace().unpack(key);
97+
final Tuple neighborValueTuple = Tuple.fromBytes(value);
98+
99+
final Tuple neighborPrimaryKey = neighborKeyTuple.getNestedTuple(2); // neighbor primary key
100+
final Vector<Half> neighborVector = StorageAdapter.vectorFromTuple(neighborValueTuple); // the entire value is the vector
101+
return new NodeReferenceWithVector(neighborPrimaryKey, neighborVector);
90102
}
91103

92104
@Override
@@ -122,4 +134,43 @@ private byte[] getNeighborKey(final int layer,
122134
@Nonnull final Tuple neighborPrimaryKey) {
123135
return getDataSubspace().pack(Tuple.from(layer, node.getPrimaryKey(), neighborPrimaryKey));
124136
}
137+
138+
@Override
139+
public Iterable<Node<NodeReferenceWithVector>> scanLayer(@Nonnull final ReadTransaction readTransaction, int layer,
140+
@Nullable final Tuple lastPrimaryKey, int maxNumRead) {
141+
final byte[] layerPrefix = getDataSubspace().pack(Tuple.from(layer));
142+
final Range range =
143+
lastPrimaryKey == null
144+
? Range.startsWith(layerPrefix)
145+
: new Range(ByteArrayUtil.strinc(getDataSubspace().pack(Tuple.from(layer, lastPrimaryKey))),
146+
ByteArrayUtil.strinc(layerPrefix));
147+
final AsyncIterable<KeyValue> itemsIterable =
148+
readTransaction.getRange(range,
149+
maxNumRead, false, StreamingMode.ITERATOR);
150+
int numRead = 0;
151+
Tuple nodePrimaryKey = null;
152+
ImmutableList.Builder<Node<NodeReferenceWithVector>> nodeBuilder = ImmutableList.builder();
153+
ImmutableList.Builder<NodeReferenceWithVector> neighborsBuilder = ImmutableList.builder();
154+
for (final KeyValue item: itemsIterable) {
155+
final NodeReferenceWithVector neighbor =
156+
neighborFromRaw(item.getKey(), item.getValue());
157+
final Tuple primaryKeyFromNodeReference = neighbor.getPrimaryKey();
158+
if (nodePrimaryKey == null) {
159+
nodePrimaryKey = primaryKeyFromNodeReference;
160+
} else {
161+
if (!nodePrimaryKey.equals(primaryKeyFromNodeReference)) {
162+
nodeBuilder.add(getNodeFactory().create(nodePrimaryKey, null, neighborsBuilder.build()));
163+
}
164+
}
165+
neighborsBuilder.add(neighbor);
166+
numRead ++;
167+
}
168+
169+
// there may be a rest
170+
if (numRead > 0 && numRead < maxNumRead) {
171+
nodeBuilder.add(getNodeFactory().create(nodePrimaryKey, null, neighborsBuilder.build()));
172+
}
173+
174+
return nodeBuilder.build();
175+
}
125176
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import com.google.common.base.Verify;
2929

3030
import javax.annotation.Nonnull;
31+
import javax.annotation.Nullable;
3132
import java.util.concurrent.CompletableFuture;
3233

3334
/**
@@ -91,6 +92,9 @@ CompletableFuture<Node<N>> fetchNode(@Nonnull ReadTransaction readTransaction,
9192
void writeNode(@Nonnull Transaction transaction, @Nonnull Node<N> node, int layer,
9293
@Nonnull NeighborsChangeSet<N> changeSet);
9394

95+
Iterable<Node<N>> scanLayer(@Nonnull ReadTransaction readTransaction, int layer, @Nullable Tuple lastPrimaryKey,
96+
int maxNumRead);
97+
9498
@Nonnull
9599
static CompletableFuture<EntryNodeReference> fetchEntryNodeReference(@Nonnull final ReadTransaction readTransaction,
96100
@Nonnull final Subspace subspace,

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,23 @@ public void testBasicInsert() {
157157
});
158158
}
159159

160+
@Test
161+
public void testBasicInsertAndScanLayer() {
162+
final Random random = new Random(0);
163+
final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool());
164+
165+
db.run(tr -> {
166+
for (int i = 0; i < 20; i ++) {
167+
hnsw.insert(tr, createRandomPrimaryKey(random), createRandomVector(random, 728)).join();
168+
}
169+
return null;
170+
});
171+
172+
hnsw.scanLayer(db, 0, 100, node -> {
173+
System.out.println(node);
174+
});
175+
}
176+
160177
private <N extends NodeReference> void writeNode(@Nonnull final Transaction transaction,
161178
@Nonnull final StorageAdapter<N> storageAdapter,
162179
@Nonnull final Node<N> node,

0 commit comments

Comments
 (0)