adding one more test

Author: normen662

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

@@ -189,7 +189,7 @@ private NodeReferenceWithVector neighborFromTuples(@Nonnull final AffineOperator
         // 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
+        // afterward, for instance 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 =

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

@@ -56,10 +56,10 @@ void testCompactNode(final long randomSeed) {
 
         final CompactNode compactNode1 = compactNode(new Random(dependentRandomSeed));
         final CompactNode compactNode1Clone = compactNode(new Random(dependentRandomSeed));
-        Assertions.assertThat(compactNode1.toString()).isEqualTo(compactNode1Clone.toString());
+        Assertions.assertThat(compactNode1).hasToString(compactNode1Clone.toString());
 
         final CompactNode compactNode2 = compactNode(random);
-        Assertions.assertThat(compactNode1.toString()).isNotEqualTo(compactNode2.toString());
+        Assertions.assertThat(compactNode1).doesNotHaveToString(compactNode2.toString());
 
         Assertions.assertThatThrownBy(compactNode1::asInliningNode).isInstanceOf(IllegalStateException.class);
     }
@@ -72,10 +72,10 @@ void testInliningNode(final long randomSeed) {
 
         final InliningNode inliningNode1 = inliningNode(new Random(dependentRandomSeed));
         final InliningNode inliningNode1Clone = inliningNode(new Random(dependentRandomSeed));
-        Assertions.assertThat(inliningNode1.toString()).isEqualTo(inliningNode1Clone.toString());
+        Assertions.assertThat(inliningNode1).hasToString(inliningNode1Clone.toString());
 
         final InliningNode inliningNode2 = inliningNode(random);
-        Assertions.assertThat(inliningNode1.toString()).isNotEqualTo(inliningNode2.toString());
+        Assertions.assertThat(inliningNode1).doesNotHaveToString(inliningNode2.toString());
 
         Assertions.assertThatThrownBy(inliningNode1::asCompactNode).isInstanceOf(IllegalStateException.class);
     }
@@ -126,11 +126,11 @@ private static <T> void assertHashCodeEqualsToString(final long randomSeed, fina
         final T t1Clone = createFunction.apply(new Random(dependentRandomSeed));
         Assertions.assertThat(t1.hashCode()).isEqualTo(t1Clone.hashCode());
         Assertions.assertThat(t1).isEqualTo(t1Clone);
-        Assertions.assertThat(t1.toString()).isEqualTo(t1Clone.toString());
+        Assertions.assertThat(t1).hasToString(t1Clone.toString());
 
         final T t2 = createFunction.apply(random);
         Assertions.assertThat(t1).isNotEqualTo(t2);
-        Assertions.assertThat(t1.toString()).isNotEqualTo(t2.toString());
+        Assertions.assertThat(t1).doesNotHaveToString(t2.toString());
     }
 
     @Nonnull

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

@@ -30,6 +30,7 @@
 import com.apple.foundationdb.linear.Quantizer;
 import com.apple.foundationdb.linear.RealVector;
 import com.apple.foundationdb.linear.StoredVecsIterator;
+import com.apple.foundationdb.rabitq.EncodedRealVector;
 import com.apple.foundationdb.test.TestDatabaseExtension;
 import com.apple.foundationdb.test.TestExecutors;
 import com.apple.foundationdb.test.TestSubspaceExtension;
@@ -83,6 +84,7 @@
 import java.util.stream.LongStream;
 import java.util.stream.Stream;
 
+import static com.apple.foundationdb.linear.RealVectorTest.createRandomDoubleVector;
 import static com.apple.foundationdb.linear.RealVectorTest.createRandomHalfVector;
 import static org.assertj.core.api.Assertions.within;
 
@@ -285,6 +287,103 @@ void testBasicInsert(final long seed, final boolean useInlining, final boolean e
         Assertions.assertThat(readIds.size()).isBetween(10, 50);
     }
 
+    @ParameterizedTest()
+    @RandomSeedSource({0x0fdbL, 0x5ca1eL, 123456L, 78910L, 1123581321345589L})
+    void testBasicInsertWithRaBitQEncodings(final long seed) {
+        final Random random = new Random(seed);
+        final Metric metric = Metric.EUCLIDEAN_METRIC;
+
+        final AtomicLong nextNodeIdAtomic = new AtomicLong(0L);
+        final int numDimensions = 128;
+        final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool(),
+                HNSW.newConfigBuilder().setMetric(metric)
+                        .setUseRaBitQ(true)
+                        .setRaBitQNumExBits(5)
+                        .setSampleVectorStatsProbability(1.0d) // every vector is sampled
+                        .setMaintainStatsProbability(1.0d) // for every vector we maintain the stats
+                        .setStatsThreshold(950) // after 950 vectors we enable RaBitQ
+                        .setM(32).setMMax(32).setMMax0(64).build(numDimensions),
+                OnWriteListener.NOOP, OnReadListener.NOOP);
+
+        final int k = 499;
+        final DoubleRealVector queryVector = createRandomDoubleVector(random, numDimensions);
+        final Map<Tuple, RealVector> dataMap = Maps.newHashMap();
+        final TreeSet<PrimaryKeyVectorAndDistance> recordsOrderedByDistance =
+                new TreeSet<>(Comparator.comparing(PrimaryKeyVectorAndDistance::getDistance));
+
+        for (int i = 0; i < 1000;) {
+            i += basicInsertBatch(hnsw, 100, nextNodeIdAtomic, new TestOnReadListener(),
+                    tr -> {
+                        final var primaryKey = createNextPrimaryKey(nextNodeIdAtomic);
+                        final DoubleRealVector dataVector = createRandomDoubleVector(random, numDimensions);
+                        final double distance = metric.distance(dataVector, queryVector);
+                        dataMap.put(primaryKey, dataVector);
+
+                        final PrimaryKeyVectorAndDistance record =
+                                new PrimaryKeyVectorAndDistance(primaryKey, dataVector, distance);
+                        recordsOrderedByDistance.add(record);
+                        if (recordsOrderedByDistance.size() > k) {
+                            recordsOrderedByDistance.pollLast();
+                        }
+                        return record;
+                    });
+        }
+
+        //
+        // If we fetch the current state back from the db some vectors are regular vectors and some vectors are
+        // RaBitQ encoded. Since that information is not surfaced through the API, we need to scan layer 0, get
+        // all vectors directly from disk (encoded/not-encoded, transformed/not-transformed) in order to check
+        // that transformations/reconstructions are applied properly.
+        //
+        final Map<Tuple, RealVector> fromDBMap = Maps.newHashMap();
+        hnsw.scanLayer(db, 0, 100,
+                node -> fromDBMap.put(node.getPrimaryKey(),
+                        node.asCompactNode().getVector().getUnderlyingVector()));
+
+        //
+        // Still run a kNN search to make sure that recall is satisfactory.
+        //
+        final List<? extends ResultEntry> results =
+                db.run(tr ->
+                        hnsw.kNearestNeighborsSearch(tr, k, 500, true, queryVector).join());
+
+        final ImmutableSet<Tuple> trueNN =
+                recordsOrderedByDistance.stream()
+                        .map(PrimaryKeyAndVector::getPrimaryKey)
+                        .collect(ImmutableSet.toImmutableSet());
+
+        int recallCount = 0;
+        int exactVectorCount = 0;
+        int encodedVectorCount = 0;
+        for (final ResultEntry resultEntry : results) {
+            if (trueNN.contains(resultEntry.getPrimaryKey())) {
+                recallCount ++;
+            }
+
+            final RealVector originalVector = dataMap.get(resultEntry.getPrimaryKey());
+            Assertions.assertThat(originalVector).isNotNull();
+            final RealVector fromDBVector = fromDBMap.get(resultEntry.getPrimaryKey());
+            Assertions.assertThat(fromDBVector).isNotNull();
+            if (!(fromDBVector instanceof EncodedRealVector)) {
+                Assertions.assertThat(originalVector).isEqualTo(fromDBVector);
+                exactVectorCount ++;
+                final double distance = metric.distance(originalVector,
+                        Objects.requireNonNull(resultEntry.getVector()));
+                Assertions.assertThat(distance).isCloseTo(0.0d, within(2E-12));
+            } else {
+                encodedVectorCount ++;
+                final double distance = metric.distance(originalVector,
+                        Objects.requireNonNull(resultEntry.getVector()).toDoubleRealVector());
+                Assertions.assertThat(distance).isCloseTo(0.0d, within(20.0d));
+            }
+        }
+        final double recall = (double)recallCount / (double)k;
+        Assertions.assertThat(recall).isGreaterThan(0.9);
+        // must have both kinds
+        Assertions.assertThat(exactVectorCount).isGreaterThan(0);
+        Assertions.assertThat(encodedVectorCount).isGreaterThan(0);
+    }
+
     private int basicInsertBatch(final HNSW hnsw, final int batchSize,
                                  @Nonnull final AtomicLong nextNodeIdAtomic, @Nonnull final TestOnReadListener onReadListener,
                                  @Nonnull final Function<Transaction, PrimaryKeyAndVector> insertFunction) {