Add support for write batch with index.

rocksdb.nim

@@ -9,9 +9,10 @@
 
 import
   ./rocksdb/[
-    backup, columnfamily, rocksdb, rocksiterator, sstfilewriter, transactiondb,
-    writebatch,
+    backup, columnfamily, optimistictxdb, rocksdb, rocksiterator, sstfilewriter,
+    transactiondb, writebatch, writebatchwi,
   ]
 
 export
-  backup, columnfamily, rocksdb, rocksiterator, sstfilewriter, transactiondb, writebatch
+  backup, columnfamily, optimistictxdb, rocksdb, rocksiterator, sstfilewriter,
+  transactiondb, writebatch, writebatchwi

rocksdb/rocksdb.nim

@@ -31,11 +31,11 @@ import
   ./options/[dbopts, readopts, writeopts],
   ./columnfamily/[cfopts, cfdescriptor, cfhandle],
   ./internal/[cftable, utils],
-  ./rocksiterator,
-  ./rocksresult,
-  ./writebatch
+  ./[rocksiterator, rocksresult, writebatch, writebatchwi]
 
-export rocksresult, dbopts, readopts, writeopts, cfdescriptor, rocksiterator, writebatch
+export
+  rocksresult, dbopts, readopts, writeopts, cfdescriptor, cfhandle, rocksiterator,
+  writebatch, writebatchwi
 
 type
   RocksDbPtr* = ptr rocksdb_t
@@ -371,6 +371,7 @@ proc openIterator*(
     db: RocksDbRef, cfHandle = db.defaultCfHandle
 ): RocksDBResult[RocksIteratorRef] =
   ## Opens an `RocksIteratorRef` for the specified column family.
+  ## The iterator should be closed using the `close` method after usage.
   doAssert not db.isClosed()
 
   let rocksIterPtr =
@@ -382,10 +383,31 @@ proc openWriteBatch*(
     db: RocksDbReadWriteRef, cfHandle = db.defaultCfHandle
 ): WriteBatchRef =
   ## Opens a `WriteBatchRef` which defaults to using the specified column family.
+  ## The write batch should be closed using the `close` method after usage.
   doAssert not db.isClosed()
 
   createWriteBatch(cfHandle)
 
+proc openWriteBatchWithIndex*(
+    db: RocksDbReadWriteRef,
+    reservedBytes = 0,
+    overwriteKey = false,
+    cfHandle = db.defaultCfHandle,
+): WriteBatchWIRef =
+  ## Opens a `WriteBatchWIRef` which defaults to using the specified column family.
+  ## The write batch should be closed using the `close` method after usage.
+  ## `WriteBatchWIRef` is similar to `WriteBatchRef` but with a binary searchable
+  ## index built for all the keys inserted which allows reading the data which has
+  ## been writen to the batch.
+  ##
+  ## Optionally set the number of bytes to be reserved for the batch by setting
+  ## `reservedBytes`. Set `overwriteKey` to true to overwrite the key in the index
+  ## when inserting a duplicate key, in this way an iterator will never show two
+  ## entries with the same key.
+  doAssert not db.isClosed()
+
+  createWriteBatch(reservedBytes, overwriteKey, db.dbOpts, cfHandle)
+
 proc write*(db: RocksDbReadWriteRef, updates: WriteBatchRef): RocksDBResult[void] =
   ## Apply the updates in the `WriteBatchRef` to the database.
   doAssert not db.isClosed()
