filip/Stalwart
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Foundation passing tests

Browse Source
This commit is contained in:
Mauro D
2023-03-30 16:52:47 +00:00
parent d74731f3c6
commit ad7cfbcee6
10 changed files with 363 additions and 677 deletions
Download Patch File Download Diff File Expand all files Collapse all files

117
src/backend/foundationdb/read.rs
View File

@@ -11,15 +11,12 @@ use futures::StreamExt;
use roaring::RoaringBitmap;
use crate::{
query::{Operator, SortedId, UnsortedIds},
query::Operator,
write::key::{DeserializeBigEndian, KeySerializer},
BitmapKey, Deserialize, IndexKey, IndexKeyPrefix, Serialize, Store, ValueKey, BM_DOCUMENT_IDS,
};
use super::{
bitmap::{DeserializeBlock, BITS_PER_BLOCK},
SUBSPACE_INDEXES,
};
use super::{bitmap::DeserializeBlock, SUBSPACE_INDEXES};
pub struct ReadTransaction<'x> {
db: &'x Database,
@@ -72,11 +69,11 @@ impl ReadTransaction<'_> {
.await
}
#[inline(always)]
pub async fn get_bitmap<T: AsRef<[u8]>>(
async fn get_bitmap_<T: AsRef<[u8]>>(
&self,
mut key: BitmapKey<T>,
) -> crate::Result<Option<RoaringBitmap>> {
bm: &mut RoaringBitmap,
) -> crate::Result<()> {
let from_key = key.serialize();
key.block_num = u32::MAX;
let to_key = key.serialize();
@@ -85,9 +82,8 @@ impl ReadTransaction<'_> {
reverse: false,
..RangeOption::from((from_key.as_ref(), to_key.as_ref()))
};
//println!("deserializing bitmap: {:?} {:?}", from_key, to_key);
let mut bm = RoaringBitmap::new();
let mut values = self.trx.get_ranges(opt, true);
while let Some(values) = values.next().await {
for value in values? {
let key = value.key();
@@ -98,23 +94,18 @@ impl ReadTransaction<'_> {
.deserialize_be_u32(key.len() - std::mem::size_of::<u32>())?,
);
}
//println!("deserializing bitmap: {:?} {:?}", value.key(), bm.len());
}
Ok(if !bm.is_empty() { Some(bm) } else { None })
Ok(())
}
#[inline(always)]
async fn get_bitmaps<T: AsRef<[u8]>>(
pub async fn get_bitmap<T: AsRef<[u8]>>(
&self,
keys: Vec<BitmapKey<T>>,
) -> crate::Result<Vec<Option<RoaringBitmap>>> {
let mut results = Vec::with_capacity(keys.len());
for key in keys {
results.push(self.get_bitmap(key).await?);
}
Ok(results)
key: BitmapKey<T>,
) -> crate::Result<Option<RoaringBitmap>> {
let mut bm = RoaringBitmap::new();
self.get_bitmap_(key, &mut bm).await?;
Ok(if !bm.is_empty() { Some(bm) } else { None })
}
pub(crate) async fn get_bitmaps_intersection<T: AsRef<[u8]>>(
@@ -122,8 +113,8 @@ impl ReadTransaction<'_> {
keys: Vec<BitmapKey<T>>,
) -> crate::Result<Option<RoaringBitmap>> {
let mut result: Option<RoaringBitmap> = None;
for bitmap in self.get_bitmaps(keys).await? {
if let Some(bitmap) = bitmap {
for key in keys {
if let Some(bitmap) = self.get_bitmap(key).await? {
if let Some(result) = &mut result {
result.bitand_assign(&bitmap);
if result.is_empty() {
@@ -143,15 +134,13 @@ impl ReadTransaction<'_> {
&self,
keys: Vec<BitmapKey<T>>,
) -> crate::Result<Option<RoaringBitmap>> {
let mut result: Option<RoaringBitmap> = None;
for bitmap in (self.get_bitmaps(keys).await?).into_iter().flatten() {
if let Some(result) = &mut result {
result.bitor_assign(&bitmap);
} else {
result = Some(bitmap);
}
let mut bm = RoaringBitmap::new();
for key in keys {
self.get_bitmap_(key, &mut bm).await?;
}
Ok(result)
Ok(if !bm.is_empty() { Some(bm) } else { None })
}
pub(crate) async fn range_to_bitmap(
@@ -180,23 +169,27 @@ impl ReadTransaction<'_> {
let (begin, end) = match op {
Operator::LowerThan => (
KeySelector::first_greater_or_equal(k1.finalize()),
KeySelector::last_less_than(k2.write(&value[..]).write(0u32).finalize()),
KeySelector::first_greater_or_equal(k2.write(&value[..]).write(0u32).finalize()),
),
Operator::LowerEqualThan => (
KeySelector::first_greater_or_equal(k1.finalize()),
KeySelector::last_less_or_equal(k2.write(&value[..]).write(u32::MAX).finalize()),
KeySelector::first_greater_or_equal(
k2.write(&value[..]).write(u32::MAX).finalize(),
),
),
Operator::GreaterThan => (
KeySelector::first_greater_than(k1.write(&value[..]).write(u32::MAX).finalize()),
KeySelector::last_less_than(k2.finalize()),
KeySelector::first_greater_or_equal(k2.finalize()),
),
Operator::GreaterEqualThan => (
KeySelector::first_greater_or_equal(k1.write(&value[..]).write(0u32).finalize()),
KeySelector::last_less_than(k2.finalize()),
KeySelector::first_greater_or_equal(k2.finalize()),
),
Operator::Equal => (
KeySelector::first_greater_or_equal(k1.write(&value[..]).write(0u32).finalize()),
KeySelector::last_less_or_equal(k2.write(&value[..]).write(u32::MAX).finalize()),
KeySelector::first_greater_or_equal(
k2.write(&value[..]).write(u32::MAX).finalize(),
),
),
};
@@ -221,15 +214,14 @@ impl ReadTransaction<'_> {
Ok(Some(bm))
}
pub(crate) async fn sort_bitmap(
pub(crate) async fn sort_index(
&self,
account_id: u32,
collection: u8,
field: u8,
documents: &impl UnsortedIds,
limit: usize,
ascending: bool,
) -> crate::Result<Vec<SortedId>> {
mut cb: impl FnMut(&[u8], u32) -> bool,
) -> crate::Result<()> {
let from_key = IndexKeyPrefix {
account_id,
collection,
@@ -242,11 +234,11 @@ impl ReadTransaction<'_> {
field: field + 1,
}
.serialize();
let mut results = Vec::with_capacity(documents.len());
let prefix_len = from_key.len();
let mut sorted_iter = self.trx.get_ranges(
RangeOption {
begin: KeySelector::first_greater_or_equal(&from_key),
end: KeySelector::last_less_than(&to_key),
end: KeySelector::first_greater_or_equal(&to_key),
mode: options::StreamingMode::Iterator,
reverse: !ascending,
..Default::default()
@@ -254,42 +246,23 @@ impl ReadTransaction<'_> {
true,
);
let mut prev_prefix = vec![];
while let Some(values) = sorted_iter.next().await {
for value in values? {
let key = value.key();
let document_id = key.deserialize_be_u32(value.key().len() - 4)?;
if documents.contains_id(document_id) {
let prefix = key
.get(..key.len() - std::mem::size_of::<u32>())
.ok_or_else(|| {
crate::Error::InternalError("Invalid key found in index".to_string())
})?;
if prefix == prev_prefix {
let last = results.last_mut().unwrap();
match last {
SortedId::Id(id) => {
*last = SortedId::GroupedId(vec![*id, document_id]);
}
SortedId::GroupedId(ids) => {
ids.push(document_id);
}
}
} else {
results.push(SortedId::Id(document_id));
prev_prefix = prefix.to_vec();
}
if results.len() == limit {
return Ok(results);
}
let id_pos = key.len() - std::mem::size_of::<u32>();
debug_assert!(key.starts_with(&from_key));
if !cb(
key.get(prefix_len..id_pos).ok_or_else(|| {
crate::Error::InternalError("Invalid key found in index".to_string())
})?,
key.deserialize_be_u32(id_pos)?,
) {
return Ok(());
}
}
}
Ok(results)
Ok(())
}
pub async fn refresh_if_old(&mut self) -> crate::Result<()> {

35
src/fts/bloom.rs
View File

@@ -7,7 +7,7 @@ use std::{
use roaring::RoaringBitmap;
use utils::codec::leb128::{Leb128Reader, Leb128Vec};
use crate::{Deserialize, Error, Serialize};
use crate::{BitmapKey, Deserialize, Error, Serialize, BLOOM_UNIGRAM, BM_BLOOM};
use super::{stemmer::StemmedToken, tokenizers::Token};
@@ -116,6 +116,39 @@ impl BloomHash {
h: [h1, h2, h3, h4, h1 ^ h2, h2 ^ h3, h3 ^ h4],
}
}
pub fn to_bitmap_key(&self, account_id: u32, collection: u8, field: u8) -> BitmapKey<Vec<u8>> {
let mut key = Vec::with_capacity(12);
key.extend_from_slice(&self.h[0].to_le_bytes()[..3]);
key.extend_from_slice(&self.h[1].to_le_bytes()[..3]);
key.extend_from_slice(&self.h[2].to_le_bytes()[..3]);
key.extend_from_slice(&self.h[3].to_le_bytes()[..3]);
BitmapKey {
account_id,
collection,
family: BM_BLOOM | BLOOM_UNIGRAM,
field,
block_num: 0,
key,
}
}
}
pub fn hash_token(item: &str) -> Vec<u8> {
let h1 = xxhash_rust::xxh3::xxh3_64(item.as_ref());
let h2 = farmhash::hash64(item.as_ref());
let h3 = AHASHER.hash_one(item);
let mut sh = *SIPHASHER;
sh.write(item.as_ref());
let h4 = sh.finish();
let mut hash = Vec::with_capacity(12);
hash.extend_from_slice(&h1.to_le_bytes()[..3]);
hash.extend_from_slice(&h2.to_le_bytes()[..3]);
hash.extend_from_slice(&h3.to_le_bytes()[..3]);
hash.extend_from_slice(&h4.to_le_bytes()[..3]);
hash
}
impl From<&str> for BloomHash {

26
src/fts/builder.rs
View File

@@ -4,11 +4,11 @@ use ahash::AHashSet;
use crate::{
write::{BatchBuilder, IntoOperations, Operation},
Serialize, BLOOM_BIGRAM, BLOOM_TRIGRAM, BLOOM_UNIGRAM, BM_BLOOM,
Serialize, BLOOM_BIGRAM, BLOOM_TRIGRAM, BLOOM_UNIGRAM, BLOOM_UNIGRAM_STEM, BM_BLOOM,
};
use super::{
bloom::{BloomFilter, BloomHash},
bloom::{hash_token, BloomFilter},
lang::{LanguageDetector, MIN_LANGUAGE_SCORE},
ngram::ToNgrams,
stemmer::Stemmer,
@@ -73,30 +73,30 @@ impl<'x> IntoOperations for FtsIndexBuilder<'x> {
let mut phrase_words = Vec::new();
for token in Stemmer::new(&part.text, language, MAX_TOKEN_LENGTH).collect::<Vec<_>>() {
unique_words.insert(token.word.to_string());
if let Some(stemmed_word) = token.stemmed_word.as_ref() {
unique_words.insert(format!("{}_", stemmed_word));
unique_words.insert((token.word.to_string(), BM_BLOOM | BLOOM_UNIGRAM));
if let Some(stemmed_word) = token.stemmed_word {
unique_words.insert((stemmed_word.into_owned(), BM_BLOOM | BLOOM_UNIGRAM_STEM));
}
phrase_words.push(token.word);
}
let mut bloom_unigram = BloomFilter::new(unique_words.len());
for word in unique_words {
let hash = BloomHash::from(word);
bloom_unigram.insert(&hash);
//let mut bloom_unigram = BloomFilter::new(unique_words.len());
for (word, family) in unique_words {
//let hash = BloomHash::from(word);
//bloom_unigram.insert(&hash);
batch.ops.push(Operation::Bitmap {
family: BM_BLOOM,
family,
field: part.field,
key: hash.as_high_rank_hash().serialize(),
key: hash_token(&word),
set: true,
});
}
batch.ops.push(Operation::Value {
/*batch.ops.push(Operation::Value {
field: part.field,
family: BM_BLOOM | BLOOM_UNIGRAM,
set: bloom_unigram.serialize().into(),
});
});*/
if phrase_words.len() > 1 {
batch.ops.push(Operation::Value {

115
src/fts/query.rs
View File

@@ -4,16 +4,17 @@ use roaring::RoaringBitmap;
use crate::{
fts::{
bloom::{BloomFilter, BloomHash, BloomHashGroup},
bloom::{hash_token, BloomFilter, BloomHash, BloomHashGroup},
builder::MAX_TOKEN_LENGTH,
ngram::ToNgrams,
stemmer::Stemmer,
tokenizers::Tokenizer,
},
BitmapKey, Serialize, Store, ValueKey, BLOOM_BIGRAM, BLOOM_TRIGRAM, BLOOM_UNIGRAM, BM_BLOOM,
BitmapKey, Store, ValueKey, BLOOM_BIGRAM, BLOOM_TRIGRAM, BLOOM_UNIGRAM, BLOOM_UNIGRAM_STEM,
BM_BLOOM,
};
use super::{Language, HIGH_RANK_MOD};
use super::Language;
impl Store {
pub(crate) async fn fts_query(
@@ -33,7 +34,7 @@ impl Store {
let mut bit_keys = Vec::new();
for token in Tokenizer::new(text, language, MAX_TOKEN_LENGTH) {
let hash = BloomHash::from(token.word.as_ref());
let key = hash.to_high_rank_key(account_id, collection, field);
let key = hash.to_bitmap_key(account_id, collection, field);
if !bit_keys.contains(&key) {
bit_keys.push(key);
}
@@ -47,11 +48,7 @@ impl Store {
match tokens.len() {
0 => return Ok(None),
1 => (
bitmaps,
vec![tokens.into_iter().next().unwrap().into()],
BM_BLOOM | BLOOM_UNIGRAM,
),
1 => return Ok(Some(bitmaps)),
2 => (
bitmaps,
<Vec<BloomHashGroup>>::to_ngrams(&tokens, 2),
@@ -64,27 +61,36 @@ impl Store {
),
}
} else {
let mut hashes = Vec::new();
let mut bitmaps = RoaringBitmap::new();
for token in Stemmer::new(text, language, MAX_TOKEN_LENGTH) {
let hash = BloomHashGroup {
h2: if let Some(stemmed_word) = token.stemmed_word {
Some(format!("{stemmed_word}_").into())
} else {
Some(format!("{}_", token.word).into())
},
h1: token.word.into(),
let token1 = hash_token(&token.word);
let token2 = if let Some(stemmed_word) = token.stemmed_word {
hash_token(&stemmed_word)
} else {
token1.clone()
};
trx.refresh_if_old().await?;
match trx
.get_bitmaps_union(vec![
hash.h1.to_high_rank_key(account_id, collection, field),
hash.h2
.as_ref()
.unwrap()
.to_high_rank_key(account_id, collection, field),
BitmapKey {
account_id,
collection,
family: BM_BLOOM | BLOOM_UNIGRAM,
field,
block_num: 0,
key: token1,
},
BitmapKey {
account_id,
collection,
family: BM_BLOOM | BLOOM_UNIGRAM_STEM,
field,
block_num: 0,
key: token2,
},
])
.await?
{
@@ -100,53 +106,13 @@ impl Store {
}
_ => return Ok(None),
};
hashes.push(hash);
}
(bitmaps, hashes, BM_BLOOM | BLOOM_UNIGRAM)
return Ok(Some(bitmaps));
};
let b_count = bitmaps.len();
/*let bm = self
.get_values::<BloomFilter>(
bitmaps
.iter()
.map(|document_id| ValueKey {
account_id,
collection,
document_id,
family,
field,
})
.collect::<Vec<_>>(),
)
.await?
.into_iter()
.zip(bitmaps)
.filter_map(|(bloom, document_id)| {
let bloom = bloom?;
if !bloom.is_empty() {
let mut matched = true;
for hash in &hashes {
if !(bloom.contains(&hash.h1)
|| hash.h2.as_ref().map_or(false, |h2| bloom.contains(h2)))
{
matched = false;
break;
}
}
if matched {
return Some(document_id);
}
}
None
})
.collect::<RoaringBitmap>();*/
let mut bm = RoaringBitmap::new();
for document_id in bitmaps {
trx.refresh_if_old().await?;
@@ -187,26 +153,3 @@ impl Store {
Ok(Some(bm))
}
}
impl BloomHash {
#[inline(always)]
pub fn as_high_rank_hash(&self) -> u16 {
(self.h[0] % HIGH_RANK_MOD) as u16
}
pub fn to_high_rank_key(
&self,
account_id: u32,
collection: u8,
field: u8,
) -> BitmapKey<Vec<u8>> {
BitmapKey {
account_id,
collection,
family: BM_BLOOM,
field,
block_num: 0,
key: self.as_high_rank_hash().serialize(),
}
}
}

5
src/lib.rs
View File

@@ -99,8 +99,9 @@ pub const BM_TAG: u8 = 0x20;
pub const BM_BLOOM: u8 = 0x40;
pub const BLOOM_UNIGRAM: u8 = 0x00;
pub const BLOOM_BIGRAM: u8 = 0x01;
pub const BLOOM_TRIGRAM: u8 = 0x02;
pub const BLOOM_UNIGRAM_STEM: u8 = 0x01;
pub const BLOOM_BIGRAM: u8 = 0x02;
pub const BLOOM_TRIGRAM: u8 = 0x04;
pub const TERM_EXACT: u8 = 0x00;
pub const TERM_STEMMED: u8 = 0x01;

33
src/query/filter.rs
View File

@@ -1,14 +1,8 @@
use std::{
ops::{BitAndAssign, BitOrAssign, BitXorAssign},
time::Instant,
};
use std::ops::{BitAndAssign, BitOrAssign, BitXorAssign};
use roaring::RoaringBitmap;
use crate::{
backend::foundationdb::read::ReadTransaction, write::Tokenize, BitmapKey, Store, BM_TERM,
TERM_EXACT,
};
use crate::{write::Tokenize, BitmapKey, Store, BM_TERM, TERM_EXACT};
use super::{Filter, ResultSet};
@@ -25,16 +19,16 @@ impl Store {
filters: Vec<Filter>,
) -> crate::Result<ResultSet> {
let mut trx = self.read_transaction().await?;
let document_ids = trx
.get_document_ids(account_id, collection)
.await?
.unwrap_or_else(RoaringBitmap::new);
let mut not_mask = RoaringBitmap::new();
let mut not_fetch = false;
if filters.is_empty() {
return Ok(ResultSet {
account_id,
collection,
results: document_ids.clone(),
document_ids,
results: trx
.get_document_ids(account_id, collection)
.await?
.unwrap_or_else(RoaringBitmap::new),
});
}
@@ -118,7 +112,15 @@ impl Store {
}
};
state.op.apply(&mut state.bm, result, &document_ids);
if matches!(state.op, Filter::Not) && !not_fetch {
not_mask = trx
.get_document_ids(account_id, collection)
.await?
.unwrap_or_else(RoaringBitmap::new);
not_fetch = true;
}
state.op.apply(&mut state.bm, result, &not_mask);
//println!("{:?}: {:?}", state.op, state.bm);
@@ -137,7 +139,6 @@ impl Store {
account_id,
collection,
results: state.bm.unwrap_or_else(RoaringBitmap::new),
document_ids,
})
}
}

32
src/query/mod.rs
View File

@@ -62,7 +62,6 @@ pub struct ResultSet {
account_id: u32,
collection: u8,
pub results: RoaringBitmap,
pub document_ids: RoaringBitmap,
}
pub struct SortedResultRet {
@@ -71,17 +70,6 @@ pub struct SortedResultRet {
pub found_anchor: bool,
}
pub enum SortedId {
Id(u32),
GroupedId(Vec<u32>),
}
#[allow(clippy::len_without_is_empty)]
pub trait UnsortedIds {
fn contains_id(&self, id: u32) -> bool;
fn len(&self) -> usize;
}
impl Filter {
pub fn cond(field: impl Into<u8>, op: Operator, value: impl Serialize) -> Self {
Filter::MatchValue {
@@ -197,23 +185,3 @@ impl Comparator {
}
}
}
impl UnsortedIds for RoaringBitmap {
fn contains_id(&self, id: u32) -> bool {
self.contains(id)
}
fn len(&self) -> usize {
self.len() as usize
}
}
impl UnsortedIds for Vec<u32> {
fn contains_id(&self, id: u32) -> bool {
self.iter().any(|&i| i == id)
}
fn len(&self) -> usize {
self.len()
}
}

666
src/query/sort.rs
View File

@@ -1,501 +1,263 @@
use std::ops::{BitAndAssign, BitXorAssign};
use ahash::AHashMap;
use foundationdb::{future::FdbValue, options, FdbResult, KeySelector, RangeOption};
use futures::{Stream, StreamExt};
use roaring::RoaringBitmap;
#[cfg(feature = "rocks")]
use rocksdb::{
DBIteratorWithThreadMode, Direction, IteratorMode, MultiThreaded, OptimisticTransactionDB,
};
#[cfg(feature = "rocks")]
use crate::backend::rocksdb::{ACCOUNT_KEY_LEN, CF_INDEXES};
use crate::{
backend::foundationdb::read::ReadTransaction, write::key::DeserializeBigEndian, Error,
IndexKeyPrefix, Serialize, Store,
};
use crate::Store;
use super::{Comparator, ResultSet, SortedResultRet};
#[cfg(feature = "rocks")]
enum IndexType<'x> {
DocumentSet {
set: RoaringBitmap,
it: Option<roaring::bitmap::IntoIter>,
},
DB {
it: Option<DBIteratorWithThreadMode<'x, OptimisticTransactionDB<MultiThreaded>>>,
prefix: Vec<u8>,
from_key: Vec<u8>,
ascending: bool,
prev_item: Option<u32>,
prev_key: Option<Box<[u8]>>,
},
}
#[cfg(feature = "rocks")]
struct IndexIterator<'x> {
index: IndexType<'x>,
remaining: RoaringBitmap,
eof: bool,
}
#[cfg(feature = "foundation")]
enum IndexType<'x, T: Stream<Item = FdbResult<FdbValue>> + Unpin + 'x> {
DocumentSet {
set: RoaringBitmap,
it: Option<roaring::bitmap::IntoIter>,
},
DB {
it: Option<T>,
from_key: Vec<u8>,
to_key: Vec<u8>,
ascending: bool,
prev_item: Option<u32>,
prev_key: Option<Box<[u8]>>,
phantom: std::marker::PhantomData<&'x ()>,
},
}
#[cfg(feature = "foundation")]
struct IndexIterator<'x, T: Stream<Item = FdbResult<FdbValue>> + Unpin + 'x> {
index: IndexType<'x, T>,
remaining: RoaringBitmap,
eof: bool,
}
impl Store {
pub async fn sort(
&self,
mut result_set: ResultSet,
comparators: Vec<Comparator>,
result_set: ResultSet,
mut comparators: Vec<Comparator>,
limit: usize,
mut position: i32,
position: i32,
anchor: Option<u32>,
mut anchor_offset: i32,
anchor_offset: i32,
) -> crate::Result<SortedResultRet> {
let has_anchor = anchor.is_some();
let mut anchor_found = false;
let requested_position = position;
let trx = self.read_transaction().await?;
let mut result = SortedResultRet {
position,
ids: Vec::with_capacity(std::cmp::min(limit, result_set.results.len() as usize)),
found_anchor: true,
let limit = match (result_set.results.len(), limit) {
(0, _) => {
return Ok(SortedResultRet {
position,
ids: vec![],
found_anchor: true,
});
}
(_, 0) => result_set.results.len() as usize,
(a, b) => std::cmp::min(a as usize, b),
};
let mut iterators = comparators
.into_iter()
.map(|comp| IndexIterator {
index: match comp {
Comparator::Field { field, ascending } => IndexType::DB {
it: None,
from_key: if !ascending {
result_set.from_key(field).serialize()
} else {
result_set.to_key(field).serialize()
},
to_key: if !ascending {
result_set.to_key(field).serialize()
} else {
result_set.from_key(field).serialize()
},
let mut paginate = Pagination::new(limit, position, anchor, anchor_offset);
if comparators.len() == 1 {
match comparators.pop().unwrap() {
Comparator::Field { field, ascending } => {
let trx = self.read_transaction().await?;
let mut results = result_set.results;
trx.sort_index(
result_set.account_id,
result_set.collection,
field,
ascending,
prev_item: None,
prev_key: None,
phantom: std::marker::PhantomData,
},
Comparator::DocumentSet { mut set, ascending } => IndexType::DocumentSet {
set: if !ascending {
if !set.is_empty() {
set.bitxor_assign(&result_set.document_ids);
set
} else {
result_set.document_ids.clone()
}
} else {
set
},
it: None,
},
},
remaining: std::mem::replace(&mut result_set.results, RoaringBitmap::new()),
eof: false,
})
.collect::<Vec<_>>();
let mut current = 0;
let iter_len = iterators.len() - 1;
'outer: loop {
let mut doc_id;
'inner: loop {
let (it_opts, mut next_it_opts) = if current < iter_len {
let (iterators_first, iterators_last) = iterators.split_at_mut(current + 1);
(
iterators_first.last_mut().unwrap(),
iterators_last.first_mut(),
|_, document_id| !results.remove(document_id) || paginate.add(document_id),
)
} else {
(&mut iterators[current], None)
};
.await?;
if !matches!(it_opts.index, IndexType::DB { prev_item,.. } if prev_item.is_some())
{
if it_opts.remaining.is_empty() {
if current > 0 {
current -= 1;
continue 'inner;
} else {
break 'outer;
}
} else if it_opts.remaining.len() == 1 || it_opts.eof {
doc_id = it_opts.remaining.min().unwrap();
it_opts.remaining.remove(doc_id);
break 'inner;
}
}
match &mut it_opts.index {
IndexType::DB {
it,
from_key,
to_key,
ascending,
prev_item,
prev_key,
..
} => {
let it = if let Some(it) = it {
it
} else {
#[cfg(feature = "foundation")]
{
*it = Some(trx.trx.get_ranges_keyvalues(
RangeOption {
begin: KeySelector::first_greater_or_equal(
from_key.clone(),
),
end: KeySelector::last_less_than(to_key.clone()),
mode: options::StreamingMode::Iterator,
reverse: !*ascending,
..Default::default()
},
true,
));
}
#[cfg(feature = "rocks")]
{
*it = Some(self.db.iterator_cf(
&self.db.cf_handle(CF_INDEXES).unwrap(),
IteratorMode::From(
from_key,
if *ascending {
Direction::Forward
} else {
Direction::Reverse
},
),
));
}
it.as_mut().unwrap()
};
let mut prev_key_prefix = prev_key
.as_ref()
.and_then(|k| k.get(..k.len() - std::mem::size_of::<u32>()))
.unwrap_or_default();
if let Some(prev_item) = prev_item.take() {
if let Some(next_it_opts) = &mut next_it_opts {
next_it_opts.remaining.insert(prev_item);
} else {
doc_id = prev_item;
break 'inner;
}
}
let mut is_eof = false;
loop {
if let Some(result) = it.next().await {
let key = result?.key().to_vec().into_boxed_slice();
/*let (key, _) = result.map_err(|e| {
Error::InternalError(format!("Iterator error: {}", e))
})?;
if !key.starts_with(prefix) {
*prev_key = None;
is_eof = true;
break;
}*/
doc_id = key
.as_ref()
.deserialize_be_u32(key.len() - std::mem::size_of::<u32>())?;
if it_opts.remaining.contains(doc_id) {
it_opts.remaining.remove(doc_id);
if let Some(next_it_opts) = &mut next_it_opts {
if let Some(prev_key_) = &*prev_key {
if key.len() != prev_key_.len()
|| !key.starts_with(prev_key_prefix)
{
*prev_item = Some(doc_id);
*prev_key = Some(key);
break;
}
} else {
*prev_key = Some(key);
prev_key_prefix = prev_key
.as_ref()
.and_then(|key| {
key.get(
..key.len() - std::mem::size_of::<u32>(),
)
})
.ok_or_else(|| {
Error::InternalError(
"Invalid index entry".to_string(),
)
})?;
}
next_it_opts.remaining.insert(doc_id);
} else {
// doc id found
break 'inner;
}
}
} else {
is_eof = true;
// Add remaining items not present in the index
if !results.is_empty() && !paginate.is_full() {
for document_id in results {
if !paginate.add(document_id) {
break;
}
}
if is_eof {
if let Some(next_it_opts) = &mut next_it_opts {
if !it_opts.remaining.is_empty() {
next_it_opts.remaining |= &it_opts.remaining;
it_opts.remaining.clear();
}
*prev_key = None;
it_opts.eof = true;
}
}
}
IndexType::DocumentSet { set, it } => {
if let Some(it) = it {
if let Some(_doc_id) = it.next() {
doc_id = _doc_id;
break 'inner;
}
Comparator::DocumentSet { set, ascending } => {
let in_set = &result_set.results & &set;
let not_in_set = &result_set.results ^ &in_set;
let sets = if ascending {
[in_set, not_in_set]
} else {
[not_in_set, in_set]
};
'outer: for set in sets {
for document_id in set {
if !paginate.add(document_id) {
break 'outer;
}
} else {
let mut set = set.clone();
set.bitand_assign(&it_opts.remaining);
let set_len = set.len();
if set_len > 0 {
it_opts.remaining.bitxor_assign(&set);
match &mut next_it_opts {
Some(next_it_opts) if set_len > 1 => {
next_it_opts.remaining = set;
}
_ if set_len == 1 => {
doc_id = set.min().unwrap();
break 'inner;
}
_ => {
let mut it_ = set.into_iter();
let result = it_.next();
*it = Some(it_);
if let Some(result) = result {
doc_id = result;
break 'inner;
} else {
break 'outer;
}
}
}
} else if !it_opts.remaining.is_empty() {
if let Some(ref mut next_it_opts) = next_it_opts {
next_it_opts.remaining = std::mem::take(&mut it_opts.remaining);
}
}
};
}
};
if let Some(next_it_opts) = next_it_opts {
if !next_it_opts.remaining.is_empty() {
if next_it_opts.remaining.len() == 1 {
doc_id = next_it_opts.remaining.min().unwrap();
next_it_opts.remaining.remove(doc_id);
break 'inner;
} else {
match &mut next_it_opts.index {
IndexType::DB {
it,
from_key,
to_key,
ascending,
prev_item,
prev_key,
..
} => {
if let Some(it) = it {
#[cfg(feature = "rocks")]
{
*it = self.db.iterator_cf(
&self.db.cf_handle(CF_INDEXES).unwrap(),
IteratorMode::From(
from_key,
if *ascending {
Direction::Forward
} else {
Direction::Reverse
},
),
);
}
#[cfg(feature = "foundation")]
{
*it = trx.trx.get_ranges_keyvalues(
RangeOption {
begin: KeySelector::first_greater_or_equal(
from_key.clone(),
),
end: KeySelector::last_less_than(
to_key.clone(),
),
mode: options::StreamingMode::Iterator,
reverse: !*ascending,
..Default::default()
},
true,
);
}
}
*prev_item = None;
*prev_key = None;
}
IndexType::DocumentSet { it, .. } => {
*it = None;
}
}
current += 1;
next_it_opts.eof = false;
continue 'inner;
}
}
}
}
} else {
let mut trx = self.read_transaction().await?;
let mut sorted_ids = AHashMap::with_capacity(paginate.limit);
it_opts.eof = true;
for (pos, comparator) in comparators.into_iter().take(4).enumerate() {
match comparator {
Comparator::Field { field, ascending } => {
let mut results = result_set.results.clone();
let mut prev_data = vec![];
let mut has_grouped_ids = false;
let mut idx = 0;
if it_opts.remaining.is_empty() {
if current > 0 {
current -= 1;
} else {
break 'outer;
trx.refresh_if_old().await?;
trx.sort_index(
result_set.account_id,
result_set.collection,
field,
ascending,
|data, document_id| {
if results.remove(document_id) {
debug_assert!(!data.is_empty());
if data != prev_data {
idx += 1;
prev_data = data.to_vec();
} else {
has_grouped_ids = true;
}
sorted_ids.entry(document_id).or_insert([0u32; 4])[pos] = idx;
!results.is_empty()
} else {
true
}
},
)
.await?;
// Add remaining items not present in the index
if !results.is_empty() {
idx += 1;
for document_id in results {
sorted_ids.entry(document_id).or_insert([0u32; 4])[pos] = idx;
}
}
if !has_grouped_ids {
// If we are sorting by multiple fields and we don't have grouped ids, we can
// stop here
break;
}
}
Comparator::DocumentSet { set, ascending } => {
let in_set = &result_set.results & &set;
let not_in_set = &result_set.results ^ &in_set;
let sets = if ascending {
[(in_set, 0), (not_in_set, 1)]
} else {
[(not_in_set, 0), (in_set, 1)]
};
for (document_ids, idx) in sets {
for document_id in document_ids {
sorted_ids.entry(document_id).or_insert([0u32; 4])[pos] = idx;
}
}
}
}
}
// Pagination
if !has_anchor {
if position >= 0 {
if position > 0 {
position -= 1;
} else {
result.ids.push(doc_id);
if limit > 0 && result.ids.len() == limit {
break 'outer;
}
}
} else {
result.ids.push(doc_id);
let mut sorted_ids = sorted_ids.into_iter().collect::<Vec<_>>();
sorted_ids.sort_by(|a, b| a.1.cmp(&b.1));
for (document_id, _) in sorted_ids {
if !paginate.add(document_id) {
break;
}
} else if anchor_offset >= 0 {
if !anchor_found {
if &doc_id != anchor.as_ref().unwrap() {
continue 'outer;
}
anchor_found = true;
}
if anchor_offset > 0 {
anchor_offset -= 1;
} else {
result.ids.push(doc_id);
if limit > 0 && result.ids.len() == limit {
break 'outer;
}
}
} else {
anchor_found = &doc_id == anchor.as_ref().unwrap();
result.ids.push(doc_id);
if !anchor_found {
continue 'outer;
}
position = anchor_offset;
break 'outer;
}
}
if !has_anchor || anchor_found {
if !has_anchor && requested_position >= 0 {
result.position = if position == 0 { requested_position } else { 0 };
} else if position >= 0 {
result.position = position;
Ok(paginate.build())
}
}
pub struct Pagination {
requested_position: i32,
position: i32,
limit: usize,
anchor: u32,
anchor_offset: i32,
has_anchor: bool,
anchor_found: bool,
ids: Vec<u32>,
}
impl Pagination {
pub fn new(limit: usize, position: i32, anchor: Option<u32>, anchor_offset: i32) -> Self {
let (has_anchor, anchor) = anchor.map(|anchor| (true, anchor)).unwrap_or((false, 0));
Self {
requested_position: position,
position,
limit,
anchor,
anchor_offset,
has_anchor,
anchor_found: false,
ids: Vec::with_capacity(limit),
}
}
pub fn add(&mut self, document_id: u32) -> bool {
// Pagination
if !self.has_anchor {
if self.position > 0 {
self.position -= 1;
} else {
let position = position.unsigned_abs() as usize;
self.ids.push(document_id);
if self.ids.len() == self.limit {
return false;
}
}
} else if self.anchor_offset >= 0 {
if !self.anchor_found {
if document_id != self.anchor {
return true;
}
self.anchor_found = true;
}
if self.anchor_offset > 0 {
self.anchor_offset -= 1;
} else {
self.ids.push(document_id);
if self.ids.len() == self.limit {
return false;
}
}
} else {
self.anchor_found = document_id == self.anchor;
self.ids.push(document_id);
if self.anchor_found {
self.position = self.anchor_offset;
return false;
}
}
true
}
pub fn is_full(&self) -> bool {
self.ids.len() == self.limit
}
pub fn build(self) -> SortedResultRet {
let mut result = SortedResultRet {
ids: self.ids,
position: 0,
found_anchor: !self.has_anchor || self.anchor_found,
};
if result.found_anchor {
if !self.has_anchor && self.requested_position >= 0 {
result.position = if self.position == 0 {
self.requested_position
} else {
0
};
} else if self.position >= 0 {
result.position = self.position;
} else {
let position = self.position.unsigned_abs() as usize;
let start_offset = if position < result.ids.len() {
result.ids.len() - position
} else {
0
};
result.position = start_offset as i32;
let end_offset = if limit > 0 {
std::cmp::min(start_offset + limit, result.ids.len())
let end_offset = if self.limit > 0 {
std::cmp::min(start_offset + self.limit, result.ids.len())
} else {
result.ids.len()
};
result.ids = result.ids[start_offset..end_offset].to_vec()
}
} else {
result.found_anchor = false;
}
Ok(result)
}
}
impl ResultSet {
pub fn from_key(&self, field: u8) -> IndexKeyPrefix {
IndexKeyPrefix {
account_id: self.account_id,
collection: self.collection,
field,
}
}
pub fn to_key(&self, field: u8) -> IndexKeyPrefix {
let (account_id, collection, field) = if field < u8::MAX {
(self.account_id, self.collection, field + 1)
} else if (self.collection) < u8::MAX {
(self.account_id, (self.collection) + 1, field)
} else {
(self.account_id + 1, self.collection, field)
};
IndexKeyPrefix {
account_id,
collection,
field,
}
result
}
}

9
src/tests/query.rs
View File

@@ -173,6 +173,7 @@ pub async fn test(db: Arc<Store>, do_insert: bool) {
}
}
}
builder.custom(fts_builder).unwrap();
documents.lock().unwrap().push(builder.build());
});
@@ -211,8 +212,8 @@ pub async fn test(db: Arc<Store>, do_insert: bool) {
println!("Insert took {} ms.", now.elapsed().as_millis());
}
//println!("Running filter tests...");
//test_filter(db.clone()).await;
println!("Running filter tests...");
test_filter(db.clone()).await;
println!("Running sort tests...");
test_sort(db).await;
@@ -420,6 +421,7 @@ pub async fn test_sort(db: Arc<Store>) {
];
for (filter, sort, expected_results) in tests {
//println!("Running test: {:?}", sort);
let mut results: Vec<String> = Vec::with_capacity(expected_results.len());
let docset = db.filter(0, COLLECTION_ID, filter).await.unwrap();
let sorted_docset = db
@@ -427,8 +429,9 @@ pub async fn test_sort(db: Arc<Store>) {
.await
.unwrap();
let db = db.read_transaction().await.unwrap();
let mut db = db.read_transaction().await.unwrap();
for document_id in sorted_docset.ids {
db.refresh_if_old().await.unwrap();
results.push(
db.get_value(ValueKey {
account_id: 0,

2
src/write/mod.rs
View File

@@ -15,10 +15,12 @@ pub struct Batch {
pub ops: Vec<Operation>,
}
#[derive(Debug)]
pub struct BatchBuilder {
pub ops: Vec<Operation>,
}
#[derive(Debug)]
pub enum Operation {
AccountId {
account_id: u32,