Add Java multiGet API for returning List<byte[]> (#1570)

Summary:
Closes https://github.com/facebook/rocksdb/pull/1570
Pull Request resolved: https://github.com/facebook/rocksdb/pull/4797

Differential Revision: D13961770

Pulled By: sagar0

fbshipit-source-id: e34fd6250d0cd3ebb0bd688e8801fe8947fd464d

java/src/main/java/org/rocksdb/RocksDB.java

@@ -1389,7 +1389,10 @@ public class RocksDB extends RocksObject {
    *
    * @throws RocksDBException thrown if error happens in underlying
    *    native library.
+   *
+   * @deprecated Consider {@link #multiGetAsList(List)} instead.
    */
+  @Deprecated
   public Map<byte[], byte[]> multiGet(final List<byte[]> keys)
       throws RocksDBException {
     assert(keys.size() != 0);
@@ -1440,7 +1443,10 @@ public class RocksDB extends RocksObject {
    *    native library.
    * @throws IllegalArgumentException thrown if the size of passed keys is not
    *    equal to the amount of passed column family handles.
+   *
+   * @deprecated Consider {@link #multiGetAsList(List, List)} instead.
    */
+  @Deprecated
   public Map<byte[], byte[]> multiGet(
       final List<ColumnFamilyHandle> columnFamilyHandleList,
       final List<byte[]> keys) throws RocksDBException,
@@ -1449,8 +1455,8 @@ public class RocksDB extends RocksObject {
     // Check if key size equals cfList size. If not a exception must be
     // thrown. If not a Segmentation fault happens.
     if (keys.size() != columnFamilyHandleList.size()) {
-        throw new IllegalArgumentException(
-            "For each key there must be a ColumnFamilyHandle.");
+      throw new IllegalArgumentException(
+          "For each key there must be a ColumnFamilyHandle.");
     }
     final long[] cfHandles = new long[columnFamilyHandleList.size()];
     for (int i = 0; i < columnFamilyHandleList.size(); i++) {
@@ -1488,7 +1494,10 @@ public class RocksDB extends RocksObject {
    *
    * @throws RocksDBException thrown if error happens in underlying
    *    native library.
+   *
+   * @deprecated Consider {@link #multiGetAsList(ReadOptions, List)} instead.
    */
+  @Deprecated
   public Map<byte[], byte[]> multiGet(final ReadOptions opt,
       final List<byte[]> keys) throws RocksDBException {
     assert(keys.size() != 0);
@@ -1534,7 +1543,11 @@ public class RocksDB extends RocksObject {
    *    native library.
    * @throws IllegalArgumentException thrown if the size of passed keys is not
    *    equal to the amount of passed column family handles.
+   *
+   * @deprecated Consider {@link #multiGetAsList(ReadOptions, List, List)}
+   *     instead.
    */
+  @Deprecated
   public Map<byte[], byte[]> multiGet(final ReadOptions opt,
       final List<ColumnFamilyHandle> columnFamilyHandleList,
       final List<byte[]> keys) throws RocksDBException {
@@ -1572,6 +1585,151 @@ public class RocksDB extends RocksObject {
     return keyValueMap;
   }
 
+  /**
+   * Takes a list of keys, and returns a list of values for the given list of
+   * keys. List will contain null for keys which could not be found.
+   *
+   * @param keys List of keys for which values need to be retrieved.
+   * @return List of values for the given list of keys. List will contain
+   * null for keys which could not be found.
+   *
+   * @throws RocksDBException thrown if error happens in underlying
+   *    native library.
+   */
+  public List<byte[]> multiGetAsList(final List<byte[]> keys)
+      throws RocksDBException {
+    assert(keys.size() != 0);
+
+    final byte[][] keysArray = keys.toArray(new byte[keys.size()][]);
+    final int keyOffsets[] = new int[keysArray.length];
+    final int keyLengths[] = new int[keysArray.length];
+    for(int i = 0; i < keyLengths.length; i++) {
+      keyLengths[i] = keysArray[i].length;
+    }
+
+    return Arrays.asList(multiGet(nativeHandle_, keysArray, keyOffsets,
+        keyLengths));
+  }
+
+  /**
+   * Returns a list of values for the given list of keys. List will contain
+   * null for keys which could not be found.
+   * <p>
+   * Note: Every key needs to have a related column family name in
+   * {@code columnFamilyHandleList}.
+   * </p>
+   *
+   * @param columnFamilyHandleList {@link java.util.List} containing
+   *     {@link org.rocksdb.ColumnFamilyHandle} instances.
+   * @param keys List of keys for which values need to be retrieved.
+   * @return List of values for the given list of keys. List will contain
+   * null for keys which could not be found.
+   *
+   * @throws RocksDBException thrown if error happens in underlying
+   *    native library.
+   * @throws IllegalArgumentException thrown if the size of passed keys is not
+   *    equal to the amount of passed column family handles.
+   */
+  public List<byte[]> multiGetAsList(
+      final List<ColumnFamilyHandle> columnFamilyHandleList,
+      final List<byte[]> keys) throws RocksDBException,
+      IllegalArgumentException {
+    assert(keys.size() != 0);
+    // Check if key size equals cfList size. If not a exception must be
+    // thrown. If not a Segmentation fault happens.
+    if (keys.size() != columnFamilyHandleList.size()) {
+        throw new IllegalArgumentException(
+            "For each key there must be a ColumnFamilyHandle.");
+    }
+    final long[] cfHandles = new long[columnFamilyHandleList.size()];
+    for (int i = 0; i < columnFamilyHandleList.size(); i++) {
+      cfHandles[i] = columnFamilyHandleList.get(i).nativeHandle_;
+    }
+
+    final byte[][] keysArray = keys.toArray(new byte[keys.size()][]);
+    final int keyOffsets[] = new int[keysArray.length];
+    final int keyLengths[] = new int[keysArray.length];
+    for(int i = 0; i < keyLengths.length; i++) {
+      keyLengths[i] = keysArray[i].length;
+    }
+
+    return Arrays.asList(multiGet(nativeHandle_, keysArray, keyOffsets,
+        keyLengths, cfHandles));
+  }
+
+  /**
+   * Returns a list of values for the given list of keys. List will contain
+   * null for keys which could not be found.
+   *
+   * @param opt Read options.
+   * @param keys of keys for which values need to be retrieved.
+   * @return List of values for the given list of keys. List will contain
+   * null for keys which could not be found.
+   *
+   * @throws RocksDBException thrown if error happens in underlying
+   *    native library.
+   */
+  public List<byte[]> multiGetAsList(final ReadOptions opt,
+      final List<byte[]> keys) throws RocksDBException {
+    assert(keys.size() != 0);
+
+    final byte[][] keysArray = keys.toArray(new byte[keys.size()][]);
+    final int keyOffsets[] = new int[keysArray.length];
+    final int keyLengths[] = new int[keysArray.length];
+    for(int i = 0; i < keyLengths.length; i++) {
+      keyLengths[i] = keysArray[i].length;
+    }
+
+    return Arrays.asList(multiGet(nativeHandle_, opt.nativeHandle_,
+        keysArray, keyOffsets, keyLengths));
+  }
+
+  /**
+   * Returns a list of values for the given list of keys. List will contain
+   * null for keys which could not be found.
+   * <p>
+   * Note: Every key needs to have a related column family name in
+   * {@code columnFamilyHandleList}.
+   * </p>
+   *
+   * @param opt Read options.
+   * @param columnFamilyHandleList {@link java.util.List} containing
+   *     {@link org.rocksdb.ColumnFamilyHandle} instances.
+   * @param keys of keys for which values need to be retrieved.
+   * @return List of values for the given list of keys. List will contain
+   * null for keys which could not be found.
+   *
+   * @throws RocksDBException thrown if error happens in underlying
+   *    native library.
+   * @throws IllegalArgumentException thrown if the size of passed keys is not
+   *    equal to the amount of passed column family handles.
+   */
+  public List<byte[]> multiGetAsList(final ReadOptions opt,
+      final List<ColumnFamilyHandle> columnFamilyHandleList,
+      final List<byte[]> keys) throws RocksDBException {
+    assert(keys.size() != 0);
+    // Check if key size equals cfList size. If not a exception must be
+    // thrown. If not a Segmentation fault happens.
+    if (keys.size()!=columnFamilyHandleList.size()){
+      throw new IllegalArgumentException(
+          "For each key there must be a ColumnFamilyHandle.");
+    }
+    final long[] cfHandles = new long[columnFamilyHandleList.size()];
+    for (int i = 0; i < columnFamilyHandleList.size(); i++) {
+      cfHandles[i] = columnFamilyHandleList.get(i).nativeHandle_;
+    }
+
+    final byte[][] keysArray = keys.toArray(new byte[keys.size()][]);
+    final int keyOffsets[] = new int[keysArray.length];
+    final int keyLengths[] = new int[keysArray.length];
+    for(int i = 0; i < keyLengths.length; i++) {
+      keyLengths[i] = keysArray[i].length;
+    }
+
+    return Arrays.asList(multiGet(nativeHandle_, opt.nativeHandle_,
+        keysArray, keyOffsets, keyLengths, cfHandles));
+  }
+
   /**
    * Remove the database entry (if any) for "key".  Returns OK on
    * success, and a non-OK status on error.  It is not an error if "key"

java/src/test/java/org/rocksdb/ColumnFamilyTest.java

@@ -419,6 +419,50 @@ public class ColumnFamilyTest {
     }
   }
 
+  @Test
+  public void multiGetAsList() throws RocksDBException {
+    final List<ColumnFamilyDescriptor> cfDescriptors = Arrays.asList(
+        new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY),
+        new ColumnFamilyDescriptor("new_cf".getBytes()));
+    final List<ColumnFamilyHandle> columnFamilyHandleList = new ArrayList<>();
+    try (final DBOptions options = new DBOptions()
+        .setCreateIfMissing(true)
+        .setCreateMissingColumnFamilies(true);
+         final RocksDB db = RocksDB.open(options,
+             dbFolder.getRoot().getAbsolutePath(),
+             cfDescriptors, columnFamilyHandleList)) {
+      try {
+        db.put(columnFamilyHandleList.get(0), "key".getBytes(),
+            "value".getBytes());
+        db.put(columnFamilyHandleList.get(1), "newcfkey".getBytes(),
+            "value".getBytes());
+
+        final List<byte[]> keys = Arrays.asList(new byte[][]{
+            "key".getBytes(), "newcfkey".getBytes()
+        });
+        List<byte[]> retValues = db.multiGetAsList(columnFamilyHandleList,
+            keys);
+        assertThat(retValues.size()).isEqualTo(2);
+        assertThat(new String(retValues.get(0)))
+            .isEqualTo("value");
+        assertThat(new String(retValues.get(1)))
+            .isEqualTo("value");
+        retValues = db.multiGetAsList(new ReadOptions(), columnFamilyHandleList,
+            keys);
+        assertThat(retValues.size()).isEqualTo(2);
+        assertThat(new String(retValues.get(0)))
+            .isEqualTo("value");
+        assertThat(new String(retValues.get(1)))
+            .isEqualTo("value");
+      } finally {
+        for (final ColumnFamilyHandle columnFamilyHandle :
+            columnFamilyHandleList) {
+          columnFamilyHandle.close();
+        }
+      }
+    }
+  }
+
   @Test
   public void properties() throws RocksDBException {
     final List<ColumnFamilyDescriptor> cfDescriptors = Arrays.asList(

java/src/test/java/org/rocksdb/RocksDBTest.java

@@ -270,6 +270,41 @@ public class RocksDBTest {
     }
   }
 
+  @Test
+  public void multiGetAsList() throws RocksDBException, InterruptedException {
+    try (final RocksDB db = RocksDB.open(dbFolder.getRoot().getAbsolutePath());
+         final ReadOptions rOpt = new ReadOptions()) {
+      db.put("key1".getBytes(), "value".getBytes());
+      db.put("key2".getBytes(), "12345678".getBytes());
+      List<byte[]> lookupKeys = new ArrayList<>();
+      lookupKeys.add("key1".getBytes());
+      lookupKeys.add("key2".getBytes());
+      List<byte[]> results = db.multiGetAsList(lookupKeys);
+      assertThat(results).isNotNull();
+      assertThat(results).hasSize(lookupKeys.size());
+      assertThat(results).
+          containsExactly("value".getBytes(), "12345678".getBytes());
+      // test same method with ReadOptions
+      results = db.multiGetAsList(rOpt, lookupKeys);
+      assertThat(results).isNotNull();
+      assertThat(results).
+          contains("value".getBytes(), "12345678".getBytes());
+
+      // remove existing key
+      lookupKeys.remove(1);
+      // add non existing key
+      lookupKeys.add("key3".getBytes());
+      results = db.multiGetAsList(lookupKeys);
+      assertThat(results).isNotNull();
+      assertThat(results).
+          containsExactly("value".getBytes(), null);
+      // test same call with readOptions
+      results = db.multiGetAsList(rOpt, lookupKeys);
+      assertThat(results).isNotNull();
+      assertThat(results).contains("value".getBytes());
+    }
+  }
+
   @Test
   public void merge() throws RocksDBException {
     try (final StringAppendOperator stringAppendOperator = new StringAppendOperator();