2024-08-28 18:45:16 +02:00
|
|
|
use std::num::NonZeroUsize;
|
2024-09-04 10:55:06 +02:00
|
|
|
use std::mem;
|
2024-08-28 18:45:16 +02:00
|
|
|
|
2024-09-03 11:02:39 +02:00
|
|
|
use grenad::{MergeFunction, Sorter};
|
2024-08-28 18:45:16 +02:00
|
|
|
use lru::LruCache;
|
|
|
|
use roaring::RoaringBitmap;
|
|
|
|
use smallvec::SmallVec;
|
|
|
|
|
2024-09-04 10:55:06 +02:00
|
|
|
use crate::update::del_add::{DelAdd, KvWriterDelAdd};
|
|
|
|
use crate::CboRoaringBitmapCodec;
|
2024-08-28 18:45:16 +02:00
|
|
|
|
|
|
|
#[derive(Debug)]
|
|
|
|
pub struct CachedSorter<MF> {
|
|
|
|
cache: lru::LruCache<SmallVec<[u8; 20]>, DelAddRoaringBitmap>,
|
|
|
|
sorter: Sorter<MF>,
|
|
|
|
deladd_buffer: Vec<u8>,
|
|
|
|
cbo_buffer: Vec<u8>,
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<MF> CachedSorter<MF> {
|
|
|
|
pub fn new(cap: NonZeroUsize, sorter: Sorter<MF>) -> Self {
|
|
|
|
CachedSorter {
|
|
|
|
cache: lru::LruCache::new(cap),
|
|
|
|
sorter,
|
|
|
|
deladd_buffer: Vec::new(),
|
|
|
|
cbo_buffer: Vec::new(),
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<MF: MergeFunction> CachedSorter<MF> {
|
|
|
|
pub fn insert_del_u32(&mut self, key: &[u8], n: u32) -> grenad::Result<(), MF::Error> {
|
|
|
|
match self.cache.get_mut(key) {
|
|
|
|
Some(DelAddRoaringBitmap { del, add: _ }) => {
|
|
|
|
del.get_or_insert_with(RoaringBitmap::new).insert(n);
|
|
|
|
}
|
|
|
|
None => {
|
|
|
|
let value = DelAddRoaringBitmap::new_del_u32(n);
|
|
|
|
if let Some((key, deladd)) = self.cache.push(key.into(), value) {
|
|
|
|
self.write_entry(key, deladd)?;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
|
|
|
pub fn insert_del(
|
|
|
|
&mut self,
|
|
|
|
key: &[u8],
|
|
|
|
bitmap: RoaringBitmap,
|
|
|
|
) -> grenad::Result<(), MF::Error> {
|
|
|
|
match self.cache.get_mut(key) {
|
|
|
|
Some(DelAddRoaringBitmap { del, add: _ }) => {
|
|
|
|
*del.get_or_insert_with(RoaringBitmap::new) |= bitmap;
|
|
|
|
}
|
|
|
|
None => {
|
|
|
|
let value = DelAddRoaringBitmap::new_del(bitmap);
|
|
|
|
if let Some((key, deladd)) = self.cache.push(key.into(), value) {
|
|
|
|
self.write_entry(key, deladd)?;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
|
|
|
pub fn insert_add_u32(&mut self, key: &[u8], n: u32) -> grenad::Result<(), MF::Error> {
|
|
|
|
match self.cache.get_mut(key) {
|
|
|
|
Some(DelAddRoaringBitmap { del: _, add }) => {
|
|
|
|
add.get_or_insert_with(RoaringBitmap::new).insert(n);
|
|
|
|
}
|
|
|
|
None => {
|
|
|
|
let value = DelAddRoaringBitmap::new_add_u32(n);
|
|
|
|
if let Some((key, deladd)) = self.cache.push(key.into(), value) {
|
|
|
|
self.write_entry(key, deladd)?;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
|
|
|
pub fn insert_add(
|
|
|
|
&mut self,
|
|
|
|
key: &[u8],
|
|
|
|
bitmap: RoaringBitmap,
|
|
|
|
) -> grenad::Result<(), MF::Error> {
|
|
|
|
match self.cache.get_mut(key) {
|
|
|
|
Some(DelAddRoaringBitmap { del: _, add }) => {
|
|
|
|
*add.get_or_insert_with(RoaringBitmap::new) |= bitmap;
|
|
|
|
}
|
|
|
|
None => {
|
|
|
|
let value = DelAddRoaringBitmap::new_add(bitmap);
|
|
|
|
if let Some((key, deladd)) = self.cache.push(key.into(), value) {
|
|
|
|
self.write_entry(key, deladd)?;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
|
|
|
pub fn insert_del_add_u32(&mut self, key: &[u8], n: u32) -> grenad::Result<(), MF::Error> {
|
|
|
|
match self.cache.get_mut(key) {
|
|
|
|
Some(DelAddRoaringBitmap { del, add }) => {
|
|
|
|
del.get_or_insert_with(RoaringBitmap::new).insert(n);
|
|
|
|
add.get_or_insert_with(RoaringBitmap::new).insert(n);
|
|
|
|
}
|
|
|
|
None => {
|
|
|
|
let value = DelAddRoaringBitmap::new_del_add_u32(n);
|
|
|
|
if let Some((key, deladd)) = self.cache.push(key.into(), value) {
|
|
|
|
self.write_entry(key, deladd)?;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
|
|
|
fn write_entry<A: AsRef<[u8]>>(
|
|
|
|
&mut self,
|
|
|
|
key: A,
|
|
|
|
deladd: DelAddRoaringBitmap,
|
|
|
|
) -> grenad::Result<(), MF::Error> {
|
2024-09-04 10:55:06 +02:00
|
|
|
/// TODO we must create a serialization trait to correctly serialize bitmaps
|
2024-08-28 18:45:16 +02:00
|
|
|
self.deladd_buffer.clear();
|
|
|
|
let mut value_writer = KvWriterDelAdd::new(&mut self.deladd_buffer);
|
|
|
|
match deladd {
|
|
|
|
DelAddRoaringBitmap { del: Some(del), add: None } => {
|
|
|
|
self.cbo_buffer.clear();
|
2024-09-04 10:55:06 +02:00
|
|
|
CboRoaringBitmapCodec::serialize_into(&del, &mut self.cbo_buffer);
|
2024-08-28 18:45:16 +02:00
|
|
|
value_writer.insert(DelAdd::Deletion, &self.cbo_buffer)?;
|
|
|
|
}
|
|
|
|
DelAddRoaringBitmap { del: None, add: Some(add) } => {
|
|
|
|
self.cbo_buffer.clear();
|
2024-09-04 10:55:06 +02:00
|
|
|
CboRoaringBitmapCodec::serialize_into(&add, &mut self.cbo_buffer);
|
2024-08-28 18:45:16 +02:00
|
|
|
value_writer.insert(DelAdd::Addition, &self.cbo_buffer)?;
|
|
|
|
}
|
|
|
|
DelAddRoaringBitmap { del: Some(del), add: Some(add) } => {
|
|
|
|
self.cbo_buffer.clear();
|
2024-09-04 10:55:06 +02:00
|
|
|
CboRoaringBitmapCodec::serialize_into(&del, &mut self.cbo_buffer);
|
2024-08-28 18:45:16 +02:00
|
|
|
value_writer.insert(DelAdd::Deletion, &self.cbo_buffer)?;
|
|
|
|
|
|
|
|
self.cbo_buffer.clear();
|
2024-09-04 10:55:06 +02:00
|
|
|
CboRoaringBitmapCodec::serialize_into(&add, &mut self.cbo_buffer);
|
2024-08-28 18:45:16 +02:00
|
|
|
value_writer.insert(DelAdd::Addition, &self.cbo_buffer)?;
|
|
|
|
}
|
|
|
|
DelAddRoaringBitmap { del: None, add: None } => return Ok(()),
|
|
|
|
}
|
|
|
|
let bytes = value_writer.into_inner().unwrap();
|
|
|
|
self.sorter.insert(key, bytes)
|
|
|
|
}
|
|
|
|
|
|
|
|
pub fn direct_insert(&mut self, key: &[u8], val: &[u8]) -> grenad::Result<(), MF::Error> {
|
|
|
|
self.sorter.insert(key, val)
|
|
|
|
}
|
|
|
|
|
|
|
|
pub fn into_sorter(mut self) -> grenad::Result<Sorter<MF>, MF::Error> {
|
|
|
|
let default_arc = LruCache::new(NonZeroUsize::MIN);
|
|
|
|
for (key, deladd) in mem::replace(&mut self.cache, default_arc) {
|
|
|
|
self.write_entry(key, deladd)?;
|
|
|
|
}
|
|
|
|
Ok(self.sorter)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Debug, Clone)]
|
|
|
|
pub struct DelAddRoaringBitmap {
|
|
|
|
pub del: Option<RoaringBitmap>,
|
|
|
|
pub add: Option<RoaringBitmap>,
|
|
|
|
}
|
|
|
|
|
|
|
|
impl DelAddRoaringBitmap {
|
|
|
|
fn new_del_add_u32(n: u32) -> Self {
|
|
|
|
DelAddRoaringBitmap {
|
|
|
|
del: Some(RoaringBitmap::from([n])),
|
|
|
|
add: Some(RoaringBitmap::from([n])),
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
fn new_del(bitmap: RoaringBitmap) -> Self {
|
|
|
|
DelAddRoaringBitmap { del: Some(bitmap), add: None }
|
|
|
|
}
|
|
|
|
|
|
|
|
fn new_del_u32(n: u32) -> Self {
|
|
|
|
DelAddRoaringBitmap { del: Some(RoaringBitmap::from([n])), add: None }
|
|
|
|
}
|
|
|
|
|
|
|
|
fn new_add(bitmap: RoaringBitmap) -> Self {
|
|
|
|
DelAddRoaringBitmap { del: None, add: Some(bitmap) }
|
|
|
|
}
|
|
|
|
|
|
|
|
fn new_add_u32(n: u32) -> Self {
|
|
|
|
DelAddRoaringBitmap { del: None, add: Some(RoaringBitmap::from([n])) }
|
|
|
|
}
|
2024-09-04 10:55:06 +02:00
|
|
|
}
|