Impl ReadableTableMetadata for untyped tables

cberner · Feb 4, 2024 · 744fd43 · 744fd43
1 parent a4e6334
commit 744fd43
Show file tree

Hide file tree

Showing 6 changed files with 175 additions and 35 deletions.
diff --git a/src/multimap_table.rs b/src/multimap_table.rs
@@ -3,9 +3,10 @@ use crate::multimap_table::DynamicCollectionType::{Inline, Subtree};
 use crate::sealed::Sealed;
 use crate::table::{ReadableTableMetadata, TableStats};
 use crate::tree_store::{
-    btree_stats, AllPageNumbersBtreeIter, BranchAccessor, Btree, BtreeMut, BtreeRangeIter,
-    BtreeStats, CachePriority, Checksum, LeafAccessor, LeafMutator, Page, PageHint, PageNumber,
-    RawBtree, RawLeafBuilder, TransactionalMemory, UntypedBtreeMut, BRANCH, LEAF, MAX_VALUE_LENGTH,
+    btree_len, btree_stats, AllPageNumbersBtreeIter, BranchAccessor, Btree, BtreeMut,
+    BtreeRangeIter, BtreeStats, CachePriority, Checksum, LeafAccessor, LeafMutator, Page, PageHint,
+    PageNumber, RawBtree, RawLeafBuilder, TransactionalMemory, UntypedBtreeMut, BRANCH, LEAF,
+    MAX_VALUE_LENGTH,
 };
 use crate::types::{Key, TypeName, Value};
 use crate::{AccessGuard, MultimapTableHandle, Result, StorageError, WriteTransaction};
@@ -18,6 +19,74 @@ use std::mem::size_of;
 use std::ops::{RangeBounds, RangeFull};
 use std::sync::{Arc, Mutex};
 
