Add release memory method

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<arrow::Buffer> 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<arrow::Buffer>(
- new arrow::Buffer(reinterpret_cast<uint8_t*>(ret_address), ret_size));
+ arrow::Buffer data(reinterpret_cast<uint8_t*>(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<arrow::Buffer> key,

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);
 
 }
 

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;
  }