Commit cda42423 authored by Adam Retter's avatar Adam Retter Committed by Zhongyi Xie
Browse files

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 https://github.com/facebook/rocksdb/pull/3334

Differential Revision: D7172605

Pulled By: miasantreble

fbshipit-source-id: e24b7eb267a3bcb6afa214e0379a1d5e8a2ceabe
parent cfbba44c
......@@ -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
......
......@@ -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\
......
......@@ -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>
// 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);
}
......@@ -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;
}
/*
......
......@@ -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);
......
......@@ -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,
......
......@@ -17,10 +17,23 @@ package org.rocksdb;
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
......
......@@ -143,7 +143,7 @@ public class ColumnFamilyOptions extends RocksObject
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 @@ public class ColumnFamilyOptions extends RocksObject
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,
......
......@@ -25,5 +25,10 @@ public abstract class Comparator extends AbstractComparator<Slice> {
return createNewComparator0(nativeParameterHandles[0]);
}
@Override
final ComparatorType getComparatorType() {
return ComparatorType.JAVA_COMPARATOR;
}
private native long createNewComparator0(final long comparatorOptionsHandle);
}
// 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).
package org.rocksdb;
enum ComparatorType {
JAVA_COMPARATOR((byte)0x0),
JAVA_DIRECT_COMPARATOR((byte)0x1),
JAVA_NATIVE_COMPARATOR_WRAPPER((byte)0x2);
private final byte value;
ComparatorType(final byte value) {
this.value = value;
}
/**
* <p>Returns the byte value of the enumerations value.</p>
*
* @return byte representation
*/
byte getValue() {
return value;
}
/**
* <p>Get the ComparatorType enumeration value by
* passing the byte identifier to this method.</p>
*
* @param byteIdentifier of ComparatorType.
*
* @return ComparatorType instance.
*
* @throws IllegalArgumentException if the comparator type for the byteIdentifier
* cannot be found
*/
static ComparatorType getComparatorType(final byte byteIdentifier) {
for (final ComparatorType comparatorType : ComparatorType.values()) {
if (comparatorType.getValue() == byteIdentifier) {
return comparatorType;
}
}
throw new IllegalArgumentException(
"Illegal value provided for ComparatorType.");
}
}
......@@ -25,6 +25,11 @@ public abstract class DirectComparator extends AbstractComparator<DirectSlice> {
return createNewDirectComparator0(nativeParameterHandles[0]);
}
@Override
final ComparatorType getComparatorType() {
return ComparatorType.JAVA_DIRECT_COMPARATOR;
}
private native long createNewDirectComparator0(
final long comparatorOptionsHandle);
}
// 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).
package org.rocksdb;
/**
* A simple abstraction to allow a Java class to wrap a custom comparator
* implemented in C++.
*
* The native comparator must directly extend rocksdb::Comparator.
*/
public abstract class NativeComparatorWrapper
extends AbstractComparator<Slice> {
@Override
final ComparatorType getComparatorType() {
return ComparatorType.JAVA_NATIVE_COMPARATOR_WRAPPER;
}
@Override
public final String name() {
throw new IllegalStateException("This should not be called. " +
"Implementation is in Native code");
}
@Override
public final int compare(final Slice s1, final Slice s2) {
throw new IllegalStateException("This should not be called. " +
"Implementation is in Native code");
}
@Override
public final String findShortestSeparator(final String start, final Slice limit) {
throw new IllegalStateException("This should not be called. " +
"Implementation is in Native code");
}
@Override
public final String findShortSuccessor(final String key) {
throw new IllegalStateException("This should not be called. " +
"Implementation is in Native code");
}
/**
* We override {@link RocksCallbackObject#disposeInternal()}
* as disposing of a native rocksd::Comparator extension requires
* a slightly different approach as it is not really a RocksCallbackObject
*/
@Override
protected void disposeInternal() {
disposeInternal(nativeHandle_);
}
private native void disposeInternal(final long handle);
}
......@@ -170,7 +170,7 @@ public class Options extends RocksObject
final AbstractComparator<? extends AbstractSlice<?>> comparator) {
assert(isOwningHandle());
setComparatorHandle(nativeHandle_, comparator.nativeHandle_,
comparator instanceof DirectComparator);
comparator.getComparatorType().getValue());
comparator_ = comparator;
return this;
}
......@@ -1734,7 +1734,7 @@ public class Options extends RocksObject
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(
......
......@@ -31,7 +31,7 @@ public class SstFileWriter extends RocksObject {
final AbstractComparator<? extends AbstractSlice<?>> comparator) {
super(newSstFileWriter(
envOptions.nativeHandle_, options.nativeHandle_, comparator.nativeHandle_,
comparator instanceof DirectComparator));
comparator.getComparatorType().getValue()));
}
/**
......@@ -225,7 +225,7 @@ public void put(final byte[] key, final byte[] value)
private native static long newSstFileWriter(
final long envOptionsHandle, final long optionsHandle,
final long userComparatorHandle, final boolean isDirect);
final long userComparatorHandle, final byte comparatorType);
private native static long newSstFileWriter(final long envOptionsHandle,
final long optionsHandle);
......
......@@ -61,7 +61,8 @@ public class WriteBatchWithIndex extends AbstractWriteBatch {
fallbackIndexComparator, final int reservedBytes,
final boolean overwriteKey) {
super(newWriteBatchWithIndex(fallbackIndexComparator.nativeHandle_,
fallbackIndexComparator instanceof DirectComparator, reservedBytes, overwriteKey));
fallbackIndexComparator.getComparatorType().getValue(), reservedBytes,
overwriteKey));
}
/**
......@@ -283,7 +284,8 @@ public class WriteBatchWithIndex extends AbstractWriteBatch {
private native static long newWriteBatchWithIndex();
private native static long newWriteBatchWithIndex(final boolean overwriteKey);
private native static long newWriteBatchWithIndex(
final long fallbackIndexComparatorHandle, final boolean isDirect, final int reservedBytes,
final long fallbackIndexComparatorHandle,
final byte comparatorType, final int reservedBytes,
final boolean overwriteKey);
private native long iterator0(final long handle);
private native long iterator1(final long handle, final long cfHandle);
......
// 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).
package org.rocksdb;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import java.util.*;
import java.util.Comparator;
import static org.junit.Assert.assertEquals;
public class NativeComparatorWrapperTest {
@Rule
public TemporaryFolder dbFolder = new TemporaryFolder();
private static final Random random = new Random();
@Test
public void rountrip() throws RocksDBException {
final String dbPath = dbFolder.getRoot().getAbsolutePath();
final int ITERATIONS = 1_000;
final String[] storedKeys = new String[ITERATIONS];
try (final NativeStringComparatorWrapper comparator = new NativeStringComparatorWrapper();
final Options opt = new Options()
.setCreateIfMissing(true)
.setComparator(comparator)) {
// store random integer keys
try (final RocksDB db = RocksDB.open(opt, dbPath)) {
for (int i = 0; i < ITERATIONS; i++) {
final String strKey = randomString();
final byte key[] = strKey.getBytes();
// does key already exist (avoid duplicates)
if (i > 0 && db.get(key) != null) {
i--; // generate a different key
} else {
db.put(key, "value".getBytes());