Added FaissHNSW and bridge to Lucene HNSW graph.

Signed-off-by: Dooyong Kim <kdooyong@amazon.com>
Co-authored-by: Dooyong Kim <kdooyong@amazon.com>

Author: Dooyong Kim

src/main/java/org/opensearch/knn/memoryoptsearch/faiss/FaissHNSW.java

@@ -0,0 +1,95 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.memoryoptsearch.faiss;
+
+import lombok.Getter;
+import org.apache.lucene.store.IndexInput;
+
+import java.io.IOException;
+
+/**
+ * Ported implementation of the FAISS HNSW graph search algorithm.
+ * While it follows the same steps as the original FAISS implementation, differences in how the JVM and C++ handle floating-point
+ * calculations can lead to slight variations in results. However, such cases are very rare, and in most instances, the results are
+ * identical to FAISS. Even when there are ranking differences, they do not impact the precision or recall of the search.
+ * For more details, refer to the [FAISS HNSW implementation](
+ * <a href="https://github.com/facebookresearch/faiss/blob/main/faiss/impl/HNSW.h">...</a>).
+ */
+@Getter
+public class FaissHNSW {
+    // Cumulative number of neighbors per each level.
+    private int[] cumNumberNeighborPerLevel;
+    // Offset to be added to cumNumberNeighborPerLevel[level] to get the actual start offset of neighbor list.
+    private long[] offsets = null;
+    // Neighbor list storage.
+    private final FaissSection neighbors = new FaissSection();
+    // Entry point in HNSW graph
+    private int entryPoint;
+    // Maximum level of HNSW graph
+    private int maxLevel = -1;
+    // Default efSearch parameter. This determines the navigation queue size.
+    // More value, algorithm will more navigate candidates.
+    private int efSearch = 16;
+    // Total number of vectors stored in graph.
+    private long totalNumberOfVectors;
+
+    /**
+     * Partially loads the FAISS HNSW graph from the provided index input stream.
+     * The graph is divided into multiple sections, and this method marks the starting offset of each section then skip to the next
+     * section instead of loading the entire graph into memory. During the search, bytes will be accessed via {@link IndexInput}.
+     *
+     * @param input An input stream for a FAISS HNSW graph file, allowing access to the neighbor list and vector locations.
+     * @param totalNumberOfVectors The total number of vectors stored in the graph.
+     * @return {@link FaissHNSW}, a graph search structure that represents the FAISS HNSW graph
+     *
+     * FYI <a href="https://github.com/facebookresearch/faiss/blob/main/faiss/impl/index_read.cpp#L363">FAISS Deserialization</a>
+     *
+     * @throws IOException
+     */
+    public static FaissHNSW load(IndexInput input, long totalNumberOfVectors) throws IOException {
+        // Total number of vectors
+        FaissHNSW faissHNSW = new FaissHNSW();
+        faissHNSW.totalNumberOfVectors = totalNumberOfVectors;
+
+        // We don't use `double[] assignProbas` for search. It is for index construction.
+        long size = input.readLong();
+        input.skipBytes(Double.BYTES * size);
+
+        // Accumulate number of neighbor per each level.
+        size = input.readLong();
+        faissHNSW.cumNumberNeighborPerLevel = new int[(int) size];
+        if (size > 0) {
+            input.readInts(faissHNSW.cumNumberNeighborPerLevel, 0, (int) size);
+        }
+
+        // We don't use `level`.
+        final FaissSection levels = new FaissSection();
+        levels.markSection(input, Integer.BYTES);
+
+        // Load `offsets` into memory.
+        size = input.readLong();
+        faissHNSW.offsets = new long[(int) size];
+        input.readLongs(faissHNSW.offsets, 0, faissHNSW.offsets.length);
+
+        // Mark neighbor list section.
+        faissHNSW.neighbors.markSection(input, Integer.BYTES);
+
+        // HNSW graph parameters
+        faissHNSW.entryPoint = input.readInt();
+
+        faissHNSW.maxLevel = input.readInt();
+
+        // We don't use this field. It's for index building.
+        final int efConstruction = input.readInt();
+
+        faissHNSW.efSearch = input.readInt();
+
+        // dummy read a deprecated field.
+        input.readInt();
+
+        return faissHNSW;
+    }
+}

src/main/java/org/opensearch/knn/memoryoptsearch/faiss/FaissHNSWFlatIndex.java

@@ -5,6 +5,7 @@
 
 package org.opensearch.knn.memoryoptsearch.faiss;
 
+import lombok.Getter;
 import org.apache.lucene.index.ByteVectorValues;
 import org.apache.lucene.index.FloatVectorValues;
 import org.apache.lucene.index.VectorEncoding;
