last round of addressing comments

Author: normen662

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

@@ -21,10 +21,10 @@
 package com.apple.foundationdb.async.hnsw;
 
 import com.apple.foundationdb.linear.RealVector;
-import com.google.common.base.Objects;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import java.util.Objects;
 
 /**
  * Class to capture the current state of this HNSW that cannot be expressed as metadata but that also is not the actual
@@ -46,17 +46,19 @@ class AccessInfo {
     private final long rotatorSeed;
 
     /**
-     * A centroid that is usually derived as an average over some vectors seen so far. It is used to create the
-     * {@link StorageTransform}.
+     * The negated centroid that is usually derived as an average over some vectors seen so far. It is used to create
+     * the {@link StorageTransform}. The centroid is stored in its negated form (i.e. {@code centroid * (-1)}) as the
+     * {@link com.apple.foundationdb.linear.AffineOperator} adds its translation vector but the centroid needs to be
+     * subtracted.
      */
     @Nullable
-    private final RealVector centroid;
+    private final RealVector negatedCentroid;
 
     public AccessInfo(@Nonnull final EntryNodeReference entryNodeReference, final long rotatorSeed,
-                      @Nullable final RealVector centroid) {
+                      @Nullable final RealVector negatedCentroid) {
         this.entryNodeReference = entryNodeReference;
         this.rotatorSeed = rotatorSeed;
-        this.centroid = centroid;
+        this.negatedCentroid = negatedCentroid;
     }
 
     @Nonnull
@@ -65,21 +67,21 @@ public EntryNodeReference getEntryNodeReference() {
     }
 
     public boolean canUseRaBitQ() {
-        return getCentroid() != null;
+        return getNegatedCentroid() != null;
     }
 
     public long getRotatorSeed() {
         return rotatorSeed;
     }
 
     @Nullable
-    public RealVector getCentroid() {
-        return centroid;
+    public RealVector getNegatedCentroid() {
+        return negatedCentroid;
     }
 
     @Nonnull
     public AccessInfo withNewEntryNodeReference(@Nonnull final EntryNodeReference entryNodeReference) {
-        return new AccessInfo(entryNodeReference, getRotatorSeed(), getCentroid());
+        return new AccessInfo(entryNodeReference, getRotatorSeed(), getNegatedCentroid());
     }
 
     @Override
@@ -89,13 +91,13 @@ public boolean equals(final Object o) {
         }
         final AccessInfo that = (AccessInfo)o;
         return rotatorSeed == that.rotatorSeed &&
-                Objects.equal(entryNodeReference, that.entryNodeReference) &&
-                Objects.equal(centroid, that.centroid);
+                Objects.equals(entryNodeReference, that.entryNodeReference) &&
+                Objects.equals(negatedCentroid, that.negatedCentroid);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hashCode(entryNodeReference, rotatorSeed, centroid);
+        return Objects.hash(entryNodeReference, rotatorSeed, negatedCentroid);
     }
 
     @Nonnull
@@ -104,6 +106,6 @@ public String toString() {
         return "AccessInfo[" +
                 "entryNodeReference=" + entryNodeReference +
                 ", rotatorSeed=" + rotatorSeed +
-                ", centroid=" + centroid + "]";
+                ", centroid=" + negatedCentroid + "]";
     }
 }

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

@@ -22,9 +22,9 @@
 
 import com.apple.foundationdb.linear.RealVector;
 import com.apple.foundationdb.linear.Transformed;
-import com.google.common.base.Objects;
 
 import javax.annotation.Nonnull;