+pub(crate) fn multimap_btree_len(
+    root: Option<PageNumber>,
+    mem: &TransactionalMemory,
+    fixed_key_size: Option<usize>,
+    fixed_value_size: Option<usize>,
+) -> Result<u64> {
+    if let Some(root) = root {
+        multimap_len_helper(root, mem, fixed_key_size, fixed_value_size)
+    } else {
+        Ok(0)
+    }
+}
+
+fn multimap_len_helper(
+    page_number: PageNumber,
+    mem: &TransactionalMemory,
+    fixed_key_size: Option<usize>,
+    fixed_value_size: Option<usize>,
+) -> Result<u64> {
+    let page = mem.get_page(page_number)?;
+    let node_mem = page.memory();
+    let mut len = 0;
+    match node_mem[0] {
+        LEAF => {
+            let accessor = LeafAccessor::new(
+                page.memory(),
+                fixed_key_size,
+                DynamicCollection::<()>::fixed_width_with(fixed_value_size),
+            );
+            for i in 0..accessor.num_pairs() {
+                let entry = accessor.entry(i).unwrap();
+                let collection: &UntypedDynamicCollection =
+                    UntypedDynamicCollection::new(entry.value());
+                match collection.collection_type() {
+                    Inline => {
+                        let inline_accessor = LeafAccessor::new(
+                            collection.as_inline(),
+                            fixed_value_size,
+                            <() as Value>::fixed_width(),
+                        );
+                        len += inline_accessor.num_pairs() as u64;
+                    }
+                    Subtree => {
+                        // this is a sub-tree, so traverse it
+                        len += btree_len(
+                            Some(collection.as_subtree().0),
+                            mem,
+                            fixed_value_size,
+                            <() as Value>::fixed_width(),
+                        )?;
+                    }
+                }
+            }
+        }
+        BRANCH => {
+            let accessor = BranchAccessor::new(&page, fixed_key_size);
+            for i in 0..accessor.count_children() {
+                if let Some(child) = accessor.child_page(i) {
+                    len += multimap_len_helper(child, mem, fixed_key_size, fixed_value_size)?;
+                }
+            }
+        }
+        _ => unreachable!(),
+    }
+
+    Ok(len)
+}
+
 pub(crate) fn multimap_btree_stats(
     root: Option<PageNumber>,
     mem: &TransactionalMemory,
@@ -1202,6 +1271,38 @@ pub struct ReadOnlyUntypedMultimapTable {
     mem: Arc<TransactionalMemory>,
 }
 
+impl Sealed for ReadOnlyUntypedMultimapTable {}
+
+impl ReadableTableMetadata for ReadOnlyUntypedMultimapTable {
+    /// Retrieves information about storage usage for the table
+    fn stats(&self) -> Result<TableStats> {
+        let tree_stats = multimap_btree_stats(
+            self.tree.get_root().map(|(p, _)| p),
+            &self.mem,
+            self.fixed_key_size,
+            self.fixed_value_size,
+        )?;
+
+        Ok(TableStats {
+            tree_height: tree_stats.tree_height,
+            leaf_pages: tree_stats.leaf_pages,
+            branch_pages: tree_stats.branch_pages,
+            stored_leaf_bytes: tree_stats.stored_leaf_bytes,
+            metadata_bytes: tree_stats.metadata_bytes,
+            fragmented_bytes: tree_stats.fragmented_bytes,
+        })
+    }
+
+    fn len(&self) -> Result<u64> {
+        multimap_btree_len(
+            self.tree.get_root().map(|(p, _)| p),
+            &self.mem,
+            self.fixed_key_size,
+            self.fixed_value_size,
+        )
+    }
+}
+
 impl ReadOnlyUntypedMultimapTable {
     pub(crate) fn new(
         root_page: Option<(PageNumber, Checksum)>,
@@ -1221,25 +1322,6 @@ impl ReadOnlyUntypedMultimapTable {
             mem,
         }
     }
-
-    /// Retrieves information about storage usage for the table
-    pub fn stats(&self) -> Result<TableStats> {
-        let tree_stats = multimap_btree_stats(
-            self.tree.get_root().map(|(p, _)| p),
-            &self.mem,
-            self.fixed_key_size,
-            self.fixed_value_size,
-        )?;
-
-        Ok(TableStats {
-            tree_height: tree_stats.tree_height,
-            leaf_pages: tree_stats.leaf_pages,
-            branch_pages: tree_stats.branch_pages,
-            stored_leaf_bytes: tree_stats.stored_leaf_bytes,
-            metadata_bytes: tree_stats.metadata_bytes,
-            fragmented_bytes: tree_stats.fragmented_bytes,
-        })
-    }
 }
 
 /// A read-only multimap table

diff --git a/src/table.rs b/src/table.rs
@@ -387,20 +387,11 @@ pub struct ReadOnlyUntypedTable {
     tree: RawBtree,
 }
 
-impl ReadOnlyUntypedTable {
-    pub(crate) fn new(
-        root_page: Option<(PageNumber, Checksum)>,
-        fixed_key_size: Option<usize>,
-        fixed_value_size: Option<usize>,
-        mem: Arc<TransactionalMemory>,
-    ) -> Self {
-        Self {
-            tree: RawBtree::new(root_page, fixed_key_size, fixed_value_size, mem),
-        }
-    }
+impl Sealed for ReadOnlyUntypedTable {}
 
+impl ReadableTableMetadata for ReadOnlyUntypedTable {
     /// Retrieves information about storage usage for the table
-    pub fn stats(&self) -> Result<TableStats> {
+    fn stats(&self) -> Result<TableStats> {
         let tree_stats = self.tree.stats()?;
 
         Ok(TableStats {
@@ -412,6 +403,23 @@ impl ReadOnlyUntypedTable {
             fragmented_bytes: tree_stats.fragmented_bytes,
         })
     }
+
+    fn len(&self) -> Result<u64> {
+        self.tree.len()
+    }
+}
+
+impl ReadOnlyUntypedTable {
+    pub(crate) fn new(
+        root_page: Option<(PageNumber, Checksum)>,
+        fixed_key_size: Option<usize>,
+        fixed_value_size: Option<usize>,
+        mem: Arc<TransactionalMemory>,
+    ) -> Self {
+        Self {
+            tree: RawBtree::new(root_page, fixed_key_size, fixed_value_size, mem),
+        }
+    }
 }
 
 /// A read-only table

diff --git a/src/tree_store/btree.rs b/src/tree_store/btree.rs
@@ -503,6 +503,15 @@ impl RawBtree {
         )
     }
 
+    pub(crate) fn len(&self) -> Result<u64> {
+        btree_len(
+            self.root.map(|(p, _)| p),
+            &self.mem,
+            self.fixed_key_size,
+            self.fixed_value_size,
+        )
+    }
+
     pub(crate) fn verify_checksum(&self) -> Result<bool> {
         if let Some((root, checksum)) = self.root {
             self.verify_checksum_helper(root, checksum)
@@ -688,6 +697,41 @@ impl<K: Key, V: Value> Btree<K, V> {
     }
 }
 
+pub(crate) fn btree_len(
+    root: Option<PageNumber>,
+    mem: &TransactionalMemory,
+    fixed_key_size: Option<usize>,
+    fixed_value_size: Option<usize>,
+) -> Result<u64> {
+    if let Some(root) = root {
+        let page = mem.get_page(root)?;
+        let node_mem = page.memory();
+        match node_mem[0] {
+            LEAF => {
+                let accessor = LeafAccessor::new(page.memory(), fixed_key_size, fixed_value_size);
+                Ok(accessor.num_pairs() as u64)
+            }
+            BRANCH => {
+                let accessor = BranchAccessor::new(&page, fixed_key_size);
+                let mut len = 0;
+                for i in 0..accessor.count_children() {
+                    len += btree_len(
+                        accessor.child_page(i),
+                        mem,
+                        fixed_key_size,
+                        fixed_value_size,
+                    )?;
+                }
+
+                Ok(len)
+            }
+            _ => unreachable!(),
+        }
+    } else {
+        Ok(0)
+    }
+}
+
 pub(crate) fn btree_stats(
     root: Option<PageNumber>,
     mem: &TransactionalMemory,

diff --git a/src/tree_store/mod.rs b/src/tree_store/mod.rs
@@ -5,7 +5,9 @@ mod btree_mutator;
 mod page_store;
 mod table_tree;
 
-pub(crate) use btree::{btree_stats, Btree, BtreeMut, BtreeStats, RawBtree, UntypedBtreeMut};
+pub(crate) use btree::{
+    btree_len, btree_stats, Btree, BtreeMut, BtreeStats, RawBtree, UntypedBtreeMut,
+};
 pub use btree_base::{AccessGuard, AccessGuardMut};
 pub(crate) use btree_base::{BranchAccessor, Checksum};
 pub(crate) use btree_base::{LeafAccessor, LeafMutator, RawLeafBuilder, BRANCH, LEAF};

diff --git a/tests/basic_tests.rs b/tests/basic_tests.rs
@@ -35,6 +35,8 @@ fn len() {
     let read_txn = db.begin_read().unwrap();
     let table = read_txn.open_table(STR_TABLE).unwrap();
     assert_eq!(table.len().unwrap(), 3);
+    let untyped_table = read_txn.open_untyped_table(STR_TABLE).unwrap();
+    assert_eq!(untyped_table.len().unwrap(), 3);
 }
 
 #[test]

diff --git a/tests/multimap_tests.rs b/tests/multimap_tests.rs
@@ -47,6 +47,8 @@ fn len() {
     let read_txn = db.begin_read().unwrap();
     let table = read_txn.open_multimap_table(STR_TABLE).unwrap();
     assert_eq!(table.len().unwrap(), 3);
+    let untyped_table = read_txn.open_untyped_multimap_table(STR_TABLE).unwrap();
+    assert_eq!(untyped_table.len().unwrap(), 3);
 }
 
 #[test]