@@ -18,30 +19,51 @@
  * For more details, please refer to <a href="https://github.com/facebookresearch/faiss/blob/main/faiss/IndexHNSW.h">...</a>
  */
 public class FaissHNSWFlatIndex extends FaissIndex {
+    public static final String IHNF = "IHNf";
+    public static final String IHNS = "IHNs";
+
+    @Getter
+    private FaissHNSW hnsw = new FaissHNSW();
+    private FaissIndex storage;
+    private VectorEncoding vectorEncoding;
+
     public FaissHNSWFlatIndex(final String indexType) {
         super(indexType);
+
+        // Set encoding
+        if (indexType.equals(IHNF)) {
+            vectorEncoding = VectorEncoding.FLOAT32;
+        } else if (indexType.equals(IHNS)) {
+            vectorEncoding = VectorEncoding.BYTE;
+        } else {
+            throw new IllegalStateException("Unsupported index type: " + indexType + " in " + FaissHNSWFlatIndex.class.getSimpleName());
+        }
     }
 
     @Override
     protected void doLoad(IndexInput input) throws IOException {
-        // TODO(KDY) : This will be covered in part-3 (FAISS HNSW).
+        // Read common header
+        readCommonHeader(input);
+
+        // Partial load HNSW graph
+        hnsw = FaissHNSW.load(input, getTotalNumberOfVectors());
+
+        // Partial load flat vector storage
+        storage = FaissIndex.load(input);
     }
 
     @Override
     public VectorEncoding getVectorEncoding() {
-        // TODO(KDY) : This will be covered in part-3 (FAISS HNSW).
-        return null;
+        return vectorEncoding;
     }
 
     @Override
     public FloatVectorValues getFloatValues(IndexInput indexInput) throws IOException {
-        // TODO(KDY) : This will be covered in part-3 (FAISS HNSW).
-        return null;
+        return storage.getFloatValues(indexInput);
     }
 
     @Override
     public ByteVectorValues getByteValues(IndexInput indexInput) throws IOException {
-        // TODO(KDY) : This will be covered in part-3 (FAISS HNSW).
-        return null;
+        return storage.getByteValues(indexInput);
     }
 }

src/main/java/org/opensearch/knn/memoryoptsearch/faiss/FaissSection.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.memoryoptsearch.faiss;
+
+import lombok.Getter;
+import org.apache.lucene.store.IndexInput;
+
+import java.io.IOException;
+
+/**
+ * This section maps to a section in FAISS index with a starting offset and the section size.
+ * A FAISS index file consists of multiple logical sections, each beginning with four bytes indicating an index type. A section may contain
+ * a nested section or vector storage, forming a tree structure with a top-level index as the starting point.
+ *
+ * Ex: FAISS index file
+ * +------------+ -> 0
+ * +            +
+ * +    IxMp    + -> FaissSection(offset=0, section_size=120)
+ * +------------+ -> 120
+ * +            +
+ * +    IHNf    + -> FaissSection(offset=120, section_size=380)
+ * +------------+ -> 500
+ * +            +
+ * +    IxF2    + -> FaissSection(offset=500, section_size=700)
+ * +------------+ -> 1200
+ *
+ */
+public class FaissSection {
+    @Getter
+    protected long baseOffset;
+    @Getter
+    protected long sectionSize;
+
+    /**
+     * Mark the starting offset and the size of section then skip to the next section.
+     *
+     * @param input Input read stream.
+     * @param singleElementSize Size of atomic element. In file, it only stores the number of elements and the size of element will be
+     *                          used to calculate the actual size of section. Ex: size=100, element=int, then the actual section size=400.
+     * @throws IOException
+     */
+    public void markSection(IndexInput input, int singleElementSize) throws IOException {
+        this.sectionSize = input.readLong() * singleElementSize;
+        this.baseOffset = input.getFilePointer();
+        // Skip the whole section and jump to the next section in the file.
+        try {
+            input.seek(baseOffset + sectionSize);
+        } catch (IOException e) {
+            throw new IOException("Failed to partial load where baseOffset=" + baseOffset + ", sectionSize=" + sectionSize, e);
+        }
+    }
+}

src/main/java/org/opensearch/knn/memoryoptsearch/faiss/IndexTypeToFaissIndexMapping.java