+import java.util.Objects;
 
 /**
  * A record-like class wrapping a {@link RealVector} and a count. This data structure is used to keep a running sum
@@ -55,12 +55,12 @@ public boolean equals(final Object o) {
             return false;
         }
         final AggregatedVector that = (AggregatedVector)o;
-        return partialCount == that.partialCount && Objects.equal(partialVector, that.partialVector);
+        return partialCount == that.partialCount && Objects.equals(partialVector, that.partialVector);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hashCode(partialCount, partialVector);
+        return Objects.hash(partialCount, partialVector);
     }
 
     @Override

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

@@ -230,7 +230,7 @@ public void writeNodeInternal(@Nonnull final Transaction transaction, @Nonnull f
         nodeItems.add(NodeKind.COMPACT.getSerialized());
         final CompactNode compactNode = node.asCompactNode();
         // getting underlying vector is okay as it is only written to the database
-        nodeItems.add(StorageAdapter.tupleFromVector(quantizer.encode(compactNode.getVector()).getUnderlyingVector()));
+        nodeItems.add(StorageAdapter.tupleFromVector(quantizer.encode(compactNode.getVector())));
 
         final Iterable<NodeReference> neighbors = neighborsChangeSet.merge();
 

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

@@ -21,10 +21,10 @@
 package com.apple.foundationdb.async.hnsw;
 
 import com.apple.foundationdb.linear.Metric;
-import com.google.common.base.Objects;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 
 import javax.annotation.Nonnull;
+import java.util.Objects;
 
 /**
  * Configuration settings for a {@link HNSW}.
@@ -292,7 +292,7 @@ public boolean equals(final Object o) {
 
     @Override
     public int hashCode() {
-        return Objects.hashCode(randomSeed, metric, numDimensions, useInlining, m, mMax, mMax0, efConstruction,
+        return Objects.hash(randomSeed, metric, numDimensions, useInlining, m, mMax, mMax0, efConstruction,
                 extendCandidates, keepPrunedConnections, sampleVectorStatsProbability, maintainStatsProbability,
                 statsThreshold, useRaBitQ, raBitQNumExBits, maxNumConcurrentNodeFetches, maxNumConcurrentNeighborhoodFetches);
     }

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

@@ -217,7 +217,7 @@ private AffineOperator storageTransform(@Nullable final AccessInfo accessInfo) {
         }
 
         return new StorageTransform(accessInfo.getRotatorSeed(),
-                getConfig().getNumDimensions(), Objects.requireNonNull(accessInfo.getCentroid()));
+                getConfig().getNumDimensions(), Objects.requireNonNull(accessInfo.getNegatedCentroid()));
     }
 
     @Nonnull
@@ -442,6 +442,8 @@ private Quantizer quantizer(@Nullable final AccessInfo accessInfo) {
                         @Nonnull final Transformed<RealVector> queryVector) {
         final Set<Tuple> visited = Sets.newConcurrentHashSet(NodeReference.primaryKeys(nodeReferences));
         final Queue<NodeReferenceWithDistance> candidates =
+                // This initial capacity is somewhat arbitrary as m is not necessarily a limit,
+                // but it gives us a number that is better than the default.
                 new PriorityQueue<>(config.getM(),
                         Comparator.comparing(NodeReferenceWithDistance::getDistance));
         candidates.addAll(nodeReferences);
@@ -1189,17 +1191,19 @@ layer, getConfig().getM(), getConfig().isExtendCandidates(), nodeCache, newVecto
                                       @Nonnull final AffineOperator storageTransform,
                                       @Nonnull final Estimator estimator,
                                       @Nonnull final NodeReferenceAndNode<N> selectedNeighbor,
-                                      int layer,
-                                      int mMax,
+                                      final int layer,
+                                      final int mMax,
                                       @Nonnull final NeighborsChangeSet<N> neighborChangeSet,
                                       @Nonnull final Map<Tuple, AbstractNode<N>> nodeCache) {
         final AbstractNode<N> selectedNeighborNode = selectedNeighbor.getNode();
-        if (selectedNeighborNode.getNeighbors().size() < mMax) {
+        final int numNeighbors =
+                Iterables.size(neighborChangeSet.merge()); // this is a view over the iterable neighbors in the set
+        if (numNeighbors < mMax) {
             return CompletableFuture.completedFuture(null);
         } else {
             if (logger.isTraceEnabled()) {
                 logger.trace("pruning neighborhood of key={} which has numNeighbors={} out of mMax={}",
-                        selectedNeighborNode.getPrimaryKey(), selectedNeighborNode.getNeighbors().size(), mMax);
+                        selectedNeighborNode.getPrimaryKey(), numNeighbors, mMax);
             }
             return fetchNeighborhood(storageAdapter, transaction, storageTransform, layer, neighborChangeSet.merge(), nodeCache)
                     .thenCompose(nodeReferenceWithVectors -> {

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

@@ -185,6 +185,13 @@ private NodeReferenceWithVector neighborFromRaw(@Nonnull final AffineOperator st
     private NodeReferenceWithVector neighborFromTuples(@Nonnull final AffineOperator storageTransform,
                                                        @Nonnull final Tuple keyTuple, @Nonnull final Tuple valueTuple) {
         final Tuple neighborPrimaryKey = keyTuple.getNestedTuple(2); // neighbor primary key
+        //
+        // Transform the raw vector that was just fetched into the internal coordinate system. If we do not have
+        // a need to transform coordinates, this transform is the identity transformation. Vectors are always stored
+        // in the internal coordinate system in use at the time the vector is written. If that coordinate system changes
+        // afterward, for instance RaBitQ by enabling RaBitQ, subsequent reads of vectors that were written prior to
+        // the coordinate system change need to be transformed when they are read back.
+        //
         final Transformed<RealVector> neighborVector =
                 storageTransform.transform(
                         StorageAdapter.vectorFromTuple(getConfig(), valueTuple)); // the entire value is the vector
@@ -255,7 +262,7 @@ public void writeNeighbor(@Nonnull final Transaction transaction, @Nonnull final
         // getting underlying vector is okay as it is only written to the database
         final byte[] value =
                 StorageAdapter.tupleFromVector(
-                        quantizer.encode(neighbor.getVector()).getUnderlyingVector()).pack();
+                        quantizer.encode(neighbor.getVector())).pack();
         transaction.set(neighborKey, value);
         getOnWriteListener().onNeighborWritten(layer, node, neighbor);
         getOnWriteListener().onKeyValueWritten(layer, neighborKey, value);

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

@@ -23,9 +23,9 @@
 import com.apple.foundationdb.linear.RealVector;
 import com.apple.foundationdb.linear.Transformed;
 import com.apple.foundationdb.tuple.Tuple;
-import com.google.common.base.Objects;
 
 import javax.annotation.Nonnull;
+import java.util.Objects;
 
 /**
  * Represents a reference to a node that includes an associated vector.
@@ -97,7 +97,7 @@ public boolean equals(final Object o) {
         if (!super.equals(o)) {
             return false;
         }
-        return Objects.equal(vector, ((NodeReferenceWithVector)o).vector);
+        return Objects.equals(vector, ((NodeReferenceWithVector)o).vector);
     }
 
     /**
@@ -106,7 +106,7 @@ public boolean equals(final Object o) {
      */
     @Override
     public int hashCode() {
-        return Objects.hashCode(super.hashCode(), vector);
+        return Objects.hash(super.hashCode(), vector);
     }
 
     /**

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

@@ -23,10 +23,10 @@
 import com.apple.foundationdb.ReadTransaction;
 import com.apple.foundationdb.linear.RealVector;
 import com.apple.foundationdb.tuple.Tuple;
-import com.google.common.base.Objects;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import java.util.Objects;
 
 /**
  * Record-like class to wrap the results of a kNN-search.
@@ -39,16 +39,16 @@ public class ResultEntry {
     private final Tuple primaryKey;
 
     /**
-     * The vector that is stored with the item in the index. This vector is expressed in the client's coordinate
-     * system and should not be of class {@link com.apple.foundationdb.linear.HalfRealVector},
+     * The vector that is stored with the item in the structure. This vector is expressed in the client's coordinate
+     * system and should be of class {@link com.apple.foundationdb.linear.HalfRealVector},
      * {@link com.apple.foundationdb.linear.FloatRealVector}, or {@link com.apple.foundationdb.linear.DoubleRealVector}.
      * This member is nullable. It is set to {@code null}, if the caller to
      * {@link HNSW#kNearestNeighborsSearch(ReadTransaction, int, int, boolean, RealVector)} requested to not return
      * vectors.
      * <p>
      * The vector, if set, may or may not be exactly equal to the vector that was originally inserted in the HNSW.
      * Depending on quantization settings (see {@link Config}, the vector that
-     * is returned may only be an approximation of the original vector..
+     * is returned may only be an approximation of the original vector.
      */
     @Nullable
     private final RealVector vector;
