redb/
table.rs

1use crate::db::TransactionGuard;
2use crate::sealed::Sealed;
3use crate::tree_store::{
4    AccessGuardMut, Btree, BtreeExtractIf, BtreeHeader, BtreeMut, BtreeRangeIter, PageHint,
5    PageNumber, RawBtree, TransactionalMemory, MAX_PAIR_LENGTH, MAX_VALUE_LENGTH,
6};
7use crate::types::{Key, MutInPlaceValue, Value};
8use crate::{AccessGuard, StorageError, WriteTransaction};
9use crate::{Result, TableHandle};
10use std::borrow::Borrow;
11use std::fmt::{Debug, Formatter};
12use std::marker::PhantomData;
13use std::ops::RangeBounds;
14use std::sync::{Arc, Mutex};
15
16/// Informational storage stats about a table
17#[derive(Debug)]
18pub struct TableStats {
19    pub(crate) tree_height: u32,
20    pub(crate) leaf_pages: u64,
21    pub(crate) branch_pages: u64,
22    pub(crate) stored_leaf_bytes: u64,
23    pub(crate) metadata_bytes: u64,
24    pub(crate) fragmented_bytes: u64,
25}
26
27impl TableStats {
28    /// Maximum traversal distance to reach the deepest (key, value) pair in the table
29    pub fn tree_height(&self) -> u32 {
30        self.tree_height
31    }
32
33    /// Number of leaf pages that store user data
34    pub fn leaf_pages(&self) -> u64 {
35        self.leaf_pages
36    }
37
38    /// Number of branch pages in the btree that store user data
39    pub fn branch_pages(&self) -> u64 {
40        self.branch_pages
41    }
42
43    /// Number of bytes consumed by keys and values that have been inserted.
44    /// Does not include indexing overhead
45    pub fn stored_bytes(&self) -> u64 {
46        self.stored_leaf_bytes
47    }
48
49    /// Number of bytes consumed by keys in internal branch pages, plus other metadata
50    pub fn metadata_bytes(&self) -> u64 {
51        self.metadata_bytes
52    }
53
54    /// Number of bytes consumed by fragmentation, both in data pages and internal metadata tables
55    pub fn fragmented_bytes(&self) -> u64 {
56        self.fragmented_bytes
57    }
58}
59
60/// A table containing key-value mappings
61pub struct Table<'txn, K: Key + 'static, V: Value + 'static> {
62    name: String,
63    transaction: &'txn WriteTransaction,
64    tree: BtreeMut<'txn, K, V>,
65}
66
67impl<K: Key + 'static, V: Value + 'static> TableHandle for Table<'_, K, V> {
68    fn name(&self) -> &str {
69        &self.name
70    }
71}
72
73impl<'txn, K: Key + 'static, V: Value + 'static> Table<'txn, K, V> {
74    pub(crate) fn new(
75        name: &str,
76        table_root: Option<BtreeHeader>,
77        freed_pages: Arc<Mutex<Vec<PageNumber>>>,
78        mem: Arc<TransactionalMemory>,
79        transaction: &'txn WriteTransaction,
80    ) -> Table<'txn, K, V> {
81        Table {
82            name: name.to_string(),
83            transaction,
84            tree: BtreeMut::new(
85                table_root,
86                transaction.transaction_guard(),
87                mem,
88                freed_pages,
89            ),
90        }
91    }
92
93    #[allow(dead_code)]
94    pub(crate) fn print_debug(&self, include_values: bool) -> Result {
95        self.tree.print_debug(include_values)
96    }
97
98    /// Removes and returns the first key-value pair in the table
99    pub fn pop_first(&mut self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
100        // TODO: optimize this
101        let first = self
102            .iter()?
103            .next()
104            .map(|x| x.map(|(key, _)| K::as_bytes(&key.value()).as_ref().to_vec()));
105        if let Some(owned_key) = first {
106            let owned_key = owned_key?;
107            let key = K::from_bytes(&owned_key);
108            let value = self.remove(&key)?.unwrap();
109            drop(key);
110            Ok(Some((AccessGuard::with_owned_value(owned_key), value)))
111        } else {
112            Ok(None)
113        }
114    }
115
116    /// Removes and returns the last key-value pair in the table
117    pub fn pop_last(&mut self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
118        // TODO: optimize this
119        let last = self
120            .iter()?
121            .next_back()
122            .map(|x| x.map(|(key, _)| K::as_bytes(&key.value()).as_ref().to_vec()));
123        if let Some(owned_key) = last {
124            let owned_key = owned_key?;
125            let key = K::from_bytes(&owned_key);
126            let value = self.remove(&key)?.unwrap();
127            drop(key);
128            Ok(Some((AccessGuard::with_owned_value(owned_key), value)))
129        } else {
130            Ok(None)
131        }
132    }
133
134    /// Applies `predicate` to all key-value pairs. All entries for which
135    /// `predicate` evaluates to `true` are returned in an iterator, and those which are read from the iterator are removed
136    ///
137    /// Note: values not read from the iterator will not be removed
138    pub fn extract_if<F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool>(
139        &mut self,
140        predicate: F,
141    ) -> Result<ExtractIf<K, V, F>> {
142        self.extract_from_if::<K::SelfType<'_>, F>(.., predicate)
143    }
144
145    /// Applies `predicate` to all key-value pairs in the specified range. All entries for which
146    /// `predicate` evaluates to `true` are returned in an iterator, and those which are read from the iterator are removed
147    ///
148    /// Note: values not read from the iterator will not be removed
149    pub fn extract_from_if<'a, KR, F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool>(
150        &mut self,
151        range: impl RangeBounds<KR> + 'a,
152        predicate: F,
153    ) -> Result<ExtractIf<K, V, F>>
154    where
155        KR: Borrow<K::SelfType<'a>> + 'a,
156    {
157        self.tree
158            .extract_from_if(&range, predicate)
159            .map(ExtractIf::new)
160    }
161
162    /// Applies `predicate` to all key-value pairs. All entries for which
163    /// `predicate` evaluates to `false` are removed.
164    ///
165    pub fn retain<F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool>(
166        &mut self,
167        predicate: F,
168    ) -> Result {
169        self.tree.retain_in::<K::SelfType<'_>, F>(predicate, ..)
170    }
171
172    /// Applies `predicate` to all key-value pairs in the range `start..end`. All entries for which
173    /// `predicate` evaluates to `false` are removed.
174    ///
175    pub fn retain_in<'a, KR, F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool>(
176        &mut self,
177        range: impl RangeBounds<KR> + 'a,
178        predicate: F,
179    ) -> Result
180    where
181        KR: Borrow<K::SelfType<'a>> + 'a,
182    {
183        self.tree.retain_in(predicate, range)
184    }
185
186    /// Insert mapping of the given key to the given value
187    ///
188    /// If key is already present it is replaced
189    ///
190    /// Returns the old value, if the key was present in the table, otherwise None is returned
191    pub fn insert<'k, 'v>(
192        &mut self,
193        key: impl Borrow<K::SelfType<'k>>,
194        value: impl Borrow<V::SelfType<'v>>,
195    ) -> Result<Option<AccessGuard<V>>> {
196        let value_len = V::as_bytes(value.borrow()).as_ref().len();
197        if value_len > MAX_VALUE_LENGTH {
198            return Err(StorageError::ValueTooLarge(value_len));
199        }
200        let key_len = K::as_bytes(key.borrow()).as_ref().len();
201        if key_len > MAX_VALUE_LENGTH {
202            return Err(StorageError::ValueTooLarge(key_len));
203        }
204        if value_len + key_len > MAX_PAIR_LENGTH {
205            return Err(StorageError::ValueTooLarge(value_len + key_len));
206        }
207        self.tree.insert(key.borrow(), value.borrow())
208    }
209
210    /// Removes the given key
211    ///
212    /// Returns the old value, if the key was present in the table
213    pub fn remove<'a>(
214        &mut self,
215        key: impl Borrow<K::SelfType<'a>>,
216    ) -> Result<Option<AccessGuard<V>>> {
217        self.tree.remove(key.borrow())
218    }
219}
220
221impl<'txn, K: Key + 'static, V: MutInPlaceValue + 'static> Table<'txn, K, V> {
222    /// Reserve space to insert a key-value pair
223    ///
224    /// If key is already present it is replaced
225    ///
226    /// The returned reference will have length equal to `value_length`
227    pub fn insert_reserve<'a>(
228        &mut self,
229        key: impl Borrow<K::SelfType<'a>>,
230        value_length: u32,
231    ) -> Result<AccessGuardMut<V>> {
232        if value_length as usize > MAX_VALUE_LENGTH {
233            return Err(StorageError::ValueTooLarge(value_length as usize));
234        }
235        let key_len = K::as_bytes(key.borrow()).as_ref().len();
236        if key_len > MAX_VALUE_LENGTH {
237            return Err(StorageError::ValueTooLarge(key_len));
238        }
239        if value_length as usize + key_len > MAX_PAIR_LENGTH {
240            return Err(StorageError::ValueTooLarge(value_length as usize + key_len));
241        }
242        self.tree.insert_reserve(key.borrow(), value_length)
243    }
244}
245
246impl<'txn, K: Key + 'static, V: Value + 'static> ReadableTableMetadata for Table<'txn, K, V> {
247    fn stats(&self) -> Result<TableStats> {
248        let tree_stats = self.tree.stats()?;
249
250        Ok(TableStats {
251            tree_height: tree_stats.tree_height,
252            leaf_pages: tree_stats.leaf_pages,
253            branch_pages: tree_stats.branch_pages,
254            stored_leaf_bytes: tree_stats.stored_leaf_bytes,
255            metadata_bytes: tree_stats.metadata_bytes,
256            fragmented_bytes: tree_stats.fragmented_bytes,
257        })
258    }
259
260    fn len(&self) -> Result<u64> {
261        self.tree.len()
262    }
263}
264
265impl<'txn, K: Key + 'static, V: Value + 'static> ReadableTable<K, V> for Table<'txn, K, V> {
266    fn get<'a>(&self, key: impl Borrow<K::SelfType<'a>>) -> Result<Option<AccessGuard<V>>> {
267        self.tree.get(key.borrow())
268    }
269
270    fn range<'a, KR>(&self, range: impl RangeBounds<KR> + 'a) -> Result<Range<K, V>>
271    where
272        KR: Borrow<K::SelfType<'a>> + 'a,
273    {
274        self.tree
275            .range(&range)
276            .map(|x| Range::new(x, self.transaction.transaction_guard()))
277    }
278
279    fn first(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
280        self.tree.first()
281    }
282
283    fn last(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
284        self.tree.last()
285    }
286}
287
288impl<K: Key, V: Value> Sealed for Table<'_, K, V> {}
289
290impl<'txn, K: Key + 'static, V: Value + 'static> Drop for Table<'txn, K, V> {
291    fn drop(&mut self) {
292        self.transaction.close_table(
293            &self.name,
294            &self.tree,
295            self.tree.get_root().map(|x| x.length).unwrap_or_default(),
296        );
297    }
298}
299
300fn debug_helper<K: Key + 'static, V: Value + 'static>(
301    f: &mut Formatter<'_>,
302    name: &str,
303    len: Result<u64>,
304    first: Result<Option<(AccessGuard<K>, AccessGuard<V>)>>,
305    last: Result<Option<(AccessGuard<K>, AccessGuard<V>)>>,
306) -> std::fmt::Result {
307    write!(f, "Table [ name: \"{name}\", ")?;
308    if let Ok(len) = len {
309        if len == 0 {
310            write!(f, "No entries")?;
311        } else if len == 1 {
312            if let Ok(first) = first {
313                let (key, value) = first.as_ref().unwrap();
314                write!(f, "One key-value: {:?} = {:?}", key.value(), value.value())?;
315            } else {
316                write!(f, "I/O Error accessing table!")?;
317            }
318        } else {
319            if let Ok(first) = first {
320                let (key, value) = first.as_ref().unwrap();
321                write!(f, "first: {:?} = {:?}, ", key.value(), value.value())?;
322            } else {
323                write!(f, "I/O Error accessing table!")?;
324            }
325            if len > 2 {
326                write!(f, "...{} more entries..., ", len - 2)?;
327            }
328            if let Ok(last) = last {
329                let (key, value) = last.as_ref().unwrap();
330                write!(f, "last: {:?} = {:?}", key.value(), value.value())?;
331            } else {
332                write!(f, "I/O Error accessing table!")?;
333            }
334        }
335    } else {
336        write!(f, "I/O Error accessing table!")?;
337    }
338    write!(f, " ]")?;
339
340    Ok(())
341}
342
343impl<K: Key + 'static, V: Value + 'static> Debug for Table<'_, K, V> {
344    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
345        debug_helper(f, &self.name, self.len(), self.first(), self.last())
346    }
347}
348
349pub trait ReadableTableMetadata {
350    /// Retrieves information about storage usage for the table
351    fn stats(&self) -> Result<TableStats>;
352
353    /// Returns the number of entries in the table
354    fn len(&self) -> Result<u64>;
355
356    /// Returns `true` if the table is empty
357    fn is_empty(&self) -> Result<bool> {
358        Ok(self.len()? == 0)
359    }
360}
361
362pub trait ReadableTable<K: Key + 'static, V: Value + 'static>: ReadableTableMetadata {
363    /// Returns the value corresponding to the given key
364    fn get<'a>(&self, key: impl Borrow<K::SelfType<'a>>) -> Result<Option<AccessGuard<V>>>;
365
366    /// Returns a double-ended iterator over a range of elements in the table
367    ///
368    /// # Examples
369    ///
370    /// Usage:
371    /// ```rust
372    /// use redb::*;
373    /// # use tempfile::NamedTempFile;
374    /// const TABLE: TableDefinition<&str, u64> = TableDefinition::new("my_data");
375    ///
376    /// # fn main() -> Result<(), Error> {
377    /// # let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
378    /// # let filename = tmpfile.path();
379    /// let db = Database::create(filename)?;
380    /// let write_txn = db.begin_write()?;
381    /// {
382    ///     let mut table = write_txn.open_table(TABLE)?;
383    ///     table.insert("a", &0)?;
384    ///     table.insert("b", &1)?;
385    ///     table.insert("c", &2)?;
386    /// }
387    /// write_txn.commit()?;
388    ///
389    /// let read_txn = db.begin_read()?;
390    /// let table = read_txn.open_table(TABLE)?;
391    /// let mut iter = table.range("a".."c")?;
392    /// let (key, value) = iter.next().unwrap()?;
393    /// assert_eq!("a", key.value());
394    /// assert_eq!(0, value.value());
395    /// # Ok(())
396    /// # }
397    /// ```
398    fn range<'a, KR>(&self, range: impl RangeBounds<KR> + 'a) -> Result<Range<K, V>>
399    where
400        KR: Borrow<K::SelfType<'a>> + 'a;
401
402    /// Returns the first key-value pair in the table, if it exists
403    fn first(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>>;
404
405    /// Returns the last key-value pair in the table, if it exists
406    fn last(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>>;
407
408    /// Returns a double-ended iterator over all elements in the table
409    fn iter(&self) -> Result<Range<K, V>> {
410        self.range::<K::SelfType<'_>>(..)
411    }
412}
413
414/// A read-only untyped table
415pub struct ReadOnlyUntypedTable {
416    tree: RawBtree,
417}
418
419impl Sealed for ReadOnlyUntypedTable {}
420
421impl ReadableTableMetadata for ReadOnlyUntypedTable {
422    /// Retrieves information about storage usage for the table
423    fn stats(&self) -> Result<TableStats> {
424        let tree_stats = self.tree.stats()?;
425
426        Ok(TableStats {
427            tree_height: tree_stats.tree_height,
428            leaf_pages: tree_stats.leaf_pages,
429            branch_pages: tree_stats.branch_pages,
430            stored_leaf_bytes: tree_stats.stored_leaf_bytes,
431            metadata_bytes: tree_stats.metadata_bytes,
432            fragmented_bytes: tree_stats.fragmented_bytes,
433        })
434    }
435
436    fn len(&self) -> Result<u64> {
437        self.tree.len()
438    }
439}
440
441impl ReadOnlyUntypedTable {
442    pub(crate) fn new(
443        root_page: Option<BtreeHeader>,
444        fixed_key_size: Option<usize>,
445        fixed_value_size: Option<usize>,
446        mem: Arc<TransactionalMemory>,
447    ) -> Self {
448        Self {
449            tree: RawBtree::new(root_page, fixed_key_size, fixed_value_size, mem),
450        }
451    }
452}
453
454/// A read-only table
455pub struct ReadOnlyTable<K: Key + 'static, V: Value + 'static> {
456    name: String,
457    tree: Btree<K, V>,
458    transaction_guard: Arc<TransactionGuard>,
459}
460
461impl<K: Key + 'static, V: Value + 'static> TableHandle for ReadOnlyTable<K, V> {
462    fn name(&self) -> &str {
463        &self.name
464    }
465}
466
467impl<K: Key + 'static, V: Value + 'static> ReadOnlyTable<K, V> {
468    pub(crate) fn new(
469        name: String,
470        root_page: Option<BtreeHeader>,
471        hint: PageHint,
472        guard: Arc<TransactionGuard>,
473        mem: Arc<TransactionalMemory>,
474    ) -> Result<ReadOnlyTable<K, V>> {
475        Ok(ReadOnlyTable {
476            name,
477            tree: Btree::new(root_page, hint, guard.clone(), mem)?,
478            transaction_guard: guard,
479        })
480    }
481
482    pub fn get<'a>(
483        &self,
484        key: impl Borrow<K::SelfType<'a>>,
485    ) -> Result<Option<AccessGuard<'static, V>>> {
486        self.tree.get(key.borrow())
487    }
488
489    /// This method is like [`ReadableTable::range()`], but the iterator is reference counted and keeps the transaction
490    /// alive until it is dropped.
491    pub fn range<'a, KR>(&self, range: impl RangeBounds<KR>) -> Result<Range<'static, K, V>>
492    where
493        KR: Borrow<K::SelfType<'a>>,
494    {
495        self.tree
496            .range(&range)
497            .map(|x| Range::new(x, self.transaction_guard.clone()))
498    }
499}
500
501impl<K: Key + 'static, V: Value + 'static> ReadableTableMetadata for ReadOnlyTable<K, V> {
502    fn stats(&self) -> Result<TableStats> {
503        let tree_stats = self.tree.stats()?;
504
505        Ok(TableStats {
506            tree_height: tree_stats.tree_height,
507            leaf_pages: tree_stats.leaf_pages,
508            branch_pages: tree_stats.branch_pages,
509            stored_leaf_bytes: tree_stats.stored_leaf_bytes,
510            metadata_bytes: tree_stats.metadata_bytes,
511            fragmented_bytes: tree_stats.fragmented_bytes,
512        })
513    }
514
515    fn len(&self) -> Result<u64> {
516        self.tree.len()
517    }
518}
519
520impl<K: Key + 'static, V: Value + 'static> ReadableTable<K, V> for ReadOnlyTable<K, V> {
521    fn get<'a>(&self, key: impl Borrow<K::SelfType<'a>>) -> Result<Option<AccessGuard<V>>> {
522        self.tree.get(key.borrow())
523    }
524
525    fn range<'a, KR>(&self, range: impl RangeBounds<KR> + 'a) -> Result<Range<K, V>>
526    where
527        KR: Borrow<K::SelfType<'a>> + 'a,
528    {
529        self.tree
530            .range(&range)
531            .map(|x| Range::new(x, self.transaction_guard.clone()))
532    }
533
534    fn first(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
535        self.tree.first()
536    }
537
538    fn last(&self) -> Result<Option<(AccessGuard<K>, AccessGuard<V>)>> {
539        self.tree.last()
540    }
541}
542
543impl<K: Key, V: Value> Sealed for ReadOnlyTable<K, V> {}
544
545impl<K: Key + 'static, V: Value + 'static> Debug for ReadOnlyTable<K, V> {
546    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
547        debug_helper(f, &self.name, self.len(), self.first(), self.last())
548    }
549}
550
551pub struct ExtractIf<
552    'a,
553    K: Key + 'static,
554    V: Value + 'static,
555    F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool,
556> {
557    inner: BtreeExtractIf<'a, K, V, F>,
558}
559
560impl<
561        'a,
562        K: Key + 'static,
563        V: Value + 'static,
564        F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool,
565    > ExtractIf<'a, K, V, F>
566{
567    fn new(inner: BtreeExtractIf<'a, K, V, F>) -> Self {
568        Self { inner }
569    }
570}
571
572impl<
573        'a,
574        K: Key + 'static,
575        V: Value + 'static,
576        F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool,
577    > Iterator for ExtractIf<'a, K, V, F>
578{
579    type Item = Result<(AccessGuard<'a, K>, AccessGuard<'a, V>)>;
580
581    fn next(&mut self) -> Option<Self::Item> {
582        let entry = self.inner.next()?;
583        Some(entry.map(|entry| {
584            let (page, key_range, value_range) = entry.into_raw();
585            let key = AccessGuard::with_page(page.clone(), key_range);
586            let value = AccessGuard::with_page(page, value_range);
587            (key, value)
588        }))
589    }
590}
591
592impl<
593        'a,
594        K: Key + 'static,
595        V: Value + 'static,
596        F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool,
597    > DoubleEndedIterator for ExtractIf<'a, K, V, F>
598{
599    fn next_back(&mut self) -> Option<Self::Item> {
600        let entry = self.inner.next_back()?;
601        Some(entry.map(|entry| {
602            let (page, key_range, value_range) = entry.into_raw();
603            let key = AccessGuard::with_page(page.clone(), key_range);
604            let value = AccessGuard::with_page(page, value_range);
605            (key, value)
606        }))
607    }
608}
609
610#[derive(Clone)]
611pub struct Range<'a, K: Key + 'static, V: Value + 'static> {
612    inner: BtreeRangeIter<K, V>,
613    _transaction_guard: Arc<TransactionGuard>,
614    // This lifetime is here so that `&` can be held on `Table` preventing concurrent mutation
615    _lifetime: PhantomData<&'a ()>,
616}
617
618impl<'a, K: Key + 'static, V: Value + 'static> Range<'a, K, V> {
619    pub(super) fn new(inner: BtreeRangeIter<K, V>, guard: Arc<TransactionGuard>) -> Self {
620        Self {
621            inner,
622            _transaction_guard: guard,
623            _lifetime: Default::default(),
624        }
625    }
626}
627
628impl<'a, K: Key + 'static, V: Value + 'static> Iterator for Range<'a, K, V> {
629    type Item = Result<(AccessGuard<'a, K>, AccessGuard<'a, V>)>;
630
631    fn next(&mut self) -> Option<Self::Item> {
632        self.inner.next().map(|x| {
633            x.map(|entry| {
634                let (page, key_range, value_range) = entry.into_raw();
635                let key = AccessGuard::with_page(page.clone(), key_range);
636                let value = AccessGuard::with_page(page, value_range);
637                (key, value)
638            })
639        })
640    }
641}
642
643impl<'a, K: Key + 'static, V: Value + 'static> DoubleEndedIterator for Range<'a, K, V> {
644    fn next_back(&mut self) -> Option<Self::Item> {
645        self.inner.next_back().map(|x| {
646            x.map(|entry| {
647                let (page, key_range, value_range) = entry.into_raw();
648                let key = AccessGuard::with_page(page.clone(), key_range);
649                let value = AccessGuard::with_page(page, value_range);
650                (key, value)
651            })
652        })
653    }
654}