diff --git a/jni/src/faiss_wrapper.cpp b/jni/src/faiss_wrapper.cpp index 03022e0d3b..4a93926207 100644 --- a/jni/src/faiss_wrapper.cpp +++ b/jni/src/faiss_wrapper.cpp @@ -22,6 +22,7 @@ #include "faiss/Index.h" #include "faiss/impl/IDSelector.h" #include "faiss/IndexIVFPQ.h" +#include "commons.h" #include #include @@ -161,6 +162,10 @@ void knn_jni::faiss_wrapper::CreateIndex(knn_jni::JNIUtilInterface * jniUtil, JN // Write the index to disk std::string indexPathCpp(jniUtil->ConvertJavaStringToCppString(env, indexPathJ)); faiss::write_index(&idMap, indexPathCpp.c_str()); + // Releasing the vectorsAddressJ memory as that is not required once we have created the index. + // This is not the ideal approach, please refer this gh issue for long term solution: + // https://github.com/opensearch-project/k-NN/issues/1600 + delete inputVectors; } void knn_jni::faiss_wrapper::CreateIndexFromTemplate(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jintArray idsJ, @@ -221,7 +226,10 @@ void knn_jni::faiss_wrapper::CreateIndexFromTemplate(knn_jni::JNIUtilInterface * auto idVector = jniUtil->ConvertJavaIntArrayToCppIntVector(env, idsJ); faiss::IndexIDMap idMap = faiss::IndexIDMap(indexWriter.get()); idMap.add_with_ids(numVectors, inputVectors->data(), idVector.data()); - + // Releasing the vectorsAddressJ memory as that is not required once we have created the index. + // This is not the ideal approach, please refer this gh issue for long term solution: + // https://github.com/opensearch-project/k-NN/issues/1600 + delete inputVectors; // Write the index to disk std::string indexPathCpp(jniUtil->ConvertJavaStringToCppString(env, indexPathJ)); faiss::write_index(&idMap, indexPathCpp.c_str()); diff --git a/jni/src/nmslib_wrapper.cpp b/jni/src/nmslib_wrapper.cpp index 00b7f4311c..311bc1e543 100644 --- a/jni/src/nmslib_wrapper.cpp +++ b/jni/src/nmslib_wrapper.cpp @@ -20,6 +20,7 @@ #include "methodfactory.h" #include "spacefactory.h" #include "space.h" +#include "commons.h" #include #include @@ -153,6 +154,12 @@ void knn_jni::nmslib_wrapper::CreateIndex(knn_jni::JNIUtilInterface * jniUtil, J } jniUtil->ReleaseIntArrayElements(env, idsJ, idsCpp, JNI_ABORT); + // Releasing the vectorsAddressJ memory as that is not required once we have created the index. + // This is not the ideal approach, please refer this gh issue for long term solution: + // https://github.com/opensearch-project/k-NN/issues/1600 + //commons::freeVectorData(vectorsAddressJ); + delete inputVectors; + std::unique_ptr> index; index.reset(similarity::MethodFactoryRegistry::Instance().CreateMethod(false, "hnsw", spaceTypeCpp, *(space), dataset)); index->CreateIndex(similarity::AnyParams(indexParameters)); diff --git a/jni/tests/faiss_wrapper_test.cpp b/jni/tests/faiss_wrapper_test.cpp index ff9f638937..05854f7edb 100644 --- a/jni/tests/faiss_wrapper_test.cpp +++ b/jni/tests/faiss_wrapper_test.cpp @@ -30,13 +30,13 @@ TEST(FaissCreateIndexTest, BasicAssertions) { // Define the data faiss::idx_t numIds = 200; std::vector ids; - std::vector vectors; + auto *vectors = new std::vector(); int dim = 2; - vectors.reserve(dim * numIds); + vectors->reserve(dim * numIds); for (int64_t i = 0; i < numIds; ++i) { ids.push_back(i); for (int j = 0; j < dim; ++j) { - vectors.push_back(test_util::RandomFloat(-500.0, 500.0)); + vectors->push_back(test_util::RandomFloat(-500.0, 500.0)); } } @@ -55,12 +55,12 @@ TEST(FaissCreateIndexTest, BasicAssertions) { EXPECT_CALL(mockJNIUtil, GetJavaObjectArrayLength( jniEnv, reinterpret_cast(&vectors))) - .WillRepeatedly(Return(vectors.size())); + .WillRepeatedly(Return(vectors->size())); // Create the index knn_jni::faiss_wrapper::CreateIndex( &mockJNIUtil, jniEnv, reinterpret_cast(&ids), - (jlong) &vectors, dim , (jstring)&indexPath, + (jlong) vectors, dim , (jstring)&indexPath, (jobject)¶metersMap); // Make sure index can be loaded @@ -74,13 +74,13 @@ TEST(FaissCreateIndexFromTemplateTest, BasicAssertions) { // Define the data faiss::idx_t numIds = 100; std::vector ids; - std::vector vectors; + auto *vectors = new std::vector(); int dim = 2; - vectors.reserve(dim * numIds); + vectors->reserve(dim * numIds); for (int64_t i = 0; i < numIds; ++i) { ids.push_back(i); for (int j = 0; j < dim; ++j) { - vectors.push_back(test_util::RandomFloat(-500.0, 500.0)); + vectors->push_back(test_util::RandomFloat(-500.0, 500.0)); } } @@ -99,7 +99,7 @@ TEST(FaissCreateIndexFromTemplateTest, BasicAssertions) { EXPECT_CALL(mockJNIUtil, GetJavaObjectArrayLength( jniEnv, reinterpret_cast(&vectors))) - .WillRepeatedly(Return(vectors.size())); + .WillRepeatedly(Return(vectors->size())); std::string spaceType = knn_jni::L2; std::unordered_map parametersMap; @@ -107,7 +107,7 @@ TEST(FaissCreateIndexFromTemplateTest, BasicAssertions) { knn_jni::faiss_wrapper::CreateIndexFromTemplate( &mockJNIUtil, jniEnv, reinterpret_cast(&ids), - (jlong)&vectors, dim, (jstring)&indexPath, + (jlong)vectors, dim, (jstring)&indexPath, reinterpret_cast(&(vectorIoWriter.data)), (jobject) ¶metersMap ); @@ -480,13 +480,13 @@ TEST(FaissCreateHnswSQfp16IndexTest, BasicAssertions) { // Define the data faiss::idx_t numIds = 200; std::vector ids; - std::vector vectors; + auto *vectors = new std::vector(); int dim = 2; - vectors.reserve(dim * numIds); + vectors->reserve(dim * numIds); for (int64_t i = 0; i < numIds; ++i) { ids.push_back(i); for (int j = 0; j < dim; ++j) { - vectors.push_back(test_util::RandomFloat(-500.0, 500.0)); + vectors->push_back(test_util::RandomFloat(-500.0, 500.0)); } } @@ -505,12 +505,12 @@ TEST(FaissCreateHnswSQfp16IndexTest, BasicAssertions) { EXPECT_CALL(mockJNIUtil, GetJavaObjectArrayLength( jniEnv, reinterpret_cast(&vectors))) - .WillRepeatedly(Return(vectors.size())); + .WillRepeatedly(Return(vectors->size())); // Create the index knn_jni::faiss_wrapper::CreateIndex( &mockJNIUtil, jniEnv, reinterpret_cast(&ids), - (jlong)&vectors, dim, (jstring)&indexPath, + (jlong)vectors, dim, (jstring)&indexPath, (jobject)¶metersMap); // Make sure index can be loaded diff --git a/jni/tests/nmslib_wrapper_test.cpp b/jni/tests/nmslib_wrapper_test.cpp index cb9c14e23a..1fd9471b0a 100644 --- a/jni/tests/nmslib_wrapper_test.cpp +++ b/jni/tests/nmslib_wrapper_test.cpp @@ -39,13 +39,13 @@ TEST(NmslibCreateIndexTest, BasicAssertions) { // Define index data int numIds = 100; std::vector ids; - std::vector vectors; + auto *vectors = new std::vector(); int dim = 2; - vectors.reserve(dim * numIds); + vectors->reserve(dim * numIds); for (int64_t i = 0; i < numIds; ++i) { ids.push_back(i); for (int j = 0; j < dim; ++j) { - vectors.push_back(test_util::RandomFloat(-500.0, 500.0)); + vectors->push_back(test_util::RandomFloat(-500.0, 500.0)); } } @@ -67,7 +67,7 @@ TEST(NmslibCreateIndexTest, BasicAssertions) { EXPECT_CALL(mockJNIUtil, GetJavaObjectArrayLength( jniEnv, reinterpret_cast(&vectors))) - .WillRepeatedly(Return(vectors.size())); + .WillRepeatedly(Return(vectors->size())); EXPECT_CALL(mockJNIUtil, GetJavaIntArrayLength(jniEnv, reinterpret_cast(&ids))) @@ -76,7 +76,7 @@ TEST(NmslibCreateIndexTest, BasicAssertions) { // Create the index knn_jni::nmslib_wrapper::CreateIndex( &mockJNIUtil, jniEnv, reinterpret_cast(&ids), - (jlong) &vectors, dim, (jstring)&indexPath, + (jlong) vectors, dim, (jstring)&indexPath, (jobject)¶metersMap); // Make sure index can be loaded diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java index 2b0e29b555..48f37681d9 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java @@ -111,67 +111,57 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer, throws IOException { // Get values to be indexed BinaryDocValues values = valuesProducer.getBinary(field); - KNNCodecUtil.Pair pair = null; - try { - pair = KNNCodecUtil.getFloats(values); - if (pair.getVectorAddress() == 0 || pair.docs.length == 0) { - logger.info("Skipping engine index creation as there are no vectors or docs in the segment"); - return; - } - long arraySize = calculateArraySize(pair.docs.length, pair.getDimension(), pair.serializationMode); - if (isMerge) { - KNNGraphValue.MERGE_CURRENT_OPERATIONS.increment(); - KNNGraphValue.MERGE_CURRENT_DOCS.incrementBy(pair.docs.length); - KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.incrementBy(arraySize); - } - // Increment counter for number of graph index requests - KNNCounter.GRAPH_INDEX_REQUESTS.increment(); - final KNNEngine knnEngine = getKNNEngine(field); - final String engineFileName = buildEngineFileName( - state.segmentInfo.name, - knnEngine.getVersion(), - field.name, - knnEngine.getExtension() - ); - final String indexPath = Paths.get( - ((FSDirectory) (FilterDirectory.unwrap(state.directory))).getDirectory().toString(), - engineFileName - ).toString(); - KNNCodecUtil.Pair finalPair = pair; - NativeIndexCreator indexCreator; - // Create library index either from model or from scratch - if (field.attributes().containsKey(MODEL_ID)) { - String modelId = field.attributes().get(MODEL_ID); - Model model = ModelCache.getInstance().get(modelId); - if (model.getModelBlob() == null) { - throw new RuntimeException(String.format("There is no trained model with id \"%s\"", modelId)); - } - indexCreator = () -> createKNNIndexFromTemplate(model.getModelBlob(), finalPair, knnEngine, indexPath); - } else { - indexCreator = () -> createKNNIndexFromScratch(field, finalPair, knnEngine, indexPath); - } - - if (isMerge) { - recordMergeStats(pair.docs.length, arraySize); + KNNCodecUtil.Pair pair = KNNCodecUtil.getFloats(values); + if (pair.getVectorAddress() == 0 || pair.docs.length == 0) { + logger.info("Skipping engine index creation as there are no vectors or docs in the segment"); + return; + } + long arraySize = calculateArraySize(pair.docs.length, pair.getDimension(), pair.serializationMode); + if (isMerge) { + KNNGraphValue.MERGE_CURRENT_OPERATIONS.increment(); + KNNGraphValue.MERGE_CURRENT_DOCS.incrementBy(pair.docs.length); + KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.incrementBy(arraySize); + } + // Increment counter for number of graph index requests + KNNCounter.GRAPH_INDEX_REQUESTS.increment(); + final KNNEngine knnEngine = getKNNEngine(field); + final String engineFileName = buildEngineFileName( + state.segmentInfo.name, + knnEngine.getVersion(), + field.name, + knnEngine.getExtension() + ); + final String indexPath = Paths.get( + ((FSDirectory) (FilterDirectory.unwrap(state.directory))).getDirectory().toString(), + engineFileName + ).toString(); + NativeIndexCreator indexCreator; + // Create library index either from model or from scratch + if (field.attributes().containsKey(MODEL_ID)) { + String modelId = field.attributes().get(MODEL_ID); + Model model = ModelCache.getInstance().get(modelId); + if (model.getModelBlob() == null) { + throw new RuntimeException(String.format("There is no trained model with id \"%s\"", modelId)); } + indexCreator = () -> createKNNIndexFromTemplate(model.getModelBlob(), pair, knnEngine, indexPath); + } else { + indexCreator = () -> createKNNIndexFromScratch(field, pair, knnEngine, indexPath); + } - if (isRefresh) { - recordRefreshStats(); - } + if (isMerge) { + recordMergeStats(pair.docs.length, arraySize); + } - // This is a bit of a hack. We have to create an output here and then immediately close it to ensure that - // engineFileName is added to the tracked files by Lucene's TrackingDirectoryWrapper. Otherwise, the file will - // not be marked as added to the directory. - state.directory.createOutput(engineFileName, state.context).close(); - indexCreator.createIndex(); - writeFooter(indexPath, engineFileName); - } finally { - // Freeing up the Native memory where vectors was stored. We added a try block here to ensure that even - // in case of exceptions we are freeing up the space to avoid memory leaks. - if (pair != null) { - JNICommons.freeVectorData(pair.getVectorAddress()); - } + if (isRefresh) { + recordRefreshStats(); } + + // This is a bit of a hack. We have to create an output here and then immediately close it to ensure that + // engineFileName is added to the tracked files by Lucene's TrackingDirectoryWrapper. Otherwise, the file will + // not be marked as added to the directory. + state.directory.createOutput(engineFileName, state.context).close(); + indexCreator.createIndex(); + writeFooter(indexPath, engineFileName); } private void recordMergeStats(int length, long arraySize) { diff --git a/src/main/java/org/opensearch/knn/jni/FaissService.java b/src/main/java/org/opensearch/knn/jni/FaissService.java index 9b7c545a54..32516ef9dd 100644 --- a/src/main/java/org/opensearch/knn/jni/FaissService.java +++ b/src/main/java/org/opensearch/knn/jni/FaissService.java @@ -50,7 +50,10 @@ class FaissService { } /** - * Create an index for the native library + * Create an index for the native library The memory occupied by the vectorsAddress will be freed up during the + * function call. So Java layer doesn't need to free up the memory. This is not an ideal behavior because Java layer + * created the memory address and that should only free up the memory. We are tracking the proper fix for this on this + * issue * * @param ids array of ids mapping to the data passed in * @param vectorsAddress address of native memory where vectors are stored diff --git a/src/main/java/org/opensearch/knn/jni/JNIService.java b/src/main/java/org/opensearch/knn/jni/JNIService.java index 70a8f5b710..5a5b6794a2 100644 --- a/src/main/java/org/opensearch/knn/jni/JNIService.java +++ b/src/main/java/org/opensearch/knn/jni/JNIService.java @@ -23,7 +23,10 @@ public class JNIService { /** - * Create an index for the native library + * Create an index for the native library. The memory occupied by the vectorsAddress will be freed up during the + * function call. So Java layer doesn't need to free up the memory. This is not an ideal behavior because Java layer + * created the memory address and that should only free up the memory. We are tracking the proper fix for this on this + * issue * * @param ids array of ids mapping to the data passed in * @param vectorsAddress address of native memory where vectors are stored diff --git a/src/main/java/org/opensearch/knn/jni/NmslibService.java b/src/main/java/org/opensearch/knn/jni/NmslibService.java index 46d9266e80..7fdc278d2d 100644 --- a/src/main/java/org/opensearch/knn/jni/NmslibService.java +++ b/src/main/java/org/opensearch/knn/jni/NmslibService.java @@ -39,7 +39,10 @@ class NmslibService { } /** - * Create an index for the native library + * Create an index for the native library. The memory occupied by the vectorsAddress will be freed up during the + * function call. So Java layer doesn't need to free up the memory. This is not an ideal behavior because Java layer + * created the memory address and that should only free up the memory. We are tracking the proper fix for this on this + * issue * * @param ids array of ids mapping to the data passed in * @param vectorsAddress address of native memory where vectors are stored