@@ -10,7 +10,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.function.Supplier;
+import java.util.function.Function;
 
 /**
  * This table maintains a mapping between FAISS index types and their corresponding index implementations.
@@ -20,12 +20,14 @@
  */
 @UtilityClass
 public class IndexTypeToFaissIndexMapping {
-    private static final Map<String, Supplier<FaissIndex>> INDEX_TYPE_TO_FAISS_INDEX;
+    private static final Map<String, Function<String, FaissIndex>> INDEX_TYPE_TO_FAISS_INDEX;
 
     static {
-        final Map<String, Supplier<FaissIndex>> mapping = new HashMap<>();
+        final Map<String, Function<String, FaissIndex>> mapping = new HashMap<>();
 
-        mapping.put(FaissIdMapIndex.IXMP, FaissIdMapIndex::new);
+        mapping.put(FaissIdMapIndex.IXMP, (indexType) -> new FaissIdMapIndex());
+        mapping.put(FaissHNSWFlatIndex.IHNF, FaissHNSWFlatIndex::new);
+        mapping.put(FaissHNSWFlatIndex.IHNS, FaissHNSWFlatIndex::new);
 
         INDEX_TYPE_TO_FAISS_INDEX = Collections.unmodifiableMap(mapping);
     }
@@ -37,9 +39,9 @@ public class IndexTypeToFaissIndexMapping {
      * @return Actual implementation that is corresponding to the given index type.
      */
     public FaissIndex getFaissIndex(final String indexType) {
-        final Supplier<FaissIndex> faissIndexSupplier = INDEX_TYPE_TO_FAISS_INDEX.get(indexType);
+        final Function<String, FaissIndex> faissIndexSupplier = INDEX_TYPE_TO_FAISS_INDEX.get(indexType);
         if (faissIndexSupplier != null) {
-            return faissIndexSupplier.get();
+            return faissIndexSupplier.apply(indexType);
         }
         throw new UnsupportedFaissIndexException("Index type [" + indexType + "] is not supported.");
     }

src/main/java/org/opensearch/knn/memoryoptsearch/faiss/LuceneFaissHnswGraph.java

@@ -0,0 +1,117 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.memoryoptsearch.faiss;
+
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.util.hnsw.HnswGraph;
+
+import java.io.IOException;
+
+import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
+
+/**
+ * This graph implements Lucene's HNSW graph interface using the FAISS HNSW graph. Conceptually, both libraries represent the graph
+ * similarly, maintaining a list of neighbor IDs. This implementation acts as a bridge, enabling Lucene's HNSW graph searcher to perform
+ * vector searches on a FAISS index.
+ */
+public class LuceneFaissHnswGraph extends HnswGraph {
+    private final FaissHNSW faissHnsw;
+    private final IndexInput indexInput;
+    private final int numVectors;
+    private int[] neighborIdList;
+    private int numNeighbors;
+    private int nextNeighborIndex;
+
+    public LuceneFaissHnswGraph(final FaissHNSW faissHNSW, final int numVectors, final IndexInput indexInput) {
+        this.faissHnsw = faissHNSW;
+        this.indexInput = indexInput;
+        this.numVectors = numVectors;
+    }
+
+    /**
+     * Seek to the starting offset of neighbor list of `internalVectorId` at the given `level`.
+     * In which, it will load all ids into a buffer array.
+     * @param level
+     * @param internalVectorId
+     */
+    @Override
+    public void seek(int level, int internalVectorId) {
+        // Get a relative starting offset of neighbor list at `level`.
+        long o = faissHnsw.getOffsets()[internalVectorId];
+
+        // `begin` and `end` represent for a pair of staring offset and end offset.
+        // But, what `end` represents is the maximum offset a neighbor list at a level can have.
+        // Therefore, it is required to traverse a list until getting a terminal `-1`.
+        // Ex: [1, 5, 20, 100, -1, -1, ..., -1]
+        final long begin = o + faissHnsw.getCumNumberNeighborPerLevel()[level];
+        final long end = o + faissHnsw.getCumNumberNeighborPerLevel()[level + 1];
+        loadNeighborIdList(begin, end);
+    }
+
+    private void loadNeighborIdList(final long begin, final long end) {
+        // Make sure we have sufficient space for neighbor list
+        final long maxLength = end - begin;
+        if (neighborIdList == null || neighborIdList.length < maxLength) {
+            neighborIdList = new int[(int) (maxLength * 1.5)];
+        }
+
+        // Seek to the first offset of neighbor list
+        try {
+            indexInput.seek(faissHnsw.getNeighbors().getBaseOffset() + Integer.BYTES * begin);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+
+        // Fill the array with neighbor ids
+        int index = 0;
+        try {
+            for (long i = begin; i < end; i++) {
+                final int neighborId = indexInput.readInt();
+                if (neighborId >= 0) {
+                    neighborIdList[index++] = neighborId;
+                } else {
+                    break;
+                }
+            }
+
+            // Set variables for navigation
+            numNeighbors = index;
+            nextNeighborIndex = 0;
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public int size() {
+        return numVectors;
+    }
+
+    @Override
+    public int nextNeighbor() {
+        if (nextNeighborIndex < numNeighbors) {
+            return neighborIdList[nextNeighborIndex++];
+        }
+
+        // Neighbor list has been exhausted.
+        return NO_MORE_DOCS;
+    }
+
+    @Override
+    public int numLevels() {
+        return faissHnsw.getMaxLevel();
+    }
+
+    @Override
+    public int entryNode() {
+        return faissHnsw.getEntryPoint();
+    }
+
+    @Override
+    public NodesIterator getNodesOnLevel(int i) {
+        throw new UnsupportedOperationException();
+    }
+}