Adding support for deleteFilesInRanges in JNI (facebook#4031)

Summary:
It is very useful method call to achieve https://github.com/facebook/rocksdb/wiki/Delete-A-Range-Of-Keys
Pull Request resolved: facebook#4031

Differential Revision: D13515418

Pulled By: vjnadimpalli

fbshipit-source-id: 930b48e0992ef07fd1edd0b0cb5f780fabb1b4b5

java/rocksjni/rocksjni.cc

@@ -18,6 +18,7 @@
 
 #include "include/org_rocksdb_RocksDB.h"
 #include "rocksdb/cache.h"
+#include "rocksdb/convenience.h"
 #include "rocksdb/db.h"
 #include "rocksdb/options.h"
 #include "rocksdb/types.h"
@@ -3044,3 +3045,73 @@ void Java_org_rocksdb_RocksDB_destroyDB(
     rocksdb::RocksDBExceptionJni::ThrowNew(env, s);
   }
 }
+
+bool get_slice_helper(JNIEnv* env, jobjectArray ranges, jsize index,
+                      std::unique_ptr<rocksdb::Slice>& slice,
+                      std::vector<std::unique_ptr<jbyte[]>>& ranges_to_free) {
+  jobject jArray = env->GetObjectArrayElement(ranges, index);
+  if (env->ExceptionCheck()) {
+    // exception thrown: ArrayIndexOutOfBoundsException
+    return false;
+  }
+
+  if (jArray == nullptr) {
+    return true;
+  }
+
+  jbyteArray jba = reinterpret_cast<jbyteArray>(jArray);
+  jsize len_ba = env->GetArrayLength(jba);
+  ranges_to_free.push_back(std::unique_ptr<jbyte[]>(new jbyte[len_ba]));
+  env->GetByteArrayRegion(jba, 0, len_ba, ranges_to_free.back().get());
+  if (env->ExceptionCheck()) {
+    // exception thrown: ArrayIndexOutOfBoundsException
+    env->DeleteLocalRef(jArray);
+    return false;
+  }
+  env->DeleteLocalRef(jArray);
+  slice.reset(new rocksdb::Slice(
+      reinterpret_cast<char*>(ranges_to_free.back().get()), len_ba));
+  return true;
+}
+/*
+ * Class:     org_rocksdb_RocksDB
+ * Method:    deleteFilesInRanges
+ * Signature: (JJLjava/util/List;Z)V
+ */
+JNIEXPORT void JNICALL Java_org_rocksdb_RocksDB_deleteFilesInRanges(
+    JNIEnv* env, jobject /*jdb*/, jlong jdb_handle, jlong jcf_handle,
+    jobjectArray ranges, jboolean include_end) {
+  jsize length = env->GetArrayLength(ranges);
+
+  std::vector<rocksdb::RangePtr> rangesVector;
+  std::vector<std::unique_ptr<rocksdb::Slice>> slices;
+  std::vector<std::unique_ptr<jbyte[]>> ranges_to_free;
+  for (jsize i = 0; (i + 1) < length; i += 2) {
+    slices.push_back(std::unique_ptr<rocksdb::Slice>());
+    if (!get_slice_helper(env, ranges, i, slices.back(), ranges_to_free)) {
+      // exception thrown
+      return;
+    }
+
+    slices.push_back(std::unique_ptr<rocksdb::Slice>());
+    if (!get_slice_helper(env, ranges, i + 1, slices.back(), ranges_to_free)) {
+      // exception thrown
+      return;
+    }
+
+    rangesVector.push_back(rocksdb::RangePtr(slices[slices.size() - 2].get(),
+                                             slices[slices.size() - 1].get()));
+  }
+
+  auto* db = reinterpret_cast<rocksdb::DB*>(jdb_handle);
+  auto* column_family =
+      reinterpret_cast<rocksdb::ColumnFamilyHandle*>(jcf_handle);
+
+  rocksdb::Status s = rocksdb::DeleteFilesInRanges(
+      db, column_family == nullptr ? db->DefaultColumnFamily() : column_family,
+      rangesVector.data(), rangesVector.size(), include_end);
+
+  if (!s.ok()) {
+    rocksdb::RocksDBExceptionJni::ThrowNew(env, s);
+  }
+}

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

@@ -3834,6 +3834,32 @@ public void endTrace() throws RocksDBException {
     endTrace(nativeHandle_);
   }
 
