Java API - Implement GetFromBatch and GetFromBatchAndDB in WBWI
Summary: Needed for working with `get` after `merge` on a WBWI. Closes https://github.com/facebook/rocksdb/pull/1093 Differential Revision: D4137978 Pulled By: yhchiang fbshipit-source-id: e18d50d
This commit is contained in:
parent
815f54afad
commit
24bceb0963
@ -900,6 +900,38 @@ class JniUtil {
|
||||
|
||||
env->ReleaseByteArrayElements(jkey, key, JNI_ABORT);
|
||||
}
|
||||
|
||||
/*
|
||||
* Helper for operations on a value
|
||||
* for example WriteBatchWithIndex->GetFromBatch
|
||||
*/
|
||||
static jbyteArray v_op(
|
||||
std::function<rocksdb::Status(rocksdb::Slice, std::string*)> op,
|
||||
JNIEnv* env, jbyteArray jkey, jint jkey_len) {
|
||||
jboolean isCopy;
|
||||
jbyte* key = env->GetByteArrayElements(jkey, &isCopy);
|
||||
rocksdb::Slice key_slice(reinterpret_cast<char*>(key), jkey_len);
|
||||
|
||||
std::string value;
|
||||
rocksdb::Status s = op(key_slice, &value);
|
||||
|
||||
env->ReleaseByteArrayElements(jkey, key, JNI_ABORT);
|
||||
|
||||
if (s.IsNotFound()) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
if (s.ok()) {
|
||||
jbyteArray jret_value =
|
||||
env->NewByteArray(static_cast<jsize>(value.size()));
|
||||
env->SetByteArrayRegion(jret_value, 0, static_cast<jsize>(value.size()),
|
||||
reinterpret_cast<const jbyte*>(value.c_str()));
|
||||
return jret_value;
|
||||
}
|
||||
rocksdb::RocksDBExceptionJni::ThrowNew(env, s);
|
||||
|
||||
return nullptr;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace rocksdb
|
||||
|
@ -273,6 +273,87 @@ jlong Java_org_rocksdb_WriteBatchWithIndex_iteratorWithBase(
|
||||
return reinterpret_cast<jlong>(iterator);
|
||||
}
|
||||
|
||||
/*
|
||||
* Class: org_rocksdb_WriteBatchWithIndex
|
||||
* Method: getFromBatch
|
||||
* Signature: (JJ[BI)[B
|
||||
*/
|
||||
jbyteArray JNICALL Java_org_rocksdb_WriteBatchWithIndex_getFromBatch__JJ_3BI(
|
||||
JNIEnv* env, jobject jobj, jlong jwbwi_handle, jlong jdbopt_handle,
|
||||
jbyteArray jkey, jint jkey_len) {
|
||||
auto* wbwi = reinterpret_cast<rocksdb::WriteBatchWithIndex*>(jwbwi_handle);
|
||||
auto* dbopt = reinterpret_cast<rocksdb::DBOptions*>(jdbopt_handle);
|
||||
|
||||
auto getter = [&wbwi, &dbopt](const rocksdb::Slice& key, std::string* value) {
|
||||
return wbwi->GetFromBatch(*dbopt, key, value);
|
||||
};
|
||||
|
||||
return rocksdb::JniUtil::v_op(getter, env, jkey, jkey_len);
|
||||
}
|
||||
|
||||
/*
|
||||
* Class: org_rocksdb_WriteBatchWithIndex
|
||||
* Method: getFromBatch
|
||||
* Signature: (JJ[BIJ)[B
|
||||
*/
|
||||
jbyteArray Java_org_rocksdb_WriteBatchWithIndex_getFromBatch__JJ_3BIJ(
|
||||
JNIEnv* env, jobject jobj, jlong jwbwi_handle, jlong jdbopt_handle,
|
||||
jbyteArray jkey, jint jkey_len, jlong jcf_handle) {
|
||||
auto* wbwi = reinterpret_cast<rocksdb::WriteBatchWithIndex*>(jwbwi_handle);
|
||||
auto* dbopt = reinterpret_cast<rocksdb::DBOptions*>(jdbopt_handle);
|
||||
auto* cf_handle = reinterpret_cast<rocksdb::ColumnFamilyHandle*>(jcf_handle);
|
||||
|
||||
auto getter =
|
||||
[&wbwi, &cf_handle, &dbopt](const rocksdb::Slice& key,
|
||||
std::string* value) {
|
||||
return wbwi->GetFromBatch(cf_handle, *dbopt, key, value);
|
||||
};
|
||||
|
||||
return rocksdb::JniUtil::v_op(getter, env, jkey, jkey_len);
|
||||
}
|
||||
|
||||
/*
|
||||
* Class: org_rocksdb_WriteBatchWithIndex
|
||||
* Method: getFromBatchAndDB
|
||||
* Signature: (JJJ[BI)[B
|
||||
*/
|
||||
jbyteArray Java_org_rocksdb_WriteBatchWithIndex_getFromBatchAndDB__JJJ_3BI(
|
||||
JNIEnv* env, jobject jobj, jlong jwbwi_handle, jlong jdb_handle,
|
||||
jlong jreadopt_handle, jbyteArray jkey, jint jkey_len) {
|
||||
auto* wbwi = reinterpret_cast<rocksdb::WriteBatchWithIndex*>(jwbwi_handle);
|
||||
auto* db = reinterpret_cast<rocksdb::DB*>(jdb_handle);
|
||||
auto* readopt = reinterpret_cast<rocksdb::ReadOptions*>(jreadopt_handle);
|
||||
|
||||
auto getter =
|
||||
[&wbwi, &db, &readopt](const rocksdb::Slice& key, std::string* value) {
|
||||
return wbwi->GetFromBatchAndDB(db, *readopt, key, value);
|
||||
};
|
||||
|
||||
return rocksdb::JniUtil::v_op(getter, env, jkey, jkey_len);
|
||||
}
|
||||
|
||||
/*
|
||||
* Class: org_rocksdb_WriteBatchWithIndex
|
||||
* Method: getFromBatchAndDB
|
||||
* Signature: (JJJ[BIJ)[B
|
||||
*/
|
||||
jbyteArray Java_org_rocksdb_WriteBatchWithIndex_getFromBatchAndDB__JJJ_3BIJ(
|
||||
JNIEnv* env, jobject jobj, jlong jwbwi_handle, jlong jdb_handle,
|
||||
jlong jreadopt_handle, jbyteArray jkey, jint jkey_len, jlong jcf_handle) {
|
||||
auto* wbwi = reinterpret_cast<rocksdb::WriteBatchWithIndex*>(jwbwi_handle);
|
||||
auto* db = reinterpret_cast<rocksdb::DB*>(jdb_handle);
|
||||
auto* readopt = reinterpret_cast<rocksdb::ReadOptions*>(jreadopt_handle);
|
||||
auto* cf_handle = reinterpret_cast<rocksdb::ColumnFamilyHandle*>(jcf_handle);
|
||||
|
||||
auto getter =
|
||||
[&wbwi, &db, &cf_handle, &readopt](const rocksdb::Slice& key,
|
||||
std::string* value) {
|
||||
return wbwi->GetFromBatchAndDB(db, *readopt, cf_handle, key, value);
|
||||
};
|
||||
|
||||
return rocksdb::JniUtil::v_op(getter, env, jkey, jkey_len);
|
||||
}
|
||||
|
||||
/*
|
||||
* Class: org_rocksdb_WriteBatchWithIndex
|
||||
* Method: disposeInternal
|
||||
|
@ -136,6 +136,88 @@ public class WriteBatchWithIndex extends AbstractWriteBatch {
|
||||
baseIterator);
|
||||
}
|
||||
|
||||
/**
|
||||
* Similar to {@link RocksDB#get(ColumnFamilyHandle, byte[])} but will only
|
||||
* read the key from this batch.
|
||||
*
|
||||
* @param columnFamilyHandle The column family to retrieve the value from
|
||||
* @param options The database options to use
|
||||
* @param key The key to read the value for
|
||||
*
|
||||
* @throws RocksDBException if the batch does not have enough data to resolve
|
||||
* Merge operations, MergeInProgress status may be returned.
|
||||
*/
|
||||
public byte[] getFromBatch(final ColumnFamilyHandle columnFamilyHandle,
|
||||
final DBOptions options, final byte[] key) throws RocksDBException {
|
||||
return getFromBatch(nativeHandle_, options.nativeHandle_,
|
||||
key, key.length, columnFamilyHandle.nativeHandle_);
|
||||
}
|
||||
|
||||
/**
|
||||
* Similar to {@link RocksDB#get(byte[])} but will only
|
||||
* read the key from this batch.
|
||||
*
|
||||
* @param options The database options to use
|
||||
* @param key The key to read the value for
|
||||
*
|
||||
* @throws RocksDBException if the batch does not have enough data to resolve
|
||||
* Merge operations, MergeInProgress status may be returned.
|
||||
*/
|
||||
public byte[] getFromBatch(final DBOptions options, final byte[] key)
|
||||
throws RocksDBException {
|
||||
return getFromBatch(nativeHandle_, options.nativeHandle_, key, key.length);
|
||||
}
|
||||
|
||||
/**
|
||||
* Similar to {@link RocksDB#get(ColumnFamilyHandle, byte[])} but will also
|
||||
* read writes from this batch.
|
||||
*
|
||||
* This function will query both this batch and the DB and then merge
|
||||
* the results using the DB's merge operator (if the batch contains any
|
||||
* merge requests).
|
||||
*
|
||||
* Setting {@link ReadOptions#setSnapshot(long, long)} will affect what is
|
||||
* read from the DB but will NOT change which keys are read from the batch
|
||||
* (the keys in this batch do not yet belong to any snapshot and will be
|
||||
* fetched regardless).
|
||||
*
|
||||
* @param columnFamilyHandle The column family to retrieve the value from
|
||||
* @param options The read options to use
|
||||
* @param key The key to read the value for
|
||||
*
|
||||
* @throws RocksDBException if the value for the key cannot be read
|
||||
*/
|
||||
public byte[] getFromBatchAndDB(final RocksDB db, final ColumnFamilyHandle columnFamilyHandle,
|
||||
final ReadOptions options, final byte[] key) throws RocksDBException {
|
||||
return getFromBatchAndDB(nativeHandle_, db.nativeHandle_,
|
||||
options.nativeHandle_, key, key.length,
|
||||
columnFamilyHandle.nativeHandle_);
|
||||
}
|
||||
|
||||
/**
|
||||
* Similar to {@link RocksDB#get(byte[])} but will also
|
||||
* read writes from this batch.
|
||||
*
|
||||
* This function will query both this batch and the DB and then merge
|
||||
* the results using the DB's merge operator (if the batch contains any
|
||||
* merge requests).
|
||||
*
|
||||
* Setting {@link ReadOptions#setSnapshot(long, long)} will affect what is
|
||||
* read from the DB but will NOT change which keys are read from the batch
|
||||
* (the keys in this batch do not yet belong to any snapshot and will be
|
||||
* fetched regardless).
|
||||
*
|
||||
* @param options The read options to use
|
||||
* @param key The key to read the value for
|
||||
*
|
||||
* @throws RocksDBException if the value for the key cannot be read
|
||||
*/
|
||||
public byte[] getFromBatchAndDB(final RocksDB db, final ReadOptions options,
|
||||
final byte[] key) throws RocksDBException {
|
||||
return getFromBatchAndDB(nativeHandle_, db.nativeHandle_,
|
||||
options.nativeHandle_, key, key.length);
|
||||
}
|
||||
|
||||
@Override protected final native void disposeInternal(final long handle);
|
||||
@Override final native int count0(final long handle);
|
||||
@Override final native void put(final long handle, final byte[] key,
|
||||
@ -167,4 +249,14 @@ public class WriteBatchWithIndex extends AbstractWriteBatch {
|
||||
private native long iterator1(final long handle, final long cfHandle);
|
||||
private native long iteratorWithBase(final long handle,
|
||||
final long baseIteratorHandle, final long cfHandle);
|
||||
private native byte[] getFromBatch(final long handle, final long optHandle,
|
||||
final byte[] key, final int keyLen);
|
||||
private native byte[] getFromBatch(final long handle, final long optHandle,
|
||||
final byte[] key, final int keyLen, final long cfHandle);
|
||||
private native byte[] getFromBatchAndDB(final long handle,
|
||||
final long dbHandle, final long readOptHandle, final byte[] key,
|
||||
final int keyLen);
|
||||
private native byte[] getFromBatchAndDB(final long handle,
|
||||
final long dbHandle, final long readOptHandle, final byte[] key,
|
||||
final int keyLen, final long cfHandle);
|
||||
}
|
||||
|
@ -307,6 +307,77 @@ public class WriteBatchWithIndexTest {
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
public void getFromBatch() throws RocksDBException {
|
||||
final byte[] k1 = "k1".getBytes();
|
||||
final byte[] k2 = "k2".getBytes();
|
||||
final byte[] k3 = "k3".getBytes();
|
||||
final byte[] k4 = "k4".getBytes();
|
||||
|
||||
final byte[] v1 = "v1".getBytes();
|
||||
final byte[] v2 = "v2".getBytes();
|
||||
final byte[] v3 = "v3".getBytes();
|
||||
|
||||
try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(true);
|
||||
final DBOptions dbOptions = new DBOptions()) {
|
||||
wbwi.put(k1, v1);
|
||||
wbwi.put(k2, v2);
|
||||
wbwi.put(k3, v3);
|
||||
|
||||
assertThat(wbwi.getFromBatch(dbOptions, k1)).isEqualTo(v1);
|
||||
assertThat(wbwi.getFromBatch(dbOptions, k2)).isEqualTo(v2);
|
||||
assertThat(wbwi.getFromBatch(dbOptions, k3)).isEqualTo(v3);
|
||||
assertThat(wbwi.getFromBatch(dbOptions, k4)).isNull();
|
||||
|
||||
wbwi.remove(k2);
|
||||
|
||||
assertThat(wbwi.getFromBatch(dbOptions, k2)).isNull();
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
public void getFromBatchAndDB() throws RocksDBException {
|
||||
final byte[] k1 = "k1".getBytes();
|
||||
final byte[] k2 = "k2".getBytes();
|
||||
final byte[] k3 = "k3".getBytes();
|
||||
final byte[] k4 = "k4".getBytes();
|
||||
|
||||
final byte[] v1 = "v1".getBytes();
|
||||
final byte[] v2 = "v2".getBytes();
|
||||
final byte[] v3 = "v3".getBytes();
|
||||
final byte[] v4 = "v4".getBytes();
|
||||
|
||||
try (final Options options = new Options().setCreateIfMissing(true);
|
||||
final RocksDB db = RocksDB.open(options,
|
||||
dbFolder.getRoot().getAbsolutePath())) {
|
||||
|
||||
db.put(k1, v1);
|
||||
db.put(k2, v2);
|
||||
db.put(k4, v4);
|
||||
|
||||
try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(true);
|
||||
final DBOptions dbOptions = new DBOptions();
|
||||
final ReadOptions readOptions = new ReadOptions()) {
|
||||
|
||||
assertThat(wbwi.getFromBatch(dbOptions, k1)).isNull();
|
||||
assertThat(wbwi.getFromBatch(dbOptions, k2)).isNull();
|
||||
assertThat(wbwi.getFromBatch(dbOptions, k4)).isNull();
|
||||
|
||||
wbwi.put(k3, v3);
|
||||
|
||||
assertThat(wbwi.getFromBatch(dbOptions, k3)).isEqualTo(v3);
|
||||
|
||||
assertThat(wbwi.getFromBatchAndDB(db, readOptions, k1)).isEqualTo(v1);
|
||||
assertThat(wbwi.getFromBatchAndDB(db, readOptions, k2)).isEqualTo(v2);
|
||||
assertThat(wbwi.getFromBatchAndDB(db, readOptions, k3)).isEqualTo(v3);
|
||||
assertThat(wbwi.getFromBatchAndDB(db, readOptions, k4)).isEqualTo(v4);
|
||||
|
||||
wbwi.remove(k4);
|
||||
|
||||
assertThat(wbwi.getFromBatchAndDB(db, readOptions, k4)).isNull();
|
||||
}
|
||||
}
|
||||
}
|
||||
private byte[] toArray(final ByteBuffer buf) {
|
||||
final byte[] ary = new byte[buf.remaining()];
|
||||
buf.get(ary);
|
||||
|
Loading…
Reference in New Issue
Block a user