From 744fd43173b78d84fa87db452e33593d0a9b8b34 Mon Sep 17 00:00:00 2001 From: Christopher Berner Date: Sun, 4 Feb 2024 13:56:27 -0800 Subject: [PATCH] Impl ReadableTableMetadata for untyped tables --- src/multimap_table.rs | 126 +++++++++++++++++++++++++++++++++------- src/table.rs | 32 ++++++---- src/tree_store/btree.rs | 44 ++++++++++++++ src/tree_store/mod.rs | 4 +- tests/basic_tests.rs | 2 + tests/multimap_tests.rs | 2 + 6 files changed, 175 insertions(+), 35 deletions(-) diff --git a/src/multimap_table.rs b/src/multimap_table.rs index bcd31e53..9bcbaccf 100644 --- 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, + mem: &TransactionalMemory, + fixed_key_size: Option, + fixed_value_size: Option, +) -> Result { + 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, + fixed_value_size: Option, +) -> Result { + 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, mem: &TransactionalMemory, @@ -1202,6 +1271,38 @@ pub struct ReadOnlyUntypedMultimapTable { mem: Arc, } +impl Sealed for ReadOnlyUntypedMultimapTable {} + +impl ReadableTableMetadata for ReadOnlyUntypedMultimapTable { + /// Retrieves information about storage usage for the table + fn stats(&self) -> Result { + 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 { + 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 { - 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 index 59eca872..385a8e16 100644 --- 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, - fixed_value_size: Option, - mem: Arc, - ) -> 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 { + fn stats(&self) -> Result { let tree_stats = self.tree.stats()?; Ok(TableStats { @@ -412,6 +403,23 @@ impl ReadOnlyUntypedTable { fragmented_bytes: tree_stats.fragmented_bytes, }) } + + fn len(&self) -> Result { + self.tree.len() + } +} + +impl ReadOnlyUntypedTable { + pub(crate) fn new( + root_page: Option<(PageNumber, Checksum)>, + fixed_key_size: Option, + fixed_value_size: Option, + mem: Arc, + ) -> 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 index 17ae9501..df30430b 100644 --- a/src/tree_store/btree.rs +++ b/src/tree_store/btree.rs @@ -503,6 +503,15 @@ impl RawBtree { ) } + pub(crate) fn len(&self) -> Result { + btree_len( + self.root.map(|(p, _)| p), + &self.mem, + self.fixed_key_size, + self.fixed_value_size, + ) + } + pub(crate) fn verify_checksum(&self) -> Result { if let Some((root, checksum)) = self.root { self.verify_checksum_helper(root, checksum) @@ -688,6 +697,41 @@ impl Btree { } } +pub(crate) fn btree_len( + root: Option, + mem: &TransactionalMemory, + fixed_key_size: Option, + fixed_value_size: Option, +) -> Result { + 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, mem: &TransactionalMemory, diff --git a/src/tree_store/mod.rs b/src/tree_store/mod.rs index a3687b57..8c740fc9 100644 --- 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 index 299162df..b9b67f45 100644 --- 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 index f29592da..567e94d5 100644 --- 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]