@@ -97,13 +97,13 @@ public boolean equals(final Object o) {
         final ResultEntry that = (ResultEntry)o;
         return Double.compare(distance, that.distance) == 0 &&
                 rankOrRowNumber == that.rankOrRowNumber &&
-                Objects.equal(primaryKey, that.primaryKey) &&
-                Objects.equal(vector, that.vector);
+                Objects.equals(primaryKey, that.primaryKey) &&
+                Objects.equals(vector, that.vector);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hashCode(primaryKey, vector, distance, rankOrRowNumber);
+        return Objects.hash(primaryKey, vector, distance, rankOrRowNumber);
     }
 
     @Override

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

@@ -38,6 +38,7 @@
 import com.apple.foundationdb.rabitq.EncodedRealVector;
 import com.apple.foundationdb.subspace.Subspace;
 import com.apple.foundationdb.tuple.Tuple;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Verify;
 import com.google.common.collect.ImmutableList;
 
@@ -63,17 +64,17 @@ interface StorageAdapter<N extends NodeReference> {
     /**
      * Subspace for data.
      */
-    byte SUBSPACE_PREFIX_DATA = 0x00;
+    long SUBSPACE_PREFIX_DATA = 0x00;
 
     /**
      * Subspace for the access info; contains entry nodes; these are kept separately from the data.
      */
-    byte SUBSPACE_PREFIX_ACCESS_INFO = 0x01;
+    long SUBSPACE_PREFIX_ACCESS_INFO = 0x01;
 
     /**
      * Subspace for (mostly) statistical analysis (like finding a centroid, etc.). Contains samples of vectors.
      */
-    byte SUBSPACE_PREFIX_SAMPLES = 0x02;
+    long SUBSPACE_PREFIX_SAMPLES = 0x02;
 
     /**
      * Returns the configuration of the HNSW graph.
@@ -175,6 +176,7 @@ void writeNode(@Nonnull Transaction transaction, @Nonnull Quantizer quantizer, @
      * @param maxNumRead the maximum number of nodes to return in this scan
      * @return an {@link AsyncIterable} that provides the nodes found in the specified layer range
      */
+    @VisibleForTesting
     Iterable<AbstractNode<N>> scanLayer(@Nonnull ReadTransaction readTransaction, int layer,
                                         @Nullable Tuple lastPrimaryKey, int maxNumRead);
 
@@ -225,6 +227,16 @@ static RealVector vectorFromBytes(@Nonnull final Config config, @Nonnull final b
         }
     }
 
