Skip to content

Commit

Permalink
Add values to Augmentation functions (#35)
Browse files Browse the repository at this point in the history
* init workable version

* loosen InnerNode Augmentation's trait req

* add V bound for Augmentation

* ready for augmentation from_leaf method change

* add values slice to augmentation from_leaf method

* add a new example to use value as count

* remove this example as new impl is more cleaner

* rename example

* update document to state that get_mut + augmentation could cause trouble
  • Loading branch information
shuoli84 authored Nov 4, 2024
1 parent 9f2e5b9 commit 94beb4a
Show file tree
Hide file tree
Showing 12 changed files with 178 additions and 66 deletions.
8 changes: 4 additions & 4 deletions examples/augment_even_count.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,8 @@ fn value_is_even(v: i64) -> bool {
v % 2 == 0
}

impl Augmentation<i64> for EvenCount {
fn from_leaf(keys: &[i64]) -> Self {
impl<V> Augmentation<i64, V> for EvenCount {
fn from_leaf(keys: &[i64], _: &[V]) -> Self {
// For leafs, we count all keys that are even
Self(keys.iter().filter(|i| value_is_even(**i)).count())
}
Expand All @@ -23,11 +23,11 @@ impl Augmentation<i64> for EvenCount {

/// This implementation enables get key by 'nth' even number. This effectively makes
/// `EvenCount` a secondary index
impl SearchAugmentation<i64> for EvenCount {
impl<V> SearchAugmentation<i64, V> for EvenCount {
/// offset of even number
type Query = usize;

fn locate_in_leaf(mut offset: usize, keys: &[i64]) -> Option<usize> {
fn locate_in_leaf(mut offset: usize, keys: &[i64], _: &[V]) -> Option<usize> {
for (idx, key) in keys.iter().enumerate() {
if value_is_even(*key) {
if offset == 0 {
Expand Down
4 changes: 2 additions & 2 deletions examples/date_statistic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,7 @@ struct DateStatistic {
}

/// How the `DateStatistic` aggregated from children
impl Augmentation<Date> for DateStatistic {
impl<V> Augmentation<Date, V> for DateStatistic {
/// aggregate for inner
fn from_inner(_keys: &[Date], augmentations: &[Self]) -> Self {
let mut augmentation_iter = augmentations.iter();
Expand Down Expand Up @@ -182,7 +182,7 @@ impl Augmentation<Date> for DateStatistic {
}

/// aggregate for leaf node
fn from_leaf(keys: &[Date]) -> Self {
fn from_leaf(keys: &[Date], _: &[V]) -> Self {
let mut keys_iter = keys.iter();
let first_key = keys_iter.next().unwrap();

Expand Down
96 changes: 96 additions & 0 deletions examples/duplicate_key.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
use sweep_bptree::augment::SearchAugmentation;
use sweep_bptree::Key;
use sweep_bptree::{augment::Augmentation, BPlusTree, NodeStoreVec};

#[derive(Debug, Copy, Clone, Default)]
pub struct Count(usize);

impl Count {
/// Get the count value
pub fn count(&self) -> usize {
self.0
}
}

impl<K: Key> Augmentation<K, usize> for Count {
fn from_leaf(_: &[K], values: &[usize]) -> Self {
Self(values.iter().sum())
}

fn from_inner(_keys: &[K], counts: &[Self]) -> Self {
Self(counts.iter().map(|a| a.0).sum())
}
}

impl<K: Key> SearchAugmentation<K, usize> for Count {
/// Query for ElementCount is index
type Query = usize;

fn locate_in_leaf(offset: usize, keys: &[K], values: &[usize]) -> Option<usize> {
let mut accum = 0usize;

for (idx, (key, count)) in keys.iter().zip(values).enumerate() {
if accum + count > offset {
return Some(idx);
}

accum += count;
}

None
}

fn locate_in_inner(mut idx: usize, _keys: &[K], counts: &[Self]) -> Option<(usize, usize)> {
for (i, a) in counts.iter().enumerate() {
if idx >= a.0 {
idx -= a.0;
} else {
return Some((i, idx));
}
}

// offset is larger than the sum of all counts
None
}
}

// insert key into tree, if it already exists, then increase count
fn insert(tree: &mut BPlusTree<NodeStoreVec<i64, usize, Count>>, key: i64) {
let prev_count = tree.get(&key).cloned().unwrap_or_default();
tree.insert(key, prev_count + 1);
}

fn delete(tree: &mut BPlusTree<NodeStoreVec<i64, usize, Count>>, key: i64) {
let prev_count = tree.get(&key).cloned().unwrap_or_default();
if prev_count == 1 {
tree.remove(&key);
} else {
tree.insert(key, prev_count - 1);
}
}

fn main() {
// create a tree with the augment
let mut tree = BPlusTree::new(NodeStoreVec::<i64, usize, Count>::new());

for _ in 0..100 {
for i in 0..100 {
insert(&mut tree, i);
}
}
assert_eq!(tree.root_augmentation().0, 10000);

for offset in 0..tree.root_augmentation().0 {
let kv = tree.get_by_augmentation::<usize>(offset).unwrap();
dbg!((offset, kv.0));
}

for _ in 0..100 {
for i in 0..100 {
delete(&mut tree, i);
}
}

assert_eq!(tree.root_augmentation().0, 0);
assert!(tree.get_by_augmentation::<usize>(0).is_none());
}
17 changes: 9 additions & 8 deletions src/augment/count.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use super::{Augmentation, RankAugmentation, SearchAugmentation};

/// This augmentation keeps track of the number of elements in the child.
/// Basicly, it turns the tree to [Order Statistic Tree](https://en.wikipedia.org/wiki/Order_statistic_tree)
#[derive(Clone, Copy, Debug, Default)]
#[derive(Debug, Copy, Clone, Default)]
pub struct Count(usize);

impl Count {
Expand All @@ -14,8 +14,8 @@ impl Count {
}
}

impl<K: Key> Augmentation<K> for Count {
fn from_leaf(keys: &[K]) -> Self {
impl<K: Key, V> Augmentation<K, V> for Count {
fn from_leaf(keys: &[K], _: &[V]) -> Self {
Self(keys.len())
}

Expand All @@ -24,11 +24,11 @@ impl<K: Key> Augmentation<K> for Count {
}
}

impl<K: Key> SearchAugmentation<K> for Count {
impl<K: Key, V> SearchAugmentation<K, V> for Count {
/// Query for ElementCount is index
type Query = usize;

fn locate_in_leaf(idx: usize, keys: &[K]) -> Option<usize> {
fn locate_in_leaf(idx: usize, keys: &[K], _: &[V]) -> Option<usize> {
if idx < keys.len() {
Some(idx)
} else {
Expand All @@ -50,7 +50,7 @@ impl<K: Key> SearchAugmentation<K> for Count {
}
}

impl<K: Key> RankAugmentation<K> for Count {
impl<K: Key, V> RankAugmentation<K, V> for Count {
/// The rank for ElementCount is index
type Rank = usize;

Expand Down Expand Up @@ -86,10 +86,11 @@ mod tests {

#[test]
fn test_element_count() {
let count = Count::from_leaf(&[1, 2, 3]);
let count = <Count as Augmentation<i32, i32>>::from_leaf(&[1, 2, 3], &[]);
assert_eq!(count.0, 3);

let count = Count::from_inner(&[1, 2, 3], &[count, count, count]);
let count =
<Count as Augmentation<i32, i32>>::from_inner(&[1, 2, 3], &[count, count, count]);
assert_eq!(count.0, 9);
}

Expand Down
10 changes: 5 additions & 5 deletions src/augment/group.rs
Original file line number Diff line number Diff line change
Expand Up @@ -148,12 +148,12 @@ pub trait FromRef<T> {
fn from_ref(input: &T) -> Self;
}

impl<K, G> Augmentation<K> for GroupCount<G>
impl<K, V, G> Augmentation<K, V> for GroupCount<G>
where
K: Key,
G: FromRef<K> + Clone + Ord + std::fmt::Debug,
{
fn from_leaf(keys: &[K]) -> Self {
fn from_leaf(keys: &[K], _: &[V]) -> Self {
let mut keys_iter = keys.iter();

let first_group = G::from_ref(match keys_iter.next() {
Expand Down Expand Up @@ -203,15 +203,15 @@ where
}
}

impl<K, G> SearchAugmentation<K> for GroupCount<G>
impl<K, V, G> SearchAugmentation<K, V> for GroupCount<G>
where
K: Key,
G: FromRef<K> + Clone + Ord + std::fmt::Debug,
{
/// It can be searched by Group and offset
type Query = (G, usize);

fn locate_in_leaf((group, offset): (G, usize), keys: &[K]) -> Option<usize> {
fn locate_in_leaf((group, offset): (G, usize), keys: &[K], _: &[V]) -> Option<usize> {
let mut in_group_offset = 0;
for (idx, k) in keys.iter().enumerate() {
let group_for_key = G::from_ref(k);
Expand Down Expand Up @@ -272,7 +272,7 @@ where
}
}

impl<K, G> RankAugmentation<K> for GroupCount<G>
impl<K, V, G> RankAugmentation<K, V> for GroupCount<G>
where
K: Key,
G: FromRef<K> + Clone + Ord + std::fmt::Debug,
Expand Down
19 changes: 10 additions & 9 deletions src/augment/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,10 @@ pub mod count;
pub mod group;

/// Augmentation trait, it is used to store augmentation, like 'size'
/// NOTE: Since the lib has no control on how value changes, so augment only calculated from keys
/// e.g: Map<i64, Arc<Mutex<i64>>
pub trait Augmentation<K: Key>: Clone + Default {
/// NOTE: if Augmentation's calculation depends on values, please DO NOT use `get_mut`
/// and then modify the value. In that case, the lib lose track of value change and
/// the augmentation is not synced.
pub trait Augmentation<K: Key, V>: Clone + Default {
fn is_zst() -> bool {
false
}
Expand All @@ -21,17 +22,17 @@ pub trait Augmentation<K: Key>: Clone + Default {
fn from_inner(keys: &[K], augmentations: &[Self]) -> Self;

/// create a new Augmentation from leaf node's key
fn from_leaf(keys: &[K]) -> Self;
fn from_leaf(keys: &[K], values: &[V]) -> Self;
}

/// Whether the augmentation able to locate element
/// `SearchAugmentation` acts like a secondary index, it is able to locate
/// the record.
pub trait SearchAugmentation<K: Key>: Augmentation<K> {
pub trait SearchAugmentation<K: Key, V>: Augmentation<K, V> {
type Query;

/// locate the offset of the element in leaf node
fn locate_in_leaf(query: Self::Query, keys: &[K]) -> Option<usize>;
fn locate_in_leaf(query: Self::Query, keys: &[K], values: &[V]) -> Option<usize>;

/// locate the child index of query in inner node, with new query
fn locate_in_inner(
Expand All @@ -42,7 +43,7 @@ pub trait SearchAugmentation<K: Key>: Augmentation<K> {
}

/// Whether the augmentation able to rank element(like the index of key)
pub trait RankAugmentation<K: Key>: Augmentation<K> {
pub trait RankAugmentation<K: Key, V>: Augmentation<K, V> {
type Rank;

/// Initial rank value, e.g: 0 for size
Expand All @@ -66,13 +67,13 @@ pub trait RankAugmentation<K: Key>: Augmentation<K> {
}

/// () is a dummy augment that turns Augmented tree to a normal tree
impl<K: Key> Augmentation<K> for () {
impl<K: Key, V> Augmentation<K, V> for () {
fn is_zst() -> bool {
true
}

#[inline(always)]
fn from_leaf(_: &[K]) -> Self {}
fn from_leaf(_: &[K], _: &[V]) -> Self {}

#[inline(always)]
fn from_inner(_: &[K], _: &[Self]) -> Self {}
Expand Down
16 changes: 8 additions & 8 deletions src/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,17 +7,17 @@ use crate::{
};

/// A B+ tree map implemented with `BPlusTree`
pub struct BPlusTreeMap<K: Key, V, A: Augmentation<K> = ()> {
pub struct BPlusTreeMap<K: Key, V, A: Augmentation<K, V> = ()> {
inner: BPlusTree<NodeStoreVec<K, V, A>>,
}

impl<K: Key, V, A: Augmentation<K>> Default for BPlusTreeMap<K, V, A> {
impl<K: Key, V, A: Augmentation<K, V>> Default for BPlusTreeMap<K, V, A> {
fn default() -> Self {
Self::new()
}
}

impl<K: Key, V, A: Augmentation<K>> BPlusTreeMap<K, V, A> {
impl<K: Key, V, A: Augmentation<K, V>> BPlusTreeMap<K, V, A> {
/// Create a new BPlusTreeMap
///
/// # Examples
Expand Down Expand Up @@ -184,7 +184,7 @@ impl<K: Key, V, A: Augmentation<K>> BPlusTreeMap<K, V, A> {
/// ```
pub fn get_by_augmentation<Q>(&self, query: Q) -> Option<(&K, &V)>
where
A: SearchAugmentation<K, Query = Q>,
A: SearchAugmentation<K, V, Query = Q>,
{
self.inner.get_by_augmentation(query)
}
Expand All @@ -209,7 +209,7 @@ impl<K: Key, V, A: Augmentation<K>> BPlusTreeMap<K, V, A> {
/// ```
pub fn get_mut_by_augmentation<Q>(&mut self, query: Q) -> Option<&mut V>
where
A: SearchAugmentation<K, Query = Q>,
A: SearchAugmentation<K, V, Query = Q>,
{
self.inner.get_mut_by_augmentation(query)
}
Expand All @@ -235,7 +235,7 @@ impl<K: Key, V, A: Augmentation<K>> BPlusTreeMap<K, V, A> {
/// ```
pub fn remove_by_augmentation<Q>(&mut self, query: Q) -> Option<(K, V)>
where
A: SearchAugmentation<K, Query = Q>,
A: SearchAugmentation<K, V, Query = Q>,
{
self.inner.remove_by_augmentation(query)
}
Expand Down Expand Up @@ -268,7 +268,7 @@ impl<K: Key, V, A: Augmentation<K>> BPlusTreeMap<K, V, A> {
/// ```
pub fn rank_by_augmentation<R>(&self, k: &K) -> Result<R, R>
where
A: RankAugmentation<K, Rank = R>,
A: RankAugmentation<K, V, Rank = R>,
{
self.inner.rank_by_augmentation(k)
}
Expand All @@ -284,7 +284,7 @@ impl<K: Key, V, A: Augmentation<K>> BPlusTreeMap<K, V, A> {
}
}

impl<K: Key, V, A: Augmentation<K>> FromIterator<(K, V)> for BPlusTreeMap<K, V, A> {
impl<K: Key, V, A: Augmentation<K, V>> FromIterator<(K, V)> for BPlusTreeMap<K, V, A> {
/// Create a BPlusTreeMap from an iterator
///
/// # Example
Expand Down
2 changes: 1 addition & 1 deletion src/tree/bulk_load.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ impl<S: NodeStore> crate::BPlusTree<S> {
nodes.push((
NodeId::Leaf(leaf_id),
leaf.key_range(),
S::Augmentation::from_leaf(leaf.keys()),
S::Augmentation::from_leaf(leaf.keys(), leaf.values()),
));
item_count += leaf.len();

Expand Down
Loading

0 comments on commit 94beb4a

Please sign in to comment.