Address Review Comments

Signed-off-by: Naveen Tatikonda <[email protected]>

Author: naveentatikonda

jni/include/faiss_index_service.h

@@ -87,6 +87,7 @@ class IndexService {
 
 protected:
     virtual void allocIndex(faiss::Index * index, size_t dim, size_t numVectors);
+    virtual jlong initAndAllocateIndex(std::unique_ptr<faiss::Index> &index, size_t threadCount, size_t dim, size_t numVectors);
 
     std::unique_ptr<FaissMethods> faissMethods;
 };  // class IndexService
@@ -157,10 +158,11 @@ class BinaryIndexService final : public IndexService {
      * @param templateIndexJ template index
      * @return memory address of the native index object
      */
-    virtual jlong initIndexFromTemplate(knn_jni::JNIUtilInterface *jniUtil, JNIEnv *env, int dim, int numVectors, int threadCount, jbyteArray templateIndexJ);
+    jlong initIndexFromTemplate(knn_jni::JNIUtilInterface *jniUtil, JNIEnv *env, int dim, int numVectors, int threadCount, jbyteArray templateIndexJ) final;
 
 protected:
     void allocIndex(faiss::Index * index, size_t dim, size_t numVectors) final;
+    jlong initAndAllocateIndex(std::unique_ptr<faiss::IndexBinary> &index, size_t threadCount, size_t dim, size_t numVectors);
 };  // class BinaryIndexService
 
 /**
@@ -229,10 +231,11 @@ class ByteIndexService final : public IndexService {
      * @param templateIndexJ template index
      * @return memory address of the native index object
      */
-    virtual jlong initIndexFromTemplate(knn_jni::JNIUtilInterface *jniUtil, JNIEnv *env, int dim, int numVectors, int threadCount, jbyteArray templateIndexJ);
+    jlong initIndexFromTemplate(knn_jni::JNIUtilInterface *jniUtil, JNIEnv *env, int dim, int numVectors, int threadCount, jbyteArray templateIndexJ) final;
 
  protected:
     void allocIndex(faiss::Index * index, size_t dim, size_t numVectors) final;
+    jlong initAndAllocateIndex(std::unique_ptr<faiss::Index> &index, size_t threadCount, size_t dim, size_t numVectors) final;
 };  // class ByteIndexService
 
 }

jni/src/faiss_index_service.cpp

@@ -68,6 +68,25 @@ void IndexService::allocIndex(faiss::Index * index, size_t dim, size_t numVector
     }
 }
 
