Skip to content

Commit 3406383

Browse files
KerollmopsKerollmops
authored
Merge pull request #252 from nolanderc/get-or-put
Add `Database::get_or_put`: insert if value does not exist, otherwise return previous value
2 parents 7e2e105 + b47b93b commit 3406383

File tree

2 files changed

+309
-6
lines changed

2 files changed

+309
-6
lines changed

‎heed/src/database.rs

Lines changed: 306 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1787,7 +1787,8 @@ impl<KC, DC, C> Database<KC, DC, C> {
17871787
RwCursor::new(txn, self.dbi).map(|cursor| RwRevPrefix::new(cursor, prefix_bytes))
17881788
}
17891789

1790-
/// Insert a key-value pair in this database. The entry is written with no specific flag.
1790+
/// Insert a key-value pair in this database, replacing any previous value. The entry is
1791+
/// written with no specific flag.
17911792
///
17921793
/// ```
17931794
/// # use std::fs;
@@ -1842,7 +1843,8 @@ impl<KC, DC, C> Database<KC, DC, C> {
18421843
Ok(())
18431844
}
18441845

1845-
/// Insert a key-value pair where the value can directly be written to disk.
1846+
/// Insert a key-value pair where the value can directly be written to disk, replacing any
1847+
/// previous value.
18461848
///
18471849
/// ```
18481850
/// # use std::fs;
@@ -1908,7 +1910,8 @@ impl<KC, DC, C> Database<KC, DC, C> {
19081910
}
19091911
}
19101912

1911-
/// Insert a key-value pair in this database. The entry is written with the specified flags.
1913+
/// Insert a key-value pair in this database, replacing any previous value. The entry is
1914+
/// written with the specified flags.
19121915
///
19131916
/// ```
19141917
/// # use std::fs;
@@ -1993,6 +1996,281 @@ impl<KC, DC, C> Database<KC, DC, C> {
19931996
Ok(())
19941997
}
19951998

