Skip to content
This repository has been archived by the owner on Aug 2, 2022. It is now read-only.

Commit

Permalink
merged master
Browse files Browse the repository at this point in the history
  • Loading branch information
@vamshin
vamshin committed Feb 6, 2020
1 parent cba3eba commit 5cbff97
Show file tree
Hide file tree
Showing 14 changed files with 283 additions and 299 deletions.
Binary file modified buildSrc/libKNNIndexV1_7_3_6.jnilib
View file
Binary file not shown.
Binary file modified buildSrc/libKNNIndexV1_7_3_6.so
View file
Binary file not shown.
123 changes: 57 additions & 66 deletions jni/src/v1736/com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,17 @@ using similarity::KNNQueue;

extern "C"

struct IndexWrapper {
IndexWrapper() {
space.reset(SpaceFactoryRegistry<float>::Instance().CreateSpace("l2", AnyParams()));
index.reset(MethodFactoryRegistry<float>::Instance().CreateMethod(false, "hnsw", "l2", *space, data));
}
std::unique_ptr<Space<float>> space;
std::unique_ptr<Index<float>> index;
// Index gets constructed with a reference to data (see above) but is otherwise unused
ObjectVector data;
};

struct JavaException {
JavaException(JNIEnv* env, const char* type = "", const char* message = "")
{
Expand Down Expand Up @@ -82,7 +93,6 @@ JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v173
int* object_ids = NULL;

try {
initLibrary();
space = SpaceFactoryRegistry<float>::Instance().CreateSpace("l2", AnyParams());
object_ids = env->GetIntArrayElements(ids, 0);
for (int i = 0; i < env->GetArrayLength(vectors); i++) {
Expand All @@ -91,8 +101,8 @@ JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v173
dataset.push_back(new Object(object_ids[i], -1, env->GetArrayLength(vectorArray)*sizeof(float), vector));
env->ReleaseFloatArrayElements(vectorArray, vector, 0);
}
// free up memory
env->ReleaseIntArrayElements(ids, object_ids, 0);

index = MethodFactoryRegistry<float>::Instance().CreateMethod(false, "hnsw", "l2", *space, dataset);

int paramsCount = env->GetArrayLength(algoParams);
Expand All @@ -106,59 +116,43 @@ JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v173

index->CreateIndex(AnyParams(paramsList));
has_exception_in_stack(env);
const char * indexString = env->GetStringUTFChars(indexPath, 0);
const char *indexString = env->GetStringUTFChars(indexPath, 0);
index->SaveIndex(indexString);
env->ReleaseStringUTFChars(indexPath, indexString);
has_exception_in_stack(env);

// Free each object in the dataset. No need to clear the vector because it goes out of scope
// immediately
for (auto it = dataset.begin(); it != dataset.end(); it++) {
delete *it;
delete *it;
}
delete index;
delete space;
}
catch (...) {
if (object_ids) { env->ReleaseIntArrayElements(ids, object_ids, 0); }
for (auto it = dataset.begin(); it != dataset.end(); it++) {
delete *it;
delete *it;
}
if (index) { delete index; }
if (space) { delete space; }
catch_cpp_exception_and_throw_java(env);
}
}

JNIEXPORT jobjectArray JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_queryIndex(JNIEnv* env, jobject indexObject, jfloatArray queryVector, jint k, jobjectArray algoParams)
JNIEXPORT jobjectArray JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_queryIndex(JNIEnv* env, jclass cls, jlong indexPointer, jfloatArray queryVector, jint k)
{
Space<float>* space = NULL;
KNNQueue<float>* result = NULL;
Object* queryObject = NULL;

try {
jclass indexClass = env->GetObjectClass(indexObject);
jmethodID getIndex = env->GetMethodID(indexClass, "getIndex", "()J");
jlong indexValue = env->CallLongMethod(indexObject, getIndex);
Index<float>* index = reinterpret_cast<Index<float>*>(indexValue);
IndexWrapper *indexWrapper = reinterpret_cast<IndexWrapper*>(indexPointer);

float* rawQueryvector = env->GetFloatArrayElements(queryVector, 0);
space = SpaceFactoryRegistry<float>::Instance().CreateSpace("l2", AnyParams());
queryObject = new Object(-1, -1, env->GetArrayLength(queryVector)*sizeof(float), rawQueryvector);
KNNQuery<float> query(*space, queryObject, k);
std::unique_ptr<const Object> queryObject(new Object(-1, -1, env->GetArrayLength(queryVector)*sizeof(float), rawQueryvector));
env->ReleaseFloatArrayElements(queryVector, rawQueryvector, 0);
has_exception_in_stack(env);

int paramsCount = env->GetArrayLength(algoParams);
vector<string> paramsList;
for (int i=0; i<paramsCount; i++) {
jstring param = (jstring) (env->GetObjectArrayElement(algoParams, i));
const char *rawString = env->GetStringUTFChars(param, 0);
paramsList.push_back(rawString);
env->ReleaseStringUTFChars(param, rawString);
}
index->SetQueryTimeParams(AnyParams(paramsList));
index->Search(&query);
result = query.Result()->Clone();
KNNQuery<float> knnQuery(*(indexWrapper->space), queryObject.get(), k);
indexWrapper->index->Search(&knnQuery);
std::unique_ptr<KNNQueue<float>> result(knnQuery.Result()->Clone());
has_exception_in_stack(env);
int resultSize = result->Size();
jclass resultClass = env->FindClass("com/amazon/opendistroforelasticsearch/knn/index/KNNQueryResult");
Expand All @@ -170,67 +164,64 @@ JNIEXPORT jobjectArray JNICALL Java_com_amazon_opendistroforelasticsearch_knn_in
env->SetObjectArrayElement(results, i, env->NewObject(resultClass, allArgs, id, distance));
}
has_exception_in_stack(env);

//free up memory
delete space;
delete result;
delete queryObject;

return results;
}
catch (...) {
if (space) { delete space; }
if (result) { delete result; }
if (queryObject) { delete queryObject; }
} catch(...) {
catch_cpp_exception_and_throw_java(env);
}
return NULL;
}

JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_init(JNIEnv* env, jobject indexObject, jstring indexPath)
JNIEXPORT jlong JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_init(JNIEnv* env, jclass cls, jstring indexPath, jobjectArray algoParams)
{
Space<float>* space = NULL;
ObjectVector* dataset = NULL;
Index<float>* index = NULL;

IndexWrapper *indexWrapper = NULL;
try {
initLibrary();
space = SpaceFactoryRegistry<float>::Instance().CreateSpace("l2", AnyParams());
dataset = new ObjectVector();
Index<float>* index = MethodFactoryRegistry<float>::Instance().CreateMethod(false, "hnsw", "l2", *space, *dataset);
const char *indexString = env->GetStringUTFChars(indexPath, 0);
index->LoadIndex(indexString);
env->ReleaseStringUTFChars(indexPath, indexString);
const char *indexPathCStr = env->GetStringUTFChars(indexPath, 0);
string indexPathString(indexPathCStr);
env->ReleaseStringUTFChars(indexPath, indexPathCStr);
has_exception_in_stack(env);
jclass indexClass = env->GetObjectClass(indexObject);
jmethodID setIndex = env->GetMethodID(indexClass, "setIndex", "(J)V");
env->CallVoidMethod(indexObject, setIndex, (jlong)index);

// Load index from file (may throw)
IndexWrapper *indexWrapper = new IndexWrapper();
indexWrapper->index->LoadIndex(indexPathString);

// Parse and set query params
int paramsCount = env->GetArrayLength(algoParams);
vector<string> paramsList;
for (int i=0; i<paramsCount; i++) {
jstring param = (jstring) (env->GetObjectArrayElement(algoParams, i));
const char *rawString = env->GetStringUTFChars(param, 0);
paramsList.push_back(rawString);
env->ReleaseStringUTFChars(param, rawString);
}
indexWrapper->index->SetQueryTimeParams(AnyParams(paramsList));
has_exception_in_stack(env);

// free up memory
delete space;
delete dataset;
return (jlong) indexWrapper;
}
// nmslib seems to throw std::runtime_error if the index cannot be read (which
// is the only known failure mode for init()).
catch (...) {
if (space) { delete space; }
if (dataset) { delete dataset; }
if (indexWrapper) delete indexWrapper;
catch_cpp_exception_and_throw_java(env);
}
return NULL;
}

JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_gc(JNIEnv* env, jobject indexObject)
JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_gc(JNIEnv* env, jclass cls, jlong indexPointer)
{
try {
jclass indexClass = env->GetObjectClass(indexObject);
jmethodID getIndex = env->GetMethodID(indexClass, "getIndex", "()J");
// index heap pointer
jlong indexValue = env->CallLongMethod(indexObject, getIndex);
Index<float>* index = reinterpret_cast<Index<float>*>(indexValue);
IndexWrapper *indexWrapper = reinterpret_cast<IndexWrapper*>(indexPointer);
has_exception_in_stack(env);
delete index;
delete indexWrapper;
has_exception_in_stack(env);
}
catch (...) {
catch_cpp_exception_and_throw_java(env);
}
}

JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_initLibrary(JNIEnv *, jclass)
{
initLibrary();

}
23 changes: 16 additions & 7 deletions jni/src/v1736/com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
* permissions and limitations under the License.
*/


/* DO NOT EDIT THIS FILE - it is machine generated */
#include <jni.h>
/* Header for class com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex */
Expand All @@ -33,26 +34,34 @@ JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v173
/*
* Class: com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex
* Method: queryIndex
* Signature: ([FI[Ljava/lang/String;)[Lcom/amazon/opendistroforelasticsearch/knn/index/KNNQueryResult;
* Signature: (J[FI[Ljava/lang/String;)[Lcom/amazon/opendistroforelasticsearch/knn/index/KNNQueryResult;
*/
JNIEXPORT jobjectArray JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_queryIndex
(JNIEnv *, jobject, jfloatArray, jint, jobjectArray);
(JNIEnv *, jclass, jlong, jfloatArray, jint);

/*
* Class: com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex
* Method: init
* Signature: (Ljava/lang/String;)V
* Signature: (Ljava/lang/String;)J
*/
JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_init
(JNIEnv *, jobject, jstring);
JNIEXPORT jlong JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_init
(JNIEnv *, jclass, jstring, jobjectArray);

/*
* Class: com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex
* Method: gc
* Signature: ()V
* Signature: (J)V
*/
JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_gc
(JNIEnv *, jobject);
(JNIEnv *, jclass, jlong);

/*
* Class: com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex
* Method: initLibrary
* Signature: ()V
*/
JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_initLibrary
(JNIEnv *, jclass);

#ifdef __cplusplus
}
Expand Down
5 changes: 1 addition & 4 deletions src/main/java/com/amazon/opendistroforelasticsearch/knn/index/KNNCircuitBreaker.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,6 @@

package com.amazon.opendistroforelasticsearch.knn.index;

import com.amazon.opendistroforelasticsearch.knn.index.v1736.KNNIndex;

import com.amazon.opendistroforelasticsearch.knn.plugin.stats.StatNames;
import com.amazon.opendistroforelasticsearch.knn.plugin.transport.KNNStatsAction;
import com.amazon.opendistroforelasticsearch.knn.plugin.transport.KNNStatsNodeResponse;
Expand All @@ -31,7 +29,6 @@

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
* Runs the circuit breaker logic and updates the settings
Expand Down Expand Up @@ -69,7 +66,7 @@ public void initialize(ThreadPool threadPool, ClusterService clusterService, Cli
this.client = client;
Runnable runnable = () -> {
if (KNNWeight.knnIndexCache.isCacheCapacityReached() && clusterService.localNode().isDataNode()) {
long currentSizeKiloBytes = KNNWeight.knnIndexCache.cache.asMap().values().stream().collect(Collectors.summingLong(KNNIndex::getIndexSize));
long currentSizeKiloBytes = KNNWeight.knnIndexCache.getWeightInKilobytes();
long circuitBreakerLimitSizeKiloBytes = KNNSettings.getCircuitBreakerLimit().getKb();
long circuitBreakerUnsetSizeKiloBytes = (long) ((KNNSettings.getCircuitBreakerUnsetPercentage()/100) * circuitBreakerLimitSizeKiloBytes);
/**
Expand Down

0 comments on commit 5cbff97

Please sign in to comment.