Address Review Comments

Signed-off-by: Naveen Tatikonda <navtat@amazon.com>

Author: naveentatikonda

jni/src/faiss_index_service.cpp

@@ -187,7 +187,7 @@ jlong IndexService::initIndexFromTemplate(
     //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
@@ -328,7 +328,7 @@ jlong BinaryIndexService::initIndexFromTemplate(
     //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
@@ -486,7 +486,7 @@ jlong ByteIndexService::initIndexFromTemplate(
     //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

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) {

src/main/java/org/opensearch/knn/index/codec/nativeindex/MemOptimizedNativeIndexBuildStrategy.java

@@ -58,31 +58,7 @@ public void buildAndWriteIndex(final BuildIndexParams indexInfo) throws IOExcept
         KNNEngine engine = indexInfo.getKnnEngine();
         Map<String, Object> indexParameters = indexInfo.getParameters();
         IndexBuildSetup indexBuildSetup = QuantizationIndexUtils.prepareIndexBuild(knnVectorValues, indexInfo);
-        long indexMemoryAddress;
-
-        if (isTemplate(indexInfo)) {
-            // Initialize the index from Template
-            indexMemoryAddress = AccessController.doPrivileged(
-                (PrivilegedAction<Long>) () -> JNIService.initIndexFromTemplate(
-                    indexInfo.getTotalLiveDocs(),
-                    indexBuildSetup.getDimensions(),
-                    indexParameters,
-                    engine,
-                    getIndexTemplate(indexInfo)
-                )
-            );
-
-        } else {
-            // Initialize the index
-            indexMemoryAddress = AccessController.doPrivileged(
-                (PrivilegedAction<Long>) () -> JNIService.initIndex(
-                    indexInfo.getTotalLiveDocs(),
-                    indexBuildSetup.getDimensions(),
-                    indexParameters,
-                    engine
-                )
-            );
-        }
+        long indexMemoryAddress = initializeIndex(indexInfo, indexBuildSetup, indexParameters, engine);
 
         try (
             final OffHeapVectorTransfer vectorTransfer = getVectorTransfer(
@@ -150,6 +126,36 @@ public void buildAndWriteIndex(final BuildIndexParams indexInfo) throws IOExcept
         }
     }
 
+    private long initializeIndex(
+        BuildIndexParams indexInfo,
+        IndexBuildSetup indexBuildSetup,
+        Map<String, Object> indexParameters,
+        KNNEngine engine
+    ) {
+        if (isTemplate(indexInfo)) {
+            // Initialize the index from Template
+            return AccessController.doPrivileged(
+                (PrivilegedAction<Long>) () -> JNIService.initIndexFromTemplate(
+                    indexInfo.getTotalLiveDocs(),
+                    indexBuildSetup.getDimensions(),
+                    indexParameters,
+                    engine,
+                    getIndexTemplate(indexInfo)
+                )
+            );
+
+        }
+        // Initialize the index
+        return AccessController.doPrivileged(
+            (PrivilegedAction<Long>) () -> JNIService.initIndex(
+                indexInfo.getTotalLiveDocs(),
+                indexBuildSetup.getDimensions(),
+                indexParameters,
+                engine
+            )
+        );
+    }
+
     private static boolean isTemplate(final BuildIndexParams indexInfo) {
         return (indexInfo.getQuantizationState() instanceof ByteScalarQuantizationState);
     }