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

result references

Browse Source
This commit is contained in:
Mauro D
2023-04-09 16:11:03 +00:00
parent 75bf0a240a
commit 633db406f7
34 changed files with 1441 additions and 253 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
crates/store/src/query/filter.rs
View File

@@ -5,10 +5,10 @@ use roaring::RoaringBitmap;
use crate::{
fts::{builder::MAX_TOKEN_LENGTH, tokenizers::space::SpaceTokenizer},
BitmapKey, ReadTransaction, Store, BM_KEYWORD,
BitmapKey, ReadTransaction, Store,
};
use super::{Filter, ResultSet};
use super::{Filter, ResultSet, TextMatch};
struct State {
op: Filter,
@@ -30,7 +30,7 @@ impl ReadTransaction<'_> {
account_id,
collection,
results: self
.get_bitmap(BitmapKey::new_document_ids(account_id, collection))
.get_bitmap(BitmapKey::document_ids(account_id, collection))
.await?
.unwrap_or_else(RoaringBitmap::new),
});
@@ -44,40 +44,32 @@ impl ReadTransaction<'_> {
self.refresh_if_old().await?;
let result = match filter {
Filter::HasKeyword { field, value } => {
self.get_bitmap(BitmapKey {
account_id,
collection,
family: BM_KEYWORD,
field,
key: value.as_bytes(),
block_num: 0,
})
.await?
}
Filter::HasKeywords { field, value } => {
self.get_bitmaps_intersection(
SpaceTokenizer::new(&value, MAX_TOKEN_LENGTH)
.collect::<HashSet<String>>()
.into_iter()
.map(|word| BitmapKey::hash(&word, account_id, collection, 0, field))
.collect(),
)
.await?
}
Filter::MatchValue { field, op, value } => {
self.range_to_bitmap(account_id, collection, field, value, op)
.await?
}
Filter::HasText {
field,
text,
language,
match_phrase,
} => {
self.fts_query(account_id, collection, field, &text, language, match_phrase)
Filter::HasText { field, text, op } => match op {
TextMatch::Exact(language) => {
self.fts_query(account_id, collection, field, &text, language, true)
.await?
}
TextMatch::Stemmed(language) => {
self.fts_query(account_id, collection, field, &text, language, false)
.await?
}
TextMatch::Tokenized => {
self.get_bitmaps_intersection(
SpaceTokenizer::new(&text, MAX_TOKEN_LENGTH)
.collect::<HashSet<String>>()
.into_iter()
.map(|word| {
BitmapKey::hash(&word, account_id, collection, 0, field)
})
.collect(),
)
.await?
}
}
},
Filter::InBitmap { family, field, key } => {
self.get_bitmap(BitmapKey {
account_id,
@@ -108,7 +100,7 @@ impl ReadTransaction<'_> {
if matches!(state.op, Filter::Not) && !not_fetch {
not_mask = self
.get_bitmap(BitmapKey::new_document_ids(account_id, collection))
.get_bitmap(BitmapKey::document_ids(account_id, collection))
.await?
.unwrap_or_else(RoaringBitmap::new);
not_fetch = true;

20
crates/store/src/query/get.rs
View File

@@ -1,4 +1,6 @@
use crate::{Deserialize, Key, Store, ValueKey};
use roaring::RoaringBitmap;
use crate::{BitmapKey, Deserialize, Key, Store, ValueKey};
impl Store {
pub async fn get_value<U>(&self, key: ValueKey) -> crate::Result<Option<U>>
@@ -72,6 +74,22 @@ impl Store {
}
}
pub async fn get_bitmap<T: AsRef<[u8]> + Send + Sync + 'static>(
&self,
key: BitmapKey<T>,
) -> crate::Result<Option<RoaringBitmap>> {
#[cfg(feature = "is_async")]
{
self.read_transaction().await?.get_bitmap(key).await
}
#[cfg(feature = "is_sync")]
{
let trx = self.read_transaction()?;
self.spawn_worker(move || trx.get_bitmap(key)).await
}
}
pub async fn iterate<T: Sync + Send + 'static>(
&self,
acc: T,

118
crates/store/src/query/mod.rs
View File

@@ -7,7 +7,8 @@ use roaring::RoaringBitmap;
use crate::{
fts::{lang::LanguageDetector, Language},
BitmapKey, Serialize, BM_DOCUMENT_IDS,
write::IntoBitmap,
BitmapKey, Serialize, BM_DOCUMENT_IDS, BM_KEYWORD,
};
#[derive(Debug, Clone, Copy)]
@@ -21,14 +22,6 @@ pub enum Operator {
#[derive(Debug)]
pub enum Filter {
HasKeyword {
field: u8,
value: String,
},
HasKeywords {
field: u8,
value: String,
},
MatchValue {
field: u8,
op: Operator,
@@ -37,8 +30,7 @@ pub enum Filter {
HasText {
field: u8,
text: String,
language: Language,
match_phrase: bool,
op: TextMatch,
},
InBitmap {
family: u8,
@@ -52,6 +44,13 @@ pub enum Filter {
End,
}
#[derive(Debug)]
pub enum TextMatch {
Exact(Language),
Stemmed(Language),
Tokenized,
}
#[derive(Debug)]
pub enum Comparator {
Field { field: u8, ascending: bool },
@@ -65,9 +64,9 @@ pub struct ResultSet {
pub results: RoaringBitmap,
}
pub struct SortedResultRet {
pub struct SortedResultSet {
pub position: i32,
pub ids: Vec<u32>,
pub ids: Vec<u64>,
pub found_anchor: bool,
}
@@ -120,57 +119,80 @@ impl Filter {
}
}
pub fn has_keyword(field: impl Into<u8>, value: impl Into<String>) -> Self {
Filter::HasKeyword {
pub fn has_keyword(field: impl Into<u8>, value: impl Serialize) -> Self {
Filter::InBitmap {
family: BM_KEYWORD,
field: field.into(),
value: value.into(),
key: value.serialize(),
}
}
pub fn has_keywords(field: impl Into<u8>, value: impl Into<String>) -> Self {
Filter::HasKeywords {
field: field.into(),
value: value.into(),
}
}
pub fn match_text(
field: impl Into<u8>,
text: impl Into<String>,
mut language: Language,
) -> Self {
pub fn has_text(field: impl Into<u8>, text: impl Into<String>, mut language: Language) -> Self {
let mut text = text.into();
let match_phrase = (text.starts_with('"') && text.ends_with('"'))
|| (text.starts_with('\'') && text.ends_with('\''));
let op = if !matches!(language, Language::None) {
let match_phrase = (text.starts_with('"') && text.ends_with('"'))
|| (text.starts_with('\'') && text.ends_with('\''));
if !match_phrase && language == Language::Unknown {
language = if let Some((l, t)) = text
.split_once(':')
.and_then(|(l, t)| (Language::from_iso_639(l)?, t.to_string()).into())
{
text = t;
l
if !match_phrase && language == Language::Unknown {
language = if let Some((l, t)) = text
.split_once(':')
.and_then(|(l, t)| (Language::from_iso_639(l)?, t.to_string()).into())
{
text = t;
l
} else {
LanguageDetector::detect_single(&text)
.and_then(|(l, c)| if c > 0.3 { Some(l) } else { None })
.unwrap_or(Language::Unknown)
};
}
if match_phrase {
TextMatch::Exact(language)
} else {
LanguageDetector::detect_single(&text)
.and_then(|(l, c)| if c > 0.3 { Some(l) } else { None })
.unwrap_or(Language::Unknown)
};
}
TextMatch::Stemmed(language)
}
} else {
TextMatch::Tokenized
};
Filter::HasText {
field: field.into(),
text,
language,
match_phrase,
op,
}
}
pub fn match_english(field: impl Into<u8>, text: impl Into<String>) -> Self {
Self::match_text(field, text, Language::English)
pub fn has_english_text(field: impl Into<u8>, text: impl Into<String>) -> Self {
Self::has_text(field, text, Language::English)
}
pub fn is_in_bitmap(field: impl Into<u8>, value: impl IntoBitmap) -> Self {
let (key, family) = value.into_bitmap();
Self::InBitmap {
family,
field: field.into(),
key,
}
}
pub fn is_in_set(set: RoaringBitmap) -> Self {
Filter::DocumentSet(set)
}
}
impl Comparator {
pub fn field(field: impl Into<u8>, ascending: bool) -> Self {
Self::Field {
field: field.into(),
ascending,
}
}
pub fn set(set: RoaringBitmap, ascending: bool) -> Self {
Self::DocumentSet { set, ascending }
}
pub fn ascending(field: impl Into<u8>) -> Self {
Self::Field {
field: field.into(),
@@ -187,10 +209,10 @@ impl Comparator {
}
impl BitmapKey<&'static [u8]> {
pub fn new_document_ids(account_id: u32, collection: u8) -> Self {
pub fn document_ids(account_id: u32, collection: impl Into<u8>) -> Self {
BitmapKey {
account_id,
collection,
collection: collection.into(),
family: BM_DOCUMENT_IDS,
field: u8::MAX,
key: b"",

129
crates/store/src/query/sort.rs
View File

@@ -1,8 +1,8 @@
use ahash::AHashMap;
use ahash::{AHashMap, AHashSet};
use crate::{ReadTransaction, Store};
use crate::{ReadTransaction, Store, ValueKey};
use super::{Comparator, ResultSet, SortedResultRet};
use super::{Comparator, ResultSet, SortedResultSet};
pub struct Pagination {
requested_position: i32,
@@ -12,7 +12,9 @@ pub struct Pagination {
anchor_offset: i32,
has_anchor: bool,
anchor_found: bool,
ids: Vec<u32>,
ids: Vec<u64>,
prefix_key: Option<ValueKey>,
prefix_unique: bool,
}
impl ReadTransaction<'_> {
@@ -21,14 +23,9 @@ impl ReadTransaction<'_> {
&mut self,
result_set: ResultSet,
mut comparators: Vec<Comparator>,
limit: usize,
position: i32,
anchor: Option<u32>,
anchor_offset: i32,
) -> crate::Result<SortedResultRet> {
let mut paginate = Pagination::new(limit, position, anchor, anchor_offset);
if comparators.len() == 1 {
mut paginate: Pagination,
) -> crate::Result<SortedResultSet> {
if comparators.len() == 1 && !paginate.prefix_unique {
match comparators.pop().unwrap() {
Comparator::Field { field, ascending } => {
let mut results = result_set.results;
@@ -38,14 +35,16 @@ impl ReadTransaction<'_> {
result_set.collection,
field,
ascending,
|_, document_id| !results.remove(document_id) || paginate.add(document_id),
|_, document_id| {
!results.remove(document_id) || paginate.add(0, document_id)
},
)
.await?;
// 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) {
if !paginate.add(0, document_id) {
break;
}
}
@@ -61,13 +60,28 @@ impl ReadTransaction<'_> {
};
'outer: for set in sets {
for document_id in set {
if !paginate.add(document_id) {
if !paginate.add(0, document_id) {
break 'outer;
}
}
}
}
}
// Obtain prefixes
let prefix_key = paginate.prefix_key.take();
let mut sorted_results = paginate.build();
if let Some(prefix_key) = prefix_key {
for id in sorted_results.ids.iter_mut() {
if let Some(prefix_id) =
self.get_value::<u32>(prefix_key.with_document_id(*id as u32))?
{
*id |= (prefix_id as u64) << 32;
}
}
}
Ok(sorted_results)
} else {
let mut sorted_ids = AHashMap::with_capacity(paginate.limit);
@@ -138,16 +152,35 @@ impl ReadTransaction<'_> {
}
}
let mut seen_prefixes = AHashSet::new();
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) {
// Obtain document prefixId
let prefix_id = if let Some(prefix_key) = &paginate.prefix_key {
if let Some(prefix_id) =
self.get_value(prefix_key.with_document_id(document_id))?
{
if paginate.prefix_unique && !seen_prefixes.insert(prefix_id) {
continue;
}
prefix_id
} else {
// Document no longer exists?
continue;
}
} else {
0
};
// Add document to results
if !paginate.add(prefix_id, document_id) {
break;
}
}
}
Ok(paginate.build())
Ok(paginate.build())
}
}
}
@@ -156,15 +189,12 @@ impl Store {
&self,
result_set: ResultSet,
comparators: Vec<Comparator>,
limit: usize,
position: i32,
anchor: Option<u32>,
anchor_offset: i32,
) -> crate::Result<SortedResultRet> {
let limit = match (result_set.results.len(), limit) {
mut paginate: Pagination,
) -> crate::Result<SortedResultSet> {
paginate.limit = match (result_set.results.len(), paginate.limit) {
(0, _) => {
return Ok(SortedResultRet {
position,
return Ok(SortedResultSet {
position: paginate.position,
ids: vec![],
found_anchor: true,
});
@@ -177,37 +207,28 @@ impl Store {
{
self.read_transaction()
.await?
.sort(
result_set,
comparators,
limit,
position,
anchor,
anchor_offset,
)
.sort(result_set, comparators, paginate)
.await
}
#[cfg(feature = "is_sync")]
{
let mut trx = self.read_transaction()?;
self.spawn_worker(move || {
trx.sort(
result_set,
comparators,
limit,
position,
anchor,
anchor_offset,
)
})
.await
self.spawn_worker(move || trx.sort(result_set, comparators, paginate))
.await
}
}
}
impl Pagination {
pub fn new(limit: usize, position: i32, anchor: Option<u32>, anchor_offset: i32) -> Self {
pub fn new(
limit: usize,
position: i32,
anchor: Option<u32>,
anchor_offset: i32,
prefix_key: Option<ValueKey>,
prefix_unique: bool,
) -> Self {
let (has_anchor, anchor) = anchor.map(|anchor| (true, anchor)).unwrap_or((false, 0));
Self {
@@ -219,16 +240,20 @@ impl Pagination {
has_anchor,
anchor_found: false,
ids: Vec::with_capacity(limit),
prefix_key,
prefix_unique,
}
}
pub fn add(&mut self, document_id: u32) -> bool {
pub fn add(&mut self, prefix_id: u32, document_id: u32) -> bool {
let id = ((prefix_id as u64) << 32) | document_id as u64;
// Pagination
if !self.has_anchor {
if self.position > 0 {
self.position -= 1;
} else {
self.ids.push(document_id);
self.ids.push(id);
if self.ids.len() == self.limit {
return false;
}
@@ -244,14 +269,14 @@ impl Pagination {
if self.anchor_offset > 0 {
self.anchor_offset -= 1;
} else {
self.ids.push(document_id);
self.ids.push(id);
if self.ids.len() == self.limit {
return false;
}
}
} else {
self.anchor_found = document_id == self.anchor;
self.ids.push(document_id);
self.ids.push(id);
if self.anchor_found {
self.position = self.anchor_offset;
@@ -266,8 +291,8 @@ impl Pagination {
self.ids.len() == self.limit
}
pub fn build(self) -> SortedResultRet {
let mut result = SortedResultRet {
pub fn build(self) -> SortedResultSet {
let mut result = SortedResultSet {
ids: self.ids,
position: 0,
found_anchor: !self.has_anchor || self.anchor_found,