1999+
/// Attempt to insert a key-value pair in this database, or if a value already exists for the
2000+
/// key, returns the previous value.
2001+
///
2002+
/// The entry is always written with the [`NO_OVERWRITE`](PutFlags::NO_OVERWRITE) flag.
2003+
///
2004+
/// ```
2005+
/// # use heed::EnvOpenOptions;
2006+
/// use heed::Database;
2007+
/// use heed::types::*;
2008+
/// use heed::byteorder::BigEndian;
2009+
///
2010+
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
2011+
/// # let dir = tempfile::tempdir()?;
2012+
/// # let env = unsafe { EnvOpenOptions::new()
2013+
/// # .map_size(10 * 1024 * 1024) // 10MB
2014+
/// # .max_dbs(3000)
2015+
/// # .open(dir.path())?
2016+
/// # };
2017+
/// type BEI32 = I32<BigEndian>;
2018+
///
2019+
/// let mut wtxn = env.write_txn()?;
2020+
/// let db: Database<BEI32, Str> = env.create_database(&mut wtxn, Some("iter-i32"))?;
2021+
///
2022+
/// # db.clear(&mut wtxn)?;
2023+
/// assert_eq!(db.get_or_put(&mut wtxn, &42, "i-am-forty-two")?, None);
2024+
/// assert_eq!(db.get_or_put(&mut wtxn, &42, "the meaning of life")?, Some("i-am-forty-two"));
2025+
///
2026+
/// let ret = db.get(&mut wtxn, &42)?;
2027+
/// assert_eq!(ret, Some("i-am-forty-two"));
2028+
///
2029+
/// wtxn.commit()?;
2030+
/// # Ok(()) }
2031+
/// ```
2032+
pub fn get_or_put<'a, 'txn>(
2033+
&'txn self,
2034+
txn: &mut RwTxn,
2035+
key: &'a KC::EItem,
2036+
data: &'a DC::EItem,
2037+
) -> Result<Option<DC::DItem>>
2038+
where
2039+
KC: BytesEncode<'a>,
2040+
DC: BytesEncode<'a> + BytesDecode<'a>,
2041+
{
2042+
self.get_or_put_with_flags(txn, PutFlags::empty(), key, data)
2043+
}
2044+
2045+
/// Attempt to insert a key-value pair in this database, or if a value already exists for the
2046+
/// key, returns the previous value.
2047+
///
2048+
/// The entry is written with the specified flags, in addition to
2049+
/// [`NO_OVERWRITE`](PutFlags::NO_OVERWRITE) which is always used.
2050+
///
2051+
/// ```
2052+
/// # use heed::EnvOpenOptions;
2053+
/// use heed::{Database, PutFlags};
2054+
/// use heed::types::*;
2055+
/// use heed::byteorder::BigEndian;
2056+
///
2057+
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
2058+
/// # let dir = tempfile::tempdir()?;
2059+
/// # let env = unsafe { EnvOpenOptions::new()
2060+
/// # .map_size(10 * 1024 * 1024) // 10MB
2061+
/// # .max_dbs(3000)
2062+
/// # .open(dir.path())?
2063+
/// # };
2064+
/// type BEI32 = I32<BigEndian>;
2065+
///
2066+
/// let mut wtxn = env.write_txn()?;
2067+
/// let db: Database<BEI32, Str> = env.create_database(&mut wtxn, Some("iter-i32"))?;
2068+
///
2069+
/// # db.clear(&mut wtxn)?;
2070+
/// assert_eq!(db.get_or_put_with_flags(&mut wtxn, PutFlags::empty(), &42, "i-am-forty-two")?, None);
2071+
/// assert_eq!(db.get_or_put_with_flags(&mut wtxn, PutFlags::empty(), &42, "the meaning of life")?, Some("i-am-forty-two"));
2072+
///
2073+
/// let ret = db.get(&mut wtxn, &42)?;
2074+
/// assert_eq!(ret, Some("i-am-forty-two"));
2075+
///
2076+
/// wtxn.commit()?;
2077+
/// # Ok(()) }
2078+
/// ```
2079+
pub fn get_or_put_with_flags<'a, 'txn>(
2080+
&'txn self,
2081+
txn: &mut RwTxn,
2082+
flags: PutFlags,
2083+
key: &'a KC::EItem,
2084+
data: &'a DC::EItem,
2085+
) -> Result<Option<DC::DItem>>
2086+
where
2087+
KC: BytesEncode<'a>,
2088+
DC: BytesEncode<'a> + BytesDecode<'a>,
2089+
{
2090+
assert_eq_env_db_txn!(self, txn);
2091+
2092+
let key_bytes: Cow<[u8]> = KC::bytes_encode(key).map_err(Error::Encoding)?;
2093+
let data_bytes: Cow<[u8]> = DC::bytes_encode(data).map_err(Error::Encoding)?;
2094+
2095+
let mut key_val = unsafe { crate::into_val(&key_bytes) };
2096+
let mut data_val = unsafe { crate::into_val(&data_bytes) };
2097+
let flags = (flags | PutFlags::NO_OVERWRITE).bits();
2098+
2099+
let result = unsafe {
2100+
mdb_result(ffi::mdb_put(txn.txn.txn, self.dbi, &mut key_val, &mut data_val, flags))
2101+
};
2102+
2103+
match result {
2104+
// the value was successfully inserted
2105+
Ok(()) => Ok(None),
2106+
// the key already exists: the previous value is stored in the data parameter
2107+
Err(MdbError::KeyExist) => {
2108+
let bytes = unsafe { crate::from_val(data_val) };
2109+
let data = DC::bytes_decode(bytes).map_err(Error::Decoding)?;
2110+
Ok(Some(data))
2111+
}
2112+
Err(error) => Err(error.into()),
2113+
}
2114+
}
2115+
2116+
/// Attempt to insert a key-value pair in this database, where the value can be directly
2117+
/// written to disk, or if a value already exists for the key, returns the previous value.
2118+
///
2119+
/// The entry is always written with the [`NO_OVERWRITE`](PutFlags::NO_OVERWRITE) and
2120+
/// [`MDB_RESERVE`](lmdb_master_sys::MDB_RESERVE) flags.
2121+
///
2122+
/// ```
2123+
/// # use heed::EnvOpenOptions;
2124+
/// use std::io::Write;
2125+
/// use heed::{Database, PutFlags};
2126+
/// use heed::types::*;
2127+
/// use heed::byteorder::BigEndian;
2128+
///
2129+
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
2130+
/// # let dir = tempfile::tempdir()?;
2131+
/// # let env = unsafe { EnvOpenOptions::new()
2132+
/// # .map_size(10 * 1024 * 1024) // 10MB
2133+
/// # .max_dbs(3000)
2134+
/// # .open(dir.path())?
2135+
/// # };
2136+
/// type BEI32 = I32<BigEndian>;
2137+
///
2138+
/// let mut wtxn = env.write_txn()?;
2139+
/// let db = env.create_database::<BEI32, Str>(&mut wtxn, Some("number-string"))?;
2140+
///
2141+
/// # db.clear(&mut wtxn)?;
2142+
/// let long = "I am a long long long value";
2143+
/// assert_eq!(
2144+
/// db.get_or_put_reserved(&mut wtxn, &42, long.len(), |reserved| {
2145+
/// reserved.write_all(long.as_bytes())
2146+
/// })?,
2147+
/// None
2148+
/// );
2149+
///
2150+
/// let longer = "I am an even longer long long long value";
2151+
/// assert_eq!(
2152+
/// db.get_or_put_reserved(&mut wtxn, &42, longer.len(), |reserved| {
2153+
/// unreachable!()
2154+
/// })?,
2155+
/// Some(long)
2156+
/// );
2157+
///
2158+
/// let ret = db.get(&mut wtxn, &42)?;
2159+
/// assert_eq!(ret, Some(long));
2160+
///
2161+
/// wtxn.commit()?;
2162+
/// # Ok(()) }
2163+
/// ```
2164+
pub fn get_or_put_reserved<'a, 'txn, F>(
2165+
&'txn self,
2166+
txn: &mut RwTxn,
2167+
key: &'a KC::EItem,
2168+
data_size: usize,
2169+
write_func: F,
2170+
) -> Result<Option<DC::DItem>>
2171+
where
2172+
KC: BytesEncode<'a>,
2173+
F: FnOnce(&mut ReservedSpace) -> io::Result<()>,
2174+
DC: BytesDecode<'a>,
2175+
{
2176+
self.get_or_put_reserved_with_flags(txn, PutFlags::empty(), key, data_size, write_func)
2177+
}
2178+
2179+
/// Attempt to insert a key-value pair in this database, where the value can be directly
2180+
/// written to disk, or if a value already exists for the key, returns the previous value.
2181+
///
2182+
/// The entry is written with the specified flags, in addition to
2183+
/// [`NO_OVERWRITE`](PutFlags::NO_OVERWRITE) and [`MDB_RESERVE`](lmdb_master_sys::MDB_RESERVE)
2184+
/// which are always used.
2185+
///
2186+
/// ```
2187+
/// # use heed::EnvOpenOptions;
2188+
/// use std::io::Write;
2189+
/// use heed::{Database, PutFlags};
2190+
/// use heed::types::*;
2191+
/// use heed::byteorder::BigEndian;
2192+
///
2193+
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
2194+
/// # let dir = tempfile::tempdir()?;
2195+
/// # let env = unsafe { EnvOpenOptions::new()
2196+
/// # .map_size(10 * 1024 * 1024) // 10MB
2197+
/// # .max_dbs(3000)
2198+
/// # .open(dir.path())?
2199+
/// # };
2200+
/// type BEI32 = I32<BigEndian>;
2201+
///
2202+
/// let mut wtxn = env.write_txn()?;
2203+
/// let db = env.create_database::<BEI32, Str>(&mut wtxn, Some("number-string"))?;
2204+
///
2205+
/// # db.clear(&mut wtxn)?;
2206+
/// let long = "I am a long long long value";
2207+
/// assert_eq!(
2208+
/// db.get_or_put_reserved_with_flags(&mut wtxn, PutFlags::empty(), &42, long.len(), |reserved| {
2209+
/// reserved.write_all(long.as_bytes())
2210+
/// })?,
2211+
/// None
2212+
/// );
2213+
///
2214+
/// let longer = "I am an even longer long long long value";
2215+
/// assert_eq!(
2216+
/// db.get_or_put_reserved_with_flags(&mut wtxn, PutFlags::empty(), &42, longer.len(), |reserved| {
2217+
/// unreachable!()
2218+
/// })?,
2219+
/// Some(long)
2220+
/// );
2221+
///
2222+
/// let ret = db.get(&mut wtxn, &42)?;
2223+
/// assert_eq!(ret, Some(long));
2224+
///
2225+
/// wtxn.commit()?;
2226+
/// # Ok(()) }
2227+
/// ```
2228+
pub fn get_or_put_reserved_with_flags<'a, 'txn, F>(
2229+
&'txn self,
2230+
txn: &mut RwTxn,
2231+
flags: PutFlags,
2232+
key: &'a KC::EItem,
2233+
data_size: usize,
2234+
write_func: F,
2235+
) -> Result<Option<DC::DItem>>
2236+
where
2237+
KC: BytesEncode<'a>,
2238+
F: FnOnce(&mut ReservedSpace) -> io::Result<()>,
2239+
DC: BytesDecode<'a>,
2240+
{
2241+
assert_eq_env_db_txn!(self, txn);
2242+
2243+
let key_bytes: Cow<[u8]> = KC::bytes_encode(key).map_err(Error::Encoding)?;
2244+
2245+
let mut key_val = unsafe { crate::into_val(&key_bytes) };
2246+
let mut reserved = ffi::reserve_size_val(data_size);
2247+
let flags = (flags | PutFlags::NO_OVERWRITE).bits() | lmdb_master_sys::MDB_RESERVE;
2248+
2249+
let result = unsafe {
2250+
mdb_result(ffi::mdb_put(txn.txn.txn, self.dbi, &mut key_val, &mut reserved, flags))
2251+
};
2252+
2253+
match result {
2254+
// value was inserted: fill the reserved space
2255+
Ok(()) => {
2256+
let mut reserved = unsafe { ReservedSpace::from_val(reserved) };
2257+
write_func(&mut reserved)?;
2258+
if reserved.remaining() == 0 {
2259+
Ok(None)
2260+
} else {
2261+
Err(io::Error::from(io::ErrorKind::UnexpectedEof).into())
2262+
}
2263+
}
2264+
// the key already exists: the previous value is stored in the data parameter
2265+
Err(MdbError::KeyExist) => {
2266+
let bytes = unsafe { crate::from_val(reserved) };
2267+
let data = DC::bytes_decode(bytes).map_err(Error::Decoding)?;
2268+
Ok(Some(data))
2269+
}
2270+
Err(error) => Err(error.into()),
2271+
}
2272+
}
2273+
19962274
/// Deletes an entry or every duplicate data items of a key
19972275
/// if the database supports duplicate data items.
19982276
///
@@ -2355,3 +2633,28 @@ pub struct DatabaseStat {
23552633
/// Number of data items.
23562634
pub entries: usize,
23572635
}
2636+
2637+
#[cfg(test)]
2638+
mod tests {
2639+
use heed_types::*;
2640+
2641+
use super::*;
2642+
2643+
#[test]
2644+
fn put_overwrite() -> Result<()> {
2645+
let dir = tempfile::tempdir()?;
2646+
let env = unsafe { EnvOpenOptions::new().open(dir.path())? };
2647+
let mut txn = env.write_txn()?;
2648+
let db = env.create_database::<Bytes, Bytes>(&mut txn, None)?;
2649+
2650+
assert_eq!(db.get(&txn, b"hello").unwrap(), None);
2651+
2652+
db.put(&mut txn, b"hello", b"hi").unwrap();
2653+
assert_eq!(db.get(&txn, b"hello").unwrap(), Some(&b"hi"[..]));
2654+
2655+
db.put(&mut txn, b"hello", b"bye").unwrap();
2656+
assert_eq!(db.get(&txn, b"hello").unwrap(), Some(&b"bye"[..]));
2657+
2658+
Ok(())
2659+
}
2660+
}

