diff --git a/java/gandiva/src/main/cpp/jni_common.cc b/java/gandiva/src/main/cpp/jni_common.cc index ef8633d703858..d5e54f38e3692 100644 --- a/java/gandiva/src/main/cpp/jni_common.cc +++ b/java/gandiva/src/main/cpp/jni_common.cc @@ -90,6 +90,7 @@ static jfieldID vector_expander_ret_capacity_; static jclass secondary_cache_class_; static jmethodID cache_get_method_; static jmethodID cache_set_method_; +static jmethodID cache_release_mem_method_; static jclass cache_buf_ret_class_; static jfieldID cache_buf_ret_address_; static jfieldID cache_buf_ret_size_; @@ -139,11 +140,12 @@ jint JNI_OnLoad(JavaVM* vm, void* reserved) { secondary_cache_class_ = (jclass)env->NewGlobalRef(local_cache_class); env->DeleteLocalRef(local_expander_ret_class); - cache_get_method_ = env->GetMethodID(secondary_cache_class_, "getSecondaryCache", + cache_get_method_ = env->GetMethodID(secondary_cache_class_, "get", "(JJ)Lorg/apache/arrow/gandiva/evaluator/" "JavaSecondaryCacheInterface$BufferResult;"); - cache_set_method_ = - env->GetMethodID(secondary_cache_class_, "setSecondaryCache", "(JJJJ)V"); + cache_set_method_ = env->GetMethodID(secondary_cache_class_, "set", "(JJJJ)V"); + cache_release_mem_method_ = + env->GetMethodID(secondary_cache_class_, "releaseBufferResult", "(J)V"); jclass local_cache_ret_class = env->FindClass( "org/apache/arrow/gandiva/evaluator/JavaSecondaryCacheInterface$BufferResult"); @@ -667,10 +669,18 @@ std::shared_ptr JavaSecondaryCache::Get( jlong ret_address = env_->GetLongField(ret, cache_buf_ret_address_); jlong ret_size = env_->GetLongField(ret, cache_buf_ret_size_); - auto data = std::shared_ptr( - new arrow::Buffer(reinterpret_cast(ret_address), ret_size)); + arrow::Buffer data(reinterpret_cast(ret_address), ret_size); - return data; + // copy the buffer and release the original memory + auto result = arrow::Buffer::CopyNonOwned(data, arrow::default_cpu_memory_manager()); + env_->CallObjectMethod(jcache_, cache_release_mem_method_, data.address()); + + if (result.ok()) { + auto buffer = std::move(result.ValueOrDie()); + return buffer; + } else { + return nullptr; + } } void JavaSecondaryCache::Set(std::shared_ptr key, diff --git a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/JavaSecondaryCacheInterface.java b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/JavaSecondaryCacheInterface.java index eedb34c6c5495..d7423dd2dd9aa 100644 --- a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/JavaSecondaryCacheInterface.java +++ b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/JavaSecondaryCacheInterface.java @@ -36,9 +36,27 @@ public BufferResult(long address, long size) { } - BufferResult getSecondaryCache(long addr, long size); + /** + * Get from the secondary cache using the key buffer. + * @param keyBufAddr address of the key buffer + * @param keyBufSize size of the key buffer + * @return + */ + BufferResult get(long keyBufAddr, long keyBufSize); - void setSecondaryCache(long addrExpr, long sizeExpr, long addr, long size); + /** + * Set the secondary cache with the value in the buffer. + * @param keyBufAddr address of the key buffer + * @param keyBufSize size of the key buffer + * @param valueBufAddr address of the value buffer + * @param valueBufSize size of the value buffer + */ + void set(long keyBufAddr, long keyBufSize, long valueBufAddr, long valueBufSize); + + /** + * release the buffer associated with the given address. + */ + void releaseBufferResult(long address); } diff --git a/java/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/ProjectorTest.java b/java/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/ProjectorTest.java index a6062b7bb36f3..8dd759ee885d2 100644 --- a/java/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/ProjectorTest.java +++ b/java/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/ProjectorTest.java @@ -2606,15 +2606,17 @@ public DummySecondaryCache() { setToCache = false; } - public BufferResult getSecondaryCache(long addrKey, long sizeKey) { + public BufferResult get(long addrKey, long sizeKey) { gotFromCache = true; return null; } - public void setSecondaryCache(long addrKey, long sizeKey, long addrValue, long sizeValue) { + public void set(long addrKey, long sizeKey, long addrValue, long sizeValue) { setToCache = true; } + public void releaseBufferResult(long address) {} + public boolean gotFromCache; public boolean setToCache; }