initial version from poc

Author: normen662

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

@@ -21,6 +21,7 @@
 package com.apple.foundationdb.async.hnsw;
 
 import com.apple.foundationdb.linear.Metric;
+import com.google.common.base.Preconditions;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 
 import javax.annotation.Nonnull;
@@ -88,19 +89,20 @@ public final class Config {
      * This attribute (named {@code M_max} by the HNSW paper) is the maximum connectivity value for nodes stored on a
      * layer greater than {@code 0}. We will never create more that {@code mMax} neighbors for a node. That means that
      * we even prune the neighbors of a node if the actual number of neighbors would otherwise exceed {@code mMax}.
+     * Note that this attribute must be greater than or equal to {@link #m}.
      */
     private final int mMax;
 
     /**
      * This attribute (named {@code M_max0} by the HNSW paper) is the maximum connectivity value for nodes stored on
      * layer {@code 0}. We will never create more that {@code mMax0} neighbors for a node that is stored on that layer.
      * That means that we even prune the neighbors of a node if the actual number of neighbors would otherwise exceed
-     * {@code mMax0}.
+     * {@code mMax0}. Note that this attribute must be greater than or equal to {@link #mMax}.
      */
     private final int mMax0;
 
     /**
-     * Maximum size of the search queues (on independent queue per layer) that are used during the insertion of a new
+     * Maximum size of the search queues (one independent queue per layer) that are used during the insertion of a new
      * node. If {@code efConstruction} is set to {@code 1}, the search naturally follows a greedy approach
      * (monotonous descent), whereas a high number for {@code efConstruction} allows for a more nuanced search that can
      * tolerate (false) local minima.
@@ -172,6 +174,8 @@ private Config(final long randomSeed, @Nonnull final Metric metric, final int nu
                    final double sampleVectorStatsProbability, final double maintainStatsProbability,
                    final int statsThreshold, final boolean useRaBitQ, final int raBitQNumExBits,
                    final int maxNumConcurrentNodeFetches, final int maxNumConcurrentNeighborhoodFetches) {
+        Preconditions.checkArgument(m <= mMax);
+        Preconditions.checkArgument(mMax <= mMax0);
         this.randomSeed = randomSeed;
         this.metric = metric;
         this.numDimensions = numDimensions;

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

@@ -29,8 +29,6 @@
 import com.apple.foundationdb.async.AsyncUtil;
 import com.apple.foundationdb.linear.AffineOperator;
 import com.apple.foundationdb.linear.DoubleRealVector;
-import com.apple.foundationdb.linear.FloatRealVector;
-import com.apple.foundationdb.linear.HalfRealVector;
 import com.apple.foundationdb.linear.Quantizer;
 import com.apple.foundationdb.linear.RealVector;
 import com.apple.foundationdb.linear.Transformed;
@@ -59,7 +57,6 @@
  * @param <N> the type of {@link NodeReference} this storage adapter manages
  */
 interface StorageAdapter<N extends NodeReference> {
-    ImmutableList<VectorType> VECTOR_TYPES = ImmutableList.copyOf(VectorType.values());
 
     /**
      * Subspace for data.
@@ -199,29 +196,24 @@ static RealVector vectorFromTuple(@Nonnull final Config config, @Nonnull final T
     /**
      * Creates a {@link RealVector} from a byte array.
      * <p>
-     * This method interprets the input byte array by interpreting the first byte of the array as the precision shift.
-     * The byte array must have the proper size, i.e. the invariant {@code (bytesLength - 1) % precision == 0} must
-     * hold.
+     * This method interprets the input byte array by interpreting the first byte of the array.
+     * It the delegates to {@link RealVector#fromBytes(VectorType, byte[])}.
      * @param config an HNSW config
      * @param vectorBytes the non-null byte array to convert.
      * @return a new {@link RealVector} instance created from the byte array.
-     * @throws com.google.common.base.VerifyException if the length of {@code vectorBytes} does not meet the invariant
-     *         {@code (bytesLength - 1) % precision == 0}
      */
     @Nonnull
     static RealVector vectorFromBytes(@Nonnull final Config config, @Nonnull final byte[] vectorBytes) {
         final byte vectorTypeOrdinal = vectorBytes[0];
-        switch (fromVectorTypeOrdinal(vectorTypeOrdinal)) {
-            case HALF:
-                return HalfRealVector.fromBytes(vectorBytes);
-            case SINGLE:
-                return FloatRealVector.fromBytes(vectorBytes);
-            case DOUBLE:
-                return DoubleRealVector.fromBytes(vectorBytes);
+        switch (RealVector.fromVectorTypeOrdinal(vectorTypeOrdinal)) {
             case RABITQ:
                 Verify.verify(config.isUseRaBitQ());
                 return EncodedRealVector.fromBytes(vectorBytes, config.getNumDimensions(),
                         config.getRaBitQNumExBits());
+            case HALF:
+            case SINGLE:
+            case DOUBLE:
+                return RealVector.fromBytes(vectorBytes);
             default:
                 throw new RuntimeException("unable to serialize vector");
         }
@@ -251,11 +243,6 @@ static Tuple tupleFromVector(@Nonnull final RealVector vector) {
         return Tuple.from(vector.getRawData());
     }
 
-    @Nonnull
-    static VectorType fromVectorTypeOrdinal(final int ordinal) {
-        return VECTOR_TYPES.get(ordinal);
-    }
-
     @Nonnull
     static CompletableFuture<AccessInfo> fetchAccessInfo(@Nonnull final Config config,
                                                          @Nonnull final ReadTransaction readTransaction,

fdb-extensions/src/main/java/com/apple/foundationdb/linear/RealVector.java

@@ -20,8 +20,9 @@
 
 package com.apple.foundationdb.linear;
 
-import com.google.common.base.Preconditions;
 import com.apple.foundationdb.half.Half;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 
 import javax.annotation.Nonnull;
 
@@ -34,6 +35,8 @@
  * data type conversions and raw data representation.
  */
 public interface RealVector {
+    ImmutableList<VectorType> VECTOR_TYPES = ImmutableList.copyOf(VectorType.values());
+
     /**
      * Returns the number of elements in the vector, i.e. the number of dimensions.
      * @return the number of dimensions
@@ -189,4 +192,47 @@ default RealVector multiply(final double scalarFactor) {
         }
         return withData(result);
     }
+
+    @Nonnull
+    static VectorType fromVectorTypeOrdinal(final int ordinal) {
+        return VECTOR_TYPES.get(ordinal);
+    }
+
+    /**
+     * Creates a {@link RealVector} from a byte array.
+     * <p>
+     * This method interprets the input byte array by interpreting the first byte of the array as the type of vector.
+     * It then delegates to {@link #fromBytes(VectorType, byte[])} to do the actual deserialization.
+     *
+     * @param vectorBytes the non-null byte array to convert.
+     * @return a new {@link RealVector} instance created from the byte array.
+     */
+    @Nonnull
+    static RealVector fromBytes(@Nonnull final byte[] vectorBytes) {
+        final byte vectorTypeOrdinal = vectorBytes[0];
+        return fromBytes(fromVectorTypeOrdinal(vectorTypeOrdinal), vectorBytes);
+    }
+
+    /**
+     * Creates a {@link RealVector} from a byte array.
+     * <p>
+     * This implementation dispatches to the actual logic that deserialize a byte array to a vector which is located in
+     * the respective implementations of {@link RealVector}.
+     * @param vectorType the vector type of the serialized vector
+     * @param vectorBytes the non-null byte array to convert.
+     * @return a new {@link RealVector} instance created from the byte array.
+     */
+    @Nonnull
+    static RealVector fromBytes(@Nonnull final VectorType vectorType, @Nonnull final byte[] vectorBytes) {
+        switch (vectorType) {
+            case HALF:
+                return HalfRealVector.fromBytes(vectorBytes);
+            case SINGLE:
+                return FloatRealVector.fromBytes(vectorBytes);
+            case DOUBLE:
+                return DoubleRealVector.fromBytes(vectorBytes);
+            default:
+                throw new RuntimeException("unable to serialize vector");
+        }
+    }
 }

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/Index.java

@@ -190,6 +190,25 @@ public Index(@Nonnull Index orig, @Nullable final IndexPredicate predicate) {
         this.lastModifiedVersion = orig.lastModifiedVersion;
     }
 
+    /**
+     * Copy constructor. This will create an index that is identical to the current <code>Index</code> with a given
+     * set of index options.
+     * @param orig original index to copy
+     * @param indexOptions the index options.
+     */
+    public Index(@Nonnull Index orig, @Nonnull final Map<String, String> indexOptions) {
+        this(orig.name, orig.rootExpression, orig.type, ImmutableMap.copyOf(indexOptions), orig.predicate);
+        if (orig.primaryKeyComponentPositions != null) {
+            this.primaryKeyComponentPositions = Arrays.copyOf(orig.primaryKeyComponentPositions, orig.primaryKeyComponentPositions.length);
+        } else {
+            this.primaryKeyComponentPositions = null;
+        }
+        this.subspaceKey = normalizeSubspaceKey(name, orig.subspaceKey);
+        this.useExplicitSubspaceKey = orig.useExplicitSubspaceKey;
+        this.addedVersion = orig.addedVersion;
+        this.lastModifiedVersion = orig.lastModifiedVersion;
+    }
+
     @SuppressWarnings({"deprecation", "squid:CallToDeprecatedMethod", "java:S3776"}) // Old (deprecated) index type needs grouping compatibility
     @SpotBugsSuppressWarnings("NP_NONNULL_FIELD_NOT_INITIALIZED_IN_CONSTRUCTOR")
     public Index(@Nonnull RecordMetaDataProto.Index proto) throws KeyExpression.DeserializationException {

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/IndexOptions.java

@@ -21,6 +21,7 @@
 package com.apple.foundationdb.record.metadata;
 
 import com.apple.foundationdb.annotation.API;
+import com.apple.foundationdb.async.hnsw.Config;
 import com.apple.foundationdb.async.rtree.RTree;
 import com.apple.foundationdb.record.provider.foundationdb.IndexMaintainer;
 
@@ -223,6 +224,143 @@ public class IndexOptions {
      */
     public static final String RTREE_USE_NODE_SLOT_INDEX = "rtreeUseNodeSlotIndex";
 
+    /**
+     * HNSW-only: The random seed that is used to probabilistically determine the highest layer of an insert into an
+     * HNSW structure. See {@link Config#getRandomSeed()}. The default random seed is
+     * {@link Config#DEFAULT_RANDOM_SEED}.
+     */
+    public static final String HNSW_RANDOM_SEED = "hnswRandomSeed";
+
+    /**
+     * HNSW-only: The metric that is used to determine distances between vectors. The default metric is
+     * {@link Config#DEFAULT_METRIC}. See {@link Config#getMetric()}.
+     */
+    public static final String HNSW_METRIC = "hnswMetric";
+
+    /**
+     * HNSW-only: The number of dimensions used. All vectors must have exactly this number of dimensions. This option
+     * must be set when interacting with a vector index as it there is no default.
+     * See {@link Config#getNumDimensions()}.
+     */
+    public static final String HNSW_NUM_DIMENSIONS = "hnswNumDimensions";
+
+    /**
+     * HNSW-only: Indicator if all layers except layer {@code 0} use inlining. If inlining is used, each node is
+     * persisted as a key/value pair per neighbor which includes the vectors of the neighbors but not for itself. If
+     * inlining is not used, each node is persisted as exactly one key/value pair per node which stores its own vector
+     * but specifically excludes the vectors of the neighbors. The default value is set to
+     * {@link Config#DEFAULT_USE_INLINING}. See {@link Config#isUseInlining()}.
+     */
+    public static final String HNSW_USE_INLINING = "hnswUseInlining";
+
+    /**
+     * HNSW-only: This option (named {@code M} by the HNSW paper) is the connectivity value for all nodes stored on
+     * any layer. While by no means enforced or even enforceable, we strive to create and maintain exactly {@code m}
+     * neighbors for a node. Due to insert/delete operations it is possible that the actual number of neighbors a node
+     * references is not exactly {@code m} at any given time. The default value is set to {@link Config#DEFAULT_M}.
+     * See {@link Config#getM()}.
+     */
+    public static final String HNSW_M = "hnswM";
+
+    /**
+     * HNSW-only: This attribute (named {@code M_max} by the HNSW paper) is the maximum connectivity value for nodes
+     * stored on a layer greater than {@code 0}. A node can never have more that {@code mMax} neighbors. That means that
+     * neighbors of a node are pruned if the actual number of neighbors would otherwise exceed {@code mMax}. Note that
+     * this option must be greater than or equal to {@link #HNSW_M}. The default value is set to
+     * {@link Config#DEFAULT_M_MAX}. See {@link Config#getMMax()}.
+     */
+    public static final String HNSW_M_MAX = "hnswMax";
+
+    /**
+     * HNSW-only: This option (named {@code M_max0} by the HNSW paper) is the maximum connectivity value for nodes
+     * stored on layer {@code 0}. We will never create more that {@code mMax0} neighbors for a node that is stored on
+     * that layer. That means that we even prune the neighbors of a node if the actual number of neighbors would
+     * otherwise exceed {@code mMax0}. Note that this option must be greater than or equal to {@link #HNSW_M_MAX}.
+     * The default value is set to {@link Config#DEFAULT_M_MAX_0}. See {@link Config#getMMax0()}.
+     */
+    public static final String HNSW_M_MAX_0 = "hnswMax0";
+
+    /**
+     * HNSW-only: Maximum size of the search queues (one independent queue per layer) that are used during the insertion
+     * of a new node. If {@code HNSW_EF_CONSTRUCTION} is set to {@code 1}, the search naturally follows a greedy
+     * approach (monotonous descent), whereas a high number for {@code HNSW_EF_CONSTRUCTION} allows for a more nuanced
+     * search that can tolerate (false) local minima. The default value is set to {@link Config#DEFAULT_EF_CONSTRUCTION}.
+     * See {@link Config#getEfConstruction()}.
+     */
+    public static final String HNSW_EF_CONSTRUCTION = "hnswEfConstruction";
+
+    /**
+     * HNSW-only: Indicator to signal if, during the insertion of a node, the set of nearest neighbors of that node is
+     * to be extended by the actual neighbors of those neighbors to form a set of candidates that the new node may be
+     * connected to during the insert operation. The default value is set to {@link Config#DEFAULT_EXTEND_CANDIDATES}.
+     * See {@link Config#isExtendCandidates()}.
+     */
+    public static final String HNSW_EXTEND_CANDIDATES = "hnswExtendCandidates";
+
+    /**
+     * HNSW-only: Indicator to signal if, during the insertion of a node, candidates that have been discarded due to not
+     * satisfying the select-neighbor heuristic may get added back in to pad the set of neighbors if the new node would
+     * otherwise have too few neighbors (see {@link Config#getM()}). The default value is set to
+     * {@link Config#DEFAULT_KEEP_PRUNED_CONNECTIONS}. See {@link Config#isKeepPrunedConnections()}.
+     */
+    public static final String HNSW_KEEP_PRUNED_CONNECTIONS = "hnswKeepPrunedConnections";
+
+    /**
+     * HNSW-only: If sampling is necessary (currently iff {@link #HNSW_USE_RABITQ} is {@code "true"}), this option
+     * represents the probability of a vector being inserted to also be written into the samples subspace of the hnsw
+     * structure. The vectors in that subspace are continuously aggregated until a total {@link #HNSW_STATS_THRESHOLD}
+     * has been reached. The default value is set to {@link Config#DEFAULT_SAMPLE_VECTOR_STATS_PROBABILITY}. See
+     * {@link Config#getSampleVectorStatsProbability()}.
+     */
+    public static final String HNSW_SAMPLE_VECTOR_STATS_PROBABILITY = "hnswSampleVectorStatsProbability";
+
+    /**
+     * HNSW-only: If sampling is necessary (currently iff {@link #HNSW_USE_RABITQ} is {@code "true"}), this option
+     * represents the probability of the samples subspace to be further aggregated (rolled-up) when a new vector is
+     * inserted. The vectors in that subspace are continuously aggregated until a total
+     * {@link #HNSW_STATS_THRESHOLD} has been reached. The default value is set to
+     * {@link Config#DEFAULT_MAINTAIN_STATS_PROBABILITY}. See {@link Config#getMaintainStatsProbability()}.
+     */
+    public static final String HNSW_MAINTAIN_STATS_PROBABILITY = "hnswMaintainStatsProbability";
+
+    /**
+     * HNSW-only: If sampling is necessary (currently iff {@link #HNSW_USE_RABITQ} is {@code "true"}), this option
+     * represents the threshold (being a number of vectors) that when reached causes the stats maintenance logic to
+     * compute the actual statistics (currently the centroid of the vectors that have been inserted to far). The result
+     * is then inserted into the access info subspace of the index. The default value is set to
+     * {@link Config#DEFAULT_STATS_THRESHOLD}. See {@link Config#getStatsThreshold()}.
+     */
+    public static final String HNSW_STATS_THRESHOLD = "hnswStatsThreshold";
+
+    /**
+     * HNSW-only: Indicator if we should RaBitQ quantization. See {@link com.apple.foundationdb.rabitq.RaBitQuantizer}
+     * for more details. The default value is set to {@link Config#DEFAULT_USE_RABITQ}.
+     * See {@link Config#isUseRaBitQ()}.
+     */
+    public static final String HNSW_USE_RABITQ = "hnswUseRaBitQ";
+
+    /**
+     * HNSW-only: Number of bits per dimensions iff {@link #HNSW_USE_RABITQ} is set to {@code "true"}, ignored
+     * otherwise. If RaBitQ encoding is used, a vector is stored using roughly
+     * {@code 25 + numDimensions * (numExBits + 1) / 8} bytes. The default value is set to
+     * {@link Config#DEFAULT_RABITQ_NUM_EX_BITS}. See {@link Config#getRaBitQNumExBits()}.
+     */
+    public static final String HNSW_RABITQ_NUM_EX_BITS = "hnswRaBitQNumExBits";
+
+    /**
+     * HNSW-only: Maximum number of concurrent node fetches during search and modification operations. The default value
+     * is set to {@link Config#DEFAULT_MAX_NUM_CONCURRENT_NODE_FETCHES}.
+     * See {@link Config#getMaxNumConcurrentNodeFetches()}.
+     */
+    public static final String HNSW_MAX_NUM_CONCURRENT_NODE_FETCHES = "hnswMaxNumConcurrentNodeFetches";
+
+    /**
+     * HNSW-only: Maximum number of concurrent neighborhood fetches during modification operations when the neighbors
+     * are pruned. The default value is set to {@link Config#DEFAULT_MAX_NUM_CONCURRENT_NEIGHBOR_FETCHES}.
+     * See {@link Config#getMaxNumConcurrentNeighborhoodFetches()}.
+     */
+    public static final String HNSW_MAX_NUM_CONCURRENT_NEIGHBORHOOD_FETCHES = "hnswMaxNumConcurrentNeighborhoodFetches";
+
     private IndexOptions() {
     }
 }

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/IndexTypes.java

@@ -164,6 +164,11 @@ public class IndexTypes {
      */
     public static final String MULTIDIMENSIONAL = "multidimensional";
 
+    /**
+     * An index using an HNSW structure.
+     */
+    public static final String VECTOR = "vector";
+
     private IndexTypes() {
     }
 }

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/FDBStoreTimer.java

@@ -761,6 +761,14 @@ public enum Counts implements Count {
         LOCKS_ATTEMPTED("number of attempts to register a lock", false),
         /** Count of the locks released. */
         LOCKS_RELEASED("number of locks released", false),
+        VECTOR_NODE_READS("intermediate nodes read", false),
+        VECTOR_NODE_READ_BYTES("intermediate node bytes read", true),
+        VECTOR_NODE0_READS("intermediate nodes read", false),
+        VECTOR_NODE0_READ_BYTES("intermediate node bytes read", true),
+        VECTOR_NODE_WRITES("intermediate nodes written", false),
+        VECTOR_NODE_WRITE_BYTES("intermediate node bytes written", true),
+        VECTOR_NODE0_WRITES("intermediate nodes written", false),
+        VECTOR_NODE0_WRITE_BYTES("intermediate node bytes written", true),
         ;
 
         private final String title;