@@ -398,6 +420,18 @@ proc write*(db: RocksDbReadWriteRef, updates: WriteBatchRef): RocksDBResult[void
 
   ok()
 
+proc write*(db: RocksDbReadWriteRef, updates: WriteBatchWIRef): RocksDBResult[void] =
+  ## Apply the updates in the `WriteBatchWIRef` to the database.
+  doAssert not db.isClosed()
+
+  var errors: cstring
+  rocksdb_write_writebatch_wi(
+    db.cPtr, db.writeOpts.cPtr, updates.cPtr, cast[cstringArray](errors.addr)
+  )
+  bailOnErrors(errors)
+
+  ok()
+
 proc ingestExternalFile*(
     db: RocksDbReadWriteRef, filePath: string, cfHandle = db.defaultCfHandle
 ): RocksDBResult[void] =

rocksdb/writebatch.nim

@@ -8,6 +8,8 @@
 # at your option. This file may not be copied, modified, or distributed except according to those terms.
 
 ## A `WriteBatchRef` holds a collection of updates to apply atomically to the database.
+## It depends on resources from an instance of `RocksDbRef' and therefore should be used
+## and closed before the `RocksDbRef` is closed.
 
 {.push raises: [].}
 

rocksdb/writebatchwi.nim

@@ -0,0 +1,153 @@
+# Nim-RocksDB
+# Copyright 2024 Status Research & Development GmbH
+# Licensed under either of
+#
+#  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or http://www.apache.org/licenses/LICENSE-2.0)
+#  * GPL license, version 2.0, ([LICENSE-GPLv2](LICENSE-GPLv2) or https://www.gnu.org/licenses/old-licenses/gpl-2.0.en.html)
+#
+# at your option. This file may not be copied, modified, or distributed except according to those terms.
+
+## A `WriteBatchWIRef` holds a collection of updates to apply atomically to the database.
+## It depends on resources from an instance of `RocksDbRef' and therefore should be used
+## and closed before the `RocksDbRef` is closed.
+##
+## `WriteBatchWIRef` is similar to `WriteBatchRef` but with a binary searchable index
+## built for all the keys inserted which allows reading the data which has been writen
+## to the batch.
+
+{.push raises: [].}
+
+import ./lib/librocksdb, ./internal/[cftable, utils], ./options/dbopts, ./rocksresult
+
+export rocksresult
+
+type
+  WriteBatchWIPtr* = ptr rocksdb_writebatch_wi_t
+
+  WriteBatchWIRef* = ref object
+    cPtr: WriteBatchWIPtr
+    dbOpts: DbOptionsRef
+    defaultCfHandle: ColFamilyHandleRef
+
+proc createWriteBatch*(
+    reservedBytes: int,
+    overwriteKey: bool,
+    dbOpts: DbOptionsRef,
+    defaultCfHandle: ColFamilyHandleRef,
+): WriteBatchWIRef =
+  WriteBatchWIRef(
+    cPtr: rocksdb_writebatch_wi_create(reservedBytes.csize_t, overwriteKey.uint8),
+    dbOpts: dbOpts,
+    defaultCfHandle: defaultCfHandle,
+  )
+
+proc isClosed*(batch: WriteBatchWIRef): bool {.inline.} =
+  ## Returns `true` if the `WriteBatchWIRef` has been closed and `false` otherwise.
+  batch.cPtr.isNil()
+
+proc cPtr*(batch: WriteBatchWIRef): WriteBatchWIPtr =
+  ## Get the underlying database pointer.
+  doAssert not batch.isClosed()
+  batch.cPtr
+
+proc clear*(batch: WriteBatchWIRef) =
+  ## Clears the write batch.
+  doAssert not batch.isClosed()
+  rocksdb_writebatch_wi_clear(batch.cPtr)
+
+proc count*(batch: WriteBatchWIRef): int =
+  ## Get the number of updates in the write batch.
+  doAssert not batch.isClosed()
+  rocksdb_writebatch_wi_count(batch.cPtr).int
+
+proc put*(
+    batch: WriteBatchWIRef, key, val: openArray[byte], cfHandle = batch.defaultCfHandle
+): RocksDBResult[void] =
+  ## Add a put operation to the write batch.
+
+  if key.len() == 0:
+    return err("rocksdb: key is empty")
+
+  rocksdb_writebatch_wi_put_cf(
+    batch.cPtr,
+    cfHandle.cPtr,
+    cast[cstring](unsafeAddr key[0]),
+    csize_t(key.len),
+    cast[cstring](if val.len > 0:
+      unsafeAddr val[0]
+    else:
+      nil
+    ),
+    csize_t(val.len),
+  )
+
+  ok()
+
+proc delete*(
+    batch: WriteBatchWIRef, key: openArray[byte], cfHandle = batch.defaultCfHandle
+): RocksDBResult[void] =
+  ## Add a delete operation to the write batch.
+
+  if key.len() == 0:
+    return err("rocksdb: key is empty")
+
+  rocksdb_writebatch_wi_delete_cf(
+    batch.cPtr, cfHandle.cPtr, cast[cstring](unsafeAddr key[0]), csize_t(key.len)
+  )
+
+  ok()
+
+proc get*(
+    batch: WriteBatchWIRef,
+    key: openArray[byte],
+    onData: DataProc,
+    cfHandle = batch.defaultCfHandle,
+): RocksDBResult[bool] =
+  ## Get the value for a given key from the batch using the provided
+  ## `onData` callback.
+
+  if key.len() == 0:
+    return err("rocksdb: key is empty")
+
+  var
+    len: csize_t
+    errors: cstring
+  let data = rocksdb_writebatch_wi_get_from_batch_cf(
+    batch.cPtr,
+    batch.dbOpts.cPtr,
+    cfHandle.cPtr,
+    cast[cstring](unsafeAddr key[0]),
+    csize_t(key.len),
+    len.addr,
+    cast[cstringArray](errors.addr),
+  )
+  bailOnErrors(errors)
+
+  if data.isNil():
+    doAssert len == 0
+    ok(false)
+  else:
+    onData(toOpenArrayByte(data, 0, len.int - 1))
+    rocksdb_free(data)
+    ok(true)
+
+proc get*(
+    batch: WriteBatchWIRef, key: openArray[byte], cfHandle = batch.defaultCfHandle
+): RocksDBResult[seq[byte]] =
+  ## Get the value for a given key from the batch.
+
+  var dataRes: RocksDBResult[seq[byte]]
+  proc onData(data: openArray[byte]) =
+    dataRes.ok(@data)
+
+  let res = batch.get(key, onData, cfHandle)
+  if res.isOk():
+    return dataRes
+
+  dataRes.err(res.error())
+
+proc close*(batch: WriteBatchWIRef) =
+  ## Close the `WriteBatchWIRef`.
+  if not batch.isClosed():
+    rocksdb_writebatch_wi_destroy(batch.cPtr)
+    batch.cPtr = nil

tests/test_all.nim

@@ -28,4 +28,5 @@ import
   ./test_rocksiterator,
   ./test_sstfilewriter,
   ./test_transactiondb,
-  ./test_writebatch
+  ./test_writebatch,
+  ./test_writebatchwi

tests/test_writebatch.nim

@@ -36,7 +36,7 @@ suite "WriteBatchRef Tests":
     removeDir($dbPath)
 
   test "Test writing batch to the default column family":
-    var batch = db.openWriteBatch()
+    let batch = db.openWriteBatch()
     defer:
       batch.close()
     check not batch.isClosed()
@@ -65,7 +65,7 @@ suite "WriteBatchRef Tests":
       not batch.isClosed()
 
   test "Test writing batch to column family":
-    var batch = db.openWriteBatch()
+    let batch = db.openWriteBatch()
     defer:
       batch.close()
     check not batch.isClosed()
@@ -93,7 +93,7 @@ suite "WriteBatchRef Tests":
       not batch.isClosed()
 
   test "Test writing to multiple column families in single batch":
-    var batch = db.openWriteBatch()
+    let batch = db.openWriteBatch()
     defer:
       batch.close()
     check not batch.isClosed()
@@ -123,17 +123,16 @@ suite "WriteBatchRef Tests":
       not batch.isClosed()
 
   test "Test writing to multiple column families in multiple batches":
-    var batch1 = db.openWriteBatch()
+    let
+      batch1 = db.openWriteBatch()
+      batch2 = db.openWriteBatch()
     defer:
       batch1.close()
-    check not batch1.isClosed()
-
-    var batch2 = db.openWriteBatch()
-    defer:
       batch2.close()
-    check not batch2.isClosed()
 
     check:
+      not batch1.isClosed()
+      not batch2.isClosed()
       batch1.put(key1, val1).isOk()
       batch1.delete(key2, otherCfHandle).isOk()
       batch1.put(key3, val3, otherCfHandle).isOk()
@@ -155,8 +154,14 @@ suite "WriteBatchRef Tests":
       db.keyExists(key2, otherCfHandle).get() == false
       db.get(key3, otherCfHandle).get() == val3
 
+      # Write batch is unchanged after write
+      batch1.count() == 3
+      batch2.count() == 3
+      not batch1.isClosed()
+      not batch2.isClosed()
+
   test "Test write empty batch":
-    var batch = db.openWriteBatch()
+    let batch = db.openWriteBatch()
     defer:
       batch.close()
     check not batch.isClosed()
@@ -166,9 +171,10 @@ suite "WriteBatchRef Tests":
     check:
       res1.isOk()
       batch.count() == 0
+      not batch.isClosed()
 
   test "Test close":
-    var batch = db.openWriteBatch()
+    let batch = db.openWriteBatch()
 
     check not batch.isClosed()
     batch.close()