feat: Introduce the new Index system

This commit is contained in:
Clément Renault 2019-04-19 18:27:57 +02:00
parent e142339106
commit 3dc057ca9c
No known key found for this signature in database
GPG key ID: 0151CDAB43460DAE
8 changed files with 235 additions and 202 deletions

View file

@ -1,175 +1,134 @@
use std::error::Error;
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use fst::{map, Map, IntoStreamer, Streamer};
use fst::raw::Fst;
use std::collections::BTreeMap;
use fst::{set, IntoStreamer, Streamer};
use sdset::{Set, SetBuf, SetOperation};
use sdset::duo::{Union, DifferenceByKey};
use sdset::{Set, SetOperation};
use crate::{DocIndex, DocumentId};
use crate::shared_data_cursor::{SharedDataCursor, FromSharedDataCursor};
use crate::write_to_bytes::WriteToBytes;
use crate::data::{DocIndexes, DocIndexesBuilder};
use crate::{DocumentId, DocIndex};
pub type Word = Vec<u8>; // TODO should be a smallvec
#[derive(Default)]
pub struct Index {
pub map: Map,
pub indexes: DocIndexes,
pub trait Store: Clone {
type Error: std::error::Error;
fn get_fst(&self) -> Result<fst::Set, Self::Error>;
fn set_fst(&self, set: &fst::Set) -> Result<(), Self::Error>;
fn get_indexes(&self, word: &[u8]) -> Result<Option<SetBuf<DocIndex>>, Self::Error>;
fn set_indexes(&self, word: &[u8], indexes: &Set<DocIndex>) -> Result<(), Self::Error>;
fn del_indexes(&self, word: &[u8]) -> Result<(), Self::Error>;
}
impl Index {
pub fn remove_documents(&self, documents: &Set<DocumentId>) -> Index {
pub struct Index<S> {
pub set: fst::Set,
pub store: S,
}
impl<S> Index<S>
where S: Store,
{
pub fn from_store(store: S) -> Result<Index<S>, S::Error> {
let set = store.get_fst()?;
Ok(Index { set, store })
}
pub fn remove_documents(&self, documents: &Set<DocumentId>) -> Result<Index<S>, S::Error> {
let mut buffer = Vec::new();
let mut builder = IndexBuilder::new();
let mut builder = fst::SetBuilder::memory();
let mut stream = self.into_stream();
while let Some((key, indexes)) = stream.next() {
buffer.clear();
while let Some((input, result)) = stream.next() {
let indexes = match result? {
Some(indexes) => indexes,
None => continue,
};
let op = DifferenceByKey::new(indexes, documents, |x| x.document_id, |x| *x);
let op = DifferenceByKey::new(&indexes, documents, |x| x.document_id, |x| *x);
buffer.clear();
op.extend_vec(&mut buffer);
if !buffer.is_empty() {
if buffer.is_empty() {
self.store.del_indexes(input)?;
} else {
builder.insert(input).unwrap();
let indexes = Set::new_unchecked(&buffer);
builder.insert(key, indexes).unwrap();
self.store.set_indexes(input, indexes)?;
}
}
builder.build()
let set = builder.into_inner().and_then(fst::Set::from_bytes).unwrap();
self.store.set_fst(&set)?;
Ok(Index { set, store: self.store.clone() })
}
pub fn union(&self, other: &Index) -> Index {
let mut builder = IndexBuilder::new();
let mut stream = map::OpBuilder::new().add(&self.map).add(&other.map).union();
pub fn insert_indexes(&self, map: BTreeMap<Word, SetBuf<DocIndex>>) -> Result<Index<S>, S::Error> {
let mut buffer = Vec::new();
while let Some((key, ivalues)) = stream.next() {
buffer.clear();
match ivalues {
[a, b] => {
let indexes = if a.index == 0 { &self.indexes } else { &other.indexes };
let indexes = &indexes[a.value as usize];
let a = Set::new_unchecked(indexes);
let mut builder = fst::SetBuilder::memory();
let set = fst::Set::from_iter(map.keys()).unwrap();
let mut union_ = self.set.op().add(&set).r#union();
let indexes = if b.index == 0 { &self.indexes } else { &other.indexes };
let indexes = &indexes[b.value as usize];
let b = Set::new_unchecked(indexes);
while let Some(input) = union_.next() {
let remote = self.store.get_indexes(input)?;
let locale = map.get(input);
let op = Union::new(a, b);
op.extend_vec(&mut buffer);
match (remote, locale) {
(Some(remote), Some(locale)) => {
buffer.clear();
Union::new(&remote, &locale).extend_vec(&mut buffer);
let indexes = Set::new_unchecked(&buffer);
if !indexes.is_empty() {
self.store.set_indexes(input, indexes)?;
builder.insert(input).unwrap();
} else {
self.store.del_indexes(input)?;
}
},
[x] => {
let indexes = if x.index == 0 { &self.indexes } else { &other.indexes };
let indexes = &indexes[x.value as usize];
buffer.extend_from_slice(indexes)
(None, Some(locale)) => {
self.store.set_indexes(input, &locale)?;
builder.insert(input).unwrap();
},
_ => continue,
}
if !buffer.is_empty() {
let indexes = Set::new_unchecked(&buffer);
builder.insert(key, indexes).unwrap();
(Some(_), None) => {
builder.insert(input).unwrap();
},
(None, None) => unreachable!(),
}
}
builder.build()
let set = builder.into_inner().and_then(fst::Set::from_bytes).unwrap();
self.store.set_fst(&set)?;
Ok(Index { set, store: self.store.clone() })
}
}
impl FromSharedDataCursor for Index {
type Error = Box<Error>;
fn from_shared_data_cursor(cursor: &mut SharedDataCursor) -> Result<Index, Self::Error> {
let len = cursor.read_u64::<LittleEndian>()? as usize;
let data = cursor.extract(len);
let fst = Fst::from_shared_bytes(data.bytes, data.offset, data.len)?;
let map = Map::from(fst);
let indexes = DocIndexes::from_shared_data_cursor(cursor)?;
Ok(Index { map, indexes})
}
pub struct Stream<'m, S> {
set_stream: set::Stream<'m>,
store: &'m S,
}
impl WriteToBytes for Index {
fn write_to_bytes(&self, bytes: &mut Vec<u8>) {
let slice = self.map.as_fst().as_bytes();
let len = slice.len() as u64;
let _ = bytes.write_u64::<LittleEndian>(len);
bytes.extend_from_slice(slice);
self.indexes.write_to_bytes(bytes);
}
}
impl<'m, 'a> IntoStreamer<'a> for &'m Index {
type Item = (&'a [u8], &'a Set<DocIndex>);
type Into = Stream<'m>;
fn into_stream(self) -> Self::Into {
Stream {
map_stream: self.map.into_stream(),
indexes: &self.indexes,
}
}
}
pub struct Stream<'m> {
map_stream: map::Stream<'m>,
indexes: &'m DocIndexes,
}
impl<'m, 'a> Streamer<'a> for Stream<'m> {
type Item = (&'a [u8], &'a Set<DocIndex>);
impl<'m, 'a, S> Streamer<'a> for Stream<'m, S>
where S: 'a + Store,
{
type Item = (&'a [u8], Result<Option<SetBuf<DocIndex>>, S::Error>);
fn next(&'a mut self) -> Option<Self::Item> {
match self.map_stream.next() {
Some((input, index)) => {
let indexes = &self.indexes[index as usize];
let indexes = Set::new_unchecked(indexes);
Some((input, indexes))
},
match self.set_stream.next() {
Some(input) => Some((input, self.store.get_indexes(input))),
None => None,
}
}
}
pub struct IndexBuilder {
map: fst::MapBuilder<Vec<u8>>,
indexes: DocIndexesBuilder<Vec<u8>>,
value: u64,
}
impl<'m, 'a, S> IntoStreamer<'a> for &'m Index<S>
where S: 'a + Store,
{
type Item = (&'a [u8], Result<Option<SetBuf<DocIndex>>, S::Error>);
type Into = Stream<'m, S>;
impl IndexBuilder {
pub fn new() -> Self {
IndexBuilder {
map: fst::MapBuilder::memory(),
indexes: DocIndexesBuilder::memory(),
value: 0,
fn into_stream(self) -> Self::Into {
Stream {
set_stream: self.set.into_stream(),
store: &self.store,
}
}
/// If a key is inserted that is less than or equal to any previous key added,
/// then an error is returned. Similarly, if there was a problem writing
/// to the underlying writer, an error is returned.
// FIXME what if one write doesn't work but the other do ?
pub fn insert<K>(&mut self, key: K, indexes: &Set<DocIndex>) -> fst::Result<()>
where K: AsRef<[u8]>,
{
self.map.insert(key, self.value)?;
self.indexes.insert(indexes);
self.value += 1;
Ok(())
}
pub fn build(self) -> Index {
let map = self.map.into_inner().unwrap();
let indexes = self.indexes.into_inner().unwrap();
let map = Map::from_bytes(map).unwrap();
let indexes = DocIndexes::from_bytes(indexes).unwrap();
Index { map, indexes }
}
}

View file

@ -13,16 +13,19 @@ use serde::{Serialize, Deserialize};
use slice_group_by::GroupBy;
use rayon::slice::ParallelSliceMut;
use zerocopy::{AsBytes, FromBytes};
pub use self::index::{Index, IndexBuilder};
pub use self::index::{Index, Store};
pub use self::query_builder::{QueryBuilder, DistinctQueryBuilder};
/// Represent an internally generated document unique identifier.
///
/// It is used to inform the database the document you want to deserialize.
/// Helpful for custom ranking.
#[derive(Serialize, Deserialize)]
#[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
#[derive(Serialize, Deserialize)]
#[derive(AsBytes, FromBytes)]
#[repr(C)]
pub struct DocumentId(pub u64);
/// This structure represent the position of a word
@ -31,6 +34,7 @@ pub struct DocumentId(pub u64);
/// This is stored in the map, generated at index time,
/// extracted and interpreted at search time.
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[derive(AsBytes, FromBytes)]
#[repr(C)]
pub struct DocIndex {
/// The document identifier where the word was found.

View file

@ -14,8 +14,8 @@ use log::info;
use crate::automaton::{self, DfaExt, AutomatonExt};
use crate::distinct_map::{DistinctMap, BufferedDistinctMap};
use crate::criterion::Criteria;
use crate::{raw_documents_from_matches, RawDocument, Document};
use crate::{Index, Match, DocumentId};
use crate::raw_documents_from_matches;
use crate::{Match, DocumentId, Index, Store, RawDocument, Document};
fn generate_automatons(query: &str) -> Vec<DfaExt> {
let has_end_whitespace = query.chars().last().map_or(false, char::is_whitespace);
@ -82,16 +82,18 @@ impl<'c, I, FI> QueryBuilder<'c, I, FI>
}
}
impl<'c, I, FI> QueryBuilder<'c, I, FI>
where I: Deref<Target=Index>,
impl<'c, I, FI, S> QueryBuilder<'c, I, FI>
where I: Deref<Target=Index<S>>,
S: Store,
{
fn query_all(&self, query: &str) -> Vec<RawDocument> {
let automatons = generate_automatons(query);
let fst = self.index.set.as_fst();
let mut stream = {
let mut op_builder = fst::map::OpBuilder::new();
let mut op_builder = fst::raw::OpBuilder::new();
for automaton in &automatons {
let stream = self.index.map.search(automaton);
let stream = fst.search(automaton);
op_builder.push(stream);
}
op_builder.r#union()
@ -105,10 +107,12 @@ where I: Deref<Target=Index>,
let distance = automaton.eval(input).to_u8();
let is_exact = distance == 0 && input.len() == automaton.query_len();
let doc_indexes = &self.index.indexes;
let doc_indexes = &doc_indexes[iv.value as usize];
// let doc_indexes = &self.index.indexes;
// let doc_indexes = &doc_indexes[iv.value as usize];
for di in doc_indexes {
let doc_indexes = self.index.store.get_indexes(input).unwrap().unwrap();
for di in doc_indexes.as_slice() {
if self.searchable_attrs.as_ref().map_or(true, |r| r.contains(&di.attribute)) {
let match_ = Match {
query_index: iv.index as u32,
@ -135,9 +139,10 @@ where I: Deref<Target=Index>,
}
}
impl<'c, I, FI> QueryBuilder<'c, I, FI>
where I: Deref<Target=Index>,
impl<'c, I, FI, S> QueryBuilder<'c, I, FI>
where I: Deref<Target=Index<S>>,
FI: Fn(DocumentId) -> bool,
S: Store,
{
pub fn query(self, query: &str, range: Range<usize>) -> Vec<Document> {
// We delegate the filter work to the distinct query builder,
@ -212,11 +217,12 @@ impl<'c, I, FI, FD> DistinctQueryBuilder<'c, I, FI, FD>
}
}
impl<'c, I, FI, FD, K> DistinctQueryBuilder<'c, I, FI, FD>
where I: Deref<Target=Index>,
impl<'c, I, FI, FD, K, S> DistinctQueryBuilder<'c, I, FI, FD>
where I: Deref<Target=Index<S>>,
FI: Fn(DocumentId) -> bool,
FD: Fn(DocumentId) -> Option<K>,
K: Hash + Eq,
S: Store,
{
pub fn query(self, query: &str, range: Range<usize>) -> Vec<Document> {
let start = Instant::now();