Skip to content

Commit 0b3fee0

Browse files
normen662normen662
committed
save point -- in the middle of just mess
1 parent 2fafc3d commit 0b3fee0

File tree

7 files changed

+154
-79
lines changed

7 files changed

+154
-79
lines changed

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

Lines changed: 28 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -75,19 +75,22 @@ public ByNodeStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final Su
7575
}
7676

7777
@Override
78-
public CompletableFuture<Node<? extends Neighbor>> fetchEntryNode(@Nonnull final Transaction transaction) {
78+
public CompletableFuture<NodeWithLayer<? extends Neighbor>> fetchEntryNode(@Nonnull final ReadTransaction readTransaction) {
7979
final byte[] key = getEntryNodeSubspace().pack();
8080

81-
return transaction.get(key)
81+
return readTransaction.get(key)
8282
.thenApply(valueBytes -> {
8383
if (valueBytes == null) {
8484
throw new IllegalStateException("cannot fetch entry point");
8585
}
86-
final Node<? extends Neighbor> node = fromTuple(Tuple.fromBytes(valueBytes));
86+
87+
final Tuple entryTuple = Tuple.fromBytes(valueBytes);
88+
final int lMax = (int)entryTuple.getLong(0);
89+
final Node<? extends Neighbor> node = nodeFromTuple(entryTuple.getNestedTuple(1));
8790
final OnReadListener onReadListener = getOnReadListener();
8891
onReadListener.onNodeRead(node);
8992
onReadListener.onKeyValueRead(node, key, valueBytes);
90-
return node;
93+
return node.withLayer(lMax);
9194
});
9295
}
9396

@@ -141,7 +144,7 @@ public CompletableFuture<Node> fetchNodeInternal(@Nonnull final ReadTransaction
141144
if (valueBytes == null) {
142145
return null;
143146
}
144-
final Node node = fromTuple(nodeId, Tuple.fromBytes(valueBytes));
147+
final Node node = nodeFromTuple(nodeId, Tuple.fromBytes(valueBytes));
145148
final OnReadListener onReadListener = getOnReadListener();
146149
onReadListener.onNodeRead(node);
147150
onReadListener.onKeyValueRead(node, key, valueBytes);
@@ -150,10 +153,17 @@ public CompletableFuture<Node> fetchNodeInternal(@Nonnull final ReadTransaction
150153
}
151154

152155
@Nonnull
153-
private Node<? extends Neighbor> fromTuple(@Nonnull final Tuple tuple) {
156+
private Node<? extends Neighbor> nodeFromTuple(@Nonnull final Tuple tuple) {
154157
final NodeKind nodeKind = NodeKind.fromSerializedNodeKind((byte)tuple.getLong(0));
155-
final Tuple neighborsTuple = tuple.getNestedTuple(1);
158+
final Tuple primaryKey = tuple.getNestedTuple(1);
159+
final Tuple vectorTuple = tuple.getNestedTuple(2);
160+
final Tuple neighborsTuple = tuple.getNestedTuple(3);
156161

162+
final Half[] vectorHalfs = new Half[vectorTuple.size()];
163+
for (int i = 0; i < vectorTuple.size(); i ++) {
164+
vectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(vectorTuple.getBytes(i)));
165+
}
166+
final Vector.HalfVector vector = new Vector.HalfVector(vectorHalfs);
157167
List<NeighborWithVector> neighborsWithVectors = null;
158168
Half[] neighborVectorHalfs = null;
159169
List<Neighbor> neighbors = null;
@@ -162,6 +172,13 @@ private Node<? extends Neighbor> fromTuple(@Nonnull final Tuple tuple) {
162172
final Tuple neighborTuple = (Tuple)neighborObject;
163173
switch (nodeKind) {
164174
case DATA:
175+
if (neighbors == null) {
176+
neighbors = Lists.newArrayListWithExpectedSize(neighborsTuple.size());
177+
}
178+
neighbors.add(new Neighbor(neighborTuple));
179+
break;
180+
181+
case INTERMEDIATE:
165182
final Tuple neighborPrimaryKey = neighborTuple.getNestedTuple(0);
166183
final Tuple neighborVectorTuple = neighborTuple.getNestedTuple(1);
167184
if (neighborsWithVectors == null) {
@@ -175,24 +192,17 @@ private Node<? extends Neighbor> fromTuple(@Nonnull final Tuple tuple) {
175192
neighborsWithVectors.add(new NeighborWithVector(neighborPrimaryKey, new Vector.HalfVector(neighborVectorHalfs)));
176193
break;
177194

178-
case INTERMEDIATE:
179-
if (neighbors == null) {
180-
neighbors = Lists.newArrayListWithExpectedSize(neighborsTuple.size());
181-
}
182-
neighbors.add(new Neighbor(neighborTuple));
183-
break;
184-
185195
default:
186196
throw new IllegalStateException("unknown node kind");
187197
}
188198
}
189199

190-
Verify.verify((nodeKind == NodeKind.DATA && neighborsWithVectors != null) ||
191-
(nodeKind == NodeKind.INTERMEDIATE && neighbors != null));
200+
Verify.verify((nodeKind == NodeKind.DATA && neighbors != null) ||
201+
(nodeKind == NodeKind.INTERMEDIATE && neighborsWithVectors != null));
192202

193203
return nodeKind == NodeKind.DATA
194-
? new DataNode(nodeId, itemSlots)
195-
: new IntermediateNode(nodeId, childSlots);
204+
? new DataNode(primaryKey, vector, neighbors)
205+
: new IntermediateNode(primaryKey, vector, neighborsWithVectors);
196206
}
197207

198208
@Nonnull

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,12 @@ public IntermediateNode asIntermediateNode() {
4949
throw new IllegalStateException("this is not a data node");
5050
}
5151

52+
@Nonnull
53+
@Override
54+
public NodeWithLayer<Neighbor> withLayer(final int layer) {
55+
return new NodeWithLayer<>(layer, this);
56+
}
57+
5258
@Nonnull
5359
@Override
5460
public NodeKind getKind() {

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

Lines changed: 67 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -526,70 +526,77 @@ private CompletableFuture<TraversalState> fetchLeftmostPathToLeaf(@Nonnull final
526526
@Nullable final Tuple lastKey,
527527
@Nonnull final Predicate<Rectangle> mbrPredicate,
528528
@Nonnull final BiPredicate<Tuple, Tuple> suffixPredicate) {
529-
final AtomicReference<byte[]> currentId = new AtomicReference<>(nodeId);
530-
final List<Deque<ChildSlot>> toBeProcessed = Lists.newArrayList();
531-
final AtomicReference<DataNode> leafNode = new AtomicReference<>(null);
532-
return AsyncUtil.whileTrue(() -> onReadListener.onAsyncRead(storageAdapter.fetchNode(readTransaction, currentId.get()))
533-
.thenApply(node -> {
534-
if (node == null) {
535-
if (Arrays.equals(currentId.get(), rootId)) {
536-
Verify.verify(leafNode.get() == null);
537-
return false;
538-
}
539-
throw new IllegalStateException("unable to fetch node for scan");
540-
}
541-
if (node.getKind() == NodeKind.INTERMEDIATE) {
542-
final Iterable<ChildSlot> childSlots = ((IntermediateNode)node).getSlots();
543-
Deque<ChildSlot> toBeProcessedThisLevel = new ArrayDeque<>();
544-
for (final Iterator<ChildSlot> iterator = childSlots.iterator(); iterator.hasNext(); ) {
545-
final ChildSlot childSlot = iterator.next();
546-
if (lastHilbertValue != null &&
547-
lastKey != null) {
548-
final int hilbertValueAndKeyCompare =
549-
childSlot.compareLargestHilbertValueAndKey(lastHilbertValue, lastKey);
550-
if (hilbertValueAndKeyCompare < 0) {
551-
//
552-
// The (lastHilbertValue, lastKey) pair is larger than the
553-
// (largestHilbertValue, largestKey) pair of the current child. Advance to the next
554-
// child.
555-
//
556-
continue;
529+
final AtomicReference<NodeWithLayer<? extends Neighbor>> currentNodeWithLayer =
530+
new AtomicReference<>();
531+
532+
storageAdapter.fetchEntryNode(readTransaction)
533+
.thenApply(nodeWithLayer -> {
534+
currentNodeWithLayer.set(nodeWithLayer);
535+
536+
final List<Deque<ChildSlot>> toBeProcessed = Lists.newArrayList();
537+
final AtomicReference<DataNode> leafNode = new AtomicReference<>(null);
538+
return AsyncUtil.whileTrue(() -> onReadListener.onAsyncRead(storageAdapter.fetchNode(readTransaction, currentId.get()))
539+
.thenApply(node -> {
540+
if (node == null) {
541+
if (Arrays.equals(currentId.get(), rootId)) {
542+
Verify.verify(leafNode.get() == null);
543+
return false;
544+
}
545+
throw new IllegalStateException("unable to fetch node for scan");
557546
}
558-
}
547+
if (node.getKind() == NodeKind.INTERMEDIATE) {
548+
final Iterable<ChildSlot> childSlots = ((IntermediateNode)node).getSlots();
549+
Deque<ChildSlot> toBeProcessedThisLevel = new ArrayDeque<>();
550+
for (final Iterator<ChildSlot> iterator = childSlots.iterator(); iterator.hasNext(); ) {
551+
final ChildSlot childSlot = iterator.next();
552+
if (lastHilbertValue != null &&
553+
lastKey != null) {
554+
final int hilbertValueAndKeyCompare =
555+
childSlot.compareLargestHilbertValueAndKey(lastHilbertValue, lastKey);
556+
if (hilbertValueAndKeyCompare < 0) {
557+
//
558+
// The (lastHilbertValue, lastKey) pair is larger than the
559+
// (largestHilbertValue, largestKey) pair of the current child. Advance to the next
560+
// child.
561+
//
562+
continue;
563+
}
564+
}
565+
566+
if (!mbrPredicate.test(childSlot.getMbr())) {
567+
onReadListener.onChildNodeDiscard(childSlot);
568+
continue;
569+
}
570+
571+
if (childSlot.suffixPredicateCanBeApplied()) {
572+
if (!suffixPredicate.test(childSlot.getSmallestKeySuffix(),
573+
childSlot.getLargestKeySuffix())) {
574+
onReadListener.onChildNodeDiscard(childSlot);
575+
continue;
576+
}
577+
}
578+
579+
toBeProcessedThisLevel.addLast(childSlot);
580+
iterator.forEachRemaining(toBeProcessedThisLevel::addLast);
581+
}
582+
toBeProcessed.add(toBeProcessedThisLevel);
559583

560-
if (!mbrPredicate.test(childSlot.getMbr())) {
561-
onReadListener.onChildNodeDiscard(childSlot);
562-
continue;
563-
}
584+
final ChildSlot nextChildSlot = resolveNextIdForFetch(toBeProcessed, mbrPredicate,
585+
suffixPredicate, onReadListener);
586+
if (nextChildSlot == null) {
587+
return false;
588+
}
564589

565-
if (childSlot.suffixPredicateCanBeApplied()) {
566-
if (!suffixPredicate.test(childSlot.getSmallestKeySuffix(),
567-
childSlot.getLargestKeySuffix())) {
568-
onReadListener.onChildNodeDiscard(childSlot);
569-
continue;
590+
currentId.set(Objects.requireNonNull(nextChildSlot.getChildId()));
591+
return true;
592+
} else {
593+
leafNode.set((DataNode)node);
594+
return false;
570595
}
571-
}
572-
573-
toBeProcessedThisLevel.addLast(childSlot);
574-
iterator.forEachRemaining(toBeProcessedThisLevel::addLast);
575-
}
576-
toBeProcessed.add(toBeProcessedThisLevel);
577-
578-
final ChildSlot nextChildSlot = resolveNextIdForFetch(toBeProcessed, mbrPredicate,
579-
suffixPredicate, onReadListener);
580-
if (nextChildSlot == null) {
581-
return false;
582-
}
583-
584-
currentId.set(Objects.requireNonNull(nextChildSlot.getChildId()));
585-
return true;
586-
} else {
587-
leafNode.set((DataNode)node);
588-
return false;
589-
}
590-
}), executor).thenApply(vignore -> leafNode.get() == null
591-
? TraversalState.end()
592-
: TraversalState.of(toBeProcessed, leafNode.get()));
596+
}), executor).thenApply(vignore -> leafNode.get() == null
597+
? TraversalState.end()
598+
: TraversalState.of(toBeProcessed, leafNode.get()));
599+
});
593600
}
594601

595602
/**

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,12 @@ public IntermediateNode asIntermediateNode() {
4949
return this;
5050
}
5151

52+
@Nonnull
53+
@Override
54+
public NodeWithLayer<NeighborWithVector> withLayer(final int layer) {
55+
return new NodeWithLayer<>(layer, this);
56+
}
57+
5258
@Nonnull
5359
@Override
5460
public NodeKind getKind() {

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,4 +67,7 @@ public interface Node<N extends Neighbor> {
6767

6868
@Nonnull
6969
IntermediateNode asIntermediateNode();
70+
71+
@Nonnull
72+
NodeWithLayer<N> withLayer(int layer);
7073
}

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

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
/*
2+
* NodeWithLayer.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+
class NodeWithLayer<N extends Neighbor> {
26+
private final int layer;
27+
@Nonnull
28+
private final Node<N> node;
29+
30+
public NodeWithLayer(final int layer, @Nonnull final Node<N> node) {
31+
this.layer = layer;
32+
this.node = node;
33+
}
34+
35+
public int getLayer() {
36+
return layer;
37+
}
38+
39+
@Nonnull
40+
public Node<N> getNode() {
41+
return node;
42+
}
43+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ interface StorageAdapter {
8282
@Nonnull
8383
OnReadListener getOnReadListener();
8484

85-
CompletableFuture<Node<? extends Neighbor>> fetchEntryNode(@Nonnull Transaction transaction);
85+
CompletableFuture<NodeWithLayer<? extends Neighbor>> fetchEntryNode(@Nonnull ReadTransaction readTransaction);
8686

8787
/**
8888
* Insert a new entry into the node index if configuration indicates we should maintain such an index.

0 commit comments

Comments
 (0)