Multi-get impl and formatting

5 files changed, 215 insertions, 1 deletions

diff --git a/java/Makefile b/java/Makefile
index 1a199e1df..1013d99e5 100644
--- a/java/Makefile
+++ b/java/Makefile
@@ -16,7 +16,7 @@ sample: java
 	javac -cp $(ROCKSDB_JAR) RocksDBSample.java
 	@rm -rf /tmp/rocksdbjni
 	@rm -rf /tmp/rocksdbjni_not_found
-	java -ea -Djava.library.path=.:../ -cp ".:./*" -Xcheck:jni RocksDBSample /tmp/rocksdbjni
+	java -ea -Djava.library.path=. -cp ".:./*" -Xcheck:jni RocksDBSample /tmp/rocksdbjni
 	@rm -rf /tmp/rocksdbjni
 	@rm -rf /tmp/rocksdbjni_not_found
 
diff --git a/java/RocksDBSample.java b/java/RocksDBSample.java
index 7e5974722..741a88422 100644
--- a/java/RocksDBSample.java
+++ b/java/RocksDBSample.java
@@ -4,6 +4,8 @@
 // of patent rights can be found in the PATENTS file in the same directory.
 
 import java.util.Arrays;
+import java.util.List;
+import java.util.ArrayList;
 import org.rocksdb.*;
 import org.rocksdb.util.SizeUnit;
 import java.io.IOException;
diff --git a/java/org/rocksdb/RocksDB.java b/java/org/rocksdb/RocksDB.java
index 0b78450d5..289638d65 100644
--- a/java/org/rocksdb/RocksDB.java
+++ b/java/org/rocksdb/RocksDB.java
@@ -5,6 +5,9 @@
 
 package org.rocksdb;
 
+import java.util.List;
+import java.util.Map;
+import java.util.HashMap;
 import java.io.Closeable;
 import java.io.IOException;
 
