diff --git a/heed/src/cursor.rs b/heed/src/cursor.rs index b5d81e6d..c08930f4 100644 --- a/heed/src/cursor.rs +++ b/heed/src/cursor.rs @@ -145,6 +145,27 @@ impl<'txn> RoCursor<'txn> { } } + pub fn move_on_key_value(&mut self, key: &[u8], data: &[u8]) -> Result> { + let mut key_val = unsafe { crate::into_val(key) }; + let mut data_val = unsafe { crate::into_val(data) }; + + // Move the cursor to the specified (key, value) + let result = unsafe { + mdb_result(ffi::mdb_cursor_get( + self.cursor, + &mut key_val, + &mut data_val, + ffi::cursor_op::MDB_GET_BOTH, + )) + }; + + match result { + Ok(()) => Ok(Some(unsafe { crate::from_val(data_val) })), + Err(e) if e.not_found() => Ok(None), + Err(e) => Err(e.into()), + } + } + pub fn move_on_key_greater_than_or_equal_to( &mut self, key: &[u8], diff --git a/heed/src/databases/database.rs b/heed/src/databases/database.rs index 46ccc902..643aeba2 100644 --- a/heed/src/databases/database.rs +++ b/heed/src/databases/database.rs @@ -386,6 +386,78 @@ impl Database { } } + /// Retrieves the value associated with a (key, value) according to the comparison function of + /// those types. If the comparison function on the value only takes into account part of the + /// value this can be used to look up the rest of the value from that part. + /// + /// If the key does not exist, then `None` is returned. + /// + /// ``` + /// # use std::fs; + /// # use std::path::Path; + /// # use heed::EnvOpenOptions; + /// use heed::Database; + /// use heed::types::*; + /// use heed::byteorder::BigEndian; + /// + /// # fn main() -> Result<(), Box> { + /// # use heed::{DatabaseFlags, IntegerComparator}; + /// # let dir = tempfile::tempdir()?; + /// # let env = unsafe { EnvOpenOptions::new() + /// # .map_size(10 * 1024 * 1024) // 10MB + /// # .max_dbs(3000) + /// # .open(dir.path())? + /// # }; + /// type BEI32 = U32; + /// + /// let mut wtxn = env.write_txn()?; + /// let db: Database = env.database_options().types().flags(DatabaseFlags::DUP_SORT | DatabaseFlags::DUP_FIXED).dup_sort_comparator::().name("test").create(&mut wtxn)?; + /// + /// # db.clear(&mut wtxn)?; + /// + /// let mut val1 = 42_usize.to_ne_bytes().to_vec(); + /// val1.extend_from_slice(&[1,2,3,4,5]); + /// db.put(&mut wtxn, &1, &val1)?; + /// + /// let mut val2 = 21_usize.to_ne_bytes().to_vec(); + /// val2.extend_from_slice(&[0,1,0,1,0]); + /// db.put(&mut wtxn, &1, &val2)?; + /// + /// let ret = db.get_duplicate(&wtxn, &1, 42_usize.to_ne_bytes().as_slice())?; + /// assert_eq!(ret, Some(val1.as_slice())); + /// + /// let ret = db.get_duplicate(&wtxn, &1, 21_usize.to_ne_bytes().as_slice())?; + /// assert_eq!(ret, Some(val2.as_slice())); + /// + /// let ret = db.get_duplicate(&wtxn, &1, 22_usize.to_ne_bytes().as_slice())?; + /// assert_eq!(ret, None); + /// + /// wtxn.commit()?; + /// # Ok(()) } + /// ``` + pub fn get_duplicate<'a, 'txn>( + &self, + txn: &'txn RoTxn, + key: &'a KC::EItem, + data: &'a DC::EItem, + ) -> Result> + where + KC: BytesEncode<'a>, + DC: BytesEncode<'a> + BytesDecode<'txn>, + { + assert_eq_env_db_txn!(self, txn); + + let key_bytes: Cow<[u8]> = KC::bytes_encode(key).map_err(Error::Encoding)?; + let data_bytes: Cow<[u8]> = DC::bytes_encode(data).map_err(Error::Decoding)?; + + let mut cursor = RoCursor::new(txn, self.dbi)?; + + cursor + .move_on_key_value(&key_bytes, &data_bytes)? + .map(|data| DC::bytes_decode(data).map_err(Error::Decoding)) + .transpose() + } + /// Returns an iterator over all of the values of a single key. /// /// You can make this iterator `Send`able between threads by opening diff --git a/heed/src/mdb/lmdb_ffi.rs b/heed/src/mdb/lmdb_ffi.rs index 3a263915..3c916dcf 100644 --- a/heed/src/mdb/lmdb_ffi.rs +++ b/heed/src/mdb/lmdb_ffi.rs @@ -33,6 +33,7 @@ pub mod cursor_op { pub const MDB_NEXT_NODUP: MDB_cursor_op = ffi::MDB_NEXT_NODUP; pub const MDB_NEXT_DUP: MDB_cursor_op = ffi::MDB_NEXT_DUP; pub const MDB_GET_CURRENT: MDB_cursor_op = ffi::MDB_GET_CURRENT; + pub const MDB_GET_BOTH: MDB_cursor_op = ffi::MDB_GET_BOTH; } pub fn reserve_size_val(size: usize) -> ffi::MDB_val {