Java wrapper for Native Comparators

Summary:
This is an abstraction for working with custom Comparators implemented in native C++ code from Java. Native code must directly extend `rocksdb::Comparator`. When the native code comparator is compiled into the RocksDB codebase, you can then create a Java Class, and JNI stub to wrap it.

Useful if the C++/JNI barrier overhead is too much for your applications comparator performance.

An example is provided in `java/rocksjni/native_comparator_wrapper_test.cc` and `java/src/main/java/org/rocksdb/NativeComparatorWrapperTest.java`.
Closes #3334

Differential Revision: D7172605

Pulled By: miasantreble

fbshipit-source-id: e24b7eb267a3bcb6afa214e0379a1d5e8a2ceabe

java/CMakeLists.txt

@@ -24,6 +24,7 @@ set(JNI_NATIVE_SOURCES
         rocksjni/lru_cache.cc
         rocksjni/memtablejni.cc
         rocksjni/merge_operator.cc
+        rocksjni/native_comparator_wrapper_test.cc
         rocksjni/options.cc
         rocksjni/options_util.cc
         rocksjni/ratelimiterjni.cc
@@ -87,6 +88,8 @@ set(NATIVE_JAVA_CLASSES
         org.rocksdb.LRUCache
         org.rocksdb.MemTableConfig
         org.rocksdb.MergeOperator
+        org.rocksdb.NativeComparatorWrapper
+        org.rocksdb.NativeComparatorWrapperTest.NativeStringComparatorWrapper
         org.rocksdb.NativeLibraryLoader
         org.rocksdb.Options
         org.rocksdb.OptionsUtil
@@ -179,6 +182,7 @@ add_jar(
   src/main/java/org/rocksdb/CompactionStyle.java
   src/main/java/org/rocksdb/Comparator.java
   src/main/java/org/rocksdb/ComparatorOptions.java
+  src/main/java/org/rocksdb/ComparatorType.java
   src/main/java/org/rocksdb/CompressionOptions.java
   src/main/java/org/rocksdb/CompressionType.java
   src/main/java/org/rocksdb/DBOptions.java
@@ -205,6 +209,7 @@ add_jar(
   src/main/java/org/rocksdb/MergeOperator.java
   src/main/java/org/rocksdb/MutableColumnFamilyOptions.java
   src/main/java/org/rocksdb/MutableColumnFamilyOptionsInterface.java
+  src/main/java/org/rocksdb/NativeComparatorWrapper.java
   src/main/java/org/rocksdb/NativeLibraryLoader.java
   src/main/java/org/rocksdb/Options.java
   src/main/java/org/rocksdb/OptionsUtil.java
@@ -246,6 +251,7 @@ add_jar(
   src/main/java/org/rocksdb/WriteOptions.java
   src/test/java/org/rocksdb/BackupEngineTest.java
   src/test/java/org/rocksdb/IngestExternalFileOptionsTest.java
+  src/test/java/org/rocksdb/NativeComparatorWrapperTest.java
   src/test/java/org/rocksdb/PlatformRandomHelper.java
   src/test/java/org/rocksdb/RocksDBExceptionTest.java
   src/test/java/org/rocksdb/RocksMemoryResource.java

java/Makefile

@@ -30,6 +30,7 @@ NATIVE_JAVA_CLASSES = org.rocksdb.AbstractCompactionFilter\
 	org.rocksdb.Logger\
 	org.rocksdb.LRUCache\
 	org.rocksdb.MergeOperator\
+	org.rocksdb.NativeComparatorWrapper\
 	org.rocksdb.OptimisticTransactionDB\
 	org.rocksdb.OptimisticTransactionOptions\
 	org.rocksdb.Options\
@@ -64,6 +65,7 @@ NATIVE_JAVA_CLASSES = org.rocksdb.AbstractCompactionFilter\
 	org.rocksdb.WBWIRocksIterator
 
 NATIVE_JAVA_TEST_CLASSES = org.rocksdb.RocksDBExceptionTest\
+    org.rocksdb.NativeComparatorWrapperTest.NativeStringComparatorWrapper\
     org.rocksdb.WriteBatchTest\
     org.rocksdb.WriteBatchTestInternalHelper
 
@@ -111,6 +113,7 @@ JAVA_TESTS = org.rocksdb.BackupableDBOptionsTest\
 	org.rocksdb.MergeTest\
 	org.rocksdb.MixedOptionsTest\
 	org.rocksdb.MutableColumnFamilyOptionsTest\
+	org.rocksdb.NativeComparatorWrapperTest\
 	org.rocksdb.NativeLibraryLoaderTest\
 	org.rocksdb.OptimisticTransactionTest\
 	org.rocksdb.OptimisticTransactionDBTest\

java/rocksjni/comparator.cc

@@ -14,6 +14,7 @@
 
 #include "include/org_rocksdb_Comparator.h"
 #include "include/org_rocksdb_DirectComparator.h"
+#include "include/org_rocksdb_NativeComparatorWrapper.h"
 #include "rocksjni/comparatorjnicallback.h"
 #include "rocksjni/portal.h"
 
@@ -49,4 +50,16 @@ jlong Java_org_rocksdb_DirectComparator_createNewDirectComparator0(
       new rocksdb::DirectComparatorJniCallback(env, jobj, copt);
   return reinterpret_cast<jlong>(c);
 }
+
+/*
+ * Class:     org_rocksdb_NativeComparatorWrapper
+ * Method:    disposeInternal
+ * Signature: (J)V
+ */
+void Java_org_rocksdb_NativeComparatorWrapper_disposeInternal(
+    JNIEnv* env, jobject jobj, jlong jcomparator_handle) {
+  auto* comparator =
+      reinterpret_cast<rocksdb::Comparator*>(jcomparator_handle);
+  delete comparator;
+}
 // </editor-fold>

java/rocksjni/native_comparator_wrapper_test.cc

@@ -0,0 +1,50 @@
+// Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#include <jni.h>
+#include <string>
+
+#include "rocksdb/comparator.h"
+#include "rocksdb/slice.h"
+
+#include "include/org_rocksdb_NativeComparatorWrapperTest_NativeStringComparatorWrapper.h"
+
+namespace rocksdb {
+
+class NativeComparatorWrapperTestStringComparator
+    : public Comparator {
+
+  const char* Name() const {
+    return "NativeComparatorWrapperTestStringComparator";
+  }
+
+  int Compare(
+      const Slice& a, const Slice& b) const {
+    return a.ToString().compare(b.ToString());
+  }
+
+  void FindShortestSeparator(
+      std::string* start, const Slice& limit) const {
+    return;
+  }
+
+  void FindShortSuccessor(
+      std::string* key) const {
+    return;
+  }
+};
+}  // end of rocksdb namespace
+
+/*
+ * Class:     org_rocksdb_NativeComparatorWrapperTest_NativeStringComparatorWrapper
+ * Method:    newStringComparator
+ * Signature: ()J
+ */
+jlong Java_org_rocksdb_NativeComparatorWrapperTest_00024NativeStringComparatorWrapper_newStringComparator(
+    JNIEnv* env , jobject jobj) {
+  auto* comparator =
+      new rocksdb::NativeComparatorWrapperTestStringComparator();
+  return reinterpret_cast<jlong>(comparator);
+}

java/rocksjni/options.cc

@@ -146,19 +146,33 @@ void Java_org_rocksdb_Options_setComparatorHandle__JI(
 /*
  * Class:     org_rocksdb_Options
  * Method:    setComparatorHandle
- * Signature: (JJZ)V
+ * Signature: (JJB)V
  */
-void Java_org_rocksdb_Options_setComparatorHandle__JJZ(
+void Java_org_rocksdb_Options_setComparatorHandle__JJB(
     JNIEnv* env, jobject jobj, jlong jopt_handle, jlong jcomparator_handle,
-    jboolean is_direct) {
-  auto* opt = reinterpret_cast<rocksdb::Options*>(jopt_handle);
-  if(is_direct) {
-    opt->comparator =
-        reinterpret_cast<rocksdb::DirectComparatorJniCallback*>(jcomparator_handle);
-  } else {
-    opt->comparator =
-        reinterpret_cast<rocksdb::ComparatorJniCallback*>(jcomparator_handle);
+    jbyte jcomparator_type) {
+  rocksdb::Comparator *comparator = nullptr;
+  switch(jcomparator_type) {
+      // JAVA_COMPARATOR
+      case 0x0:
+        comparator =
+            reinterpret_cast<rocksdb::ComparatorJniCallback*>(jcomparator_handle);
+        break;
+
+      // JAVA_DIRECT_COMPARATOR
+      case 0x1:
+        comparator =
+            reinterpret_cast<rocksdb::DirectComparatorJniCallback*>(jcomparator_handle);
+        break;
+
+      // JAVA_NATIVE_COMPARATOR_WRAPPER
+      case 0x2:
+        comparator =
+            reinterpret_cast<rocksdb::Comparator*>(jcomparator_handle);
+        break;
   }
+  auto* opt = reinterpret_cast<rocksdb::Options*>(jopt_handle);
+  opt->comparator = comparator;
 }
 
 /*
@@ -2960,19 +2974,33 @@ void Java_org_rocksdb_ColumnFamilyOptions_setComparatorHandle__JI(
 /*
  * Class:     org_rocksdb_ColumnFamilyOptions
  * Method:    setComparatorHandle
- * Signature: (JJZ)V
+ * Signature: (JJB)V
  */
-void Java_org_rocksdb_ColumnFamilyOptions_setComparatorHandle__JJZ(
+void Java_org_rocksdb_ColumnFamilyOptions_setComparatorHandle__JJB(
     JNIEnv* env, jobject jobj, jlong jopt_handle, jlong jcomparator_handle,
-    jboolean is_direct) {
-  auto* opt = reinterpret_cast<rocksdb::ColumnFamilyOptions*>(jopt_handle);
-  if(is_direct) {
-    opt->comparator =
-        reinterpret_cast<rocksdb::DirectComparatorJniCallback*>(jcomparator_handle);
-  } else {
-    opt->comparator =
-        reinterpret_cast<rocksdb::ComparatorJniCallback*>(jcomparator_handle);
+    jbyte jcomparator_type) {
+  rocksdb::Comparator *comparator = nullptr;
+  switch(jcomparator_type) {
+      // JAVA_COMPARATOR
+      case 0x0:
+        comparator =
+            reinterpret_cast<rocksdb::ComparatorJniCallback*>(jcomparator_handle);
+        break;
+
+      // JAVA_DIRECT_COMPARATOR
+      case 0x1:
+        comparator =
+            reinterpret_cast<rocksdb::DirectComparatorJniCallback*>(jcomparator_handle);
+        break;
+
+      // JAVA_NATIVE_COMPARATOR_WRAPPER
+      case 0x2:
+        comparator =
+            reinterpret_cast<rocksdb::Comparator*>(jcomparator_handle);
+        break;
   }
+  auto* opt = reinterpret_cast<rocksdb::ColumnFamilyOptions*>(jopt_handle);
+  opt->comparator = comparator;
 }
 
 /*

java/rocksjni/sst_file_writerjni.cc

@@ -20,24 +20,34 @@
 /*
  * Class:     org_rocksdb_SstFileWriter
  * Method:    newSstFileWriter
- * Signature: (JJJZ)J
+ * Signature: (JJJB)J
  */
-jlong Java_org_rocksdb_SstFileWriter_newSstFileWriter__JJJZ(JNIEnv *env,
-    jclass jcls, jlong jenvoptions,  jlong joptions, jlong jcomparator,
-    jboolean is_direct) {
+jlong Java_org_rocksdb_SstFileWriter_newSstFileWriter__JJJB(JNIEnv *env,
+    jclass jcls, jlong jenvoptions,  jlong joptions, jlong jcomparator_handle,
+    jbyte jcomparator_type) {
+  rocksdb::Comparator *comparator = nullptr;
+  switch(jcomparator_type) {
+      // JAVA_COMPARATOR
+      case 0x0:
+        comparator =
+            reinterpret_cast<rocksdb::ComparatorJniCallback*>(jcomparator_handle);
+        break;
+
+      // JAVA_DIRECT_COMPARATOR
+      case 0x1:
+        comparator =
+            reinterpret_cast<rocksdb::DirectComparatorJniCallback*>(jcomparator_handle);
+        break;
+
+      // JAVA_NATIVE_COMPARATOR_WRAPPER
+      case 0x2:
+        comparator =
+            reinterpret_cast<rocksdb::Comparator*>(jcomparator_handle);
+        break;
+  }
   auto *env_options =
       reinterpret_cast<const rocksdb::EnvOptions *>(jenvoptions);
   auto *options = reinterpret_cast<const rocksdb::Options *>(joptions);
-
-  rocksdb::Comparator *comparator = nullptr;
-  if(is_direct) {
-    comparator =
-        reinterpret_cast<rocksdb::DirectComparatorJniCallback*>(jcomparator);
-  } else {
-    comparator =
-        reinterpret_cast<rocksdb::ComparatorJniCallback*>(jcomparator);
-  }
-
   rocksdb::SstFileWriter *sst_file_writer =
       new rocksdb::SstFileWriter(*env_options, *options, comparator);
   return reinterpret_cast<jlong>(sst_file_writer);

java/rocksjni/write_batch_with_index.cc

@@ -39,19 +39,31 @@ jlong Java_org_rocksdb_WriteBatchWithIndex_newWriteBatchWithIndex__Z(
 /*
  * Class:     org_rocksdb_WriteBatchWithIndex
  * Method:    newWriteBatchWithIndex
- * Signature: (JZIZ)J
+ * Signature: (JBIZ)J
  */
-jlong Java_org_rocksdb_WriteBatchWithIndex_newWriteBatchWithIndex__JZIZ(
+jlong Java_org_rocksdb_WriteBatchWithIndex_newWriteBatchWithIndex__JBIZ(
     JNIEnv* env, jclass jcls, jlong jfallback_index_comparator_handle,
-    jboolean is_direct, jint jreserved_bytes, jboolean joverwrite_key) {
-rocksdb::Comparator *fallback_comparator = nullptr;
-if(is_direct) {
-  fallback_comparator =
-      reinterpret_cast<rocksdb::DirectComparatorJniCallback*>(jfallback_index_comparator_handle);
-} else {
-  fallback_comparator =
-      reinterpret_cast<rocksdb::ComparatorJniCallback*>(jfallback_index_comparator_handle);
-}
+    jbyte jcomparator_type, jint jreserved_bytes, jboolean joverwrite_key) {
+  rocksdb::Comparator *fallback_comparator = nullptr;
+  switch(jcomparator_type) {
+      // JAVA_COMPARATOR
+      case 0x0:
+        fallback_comparator =
+            reinterpret_cast<rocksdb::ComparatorJniCallback*>(jfallback_index_comparator_handle);
+        break;
+
+      // JAVA_DIRECT_COMPARATOR
+      case 0x1:
+        fallback_comparator =
+            reinterpret_cast<rocksdb::DirectComparatorJniCallback*>(jfallback_index_comparator_handle);
+        break;
+
+      // JAVA_NATIVE_COMPARATOR_WRAPPER
+      case 0x2:
+        fallback_comparator =
+            reinterpret_cast<rocksdb::Comparator*>(jfallback_index_comparator_handle);
+        break;
+  }
   auto* wbwi =
       new rocksdb::WriteBatchWithIndex(
           fallback_comparator,

java/src/main/java/org/rocksdb/AbstractComparator.java

@@ -17,10 +17,23 @@
 public abstract class AbstractComparator<T extends AbstractSlice<?>>
     extends RocksCallbackObject {
 
+  protected AbstractComparator() {
+    super();
+  }
+
   protected AbstractComparator(final ComparatorOptions copt) {
     super(copt.nativeHandle_);
   }
 
+  /**
+   * Get the type of this comparator.
+   *
+   * Used for determining the correct C++ cast in native code.
+   *
+   * @return The type of the comparator.
+   */
+  abstract ComparatorType getComparatorType();
+
   /**
    * The name of the comparator.  Used to check for comparator
    * mismatches (i.e., a DB created with one comparator is

java/src/main/java/org/rocksdb/ColumnFamilyOptions.java

@@ -143,7 +143,7 @@ public ColumnFamilyOptions setComparator(
       final AbstractComparator<? extends AbstractSlice<?>> comparator) {
     assert (isOwningHandle());
     setComparatorHandle(nativeHandle_, comparator.nativeHandle_,
-            comparator instanceof DirectComparator);
+            comparator.getComparatorType().getValue());
     comparator_ = comparator;
     return this;
   }
@@ -795,7 +795,7 @@ private native void optimizeUniversalStyleCompaction(long handle,
       long memtableMemoryBudget);
   private native void setComparatorHandle(long handle, int builtinComparator);
   private native void setComparatorHandle(long optHandle,
-      long comparatorHandle, boolean isDirect);
+      long comparatorHandle, byte comparatorType);
   private native void setMergeOperatorName(long handle, String name);
   private native void setMergeOperator(long handle, long mergeOperatorHandle);
   private native void setCompactionFilterHandle(long handle,

java/src/main/java/org/rocksdb/Comparator.java

@@ -25,5 +25,10 @@ protected long initializeNative(final long... nativeParameterHandles) {
     return createNewComparator0(nativeParameterHandles[0]);
   }
 
+  @Override
+  final ComparatorType getComparatorType() {
+    return ComparatorType.JAVA_COMPARATOR;
+  }
+
   private native long createNewComparator0(final long comparatorOptionsHandle);
 }