+jlong IndexService::initAndAllocateIndex(std::unique_ptr<faiss::Index> &index, size_t threadCount, size_t dim, size_t numVectors) {
+    // Set thread count if it is passed in as a parameter. Setting this variable will only impact the current thread
+    if (threadCount != 0) {
+        omp_set_num_threads(threadCount);
+    }
+
+    std::unique_ptr<faiss::IndexIDMap> idMap (faissMethods->indexIdMap(index.get()));
+    //Makes sure the index is deleted when the destructor is called, this cannot be passed in the constructor
+    idMap->own_fields = true;
+
+    // TODO: allocIndex for IVF
+    allocIndex(dynamic_cast<faiss::Index *>(idMap->index), dim, numVectors);
+
+    //Release the ownership so as to make sure not delete the underlying index that is created. The index is needed later
+    //in insert and write operations
+    index.release();
+    return reinterpret_cast<jlong>(idMap.release());
+}
+
 jlong IndexService::initIndex(
         knn_jni::JNIUtilInterface * jniUtil,
         JNIEnv * env,
@@ -81,11 +100,6 @@ jlong IndexService::initIndex(
     // Create index using Faiss factory method
     std::unique_ptr<faiss::Index> index(faissMethods->indexFactory(dim, indexDescription.c_str(), metric));
 
-    // Set thread count if it is passed in as a parameter. Setting this variable will only impact the current thread
-    if (threadCount != 0) {
-        omp_set_num_threads(threadCount);
-    }
-
     // Add extra parameters that cant be configured with the index factory
     SetExtraParameters<faiss::Index, faiss::IndexIVF, faiss::IndexHNSW>(jniUtil, env, parameters, index.get());
 
@@ -94,16 +108,7 @@ jlong IndexService::initIndex(
         throw std::runtime_error("Index is not trained");
     }
 
-    std::unique_ptr<faiss::IndexIDMap> idMap (faissMethods->indexIdMap(index.get()));
-    //Makes sure the index is deleted when the destructor is called, this cannot be passed in the constructor
-    idMap->own_fields = true;
-
-    allocIndex(dynamic_cast<faiss::Index *>(idMap->index), dim, numVectors);
-
-    //Release the ownership so as to make sure not delete the underlying index that is created. The index is needed later
-    //in insert and write operations
-    index.release();
-    return reinterpret_cast<jlong>(idMap.release());
+    return initAndAllocateIndex(index, threadCount, dim, numVectors);
 }
 
 void IndexService::insertToIndex(
@@ -178,16 +183,31 @@ jlong IndexService::initIndexFromTemplate(
     std::unique_ptr<faiss::Index> index;
     index.reset(faiss::read_index(&vectorIoReader, 0));
 
+    return initAndAllocateIndex(index, threadCount, dim, numVectors);
+}
+
+BinaryIndexService::BinaryIndexService(std::unique_ptr<FaissMethods> _faissMethods)
+  : IndexService(std::move(_faissMethods)) {
+}
+
+void BinaryIndexService::allocIndex(faiss::Index * index, size_t dim, size_t numVectors) {
+    if (auto * indexBinaryHNSW = dynamic_cast<faiss::IndexBinaryHNSW *>(index)) {
+        auto * indexBinaryFlat = dynamic_cast<faiss::IndexBinaryFlat *>(indexBinaryHNSW->storage);
+        indexBinaryFlat->xb.reserve(dim * numVectors / 8);
+    }
+}
+
+jlong BinaryIndexService::initAndAllocateIndex(std::unique_ptr<faiss::IndexBinary> &index, size_t threadCount, size_t dim, size_t numVectors) {
     // Set thread count if it is passed in as a parameter. Setting this variable will only impact the current thread
     if (threadCount != 0) {
         omp_set_num_threads(threadCount);
     }
 
-    std::unique_ptr<faiss::IndexIDMap> idMap (faissMethods->indexIdMap(index.get()));
+    std::unique_ptr<faiss::IndexBinaryIDMap> idMap(faissMethods->indexBinaryIdMap(index.get()));
     //Makes sure the index is deleted when the destructor is called, this cannot be passed in the constructor
     idMap->own_fields = true;
 
-    // TODO: allocIndex
+    // TODO: allocIndex for IVF
     allocIndex(dynamic_cast<faiss::Index *>(idMap->index), dim, numVectors);
 
     //Release the ownership so as to make sure not delete the underlying index that is created. The index is needed later
@@ -196,17 +216,6 @@ jlong IndexService::initIndexFromTemplate(
     return reinterpret_cast<jlong>(idMap.release());
 }
 
-BinaryIndexService::BinaryIndexService(std::unique_ptr<FaissMethods> _faissMethods)
-  : IndexService(std::move(_faissMethods)) {
-}
-
-void BinaryIndexService::allocIndex(faiss::Index * index, size_t dim, size_t numVectors) {
-    if (auto * indexBinaryHNSW = dynamic_cast<faiss::IndexBinaryHNSW *>(index)) {
-        auto * indexBinaryFlat = dynamic_cast<faiss::IndexBinaryFlat *>(indexBinaryHNSW->storage);
-        indexBinaryFlat->xb.reserve(dim * numVectors / 8);
-    }
-}
-
 jlong BinaryIndexService::initIndex(
         knn_jni::JNIUtilInterface * jniUtil,
         JNIEnv * env,
@@ -219,10 +228,6 @@ jlong BinaryIndexService::initIndex(
     ) {
     // Create index using Faiss factory method
     std::unique_ptr<faiss::IndexBinary> index(faissMethods->indexBinaryFactory(dim, indexDescription.c_str()));
-    // Set thread count if it is passed in as a parameter. Setting this variable will only impact the current thread
-    if (threadCount != 0) {
-       omp_set_num_threads(threadCount);
-    }
 
     // Add extra parameters that cant be configured with the index factory
     SetExtraParameters<faiss::IndexBinary, faiss::IndexBinaryIVF, faiss::IndexBinaryHNSW>(jniUtil, env, parameters, index.get());
@@ -232,16 +237,7 @@ jlong BinaryIndexService::initIndex(
         throw std::runtime_error("Index is not trained");
     }
 
-    std::unique_ptr<faiss::IndexBinaryIDMap> idMap(faissMethods->indexBinaryIdMap(index.get()));
-    //Makes sure the index is deleted when the destructor is called
-    idMap->own_fields = true;
-
-    allocIndex(dynamic_cast<faiss::Index *>(idMap->index), dim, numVectors);
-
-    //Release the ownership so as to make sure not delete the underlying index that is created. The index is needed later
-    //in insert and write operations
-    index.release();
-    return reinterpret_cast<jlong>(idMap.release());
+    return initAndAllocateIndex(index, threadCount, dim, numVectors);
 }
 
 void BinaryIndexService::insertToIndex(
@@ -319,16 +315,32 @@ jlong BinaryIndexService::initIndexFromTemplate(
     std::unique_ptr<faiss::IndexBinary> index;
     index.reset(faiss::read_index_binary(&vectorIoReader, 0));
 
+    return initAndAllocateIndex(index, threadCount, dim, numVectors);
+}
+
+ByteIndexService::ByteIndexService(std::unique_ptr<FaissMethods> _faissMethods)
+  : IndexService(std::move(_faissMethods)) {
+}
+
+void ByteIndexService::allocIndex(faiss::Index * index, size_t dim, size_t numVectors) {
+    if (auto * indexHNSWSQ = dynamic_cast<faiss::IndexHNSWSQ *>(index)) {
+        if(auto * indexScalarQuantizer = dynamic_cast<faiss::IndexScalarQuantizer *>(indexHNSWSQ->storage)) {
+            indexScalarQuantizer->codes.reserve(indexScalarQuantizer->code_size * numVectors);
+        }
+    }
+}
+
+jlong ByteIndexService::initAndAllocateIndex(std::unique_ptr<faiss::Index> &index, size_t threadCount, size_t dim, size_t numVectors) {
     // Set thread count if it is passed in as a parameter. Setting this variable will only impact the current thread
     if (threadCount != 0) {
         omp_set_num_threads(threadCount);
     }
 
-    std::unique_ptr<faiss::IndexBinaryIDMap> idMap (faissMethods->indexBinaryIdMap(index.get()));
+    std::unique_ptr<faiss::IndexIDMap> idMap (faissMethods->indexIdMap(index.get()));
     //Makes sure the index is deleted when the destructor is called, this cannot be passed in the constructor
     idMap->own_fields = true;
 
-    // TODO: allocIndex
+    // TODO: allocIndex for IVF
     allocIndex(dynamic_cast<faiss::Index *>(idMap->index), dim, numVectors);
 
     //Release the ownership so as to make sure not delete the underlying index that is created. The index is needed later
@@ -337,18 +349,6 @@ jlong BinaryIndexService::initIndexFromTemplate(
     return reinterpret_cast<jlong>(idMap.release());
 }
 
-ByteIndexService::ByteIndexService(std::unique_ptr<FaissMethods> _faissMethods)
-  : IndexService(std::move(_faissMethods)) {
-}
-
-void ByteIndexService::allocIndex(faiss::Index * index, size_t dim, size_t numVectors) {
-    if (auto * indexHNSWSQ = dynamic_cast<faiss::IndexHNSWSQ *>(index)) {
-        if(auto * indexScalarQuantizer = dynamic_cast<faiss::IndexScalarQuantizer *>(indexHNSWSQ->storage)) {
-            indexScalarQuantizer->codes.reserve(indexScalarQuantizer->code_size * numVectors);
-        }
-    }
-}
-
 jlong ByteIndexService::initIndex(
         knn_jni::JNIUtilInterface * jniUtil,
         JNIEnv * env,
@@ -362,11 +362,6 @@ jlong ByteIndexService::initIndex(
     // Create index using Faiss factory method
     std::unique_ptr<faiss::Index> index(faissMethods->indexFactory(dim, indexDescription.c_str(), metric));
 
-    // Set thread count if it is passed in as a parameter. Setting this variable will only impact the current thread
-    if (threadCount != 0) {
-        omp_set_num_threads(threadCount);
-    }
-
     // Add extra parameters that cant be configured with the index factory
     SetExtraParameters<faiss::Index, faiss::IndexIVF, faiss::IndexHNSW>(jniUtil, env, parameters, index.get());
 
@@ -375,16 +370,7 @@ jlong ByteIndexService::initIndex(
         throw std::runtime_error("Index is not trained");
     }
 
-    std::unique_ptr<faiss::IndexIDMap> idMap (faissMethods->indexIdMap(index.get()));
-    //Makes sure the index is deleted when the destructor is called, this cannot be passed in the constructor
-    idMap->own_fields = true;
-
-    allocIndex(dynamic_cast<faiss::Index *>(idMap->index), dim, numVectors);
-
-    //Release the ownership so as to make sure not delete the underlying index that is created. The index is needed later
-    //in insert and write operations
-    index.release();
-    return reinterpret_cast<jlong>(idMap.release());
+    return initAndAllocateIndex(index, threadCount, dim, numVectors);
 }
 
 void ByteIndexService::insertToIndex(
@@ -477,22 +463,7 @@ jlong ByteIndexService::initIndexFromTemplate(
     std::unique_ptr<faiss::Index> index;
     index.reset(faiss::read_index(&vectorIoReader, 0));
 
-    // Set thread count if it is passed in as a parameter. Setting this variable will only impact the current thread
-    if (threadCount != 0) {
-        omp_set_num_threads(threadCount);
-    }
-
-    std::unique_ptr<faiss::IndexIDMap> idMap (faissMethods->indexIdMap(index.get()));
-    //Makes sure the index is deleted when the destructor is called, this cannot be passed in the constructor
-    idMap->own_fields = true;
-
-    // TODO: allocIndex
-    allocIndex(dynamic_cast<faiss::Index *>(idMap->index), dim, numVectors);
-
-    //Release the ownership so as to make sure not delete the underlying index that is created. The index is needed later
-    //in insert and write operations
-    index.release();
-    return reinterpret_cast<jlong>(idMap.release());
+    return initAndAllocateIndex(index, threadCount, dim, numVectors);
 }
 } // namespace faiss_wrapper
 } // namesapce knn_jni

src/main/java/org/opensearch/knn/index/KNNIndexShard.java

@@ -193,12 +193,7 @@ List<EngineFileContext> getEngineFileContexts(IndexReader indexReader, KNNEngine
                             fileExtension,
                             spaceType,
                             modelId,
-                            quantizationConfig == QuantizationConfig.EMPTY
-                                ? VectorDataType.get(
-                                    fieldInfo.attributes().getOrDefault(VECTOR_DATA_TYPE_FIELD, VectorDataType.FLOAT.getValue())
-                                )
-                                : quantizationConfig.getQuantizationType() == ScalarQuantizationType.EIGHT_BIT ? VectorDataType.BYTE
-                                : VectorDataType.BINARY
+                            getVectorDataType(quantizationConfig, fieldInfo)
                         )
                     );
                 }
@@ -228,6 +223,16 @@ List<EngineFileContext> getEngineFileContexts(
             .collect(Collectors.toList());
     }
 
+    private VectorDataType getVectorDataType(QuantizationConfig quantizationConfig, FieldInfo fieldInfo) {
+        if (quantizationConfig == QuantizationConfig.EMPTY) {
+            return VectorDataType.get(fieldInfo.attributes().getOrDefault(VECTOR_DATA_TYPE_FIELD, VectorDataType.FLOAT.getValue()));
+        }
+        if (quantizationConfig.getQuantizationType() == ScalarQuantizationType.EIGHT_BIT) {
+            return VectorDataType.BYTE;
+        }
+        return VectorDataType.BINARY;
+    }
+
     @AllArgsConstructor
     @Getter
     @VisibleForTesting

src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsWriter.java

@@ -257,18 +257,8 @@ private QuantizationState train(
         if (quantizationParams != null && totalLiveDocs > 0) {
             KNNVectorValues<?> knnVectorValues = knnVectorValuesSupplier.get();
 
-            // We will not be writing quantization state for 8 bits into a segment file because we are not quantizing the query vectors and
-            // we are storing the template index after training in the quantization state to use it later in the index build strategy for
-            // ingesting data
-            if ((quantizationParams.getTypeIdentifier()).equals(
-                ScalarQuantizationParams.generateTypeIdentifier(ScalarQuantizationType.EIGHT_BIT)
-            )) {
-                quantizationState = quantizationService.train(quantizationParams, knnVectorValues, totalLiveDocs, fieldInfo);
-            } else {
-                initQuantizationStateWriterIfNecessary();
-                quantizationState = quantizationService.train(quantizationParams, knnVectorValues, totalLiveDocs, fieldInfo);
-                quantizationStateWriter.writeState(fieldInfo.getFieldNumber(), quantizationState);
-            }
+            quantizationState = quantizationService.train(quantizationParams, knnVectorValues, totalLiveDocs, fieldInfo);
+            writeQuantizationState(quantizationParams, quantizationState, fieldInfo.getFieldNumber());
         }
 
         return quantizationState;
@@ -289,11 +279,24 @@ private int getLiveDocs(KNNVectorValues<?> vectorValues) throws IOException {
         return liveDocs;
     }
 
-    private void initQuantizationStateWriterIfNecessary() throws IOException {
+    private void writeQuantizationState(QuantizationParams quantizationParams, QuantizationState quantizationState, int fieldNumber)
+        throws IOException {
+
+        // We will not be writing quantization state for 8 bits into a segment file because we are not quantizing the query vectors and
+        // we are storing the template index after training in the quantization state to use it later in the index build strategy for
+        // ingesting data
+        if ((quantizationParams.getTypeIdentifier()).equals(
+            ScalarQuantizationParams.generateTypeIdentifier(ScalarQuantizationType.EIGHT_BIT)
+        )) {
+            return;
+        }
+
+        // Initialize quantization state writer if required
         if (quantizationStateWriter == null) {
             quantizationStateWriter = new KNN990QuantizationStateWriter(segmentWriteState);
             quantizationStateWriter.writeHeader(segmentWriteState);
         }
+        quantizationStateWriter.writeState(fieldNumber, quantizationState);
     }
 
     private boolean shouldSkipBuildingVectorDataStructure(final long docCount) {