@@ -142,6 +145,7 @@ public class RocksDB {
    * returned if the specified key is not found.
    *
    * @param key the key retrieve the value.
+   * @param opt Read options.
    * @return a byte array storing the value associated with the input key if
    *     any.  null if it does not find the specified key.
    *
@@ -150,6 +154,60 @@ public class RocksDB {
   public byte[] get(ReadOptions opt, byte[] key) throws RocksDBException {
     return get(nativeHandle_, opt.nativeHandle_, key, key.length);
   }
+  
+  /**
+   * Returns a map of keys for which values were found in DB.
+   * 
+   * @param keys List of keys for which values need to be retrieved.
+   * @return Map where key of map is the key passed by user and value for map 
+   * entry is the corresponding value in DB.
+   * 
+   * @see RocksDBException 
+   */
+  public Map<byte[], byte[]> multiGet(List<byte[]> keys)
+      throws RocksDBException {
+    List<byte[]> values = multiGet(
+        nativeHandle_, keys, keys.size());
+    
+    Map<byte[], byte[]> keyValueMap = new HashMap<byte[], byte[]>();    
+    for(int i = 0; i < values.size(); i++) {
+      if(values.get(i) == null) {
+        continue;
+      }
+      
+      keyValueMap.put(keys.get(i), values.get(i));
+    }
+    
+    return keyValueMap;
+  }
+  
+  
+  /**
+   * Returns a map of keys for which values were found in DB.
+   * 
+   * @param List of keys for which values need to be retrieved.
+   * @param opt Read options.
+   * @return Map where key of map is the key passed by user and value for map 
+   * entry is the corresponding value in DB.
+   * 
+   * @see RocksDBException 
+   */
+  public Map<byte[], byte[]> multiGet(ReadOptions opt, List<byte[]> keys)
+      throws RocksDBException {
+    List<byte[]> values = multiGet(
+        nativeHandle_, opt.nativeHandle_, keys, keys.size());
+    
+    Map<byte[], byte[]> keyValueMap = new HashMap<byte[], byte[]>();    
+    for(int i = 0; i < values.size(); i++) {
+      if(values.get(i) == null) {
+        continue;
+      }
+      
+      keyValueMap.put(keys.get(i), values.get(i));
+    }
+    
+    return keyValueMap;
+  }
 
   /**
    * Remove the database entry (if any) for "key".  Returns OK on
@@ -213,6 +271,10 @@ public class RocksDB {
   protected native int get(
       long handle, long readOptHandle, byte[] key, int keyLen,
       byte[] value, int valueLen) throws RocksDBException;
+  protected native List<byte[]> multiGet(
+      long dbHandle, List<byte[]> keys, int keysCount);
+  protected native List<byte[]> multiGet(
+      long dbHandle, long rOptHandle, List<byte[]> keys, int keysCount);
   protected native byte[] get(
       long handle, byte[] key, int keyLen) throws RocksDBException;
   protected native byte[] get(
diff --git a/java/rocksjni/portal.h b/java/rocksjni/portal.h
index 4c4444329..11764eee8 100644
--- a/java/rocksjni/portal.h
+++ b/java/rocksjni/portal.h
@@ -315,5 +315,69 @@ class FilterJni {
         reinterpret_cast<jlong>(op));
   }
 };
+
+class ListJni {
+ public:
+  // Get the java class id of java.util.List.
+  static jclass getListClass(JNIEnv* env) {
+    static jclass jclazz = env->FindClass("java/util/List");
+    assert(jclazz != nullptr);
+    return jclazz;
+  }
+  
+  // Get the java class id of java.util.ArrayList.
+  static jclass getArrayListClass(JNIEnv* env) {
+    static jclass jclazz = env->FindClass("java/util/ArrayList");
+    assert(jclazz != nullptr);
+    return jclazz;
+  }
+
+  // Get the java class id of java.util.Iterator.
+  static jclass getIteratorClass(JNIEnv* env) {
+    static jclass jclazz = env->FindClass("java/util/Iterator");
+    assert(jclazz != nullptr);
+    return jclazz;
+  }
+  
+  // Get the java method id of java.util.List.iterator().
+  static jmethodID getIteratorMethod(JNIEnv* env) {
+    static jmethodID mid = env->GetMethodID(
+        getListClass(env), "iterator", "()Ljava/util/Iterator;");
+    assert(mid != nullptr);
+    return mid;
+  }
+  
+  // Get the java method id of java.util.Iterator.hasNext().
+  static jmethodID getHasNextMethod(JNIEnv* env) {
+    static jmethodID mid = env->GetMethodID(
+        getIteratorClass(env), "hasNext", "()Z");
+    assert(mid != nullptr);
+    return mid;
+  }
+  
+  // Get the java method id of java.util.Iterator.next().
+  static jmethodID getNextMethod(JNIEnv* env) {
+    static jmethodID mid = env->GetMethodID(
+        getIteratorClass(env), "next", "()Ljava/lang/Object;");
+    assert(mid != nullptr);
+    return mid;
+  }
+  
+  // Get the java method id of arrayList constructor.
+  static jmethodID getArrayListConstructorMethodId(JNIEnv* env, jclass jclazz) {
+    static jmethodID mid = env->GetMethodID(
+        jclazz, "<init>", "(I)V");
+    assert(mid != nullptr);
+    return mid;
+  }
+  
+  // Get the java method id of java.util.List.add().
+  static jmethodID getListAddMethodId(JNIEnv* env) {
+    static jmethodID mid = env->GetMethodID(
+        getListClass(env), "add", "(Ljava/lang/Object;)Z");
+    assert(mid != nullptr);
+    return mid;
+  }
+};
 }  // namespace rocksdb
 #endif  // JAVA_ROCKSJNI_PORTAL_H_
diff --git a/java/rocksjni/rocksjni.cc b/java/rocksjni/rocksjni.cc
index 17c7b8b10..a8e6dfbb0 100644
--- a/java/rocksjni/rocksjni.cc
+++ b/java/rocksjni/rocksjni.cc
@@ -10,6 +10,7 @@
 #include <stdlib.h>
 #include <jni.h>
 #include <string>
+#include <vector>
 
 #include "include/org_rocksdb_RocksDB.h"
 #include "rocksjni/portal.h"
@@ -244,6 +245,91 @@ jint rocksdb_get_helper(
   return cvalue_len;
 }
 
+jobject multi_get_helper(JNIEnv* env, jobject jdb, rocksdb::DB* db,
+    const rocksdb::ReadOptions& rOpt, jobject jkey_list, jint jkeys_count) {
+  std::vector<rocksdb::Slice> keys;
+  std::vector<jbyte*> keys_to_free;
+  
+  // get iterator
+  jobject iteratorObj = env->CallObjectMethod(
+      jkey_list, rocksdb::ListJni::getIteratorMethod(env));
+  
+  // iterate over keys and convert java byte array to slice
+  while(env->CallBooleanMethod(
+      iteratorObj, rocksdb::ListJni::getHasNextMethod(env)) == JNI_TRUE) {
+    jbyteArray jkey = (jbyteArray) env->CallObjectMethod(
+       iteratorObj, rocksdb::ListJni::getNextMethod(env));
+    jint key_length = env->GetArrayLength(jkey);
+    
+    jbyte* key = new jbyte[key_length];
+    env->GetByteArrayRegion(jkey, 0, key_length, key);
+    // store allocated jbyte to free it after multiGet call
+    keys_to_free.push_back(key); 
+    
+    rocksdb::Slice key_slice(
+      reinterpret_cast<char*>(key), key_length);
+    keys.push_back(key_slice);
+  }
+  
+  std::vector<std::string> values;
+  std::vector<rocksdb::Status> s = db->MultiGet(rOpt, keys, &values);
+  
+  // Don't reuse class pointer
+  jclass jclazz = env->FindClass("java/util/ArrayList");  
+  jmethodID mid = rocksdb::ListJni::getArrayListConstructorMethodId(
+      env, jclazz);  
+  jobject jvalue_list = env->NewObject(jclazz, mid, jkeys_count);
+  
+  // insert in java list
+  for(std::vector<rocksdb::Status>::size_type i = 0; i != s.size(); i++) {
+    if(s[i].ok()) {
+      jbyteArray jvalue = env->NewByteArray(values[i].size());
+      env->SetByteArrayRegion(
+          jvalue, 0, values[i].size(),
+          reinterpret_cast<const jbyte*>(values[i].c_str()));
+      env->CallBooleanMethod(
+          jvalue_list, rocksdb::ListJni::getListAddMethodId(env), jvalue);
+    }
+    else {
+      env->CallBooleanMethod(
+          jvalue_list, rocksdb::ListJni::getListAddMethodId(env), nullptr);
+    }
+  }
+  
+  // free up allocated byte arrays
+  for(std::vector<jbyte*>::size_type i = 0; i != keys_to_free.size(); i++) {
+    delete[] keys_to_free[i];
+  }
+  keys_to_free.clear();
+  
+  return jvalue_list;
+}
+
+/*
+ * Class:     org_rocksdb_RocksDB
+ * Method:    multiGet
+ * Signature: (JLjava/util/List;I)Ljava/util/List;
+ */
+jobject Java_org_rocksdb_RocksDB_multiGet(
+    JNIEnv* env, jobject jdb, jlong jdb_handle,
+    jobject jkey_list, jint jkeys_count) {
+  return multi_get_helper(env, jdb, reinterpret_cast<rocksdb::DB*>(jdb_handle),
+      rocksdb::ReadOptions(), jkey_list, jkeys_count);
+}
+
+/*
+ * Class:     org_rocksdb_RocksDB
+ * Method:    multiGet
+ * Signature: (JJLjava/util/List;I)Ljava/util/List;
+ */
+jobject Java_org_rocksdb_RocksDB_multiGet(
+    JNIEnv* env, jobject jdb, jlong jdb_handle,
+    jlong jropt_handle, jobject jkey_list, jint jkeys_count) {
+  return multi_get_helper(env, jdb, reinterpret_cast<rocksdb::DB*>(jdb_handle),
+      *reinterpret_cast<rocksdb::ReadOptions*>(jropt_handle), jkey_list,
+      jkeys_count);
+}
+
 /*
  * Class:     org_rocksdb_RocksDB
  * Method:    get