From ad7cfbcee69e2b92ea4f2032547866ed2a7d4c97 Mon Sep 17 00:00:00 2001 From: Mauro D Date: Thu, 30 Mar 2023 16:52:47 +0000 Subject: [PATCH] Foundation passing tests --- src/backend/foundationdb/read.rs | 117 +++--- src/fts/bloom.rs | 35 +- src/fts/builder.rs | 26 +- src/fts/query.rs | 115 ++---- src/lib.rs | 5 +- src/query/filter.rs | 33 +- src/query/mod.rs | 32 -- src/query/sort.rs | 666 ++++++++++--------------------- src/tests/query.rs | 9 +- src/write/mod.rs | 2 + 10 files changed, 363 insertions(+), 677 deletions(-) diff --git a/src/backend/foundationdb/read.rs b/src/backend/foundationdb/read.rs index 6f2542e5..efba51f1 100644 --- a/src/backend/foundationdb/read.rs +++ b/src/backend/foundationdb/read.rs @@ -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>( + async fn get_bitmap_>( &self, mut key: BitmapKey, - ) -> crate::Result> { + 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::())?, ); } - //println!("deserializing bitmap: {:?} {:?}", value.key(), bm.len()); } - Ok(if !bm.is_empty() { Some(bm) } else { None }) + Ok(()) } - #[inline(always)] - async fn get_bitmaps>( + pub async fn get_bitmap>( &self, - keys: Vec>, - ) -> crate::Result>> { - let mut results = Vec::with_capacity(keys.len()); - for key in keys { - results.push(self.get_bitmap(key).await?); - } - - Ok(results) + key: BitmapKey, + ) -> crate::Result> { + 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>( @@ -122,8 +113,8 @@ impl ReadTransaction<'_> { keys: Vec>, ) -> crate::Result> { let mut result: Option = 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>, ) -> crate::Result> { - let mut result: Option = 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> { + 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::()) - .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::(); + 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<()> { diff --git a/src/fts/bloom.rs b/src/fts/bloom.rs index 60dfb397..6e35ffde 100644 --- a/src/fts/bloom.rs +++ b/src/fts/bloom.rs @@ -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> { + 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 { + 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 { diff --git a/src/fts/builder.rs b/src/fts/builder.rs index f6487f53..14d84a4b 100644 --- a/src/fts/builder.rs +++ b/src/fts/builder.rs @@ -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::>() { - 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 { diff --git a/src/fts/query.rs b/src/fts/query.rs index 164a0082..574a9073 100644 --- a/src/fts/query.rs +++ b/src/fts/query.rs @@ -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, >::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::( - bitmaps - .iter() - .map(|document_id| ValueKey { - account_id, - collection, - document_id, - family, - field, - }) - .collect::>(), - ) - .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::();*/ - 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> { - BitmapKey { - account_id, - collection, - family: BM_BLOOM, - field, - block_num: 0, - key: self.as_high_rank_hash().serialize(), - } - } -} diff --git a/src/lib.rs b/src/lib.rs index a11c6f1d..cccdb536 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -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; diff --git a/src/query/filter.rs b/src/query/filter.rs index 58d57cc1..b037c374 100644 --- a/src/query/filter.rs +++ b/src/query/filter.rs @@ -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, ) -> crate::Result { 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, ¬_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, }) } } diff --git a/src/query/mod.rs b/src/query/mod.rs index 1117c40b..09571c00 100644 --- a/src/query/mod.rs +++ b/src/query/mod.rs @@ -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), -} - -#[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, 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 { - fn contains_id(&self, id: u32) -> bool { - self.iter().any(|&i| i == id) - } - - fn len(&self) -> usize { - self.len() - } -} diff --git a/src/query/sort.rs b/src/query/sort.rs index 3eeb3bc8..5563e08b 100644 --- a/src/query/sort.rs +++ b/src/query/sort.rs @@ -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, - }, - DB { - it: Option>>, - prefix: Vec, - from_key: Vec, - ascending: bool, - prev_item: Option, - prev_key: Option>, - }, -} - -#[cfg(feature = "rocks")] -struct IndexIterator<'x> { - index: IndexType<'x>, - remaining: RoaringBitmap, - eof: bool, -} - -#[cfg(feature = "foundation")] -enum IndexType<'x, T: Stream> + Unpin + 'x> { - DocumentSet { - set: RoaringBitmap, - it: Option, - }, - DB { - it: Option, - from_key: Vec, - to_key: Vec, - ascending: bool, - prev_item: Option, - prev_key: Option>, - phantom: std::marker::PhantomData<&'x ()>, - }, -} - -#[cfg(feature = "foundation")] -struct IndexIterator<'x, T: Stream> + Unpin + 'x> { - index: IndexType<'x, T>, - remaining: RoaringBitmap, - eof: bool, -} - impl Store { pub async fn sort( &self, - mut result_set: ResultSet, - comparators: Vec, + result_set: ResultSet, + mut comparators: Vec, limit: usize, - mut position: i32, + position: i32, anchor: Option, - mut anchor_offset: i32, + anchor_offset: i32, ) -> crate::Result { - 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::>(); - - 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::())) - .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::())?; - 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::(), - ) - }) - .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::>(); + 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, +} + +impl Pagination { + pub fn new(limit: usize, position: i32, anchor: Option, 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 } } diff --git a/src/tests/query.rs b/src/tests/query.rs index a1f526cc..ebc6f789 100644 --- a/src/tests/query.rs +++ b/src/tests/query.rs @@ -173,6 +173,7 @@ pub async fn test(db: Arc, do_insert: bool) { } } } + builder.custom(fts_builder).unwrap(); documents.lock().unwrap().push(builder.build()); }); @@ -211,8 +212,8 @@ pub async fn test(db: Arc, 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) { ]; for (filter, sort, expected_results) in tests { + //println!("Running test: {:?}", sort); let mut results: Vec = 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) { .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, diff --git a/src/write/mod.rs b/src/write/mod.rs index 56fb7600..97662bbe 100644 --- a/src/write/mod.rs +++ b/src/write/mod.rs @@ -15,10 +15,12 @@ pub struct Batch { pub ops: Vec, } +#[derive(Debug)] pub struct BatchBuilder { pub ops: Vec, } +#[derive(Debug)] pub enum Operation { AccountId { account_id: u32,