‎heed/src/iterator/prefix.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ use crate::*;
1313
/// defined by the `C` comparator. If no successor exists (i.e. `bytes` is the maximal
1414
/// value), it remains unchanged and the function returns `false`. Otherwise, updates
1515
/// `bytes` and returns `true`.
16-
fn advance_prefix<C: LexicographicComparator>(bytes: &mut Vec<u8>) -> bool {
16+
fn advance_prefix<C: LexicographicComparator>(bytes: &mut [u8]) -> bool {
1717
let mut idx = bytes.len();
1818
while idx > 0 && bytes[idx - 1] == C::max_elem() {
1919
idx -= 1;
@@ -32,7 +32,7 @@ fn advance_prefix<C: LexicographicComparator>(bytes: &mut Vec<u8>) -> bool {
3232
/// defined by the `C` comparator. If no predecessor exists (i.e. `bytes` is the minimum
3333
/// value), it remains unchanged and the function returns `false`. Otherwise, updates
3434
/// `bytes` and returns `true`.
35-
fn retreat_prefix<C: LexicographicComparator>(bytes: &mut Vec<u8>) -> bool {
35+
fn retreat_prefix<C: LexicographicComparator>(bytes: &mut [u8]) -> bool {
3636
let mut idx = bytes.len();
3737
while idx > 0 && bytes[idx - 1] == C::min_elem() {
3838
idx -= 1;
@@ -49,7 +49,7 @@ fn retreat_prefix<C: LexicographicComparator>(bytes: &mut Vec<u8>) -> bool {
4949

5050
fn move_on_prefix_end<'txn, C: LexicographicComparator>(
5151
cursor: &mut RoCursor<'txn>,
52-
prefix: &mut Vec<u8>,
52+
prefix: &mut [u8],
5353
) -> Result<Option<(&'txn [u8], &'txn [u8])>> {
5454
if advance_prefix::<C>(prefix) {
5555
let result = cursor

0 commit comments

Comments
 (0)