+  /*
+   * Delete files in multiple ranges at once
+   * Delete files in a lot of ranges one at a time can be slow, use this API for
+   * better performance in that case.
+   * @param columnFamily - The column family for operation (null for default)
+   * @param includeEnd - Whether ranges should include end
+   * @param ranges - pairs of ranges (from1, to1, from2, to2, ...)
+   * @throws RocksDBException thrown if error happens in underlying
+   *     native library.
+   */
+  public void deleteFilesInRanges(final ColumnFamilyHandle columnFamily, final List<byte[]> ranges,
+      final boolean includeEnd) throws RocksDBException {
+    if (ranges.size() == 0) {
+      return;
+    }
+    if ((ranges.size() % 2) != 0) {
+      throw new IllegalArgumentException("Ranges size needs to be multiple of 2 "
+          + "(from1, to1, from2, to2, ...), but is " + ranges.size());
+    }
+
+    final byte[][] rangesArray = ranges.toArray(new byte[ranges.size()][]);
+
+    deleteFilesInRanges(nativeHandle_, columnFamily == null ? 0 : columnFamily.nativeHandle_,
+        rangesArray, includeEnd);
+  }
+
   /**
    * Static method to destroy the contents of the specified database.
    * Be very careful using this method.
@@ -4171,7 +4197,8 @@ private native void promoteL0(final long handle,
   private native void startTrace(final long handle, final long maxTraceFileSize,
       final long traceWriterHandle) throws RocksDBException;
   private native void endTrace(final long handle) throws RocksDBException;
-
+  private native void deleteFilesInRanges(long handle, long cfHandle, final byte[][] ranges,
+      boolean include_end) throws RocksDBException;
 
   private native static void destroyDB(final String path,
       final long optionsHandle) throws RocksDBException;

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

@@ -869,6 +869,62 @@ public void compactRangeToLevel()
     }
   }
 
+  @Test
+  public void deleteFilesInRange() throws RocksDBException, InterruptedException {
+    final int KEY_SIZE = 20;
+    final int VALUE_SIZE = 1000;
+    final int FILE_SIZE = 64000;
+    final int NUM_FILES = 10;
+
+    final int KEY_INTERVAL = 10000;
+    /*
+     * Intention of these options is to end up reliably with 10 files
+     * we will be deleting using deleteFilesInRange.
+     * It is writing roughly number of keys that will fit in 10 files (target size)
+     * It is writing interleaved so that files from memory on L0 will overlap
+     * Then compaction cleans everything and we should end up with 10 files
+     */
+    try (final Options opt = new Options()
+                                 .setCreateIfMissing(true)
+                                 .setCompressionType(CompressionType.NO_COMPRESSION)
+                                 .setTargetFileSizeBase(FILE_SIZE)
+                                 .setWriteBufferSize(FILE_SIZE / 2)
+                                 .setDisableAutoCompactions(true);
+         final RocksDB db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath())) {
+      int records = FILE_SIZE / (KEY_SIZE + VALUE_SIZE);
+
+      // fill database with key/value pairs
+      byte[] value = new byte[VALUE_SIZE];
+      int key_init = 0;
+      for (int o = 0; o < NUM_FILES; ++o) {
+        int int_key = key_init++;
+        for (int i = 0; i < records; ++i) {
+          int_key += KEY_INTERVAL;
+          rand.nextBytes(value);
+
+          db.put(String.format("%020d", int_key).getBytes(), value);
+        }
+      }
+      db.flush(new FlushOptions().setWaitForFlush(true));
+      db.compactRange();
+      // Make sure we do create one more L0 files.
+      assertThat(db.getProperty("rocksdb.num-files-at-level0")).isEqualTo("0");
+
+      // Should be 10, but we are OK with asserting +- 2
+      int files = Integer.parseInt(db.getProperty("rocksdb.num-files-at-level1"));
+      assertThat(files).isBetween(8, 12);
+
+      // Delete lower 60% (roughly). Result should be 5, but we are OK with asserting +- 2
+      // Important is that we know something was deleted (JNI call did something)
+      // Exact assertions are done in C++ unit tests
+      db.deleteFilesInRanges(null,
+          Arrays.asList(null, String.format("%020d", records * KEY_INTERVAL * 6 / 10).getBytes()),
+          false);
+      files = Integer.parseInt(db.getProperty("rocksdb.num-files-at-level1"));
+      assertThat(files).isBetween(3, 7);
+    }
+  }
+
   @Test
   public void compactRangeToLevelColumnFamily()
       throws RocksDBException {