+    /**
+     * Converts a transformed vector into a tuple.
+     * @param vector a transformed vector
+     * @return a new, non-null {@code Tuple} instance representing the contents of the underlying vector.
+     */
+    @Nonnull
+    static Tuple tupleFromVector(@Nonnull final Transformed<RealVector> vector) {
+        return tupleFromVector(vector.getUnderlyingVector());
+    }
+
     /**
      * Converts a {@link RealVector} into a {@link Tuple}.
      * <p>
@@ -295,12 +307,12 @@ static void writeAccessInfo(@Nonnull final Transaction transaction,
                                 @Nonnull final OnWriteListener onWriteListener) {
         final Subspace entryNodeSubspace = accessInfoSubspace(subspace);
         final EntryNodeReference entryNodeReference = accessInfo.getEntryNodeReference();
-        final RealVector centroid = accessInfo.getCentroid();
+        final RealVector centroid = accessInfo.getNegatedCentroid();
         final byte[] key = entryNodeSubspace.pack();
         final byte[] value = Tuple.from(entryNodeReference.getLayer(),
                 entryNodeReference.getPrimaryKey(),
                 // getting underlying is okay as it is only written to the database
-                StorageAdapter.tupleFromVector(entryNodeReference.getVector().getUnderlyingVector()),
+                StorageAdapter.tupleFromVector(entryNodeReference.getVector()),
                 accessInfo.getRotatorSeed(),
                 centroid == null ? null : StorageAdapter.tupleFromVector(centroid)).pack();
         transaction.set(key, value);
@@ -313,7 +325,6 @@ static CompletableFuture<List<AggregatedVector>> consumeSampledVectors(@Nonnull
                                                                            final int numMaxVectors,
                                                                            @Nonnull final OnReadListener onReadListener) {
         final Subspace prefixSubspace = samplesSubspace(subspace);
-
         final byte[] prefixKey = prefixSubspace.pack();
         final ReadTransaction snapshot = transaction.snapshot();
         final Range range = Range.startsWith(prefixKey);

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

@@ -65,17 +65,6 @@ public RealVector apply(@Nonnull final RealVector vector) {
     @Nonnull
     @Override
     public RealVector invertedApply(@Nonnull final RealVector vector) {
-        //
-        // Only transform the vector if it is needed. We make the decision based on whether the vector is encoded or
-        // not. When we switch on encoding, we apply the new coordinate system from that point onwards meaning that all
-        // vectors inserted before use the client coordinate system. For the inverted case, we only have to transform
-        // the encoded vectors as they are expressed in the internal coordinate system, while regular non-encoded
-        // vectors are already expressed in the client system.
-        //
-        if (!(vector instanceof EncodedRealVector)) {
-            return vector;
-        }
-
         return super.invertedApply(vector);
     }
 }