Add keyExists Java API (#11705)

Summary:
Add a new method to check if a key exists in the database. It avoids copying data between C++ and Java code.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/11705

Reviewed By: ajkr

Differential Revision: D50370934

Pulled By: akankshamahajan15

fbshipit-source-id: ab2d42213fbebcaff919b0ffbbef9d45e88ca365
This commit is contained in:
Radek Hubner 2023-10-18 12:46:35 -07:00 committed by Facebook GitHub Bot
parent 0bb3a26d89
commit a80e3f6c57
4 changed files with 587 additions and 0 deletions

View File

@ -142,6 +142,7 @@ JAVA_TESTS = \
org.rocksdb.FilterTest\
org.rocksdb.FlushTest\
org.rocksdb.InfoLogLevelTest\
org.rocksdb.KeyExistsTest \
org.rocksdb.KeyMayExistTest\
org.rocksdb.ConcurrentTaskLimiterTest\
org.rocksdb.LoggerTest\

View File

@ -2215,6 +2215,108 @@ bool key_may_exist_direct_helper(JNIEnv* env, jlong jdb_handle,
return exists;
}
jboolean key_exists_helper(JNIEnv* env, jlong jdb_handle, jlong jcf_handle,
jlong jread_opts_handle, char* key, jint jkey_len) {
std::string value;
bool value_found = false;
auto* db = reinterpret_cast<ROCKSDB_NAMESPACE::DB*>(jdb_handle);
ROCKSDB_NAMESPACE::ColumnFamilyHandle* cf_handle;
if (jcf_handle == 0) {
cf_handle = db->DefaultColumnFamily();
} else {
cf_handle =
reinterpret_cast<ROCKSDB_NAMESPACE::ColumnFamilyHandle*>(jcf_handle);
}
ROCKSDB_NAMESPACE::ReadOptions read_opts =
jread_opts_handle == 0
? ROCKSDB_NAMESPACE::ReadOptions()
: *(reinterpret_cast<ROCKSDB_NAMESPACE::ReadOptions*>(
jread_opts_handle));
ROCKSDB_NAMESPACE::Slice key_slice(key, jkey_len);
const bool may_exist =
db->KeyMayExist(read_opts, cf_handle, key_slice, &value, &value_found);
if (may_exist) {
ROCKSDB_NAMESPACE::Status s;
{
ROCKSDB_NAMESPACE::PinnableSlice pinnable_val;
s = db->Get(read_opts, cf_handle, key_slice, &pinnable_val);
}
if (s.IsNotFound()) {
return JNI_FALSE;
} else if (s.ok()) {
return JNI_TRUE;
} else {
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(env, s);
return JNI_FALSE;
}
} else {
return JNI_FALSE;
}
}
/*
* Class: org_rocksdb_RocksDB
* Method: keyExist
* Signature: (JJJ[BII)Z
*/
jboolean Java_org_rocksdb_RocksDB_keyExists(JNIEnv* env, jobject,
jlong jdb_handle, jlong jcf_handle,
jlong jread_opts_handle,
jbyteArray jkey, jint jkey_offset,
jint jkey_len) {
jbyte* key = new jbyte[jkey_len];
env->GetByteArrayRegion(jkey, jkey_offset, jkey_len, key);
if (env->ExceptionCheck()) {
// exception thrown: ArrayIndexOutOfBoundsException
delete[] key;
return JNI_FALSE;
} else {
jboolean key_exists =
key_exists_helper(env, jdb_handle, jcf_handle, jread_opts_handle,
reinterpret_cast<char*>(key), jkey_len);
delete[] key;
return key_exists;
}
}
/*
private native boolean keyExistDirect(final long handle, final long
cfHandle, final long readOptHandle, final ByteBuffer key, final int keyOffset,
final int keyLength);
* Class: org_rocksdb_RocksDB
* Method: keyExistDirect
* Signature: (JJJLjava/nio/ByteBuffer;II)Z
*/
jboolean Java_org_rocksdb_RocksDB_keyExistsDirect(
JNIEnv* env, jobject, jlong jdb_handle, jlong jcf_handle,
jlong jread_opts_handle, jobject jkey, jint jkey_offset, jint jkey_len) {
char* key = reinterpret_cast<char*>(env->GetDirectBufferAddress(jkey));
if (key == nullptr) {
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(
env,
"Invalid key argument (argument is not a valid direct ByteBuffer)");
return JNI_FALSE;
}
if (env->GetDirectBufferCapacity(jkey) < (jkey_offset + jkey_len)) {
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(
env,
"Invalid key argument. Capacity is less than requested region (offset "
"+ length).");
return JNI_FALSE;
}
return key_exists_helper(env, jdb_handle, jcf_handle, jread_opts_handle, key,
jkey_len);
}
/*
* Class: org_rocksdb_RocksDB
* Method: keyMayExist

View File

@ -2435,6 +2435,259 @@ public class RocksDB extends RocksObject {
return results;
}
/**
* Check if a key exists in the database.
* This method is not as lightweight as {@code keyMayExist} but it gives a 100% guarantee
* of a correct result, whether the key exists or not.
*
* Internally it checks if the key may exist and then double checks with read operation
* that confirms the key exists. This deals with the case where {@code keyMayExist} may return
* a false positive.
*
* The code crosses the Java/JNI boundary only once.
* @param key byte array of a key to search for*
* @return true if key exist in database, otherwise false.
*/
public boolean keyExists(final byte[] key) {
return keyExists(key, 0, key.length);
}
/**
* Check if a key exists in the database.
* This method is not as lightweight as {@code keyMayExist} but it gives a 100% guarantee
* of a correct result, whether the key exists or not.
*
* Internally it checks if the key may exist and then double checks with read operation
* that confirms the key exists. This deals with the case where {@code keyMayExist} may return
* a false positive.
*
* The code crosses the Java/JNI boundary only once.
* @param key byte array of a key to search for
* @param offset the offset of the "key" array to be used, must be
* non-negative and no larger than "key".length
* @param len the length of the "key" array to be used, must be non-negative
* and no larger than "key".length
* @return true if key exist in database, otherwise false.
*/
public boolean keyExists(final byte[] key, final int offset, final int len) {
return keyExists(null, null, key, offset, len);
}
/**
* Check if a key exists in the database.
* This method is not as lightweight as {@code keyMayExist} but it gives a 100% guarantee
* of a correct result, whether the key exists or not.
*
* Internally it checks if the key may exist and then double checks with read operation
* that confirms the key exists. This deals with the case where {@code keyMayExist} may return
* a false positive.
*
* The code crosses the Java/JNI boundary only once.
*
* @param columnFamilyHandle {@link ColumnFamilyHandle} instance
* @param key byte array of a key to search for
* @return true if key exist in database, otherwise false.
*/
public boolean keyExists(final ColumnFamilyHandle columnFamilyHandle, final byte[] key) {
return keyExists(columnFamilyHandle, key, 0, key.length);
}
/**
* Check if a key exists in the database.
* This method is not as lightweight as {@code keyMayExist} but it gives a 100% guarantee
* of a correct result, whether the key exists or not.
*
* Internally it checks if the key may exist and then double checks with read operation
* that confirms the key exists. This deals with the case where {@code keyMayExist} may return
* a false positive.
*
* The code crosses the Java/JNI boundary only once.
*
* @param columnFamilyHandle {@link ColumnFamilyHandle} instance
* @param key byte array of a key to search for
* @param offset the offset of the "key" array to be used, must be
* non-negative and no larger than "key".length
* @param len the length of the "key" array to be used, must be non-negative
* and no larger than "key".length
* @return true if key exist in database, otherwise false.
*/
public boolean keyExists(final ColumnFamilyHandle columnFamilyHandle, final byte[] key,
final int offset, final int len) {
return keyExists(columnFamilyHandle, null, key, offset, len);
}
/**
* Check if a key exists in the database.
* This method is not as lightweight as {@code keyMayExist} but it gives a 100% guarantee
* of a correct result, whether the key exists or not.
*
* Internally it checks if the key may exist and then double checks with read operation
* that confirms the key exists. This deals with the case where {@code keyMayExist} may return
* a false positive.
*
* The code crosses the Java/JNI boundary only once.
*
* @param readOptions {@link ReadOptions} instance
* @param key byte array of a key to search for
* @return true if key exist in database, otherwise false.
*/
public boolean keyExists(final ReadOptions readOptions, final byte[] key) {
return keyExists(readOptions, key, 0, key.length);
}
/**
* Check if a key exists in the database.
* This method is not as lightweight as {@code keyMayExist} but it gives a 100% guarantee
* of a correct result, whether the key exists or not.
*
* Internally it checks if the key may exist and then double checks with read operation
* that confirms the key exists. This deals with the case where {@code keyMayExist} may return
* a false positive.
*
* The code crosses the Java/JNI boundary only once.
*
* @param readOptions {@link ReadOptions} instance
* @param key byte array of a key to search for
* @param offset the offset of the "key" array to be used, must be
* non-negative and no larger than "key".length
* @param len the length of the "key" array to be used, must be non-negative
* and no larger than "key".length
* @return true if key exist in database, otherwise false.
*/
public boolean keyExists(
final ReadOptions readOptions, final byte[] key, final int offset, final int len) {
return keyExists(null, readOptions, key, offset, len);
}
/**
* Check if a key exists in the database.
* This method is not as lightweight as {@code keyMayExist} but it gives a 100% guarantee
* of a correct result, whether the key exists or not.
*
* Internally it checks if the key may exist and then double checks with read operation
* that confirms the key exists. This deals with the case where {@code keyMayExist} may return
* a false positive.
*
* The code crosses the Java/JNI boundary only once.
*
* @param columnFamilyHandle {@link ColumnFamilyHandle} instance
* @param readOptions {@link ReadOptions} instance
* @param key byte array of a key to search for
* @return true if key exist in database, otherwise false.
*/
public boolean keyExists(final ColumnFamilyHandle columnFamilyHandle,
final ReadOptions readOptions, final byte[] key) {
return keyExists(columnFamilyHandle, readOptions, key, 0, key.length);
}
/**
* Check if a key exists in the database.
* This method is not as lightweight as {@code keyMayExist} but it gives a 100% guarantee
* of a correct result, whether the key exists or not.
*
* Internally it checks if the key may exist and then double checks with read operation
* that confirms the key exists. This deals with the case where {@code keyMayExist} may return
* a false positive.
*
* The code crosses the Java/JNI boundary only once.
*
* @param columnFamilyHandle {@link ColumnFamilyHandle} instance
* @param readOptions {@link ReadOptions} instance
* @param key byte array of a key to search for
* @param offset the offset of the "key" array to be used, must be
* non-negative and no larger than "key".length
* @param len the length of the "key" array to be used, must be non-negative
* and no larger than "key".length
* @return true if key exist in database, otherwise false.
*/
public boolean keyExists(final ColumnFamilyHandle columnFamilyHandle,
final ReadOptions readOptions, final byte[] key, final int offset, final int len) {
checkBounds(offset, len, key.length);
return keyExists(nativeHandle_,
columnFamilyHandle == null ? 0 : columnFamilyHandle.nativeHandle_,
readOptions == null ? 0 : readOptions.nativeHandle_, key, offset, len);
}
/**
* Check if a key exists in the database.
* This method is not as lightweight as {@code keyMayExist} but it gives a 100% guarantee
* of a correct result, whether the key exists or not.
*
* Internally it checks if the key may exist and then double checks with read operation
* that confirms the key exists. This deals with the case where {@code keyMayExist} may return
* a false positive.
*
* The code crosses the Java/JNI boundary only once.
*
* @param key ByteBuffer with key. Must be allocated as direct.
* @return true if key exist in database, otherwise false.
*/
public boolean keyExists(final ByteBuffer key) {
return keyExists(null, null, key);
}
/**
* Check if a key exists in the database.
* This method is not as lightweight as {@code keyMayExist} but it gives a 100% guarantee
* of a correct result, whether the key exists or not.
*
* Internally it checks if the key may exist and then double checks with read operation
* that confirms the key exists. This deals with the case where {@code keyMayExist} may return
* a false positive.
*
* The code crosses the Java/JNI boundary only once.
*
* @param columnFamilyHandle {@link ColumnFamilyHandle} instance
* @param key ByteBuffer with key. Must be allocated as direct.
* @return true if key exist in database, otherwise false.
*/
public boolean keyExists(final ColumnFamilyHandle columnFamilyHandle, final ByteBuffer key) {
return keyExists(columnFamilyHandle, null, key);
}
/**
* Check if a key exists in the database.
* This method is not as lightweight as {@code keyMayExist} but it gives a 100% guarantee
* of a correct result, whether the key exists or not.
*
* Internally it checks if the key may exist and then double checks with read operation
* that confirms the key exists. This deals with the case where {@code keyMayExist} may return
* a false positive.
*
* The code crosses the Java/JNI boundary only once.
*
* @param readOptions {@link ReadOptions} instance
* @param key ByteBuffer with key. Must be allocated as direct.
* @return true if key exist in database, otherwise false.
*/
public boolean keyExists(final ReadOptions readOptions, final ByteBuffer key) {
return keyExists(null, readOptions, key);
}
/**
* Check if a key exists in the database.
* This method is not as lightweight as {@code keyMayExist} but it gives a 100% guarantee
* of a correct result, whether the key exists or not.
*
* Internally it checks if the key may exist and then double checks with read operation
* that confirms the key exists. This deals with the case where {@code keyMayExist} may return
* a false positive.
*
* The code crosses the Java/JNI boundary only once.
*
* @param columnFamilyHandle {@link ColumnFamilyHandle} instance
* @param readOptions {@link ReadOptions} instance
* @param key ByteBuffer with key. Must be allocated as direct.
* @return true if key exist in database, otherwise false.
*/
public boolean keyExists(final ColumnFamilyHandle columnFamilyHandle,
final ReadOptions readOptions, final ByteBuffer key) {
assert key != null : "key ByteBuffer parameter cannot be null";
assert key.isDirect() : "key parameter must be a direct ByteBuffer";
return keyExistsDirect(nativeHandle_,
columnFamilyHandle == null ? 0 : columnFamilyHandle.nativeHandle_,
readOptions == null ? 0 : readOptions.nativeHandle_, key, key.position(), key.limit());
}
/**
* If the key definitely does not exist in the database, then this method
* returns false, otherwise it returns true if the key might exist.
@ -4559,6 +4812,12 @@ public class RocksDB extends RocksObject {
final int[] keyLengths, final ByteBuffer[] valuesArray, final int[] valuesSizeArray,
final Status[] statusArray);
private native boolean keyExists(final long handle, final long cfHandle, final long readOptHandle,
final byte[] key, final int keyOffset, final int keyLength);
private native boolean keyExistsDirect(final long handle, final long cfHandle,
final long readOptHandle, final ByteBuffer key, final int keyOffset, final int keyLength);
private native boolean keyMayExist(
final long handle, final long cfHandle, final long readOptHandle,
final byte[] key, final int keyOffset, final int keyLength);

View File

@ -0,0 +1,225 @@
package org.rocksdb;
import static java.nio.charset.StandardCharsets.UTF_8;
import static org.assertj.core.api.Assertions.assertThat;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import org.junit.*;
import org.junit.rules.ExpectedException;
import org.junit.rules.TemporaryFolder;
public class KeyExistsTest {
@ClassRule
public static final RocksNativeLibraryResource ROCKS_NATIVE_LIBRARY_RESOURCE =
new RocksNativeLibraryResource();
@Rule public TemporaryFolder dbFolder = new TemporaryFolder();
@Rule public ExpectedException exceptionRule = ExpectedException.none();
List<ColumnFamilyDescriptor> cfDescriptors;
List<ColumnFamilyHandle> columnFamilyHandleList = new ArrayList<>();
RocksDB db;
@Before
public void before() throws RocksDBException {
cfDescriptors = Arrays.asList(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY),
new ColumnFamilyDescriptor("new_cf".getBytes()));
final DBOptions options =
new DBOptions().setCreateIfMissing(true).setCreateMissingColumnFamilies(true);
db = RocksDB.open(
options, dbFolder.getRoot().getAbsolutePath(), cfDescriptors, columnFamilyHandleList);
}
@After
public void after() {
for (final ColumnFamilyHandle columnFamilyHandle : columnFamilyHandleList) {
columnFamilyHandle.close();
}
db.close();
}
@Test
public void keyExists() throws RocksDBException {
db.put("key".getBytes(UTF_8), "value".getBytes(UTF_8));
boolean exists = db.keyExists("key".getBytes(UTF_8));
assertThat(exists).isTrue();
exists = db.keyExists("key2".getBytes(UTF_8));
assertThat(exists).isFalse();
}
@Test
public void keyExistsColumnFamily() throws RocksDBException {
byte[] key1 = "keyBBCF0".getBytes(UTF_8);
byte[] key2 = "keyBBCF1".getBytes(UTF_8);
db.put(columnFamilyHandleList.get(0), key1, "valueBBCF0".getBytes(UTF_8));
db.put(columnFamilyHandleList.get(1), key2, "valueBBCF1".getBytes(UTF_8));
assertThat(db.keyExists(columnFamilyHandleList.get(0), key1)).isTrue();
assertThat(db.keyExists(columnFamilyHandleList.get(0), key2)).isFalse();
assertThat(db.keyExists(columnFamilyHandleList.get(1), key1)).isFalse();
assertThat(db.keyExists(columnFamilyHandleList.get(1), key2)).isTrue();
}
@Test
public void keyExistsColumnFamilyReadOptions() throws RocksDBException {
try (final ReadOptions readOptions = new ReadOptions()) {
byte[] key1 = "keyBBCF0".getBytes(UTF_8);
byte[] key2 = "keyBBCF1".getBytes(UTF_8);
db.put(columnFamilyHandleList.get(0), key1, "valueBBCF0".getBytes(UTF_8));
db.put(columnFamilyHandleList.get(1), key2, "valueBBCF1".getBytes(UTF_8));
assertThat(db.keyExists(columnFamilyHandleList.get(0), readOptions, key1)).isTrue();
assertThat(db.keyExists(columnFamilyHandleList.get(0), readOptions, key2)).isFalse();
assertThat(db.keyExists(columnFamilyHandleList.get(1), readOptions, key1)).isFalse();
assertThat(db.keyExists(columnFamilyHandleList.get(1), readOptions, key2)).isTrue();
}
}
@Test
public void keyExistsReadOptions() throws RocksDBException {
try (final ReadOptions readOptions = new ReadOptions()) {
db.put("key".getBytes(UTF_8), "value".getBytes(UTF_8));
boolean exists = db.keyExists(readOptions, "key".getBytes(UTF_8));
assertThat(exists).isTrue();
exists = db.keyExists("key2".getBytes(UTF_8));
assertThat(exists).isFalse();
}
}
@Test
public void keyExistsAfterDelete() throws RocksDBException {
db.put("key".getBytes(UTF_8), "value".getBytes(UTF_8));
boolean exists = db.keyExists(null, null, "key".getBytes(UTF_8), 0, 3);
assertThat(exists).isTrue();
db.delete("key".getBytes(UTF_8));
exists = db.keyExists(null, null, "key".getBytes(UTF_8), 0, 3);
assertThat(exists).isFalse();
}
@Test
public void keyExistsArrayIndexOutOfBoundsException() throws RocksDBException {
db.put("key".getBytes(UTF_8), "value".getBytes(UTF_8));
exceptionRule.expect(IndexOutOfBoundsException.class);
db.keyExists(null, null, "key".getBytes(UTF_8), 0, 5);
}
@Test()
public void keyExistsArrayIndexOutOfBoundsExceptionWrongOffset() throws RocksDBException {
db.put("key".getBytes(UTF_8), "value".getBytes(UTF_8));
exceptionRule.expect(IndexOutOfBoundsException.class);
db.keyExists(null, null, "key".getBytes(UTF_8), 6, 2);
}
@Test
public void keyExistsDirectByteBuffer() throws RocksDBException {
byte[] key = "key".getBytes(UTF_8);
db.put(key, "value".getBytes(UTF_8));
ByteBuffer buff = ByteBuffer.allocateDirect(key.length);
buff.put(key);
buff.flip();
boolean exists = db.keyExists(buff);
assertThat(exists).isTrue();
}
@Test
public void keyExistsDirectByteBufferReadOptions() throws RocksDBException {
try (final ReadOptions readOptions = new ReadOptions()) {
byte[] key = "key".getBytes(UTF_8);
db.put(key, "value".getBytes(UTF_8));
ByteBuffer buff = ByteBuffer.allocateDirect(key.length);
buff.put(key);
buff.flip();
boolean exists = db.keyExists(buff);
assertThat(exists).isTrue();
}
}
@Test
public void keyExistsDirectByteBufferAfterDelete() throws RocksDBException {
byte[] key = "key".getBytes(UTF_8);
db.put(key, "value".getBytes(UTF_8));
ByteBuffer buff = ByteBuffer.allocateDirect(key.length);
buff.put(key);
buff.flip();
boolean exists = db.keyExists(buff);
assertThat(exists).isTrue();
db.delete(key);
exists = db.keyExists(buff);
assertThat(exists).isFalse();
}
@Test
public void keyExistsDirectByteBufferColumnFamily() throws RocksDBException {
byte[] key1 = "keyBBCF0".getBytes(UTF_8);
byte[] key2 = "keyBBCF1".getBytes(UTF_8);
db.put(columnFamilyHandleList.get(0), key1, "valueBBCF0".getBytes(UTF_8));
db.put(columnFamilyHandleList.get(1), key2, "valueBBCF1".getBytes(UTF_8));
ByteBuffer key1Buff = ByteBuffer.allocateDirect(key1.length);
key1Buff.put(key1);
key1Buff.flip();
ByteBuffer key2Buff = ByteBuffer.allocateDirect(key2.length);
key2Buff.put(key2);
key2Buff.flip();
assertThat(db.keyExists(columnFamilyHandleList.get(0), key1Buff)).isTrue();
assertThat(db.keyExists(columnFamilyHandleList.get(0), key2Buff)).isFalse();
assertThat(db.keyExists(columnFamilyHandleList.get(1), key1Buff)).isFalse();
assertThat(db.keyExists(columnFamilyHandleList.get(1), key2Buff)).isTrue();
}
@Test
public void keyExistsDirectByteBufferColumnFamilyReadOptions() throws RocksDBException {
try (final ReadOptions readOptions = new ReadOptions()) {
byte[] key1 = "keyBBCF0".getBytes(UTF_8);
byte[] key2 = "keyBBCF1".getBytes(UTF_8);
db.put(columnFamilyHandleList.get(0), key1, "valueBBCF0".getBytes(UTF_8));
db.put(columnFamilyHandleList.get(1), key2, "valueBBCF1".getBytes(UTF_8));
ByteBuffer key1Buff = ByteBuffer.allocateDirect(key1.length);
key1Buff.put(key1);
key1Buff.flip();
ByteBuffer key2Buff = ByteBuffer.allocateDirect(key2.length);
key2Buff.put(key2);
key2Buff.flip();
assertThat(db.keyExists(columnFamilyHandleList.get(0), readOptions, key1Buff)).isTrue();
assertThat(db.keyExists(columnFamilyHandleList.get(0), readOptions, key2Buff)).isFalse();
assertThat(db.keyExists(columnFamilyHandleList.get(1), readOptions, key1Buff)).isFalse();
assertThat(db.keyExists(columnFamilyHandleList.get(1), readOptions, key2Buff)).isTrue();
}
}
@Test
public void keyExistsDirectReadOptions() throws RocksDBException {
try (final ReadOptions readOptions = new ReadOptions()) {
byte[] key = "key1".getBytes(UTF_8);
db.put(key, "value".getBytes(UTF_8));
ByteBuffer buff = ByteBuffer.allocateDirect(key.length);
buff.put(key);
buff.flip();
boolean exists = db.keyExists(readOptions, key);
assertThat(exists).isTrue();
buff.clear();
buff.put("key2".getBytes(UTF_8));
buff.flip();
exists = db.keyExists("key2".getBytes(UTF_8));
assertThat(exists).isFalse();
}
}
}