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