save point

Author: normen662

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

@@ -61,7 +61,6 @@
 import java.util.function.BiPredicate;
 import java.util.function.Function;
 import java.util.function.Predicate;
-import java.util.function.Supplier;
 
 /**
  * TODO.
@@ -546,10 +545,10 @@ private <N extends NodeReference> CompletableFuture<SearchResult<N>> searchLayer
                 new PriorityBlockingQueue<>(config.getM(),
                         Comparator.comparing(NodeReferenceWithDistance::getDistance));
         candidates.addAll(entryNeighbors);
-        final Queue<NodeReferenceWithDistance> furthestNeighbors =
+        final Queue<NodeReferenceWithDistance> nearestNeighbors =
                 new PriorityBlockingQueue<>(config.getM(),
                         Comparator.comparing(NodeReferenceWithDistance::getDistance).reversed());
-        furthestNeighbors.addAll(entryNeighbors);
+        nearestNeighbors.addAll(entryNeighbors);
         final Map<Tuple, Node<N>> nodeCache = Maps.newConcurrentMap();
         final Metric metric = getConfig().getMetric();
 
@@ -559,7 +558,7 @@ private <N extends NodeReference> CompletableFuture<SearchResult<N>> searchLayer
             }
 
             final NodeReferenceWithDistance candidate = candidates.poll();
-            final NodeReferenceWithDistance furthestNeighbor = Objects.requireNonNull(furthestNeighbors.peek());
+            final NodeReferenceWithDistance furthestNeighbor = Objects.requireNonNull(nearestNeighbors.peek());
 
             if (candidate.getDistance() > furthestNeighbor.getDistance()) {
                 return AsyncUtil.READY_FALSE;
@@ -575,23 +574,23 @@ private <N extends NodeReference> CompletableFuture<SearchResult<N>> searchLayer
                         for (final NodeReferenceWithVector current : neighborReferences) {
                             visited.add(current.getPrimaryKey());
                             final double furthestDistance =
-                                    Objects.requireNonNull(furthestNeighbors.peek()).getDistance();
+                                    Objects.requireNonNull(nearestNeighbors.peek()).getDistance();
 
                             final double currentDistance =
                                     Vector.comparativeDistance(metric, current.getVector(), queryVector);
-                            if (currentDistance < furthestDistance || furthestNeighbors.size() < efSearch) {
+                            if (currentDistance < furthestDistance || nearestNeighbors.size() < efSearch) {
                                 final NodeReferenceWithDistance currentWithDistance =
                                         new NodeReferenceWithDistance(current.getPrimaryKey(), currentDistance);
                                 candidates.add(currentWithDistance);
-                                furthestNeighbors.add(currentWithDistance);
-                                if (furthestNeighbors.size() > efSearch) {
-                                    furthestNeighbors.poll();
+                                nearestNeighbors.add(currentWithDistance);
+                                if (nearestNeighbors.size() > efSearch) {
+                                    nearestNeighbors.poll();
                                 }
                             }
                         }
                         return true;
                     });
-        }).thenCompose(ignored -> fetchResultsIfNecessary(nodeFactory, readTransaction, layer, furthestNeighbors,
+        }).thenCompose(ignored -> fetchResultsIfNecessary(nodeFactory, readTransaction, layer, nearestNeighbors,
                 nodeCache));
     }
 
@@ -606,7 +605,10 @@ private <N extends NodeReference> CompletableFuture<Node<N>> fetchNodeIfNotCache
                                                                                       @Nonnull final Map<Tuple, Node<N>> nodeCache) {
         return fetchNodeIfNecessaryAndApply(nodeFactory, readTransaction, layer, nodeReference,
                 nR -> nodeCache.get(nR.getPrimaryKey()),
-                (ignored, node) -> node);
+                (nR, node) -> {
+                    nodeCache.put(nR.getPrimaryKey(), node);
+                    return node;
+                });
     }
 
     /**
@@ -649,8 +651,10 @@ private <N extends NodeReference> CompletableFuture<List<NodeReferenceWithVector
                     }
                     return new NodeReferenceWithVector(neighborReference.getPrimaryKey(), neighborNode.asCompactNode().getVector());
                 },
-                (neighborReference, neighborNode) ->
-                        new NodeReferenceWithVector(neighborReference.getPrimaryKey(), neighborNode.asCompactNode().getVector()));
+                (neighborReference, neighborNode) -> {
+                    nodeCache.put(neighborReference.getPrimaryKey(), neighborNode);
+                    return new NodeReferenceWithVector(neighborReference.getPrimaryKey(), neighborNode.asCompactNode().getVector());
+                });
     }
 
     /**
@@ -670,7 +674,10 @@ private <N extends NodeReference> CompletableFuture<SearchResult<N>> fetchResult
                     }
                     return new SearchResult.NodeReferenceWithNode<>(nodeReference, node);
                 },
-                SearchResult.NodeReferenceWithNode::new)
+                (nodeReferenceWithDistance, node) -> {
+                    nodeCache.put(nodeReferenceWithDistance.getPrimaryKey(), node);
+                    return new SearchResult.NodeReferenceWithNode<N>(nodeReferenceWithDistance, node);
+                })
                 .thenApply(nodeReferencesWithNodes -> {
                     final ImmutableMap.Builder<NodeReferenceWithDistance, Node<N>> nodeMapBuilder =
                             ImmutableMap.builder();

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

@@ -35,9 +35,12 @@
 public abstract class Vector<R extends Number> {
     @Nonnull
     protected R[] data;
+    @Nonnull
+    protected Supplier<Integer> hashCodeSupplier;
 
     public Vector(@Nonnull final R[] data) {
         this.data = data;
+        this.hashCodeSupplier = Suppliers.memoize(this::computeHashCode);
     }
 
     public int size() {
@@ -71,6 +74,10 @@ public boolean equals(final Object o) {
 
     @Override
     public int hashCode() {
+        return hashCodeSupplier.get();
+    }
+
+    private int computeHashCode() {
         return Arrays.hashCode(data);
     }