mirror of
https://github.com/meilisearch/MeiliSearch
synced 2024-11-27 07:14:26 +01:00
Introduce a function to retrieve the facet level range docids
This commit is contained in:
parent
57d253aeda
commit
9ec95679e1
62
src/heed_codec/facet/facet_level_value_f64_codec.rs
Normal file
62
src/heed_codec/facet/facet_level_value_f64_codec.rs
Normal file
@ -0,0 +1,62 @@
|
||||
use std::borrow::Cow;
|
||||
use std::convert::TryInto;
|
||||
|
||||
use crate::facet::value_encoding::f64_into_bytes;
|
||||
|
||||
// TODO do not de/serialize right bound when level = 0
|
||||
pub struct FacetLevelValueF64Codec;
|
||||
|
||||
impl<'a> heed::BytesDecode<'a> for FacetLevelValueF64Codec {
|
||||
type DItem = (u8, u8, f64, f64);
|
||||
|
||||
fn bytes_decode(bytes: &'a [u8]) -> Option<Self::DItem> {
|
||||
let (field_id, bytes) = bytes.split_first()?;
|
||||
let (level, bytes) = bytes.split_first()?;
|
||||
|
||||
let left = bytes[16..24].try_into().ok().map(f64::from_be_bytes)?;
|
||||
let right = bytes[24..].try_into().ok().map(f64::from_be_bytes)?;
|
||||
|
||||
Some((*field_id, *level, left, right))
|
||||
}
|
||||
}
|
||||
|
||||
impl heed::BytesEncode<'_> for FacetLevelValueF64Codec {
|
||||
type EItem = (u8, u8, f64, f64);
|
||||
|
||||
fn bytes_encode((field_id, level, left, right): &Self::EItem) -> Option<Cow<[u8]>> {
|
||||
let mut buffer = [0u8; 32];
|
||||
|
||||
// Write the globally ordered floats.
|
||||
let bytes = f64_into_bytes(*left)?;
|
||||
buffer[..8].copy_from_slice(&bytes[..]);
|
||||
|
||||
let bytes = f64_into_bytes(*right)?;
|
||||
buffer[8..16].copy_from_slice(&bytes[..]);
|
||||
|
||||
// Then the f64 values just to be able to read them back.
|
||||
let bytes = left.to_be_bytes();
|
||||
buffer[16..24].copy_from_slice(&bytes[..]);
|
||||
|
||||
let bytes = right.to_be_bytes();
|
||||
buffer[24..].copy_from_slice(&bytes[..]);
|
||||
|
||||
let mut bytes = Vec::with_capacity(buffer.len() + 2);
|
||||
bytes.push(*field_id);
|
||||
bytes.push(*level);
|
||||
bytes.extend_from_slice(&buffer[..]);
|
||||
Some(Cow::Owned(bytes))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use heed::{BytesEncode, BytesDecode};
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn globally_ordered_f64() {
|
||||
let bytes = FacetLevelValueF64Codec::bytes_encode(&(3, 0, -32.0, 32.0)).unwrap();
|
||||
let (name, level, left, right) = FacetLevelValueF64Codec::bytes_decode(&bytes).unwrap();
|
||||
assert_eq!((name, level, left, right), (3, 0, -32.0, 32.0));
|
||||
}
|
||||
}
|
43
src/heed_codec/facet/facet_level_value_i64_codec.rs
Normal file
43
src/heed_codec/facet/facet_level_value_i64_codec.rs
Normal file
@ -0,0 +1,43 @@
|
||||
use std::borrow::Cow;
|
||||
use std::convert::TryInto;
|
||||
|
||||
use crate::facet::value_encoding::{i64_from_bytes, i64_into_bytes};
|
||||
|
||||
pub struct FacetLevelValueI64Codec;
|
||||
|
||||
impl<'a> heed::BytesDecode<'a> for FacetLevelValueI64Codec {
|
||||
type DItem = (u8, u8, i64, i64);
|
||||
|
||||
fn bytes_decode(bytes: &'a [u8]) -> Option<Self::DItem> {
|
||||
let (field_id, bytes) = bytes.split_first()?;
|
||||
let (level, bytes) = bytes.split_first()?;
|
||||
|
||||
let left = bytes[..8].try_into().map(i64_from_bytes).ok()?;
|
||||
let right = if *level != 0 {
|
||||
bytes[8..].try_into().map(i64_from_bytes).ok()?
|
||||
} else {
|
||||
left
|
||||
};
|
||||
|
||||
Some((*field_id, *level, left, right))
|
||||
}
|
||||
}
|
||||
|
||||
impl heed::BytesEncode<'_> for FacetLevelValueI64Codec {
|
||||
type EItem = (u8, u8, i64, i64);
|
||||
|
||||
fn bytes_encode((field_id, level, left, right): &Self::EItem) -> Option<Cow<[u8]>> {
|
||||
let left = i64_into_bytes(*left);
|
||||
let right = i64_into_bytes(*right);
|
||||
|
||||
let mut bytes = Vec::with_capacity(2 + left.len() + right.len());
|
||||
bytes.push(*field_id);
|
||||
bytes.push(*level);
|
||||
bytes.extend_from_slice(&left[..]);
|
||||
if *level != 0 {
|
||||
bytes.extend_from_slice(&right[..]);
|
||||
}
|
||||
|
||||
Some(Cow::Owned(bytes))
|
||||
}
|
||||
}
|
@ -1,50 +0,0 @@
|
||||
use std::borrow::Cow;
|
||||
use std::convert::TryInto;
|
||||
|
||||
use crate::facet::value_encoding::f64_into_bytes;
|
||||
|
||||
pub struct FacetValueF64Codec;
|
||||
|
||||
impl<'a> heed::BytesDecode<'a> for FacetValueF64Codec {
|
||||
type DItem = (u8, f64);
|
||||
|
||||
fn bytes_decode(bytes: &'a [u8]) -> Option<Self::DItem> {
|
||||
let (field_id, buffer) = bytes.split_first()?;
|
||||
let value = buffer[8..].try_into().ok().map(f64::from_be_bytes)?;
|
||||
Some((*field_id, value))
|
||||
}
|
||||
}
|
||||
|
||||
impl heed::BytesEncode<'_> for FacetValueF64Codec {
|
||||
type EItem = (u8, f64);
|
||||
|
||||
fn bytes_encode((field_id, value): &Self::EItem) -> Option<Cow<[u8]>> {
|
||||
let mut buffer = [0u8; 16];
|
||||
|
||||
// Write the globally ordered float.
|
||||
let bytes = f64_into_bytes(*value)?;
|
||||
buffer[..8].copy_from_slice(&bytes[..]);
|
||||
|
||||
// Then the f64 value just to be able to read it back.
|
||||
let bytes = value.to_be_bytes();
|
||||
buffer[8..].copy_from_slice(&bytes[..]);
|
||||
|
||||
let mut bytes = Vec::with_capacity(buffer.len() + 1);
|
||||
bytes.push(*field_id);
|
||||
bytes.extend_from_slice(&buffer[..]);
|
||||
Some(Cow::Owned(bytes))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use heed::{BytesEncode, BytesDecode};
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn globally_ordered_f64() {
|
||||
let bytes = FacetValueF64Codec::bytes_encode(&(3, -32.0)).unwrap();
|
||||
let (name, value) = FacetValueF64Codec::bytes_decode(&bytes).unwrap();
|
||||
assert_eq!((name, value), (3, -32.0));
|
||||
}
|
||||
}
|
@ -1,28 +0,0 @@
|
||||
use std::borrow::Cow;
|
||||
use std::convert::TryInto;
|
||||
|
||||
use crate::facet::value_encoding::{i64_from_bytes, i64_into_bytes};
|
||||
|
||||
pub struct FacetValueI64Codec;
|
||||
|
||||
impl<'a> heed::BytesDecode<'a> for FacetValueI64Codec {
|
||||
type DItem = (u8, i64);
|
||||
|
||||
fn bytes_decode(bytes: &'a [u8]) -> Option<Self::DItem> {
|
||||
let (field_id, buffer) = bytes.split_first()?;
|
||||
let value = buffer.try_into().map(i64_from_bytes).ok()?;
|
||||
Some((*field_id, value))
|
||||
}
|
||||
}
|
||||
|
||||
impl heed::BytesEncode<'_> for FacetValueI64Codec {
|
||||
type EItem = (u8, i64);
|
||||
|
||||
fn bytes_encode((field_id, value): &Self::EItem) -> Option<Cow<[u8]>> {
|
||||
let value = i64_into_bytes(*value);
|
||||
let mut bytes = Vec::with_capacity(value.len() + 1);
|
||||
bytes.push(*field_id);
|
||||
bytes.extend_from_slice(&value[..]);
|
||||
Some(Cow::Owned(bytes))
|
||||
}
|
||||
}
|
@ -1,7 +1,7 @@
|
||||
mod facet_value_f64_codec;
|
||||
mod facet_value_i64_codec;
|
||||
mod facet_level_value_f64_codec;
|
||||
mod facet_level_value_i64_codec;
|
||||
mod facet_value_string_codec;
|
||||
|
||||
pub use self::facet_value_f64_codec::FacetValueF64Codec;
|
||||
pub use self::facet_value_i64_codec::FacetValueI64Codec;
|
||||
pub use self::facet_level_value_f64_codec::FacetLevelValueF64Codec;
|
||||
pub use self::facet_level_value_i64_codec::FacetLevelValueI64Codec;
|
||||
pub use self::facet_value_string_codec::FacetValueStringCodec;
|
||||
|
113
src/search.rs
113
src/search.rs
@ -4,6 +4,7 @@ use std::fmt;
|
||||
|
||||
use anyhow::{bail, ensure, Context};
|
||||
use fst::{IntoStreamer, Streamer};
|
||||
use heed::types::DecodeIgnore;
|
||||
use levenshtein_automata::DFA;
|
||||
use levenshtein_automata::LevenshteinAutomatonBuilder as LevBuilder;
|
||||
use log::debug;
|
||||
@ -11,7 +12,7 @@ use once_cell::sync::Lazy;
|
||||
use roaring::bitmap::RoaringBitmap;
|
||||
|
||||
use crate::facet::FacetType;
|
||||
use crate::heed_codec::{CboRoaringBitmapCodec, facet::FacetValueI64Codec};
|
||||
use crate::heed_codec::{facet::FacetLevelValueI64Codec, CboRoaringBitmapCodec};
|
||||
use crate::mdfs::Mdfs;
|
||||
use crate::query_tokens::{QueryTokens, QueryToken};
|
||||
use crate::{Index, DocumentId};
|
||||
@ -238,31 +239,107 @@ impl<'a> Search<'a> {
|
||||
// We create the original candidates with the facet conditions results.
|
||||
let facet_candidates = match self.facet_condition {
|
||||
Some(FacetCondition::Operator(fid, operator)) => {
|
||||
use std::ops::Bound::{Included, Excluded};
|
||||
use std::ops::Bound::{self, Included, Excluded};
|
||||
use FacetOperator::*;
|
||||
// Make sure we always bound the ranges with the field id, as the facets
|
||||
// values are all in the same database and prefixed by the field id.
|
||||
let range = match operator {
|
||||
GreaterThan(val) => (Excluded((fid, val)), Included((fid, i64::MAX))),
|
||||
GreaterThanOrEqual(val) => (Included((fid, val)), Included((fid, i64::MAX))),
|
||||
LowerThan(val) => (Included((fid, i64::MIN)), Excluded((fid, val))),
|
||||
LowerThanOrEqual(val) => (Included((fid, i64::MIN)), Included((fid, val))),
|
||||
Equal(val) => (Included((fid, val)), Included((fid, val))),
|
||||
Between(left, right) => (Included((fid, left)), Included((fid, right))),
|
||||
|
||||
fn explore_facet_levels(
|
||||
rtxn: &heed::RoTxn,
|
||||
db: &heed::Database<FacetLevelValueI64Codec, CboRoaringBitmapCodec>,
|
||||
field_id: u8,
|
||||
level: u8,
|
||||
left: Bound<i64>,
|
||||
right: Bound<i64>,
|
||||
candidates: &mut RoaringBitmap,
|
||||
) -> anyhow::Result<()>
|
||||
{
|
||||
let mut left_found = left;
|
||||
let mut right_found = right;
|
||||
|
||||
let range = {
|
||||
let left = match left {
|
||||
Included(left) => Included((field_id, level, left, i64::MIN)),
|
||||
Excluded(left) => Excluded((field_id, level, left, i64::MIN)),
|
||||
Bound::Unbounded => Bound::Unbounded,
|
||||
};
|
||||
let right = Included((field_id, level, i64::MAX, i64::MAX));
|
||||
(left, right)
|
||||
};
|
||||
|
||||
let mut candidates = RoaringBitmap::new();
|
||||
|
||||
let db = self.index.facet_field_id_value_docids;
|
||||
let db = db.remap_types::<FacetValueI64Codec, CboRoaringBitmapCodec>();
|
||||
for result in db.range(self.rtxn, &range)? {
|
||||
let ((_fid, _value), docids) = result?;
|
||||
candidates.union_with(&docids);
|
||||
for (i, result) in db.range(rtxn, &range)?.enumerate() {
|
||||
let ((_fid, _level, l, r), docids) = result?;
|
||||
match right {
|
||||
Included(right) if r > right => break,
|
||||
Excluded(right) if r >= right => break,
|
||||
_ => (),
|
||||
}
|
||||
|
||||
eprintln!("{} to {} (level {})", l, r, _level);
|
||||
candidates.union_with(&docids);
|
||||
// We save the leftest and rightest bounds we actually found at this level.
|
||||
if i == 0 { left_found = Excluded(l); }
|
||||
right_found = Excluded(r);
|
||||
}
|
||||
|
||||
// Can we go deeper?
|
||||
let deeper_level = match level.checked_sub(1) {
|
||||
Some(level) => level,
|
||||
None => return Ok(()),
|
||||
};
|
||||
|
||||
// We must refine the left and right bounds of this range by retrieving the
|
||||
// missing part in a deeper level.
|
||||
// TODO we must avoid going at deeper when the bounds are already satisfied.
|
||||
explore_facet_levels(rtxn, db, field_id, deeper_level, left, left_found, candidates)?;
|
||||
explore_facet_levels(rtxn, db, field_id, deeper_level, right_found, right, candidates)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Make sure we always bound the ranges with the field id, as the facets
|
||||
// values are all in the same database and prefixed by the field id.
|
||||
let (left, right) = match operator {
|
||||
GreaterThan(val) => (Excluded(val), Included(i64::MAX)),
|
||||
GreaterThanOrEqual(val) => (Included(val), Included(i64::MAX)),
|
||||
LowerThan(val) => (Included(i64::MIN), Excluded(val)),
|
||||
LowerThanOrEqual(val) => (Included(i64::MIN), Included(val)),
|
||||
Equal(val) => (Included(val), Included(val)),
|
||||
Between(left, right) => (Included(left), Included(right)),
|
||||
};
|
||||
|
||||
let db = self.index
|
||||
.facet_field_id_value_docids
|
||||
.remap_key_type::<FacetLevelValueI64Codec>();
|
||||
|
||||
let biggest_level = match fid.checked_add(1) {
|
||||
Some(next_fid) => {
|
||||
// If we are able to find the next field id we ask the key that is before
|
||||
// the first entry of it which corresponds to the last key of our field id.
|
||||
let db = db.remap_data_type::<DecodeIgnore>();
|
||||
match db.get_lower_than(self.rtxn, &(next_fid, 0, i64::MIN, i64::MIN))? {
|
||||
Some(((id, level, _left, _right), _docids)) if fid == id => Some(level),
|
||||
_ => None,
|
||||
}
|
||||
},
|
||||
None => {
|
||||
// If we can't generate a bigger field id, it must be equal to 255 and
|
||||
// therefore the last key of the database must be the one we want.
|
||||
match db.remap_data_type::<DecodeIgnore>().last(self.rtxn)? {
|
||||
Some(((id, level, _left, _right), _docids)) if fid == id => Some(level),
|
||||
_ => None,
|
||||
}
|
||||
},
|
||||
};
|
||||
|
||||
match biggest_level {
|
||||
Some(level) => {
|
||||
let mut candidates = RoaringBitmap::new();
|
||||
explore_facet_levels(self.rtxn, &db, fid, level, left, right, &mut candidates)?;
|
||||
Some(candidates)
|
||||
},
|
||||
None => None,
|
||||
}
|
||||
},
|
||||
None => None,
|
||||
};
|
||||
|
||||
let (candidates, derived_words) = match (facet_candidates, derived_words) {
|
||||
|
@ -2,9 +2,9 @@ use std::path::PathBuf;
|
||||
use std::{str, io};
|
||||
|
||||
use anyhow::Context;
|
||||
use crate::Index;
|
||||
use heed::EnvOpenOptions;
|
||||
use structopt::StructOpt;
|
||||
use crate::Index;
|
||||
|
||||
use Command::*;
|
||||
|
||||
@ -225,12 +225,52 @@ fn most_common_words(index: &Index, rtxn: &heed::RoTxn, limit: usize) -> anyhow:
|
||||
Ok(wtr.flush()?)
|
||||
}
|
||||
|
||||
fn facet_values_iter<'txn, DC: 'txn>(
|
||||
rtxn: &'txn heed::RoTxn,
|
||||
db: heed::Database<heed::types::ByteSlice, DC>,
|
||||
field_id: u8,
|
||||
facet_type: crate::facet::FacetType,
|
||||
) -> heed::Result<Box<dyn Iterator<Item=heed::Result<(String, DC::DItem)>> + 'txn>>
|
||||
where
|
||||
DC: heed::BytesDecode<'txn>,
|
||||
{
|
||||
use crate::facet::FacetType;
|
||||
use crate::heed_codec::facet::{
|
||||
FacetValueStringCodec, FacetLevelValueF64Codec, FacetLevelValueI64Codec,
|
||||
};
|
||||
|
||||
let iter = db.prefix_iter(&rtxn, &[field_id])?;
|
||||
match facet_type {
|
||||
FacetType::String => {
|
||||
let iter = iter.remap_key_type::<FacetValueStringCodec>()
|
||||
.map(|r| r.map(|((_, key), value)| (key.to_string(), value)));
|
||||
Ok(Box::new(iter) as Box<dyn Iterator<Item=_>>)
|
||||
},
|
||||
FacetType::Float => {
|
||||
let iter = iter.remap_key_type::<FacetLevelValueF64Codec>()
|
||||
.map(|r| r.map(|((_, level, left, right), value)| if level == 0 {
|
||||
(format!("{} (level {})", left, level), value)
|
||||
} else {
|
||||
(format!("{} to {} (level {})", left, right, level), value)
|
||||
}));
|
||||
Ok(Box::new(iter))
|
||||
},
|
||||
FacetType::Integer => {
|
||||
let iter = iter.remap_key_type::<FacetLevelValueI64Codec>()
|
||||
.map(|r| r.map(|((_, level, left, right), value)| if level == 0 {
|
||||
(format!("{} (level {})", left, level), value)
|
||||
} else {
|
||||
(format!("{} to {} (level {})", left, right, level), value)
|
||||
}));
|
||||
Ok(Box::new(iter))
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn biggest_value_sizes(index: &Index, rtxn: &heed::RoTxn, limit: usize) -> anyhow::Result<()> {
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::BinaryHeap;
|
||||
use heed::types::{Str, ByteSlice};
|
||||
use crate::facet::FacetType;
|
||||
use crate::heed_codec::facet::{FacetValueStringCodec, FacetValueF64Codec, FacetValueI64Codec};
|
||||
|
||||
let Index {
|
||||
env: _env,
|
||||
@ -285,25 +325,9 @@ fn biggest_value_sizes(index: &Index, rtxn: &heed::RoTxn, limit: usize) -> anyho
|
||||
let fields_ids_map = index.fields_ids_map(rtxn)?;
|
||||
for (field_id, field_type) in faceted_fields {
|
||||
let facet_name = fields_ids_map.name(field_id).unwrap();
|
||||
let iter = facet_field_id_value_docids.prefix_iter(&rtxn, &[field_id])?;
|
||||
let iter = match field_type {
|
||||
FacetType::String => {
|
||||
let iter = iter.remap_types::<FacetValueStringCodec, ByteSlice>()
|
||||
.map(|r| r.map(|((_, k), v)| (k.to_string(), v)));
|
||||
Box::new(iter) as Box<dyn Iterator<Item=_>>
|
||||
},
|
||||
FacetType::Float => {
|
||||
let iter = iter.remap_types::<FacetValueF64Codec, ByteSlice>()
|
||||
.map(|r| r.map(|((_, k), v)| (k.to_string(), v)));
|
||||
Box::new(iter)
|
||||
},
|
||||
FacetType::Integer => {
|
||||
let iter = iter.remap_types::<FacetValueI64Codec, ByteSlice>()
|
||||
.map(|r| r.map(|((_, k), v)| (k.to_string(), v)));
|
||||
Box::new(iter)
|
||||
},
|
||||
};
|
||||
for result in iter {
|
||||
|
||||
let db = facet_field_id_value_docids.remap_data_type::<ByteSlice>();
|
||||
for result in facet_values_iter(rtxn, db, field_id, field_type)? {
|
||||
let (fvalue, value) = result?;
|
||||
let key = format!("{} {}", facet_name, fvalue);
|
||||
heap.push(Reverse((value.len(), key, facet_field_id_value_docids_name)));
|
||||
@ -349,10 +373,6 @@ fn words_docids(index: &Index, rtxn: &heed::RoTxn, debug: bool, words: Vec<Strin
|
||||
}
|
||||
|
||||
fn facet_values_docids(index: &Index, rtxn: &heed::RoTxn, debug: bool, field_name: String) -> anyhow::Result<()> {
|
||||
use crate::facet::FacetType;
|
||||
use crate::heed_codec::facet::{FacetValueStringCodec, FacetValueF64Codec, FacetValueI64Codec};
|
||||
use heed::{BytesDecode, Error::Decoding};
|
||||
|
||||
let fields_ids_map = index.fields_ids_map(&rtxn)?;
|
||||
let faceted_fields = index.faceted_fields(&rtxn)?;
|
||||
|
||||
@ -361,39 +381,12 @@ fn facet_values_docids(index: &Index, rtxn: &heed::RoTxn, debug: bool, field_nam
|
||||
let field_type = faceted_fields.get(&field_id)
|
||||
.with_context(|| format!("field {} is not faceted", field_name))?;
|
||||
|
||||
let iter = index.facet_field_id_value_docids.prefix_iter(&rtxn, &[field_id])?;
|
||||
let iter = match field_type {
|
||||
FacetType::String => {
|
||||
let iter = iter
|
||||
.map(|result| result.and_then(|(key, value)| {
|
||||
let (_, key) = FacetValueStringCodec::bytes_decode(key).ok_or(Decoding)?;
|
||||
Ok((key.to_string(), value))
|
||||
}));
|
||||
Box::new(iter) as Box<dyn Iterator<Item=_>>
|
||||
},
|
||||
FacetType::Float => {
|
||||
let iter = iter
|
||||
.map(|result| result.and_then(|(key, value)| {
|
||||
let (_, key) = FacetValueF64Codec::bytes_decode(key).ok_or(Decoding)?;
|
||||
Ok((key.to_string(), value))
|
||||
}));
|
||||
Box::new(iter)
|
||||
},
|
||||
FacetType::Integer => {
|
||||
let iter = iter
|
||||
.map(|result| result.and_then(|(key, value)| {
|
||||
let (_, key) = FacetValueI64Codec::bytes_decode(key).ok_or(Decoding)?;
|
||||
Ok((key.to_string(), value))
|
||||
}));
|
||||
Box::new(iter)
|
||||
},
|
||||
};
|
||||
|
||||
let stdout = io::stdout();
|
||||
let mut wtr = csv::Writer::from_writer(stdout.lock());
|
||||
wtr.write_record(&["facet_value", "documents_ids"])?;
|
||||
|
||||
for result in iter {
|
||||
let db = index.facet_field_id_value_docids;
|
||||
for result in facet_values_iter(rtxn, db, field_id, *field_type)? {
|
||||
let (value, docids) = result?;
|
||||
let docids = if debug {
|
||||
format!("{:?}", docids)
|
||||
|
125
src/update/index_documents/facet_level.rs
Normal file
125
src/update/index_documents/facet_level.rs
Normal file
@ -0,0 +1,125 @@
|
||||
use std::fs::File;
|
||||
|
||||
use grenad::{CompressionType, Reader, Writer, FileFuse};
|
||||
use heed::types::{ByteSlice, DecodeIgnore};
|
||||
use heed::{BytesEncode, Error};
|
||||
use roaring::RoaringBitmap;
|
||||
|
||||
use crate::facet::FacetType;
|
||||
use crate::heed_codec::{facet::FacetLevelValueI64Codec, CboRoaringBitmapCodec};
|
||||
use crate::update::index_documents::{create_writer, writer_into_reader};
|
||||
|
||||
pub fn clear_field_levels(
|
||||
wtxn: &mut heed::RwTxn,
|
||||
db: heed::Database<ByteSlice, CboRoaringBitmapCodec>,
|
||||
field_id: u8,
|
||||
) -> heed::Result<()>
|
||||
{
|
||||
let range = (field_id, 1, i64::MIN, i64::MIN)..=(field_id, u8::MAX, i64::MAX, i64::MAX);
|
||||
db.remap_key_type::<FacetLevelValueI64Codec>()
|
||||
.delete_range(wtxn, &range)
|
||||
.map(drop)
|
||||
}
|
||||
|
||||
pub fn compute_facet_levels(
|
||||
rtxn: &heed::RoTxn,
|
||||
db: heed::Database<ByteSlice, CboRoaringBitmapCodec>,
|
||||
compression_type: CompressionType,
|
||||
compression_level: Option<u32>,
|
||||
shrink_size: Option<u64>,
|
||||
field_id: u8,
|
||||
facet_type: FacetType,
|
||||
) -> anyhow::Result<Reader<FileFuse>>
|
||||
{
|
||||
let last_level_size = 5;
|
||||
let number_of_levels = 5;
|
||||
let first_level_size = db.prefix_iter(rtxn, &[field_id])?
|
||||
.remap_types::<DecodeIgnore, DecodeIgnore>()
|
||||
.fold(Ok(0u64), |count, result| result.and(count).map(|c| c + 1))?;
|
||||
|
||||
// It is forbidden to keep a cursor and write in a database at the same time with LMDB
|
||||
// therefore we write the facet levels entries into a grenad file before transfering them.
|
||||
let mut writer = tempfile::tempfile().and_then(|file| {
|
||||
create_writer(compression_type, compression_level, file)
|
||||
})?;
|
||||
|
||||
let level_0_range = (field_id, 0, i64::MIN, i64::MIN)..=(field_id, 0, i64::MAX, i64::MAX);
|
||||
let level_sizes_iter = levels_iterator(first_level_size, last_level_size, number_of_levels)
|
||||
.enumerate()
|
||||
.skip(1);
|
||||
|
||||
// TODO we must not create levels with identical group sizes.
|
||||
for (level, size) in level_sizes_iter {
|
||||
let level_entry_sizes = (first_level_size as f64 / size as f64).ceil() as usize;
|
||||
let mut left = 0;
|
||||
let mut right = 0;
|
||||
let mut group_docids = RoaringBitmap::new();
|
||||
|
||||
let db = db.remap_key_type::<FacetLevelValueI64Codec>();
|
||||
for (i, result) in db.range(rtxn, &level_0_range)?.enumerate() {
|
||||
let ((_field_id, _level, value, _right), docids) = result?;
|
||||
|
||||
if i == 0 {
|
||||
left = value;
|
||||
} else if i % level_entry_sizes == 0 {
|
||||
// we found the first bound of the next group, we must store the left
|
||||
// and right bounds associated with the docids.
|
||||
write_entry(&mut writer, field_id, level as u8, left, right, &group_docids)?;
|
||||
|
||||
// We save the left bound for the new group and also reset the docids.
|
||||
group_docids = RoaringBitmap::new();
|
||||
left = value;
|
||||
}
|
||||
|
||||
// The right bound is always the bound we run through.
|
||||
group_docids.union_with(&docids);
|
||||
right = value;
|
||||
}
|
||||
|
||||
if !group_docids.is_empty() {
|
||||
write_entry(&mut writer, field_id, level as u8, left, right, &group_docids)?;
|
||||
}
|
||||
}
|
||||
|
||||
writer_into_reader(writer, shrink_size)
|
||||
}
|
||||
|
||||
fn write_entry(
|
||||
writer: &mut Writer<File>,
|
||||
field_id: u8,
|
||||
level: u8,
|
||||
left: i64,
|
||||
right: i64,
|
||||
ids: &RoaringBitmap,
|
||||
) -> anyhow::Result<()>
|
||||
{
|
||||
let key = (field_id, level, left, right);
|
||||
let key = FacetLevelValueI64Codec::bytes_encode(&key).ok_or(Error::Encoding)?;
|
||||
let data = CboRoaringBitmapCodec::bytes_encode(&ids).ok_or(Error::Encoding)?;
|
||||
writer.insert(&key, &data)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn levels_iterator(
|
||||
first_level_size: u64, // biggest level
|
||||
last_level_size: u64, // smallest level
|
||||
number_of_levels: u64,
|
||||
) -> impl Iterator<Item=u64>
|
||||
{
|
||||
// Go look at the function definitions here:
|
||||
// https://docs.rs/easer/0.2.1/easer/index.html
|
||||
// https://easings.net/#easeOutExpo
|
||||
fn ease_out_expo(t: f64, b: f64, c: f64, d: f64) -> f64 {
|
||||
if t == d {
|
||||
b + c
|
||||
} else {
|
||||
c * (-2.0_f64.powf(-10.0 * t / d) + 1.0) + b
|
||||
}
|
||||
}
|
||||
|
||||
let b = last_level_size as f64;
|
||||
let end = first_level_size as f64;
|
||||
let c = end - b;
|
||||
let d = number_of_levels;
|
||||
(0..=d).map(move |t| ((end + b) - ease_out_expo(t as f64, b, c, d as f64)) as u64)
|
||||
}
|
@ -14,6 +14,7 @@ use memmap::Mmap;
|
||||
use rayon::prelude::*;
|
||||
use rayon::ThreadPool;
|
||||
|
||||
use crate::facet::FacetType;
|
||||
use crate::index::Index;
|
||||
use crate::update::UpdateIndexingStep;
|
||||
use self::store::{Store, Readers};
|
||||
@ -22,10 +23,12 @@ use self::merge_function::{
|
||||
docid_word_positions_merge, documents_merge, facet_field_value_docids_merge,
|
||||
};
|
||||
pub use self::transform::{Transform, TransformOutput};
|
||||
pub use self::facet_level::{clear_field_levels, compute_facet_levels};
|
||||
|
||||
use crate::MergeFn;
|
||||
use super::UpdateBuilder;
|
||||
|
||||
mod facet_level;
|
||||
mod merge_function;
|
||||
mod store;
|
||||
mod transform;
|
||||
@ -327,7 +330,7 @@ impl<'t, 'u, 'i, 'a> IndexDocuments<'t, 'u, 'i, 'a> {
|
||||
enum DatabaseType {
|
||||
Main,
|
||||
WordDocids,
|
||||
FacetValuesDocids,
|
||||
FacetLevel0ValuesDocids,
|
||||
}
|
||||
|
||||
let faceted_fields = self.index.faceted_fields(self.wtxn)?;
|
||||
@ -427,7 +430,7 @@ impl<'t, 'u, 'i, 'a> IndexDocuments<'t, 'u, 'i, 'a> {
|
||||
(DatabaseType::Main, main_readers, main_merge as MergeFn),
|
||||
(DatabaseType::WordDocids, word_docids_readers, word_docids_merge),
|
||||
(
|
||||
DatabaseType::FacetValuesDocids,
|
||||
DatabaseType::FacetLevel0ValuesDocids,
|
||||
facet_field_value_docids_readers,
|
||||
facet_field_value_docids_merge,
|
||||
),
|
||||
@ -475,6 +478,9 @@ impl<'t, 'u, 'i, 'a> IndexDocuments<'t, 'u, 'i, 'a> {
|
||||
// We write the external documents ids into the main database.
|
||||
self.index.put_external_documents_ids(self.wtxn, &external_documents_ids)?;
|
||||
|
||||
// We get the faceted fields to be able to create the facet levels.
|
||||
let faceted_fields = self.index.faceted_fields(self.wtxn)?;
|
||||
|
||||
// We merge the new documents ids with the existing ones.
|
||||
documents_ids.union_with(&new_documents_ids);
|
||||
documents_ids.union_with(&replaced_documents_ids);
|
||||
@ -557,7 +563,7 @@ impl<'t, 'u, 'i, 'a> IndexDocuments<'t, 'u, 'i, 'a> {
|
||||
write_method,
|
||||
)?;
|
||||
},
|
||||
DatabaseType::FacetValuesDocids => {
|
||||
DatabaseType::FacetLevel0ValuesDocids => {
|
||||
debug!("Writing the facet values docids into LMDB on disk...");
|
||||
let db = *self.index.facet_field_id_value_docids.as_polymorph();
|
||||
write_into_lmdb_database(
|
||||
@ -577,6 +583,35 @@ impl<'t, 'u, 'i, 'a> IndexDocuments<'t, 'u, 'i, 'a> {
|
||||
});
|
||||
}
|
||||
|
||||
debug!("Computing and writing the facet values levels docids into LMDB on disk...");
|
||||
for (field_id, facet_type) in faceted_fields {
|
||||
if facet_type == FacetType::String { continue }
|
||||
|
||||
clear_field_levels(
|
||||
self.wtxn,
|
||||
self.index.facet_field_id_value_docids,
|
||||
field_id,
|
||||
)?;
|
||||
|
||||
let content = compute_facet_levels(
|
||||
self.wtxn,
|
||||
self.index.facet_field_id_value_docids,
|
||||
chunk_compression_type,
|
||||
chunk_compression_level,
|
||||
chunk_fusing_shrink_size,
|
||||
field_id,
|
||||
facet_type,
|
||||
)?;
|
||||
|
||||
write_into_lmdb_database(
|
||||
self.wtxn,
|
||||
*self.index.facet_field_id_value_docids.as_polymorph(),
|
||||
content,
|
||||
|_, _| anyhow::bail!("invalid facet level merging"),
|
||||
WriteMethod::GetMergePut,
|
||||
)?;
|
||||
}
|
||||
|
||||
debug_assert_eq!(database_count, total_databases);
|
||||
|
||||
info!("Transform output indexed in {:.02?}", before_indexing.elapsed());
|
||||
|
@ -19,7 +19,7 @@ use tempfile::tempfile;
|
||||
|
||||
use crate::facet::FacetType;
|
||||
use crate::heed_codec::{BoRoaringBitmapCodec, CboRoaringBitmapCodec};
|
||||
use crate::heed_codec::facet::{FacetValueStringCodec, FacetValueF64Codec, FacetValueI64Codec};
|
||||
use crate::heed_codec::facet::{FacetValueStringCodec, FacetLevelValueF64Codec, FacetLevelValueI64Codec};
|
||||
use crate::tokenizer::{simple_tokenizer, only_token};
|
||||
use crate::update::UpdateIndexingStep;
|
||||
use crate::{json_to_string, SmallVec8, SmallVec32, SmallString32, Position, DocumentId};
|
||||
@ -337,8 +337,8 @@ impl Store {
|
||||
for ((field_id, value), docids) in iter {
|
||||
let result = match value {
|
||||
String(s) => FacetValueStringCodec::bytes_encode(&(field_id, &s)).map(Cow::into_owned),
|
||||
Float(f) => FacetValueF64Codec::bytes_encode(&(field_id, *f)).map(Cow::into_owned),
|
||||
Integer(i) => FacetValueI64Codec::bytes_encode(&(field_id, i)).map(Cow::into_owned),
|
||||
Float(f) => FacetLevelValueF64Codec::bytes_encode(&(field_id, 0, *f, *f)).map(Cow::into_owned),
|
||||
Integer(i) => FacetLevelValueI64Codec::bytes_encode(&(field_id, 0, i, i)).map(Cow::into_owned),
|
||||
};
|
||||
let key = result.context("could not serialize facet key")?;
|
||||
let bytes = CboRoaringBitmapCodec::bytes_encode(&docids)
|
||||
|
@ -412,7 +412,8 @@ mod tests {
|
||||
let rtxn = index.read_txn().unwrap();
|
||||
let fields_ids = index.faceted_fields(&rtxn).unwrap();
|
||||
assert_eq!(fields_ids, hashmap!{ 1 => FacetType::Integer });
|
||||
let count = index.facet_field_id_value_docids.len(&rtxn).unwrap();
|
||||
// Only count the field_id 0 and level 0 facet values.
|
||||
let count = index.facet_field_id_value_docids.prefix_iter(&rtxn, &[1, 0]).unwrap().count();
|
||||
assert_eq!(count, 3);
|
||||
drop(rtxn);
|
||||
|
||||
@ -425,7 +426,8 @@ mod tests {
|
||||
wtxn.commit().unwrap();
|
||||
|
||||
let rtxn = index.read_txn().unwrap();
|
||||
let count = index.facet_field_id_value_docids.len(&rtxn).unwrap();
|
||||
// Only count the field_id 0 and level 0 facet values.
|
||||
let count = index.facet_field_id_value_docids.prefix_iter(&rtxn, &[1, 0]).unwrap().count();
|
||||
assert_eq!(count, 4);
|
||||
drop(rtxn);
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user