From 9ec95679e17723d9cffb774a5bf16215723a78ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= Date: Wed, 18 Nov 2020 16:29:07 +0100 Subject: [PATCH] Introduce a function to retrieve the facet level range docids --- .../facet/facet_level_value_f64_codec.rs | 62 +++++++++ .../facet/facet_level_value_i64_codec.rs | 43 ++++++ src/heed_codec/facet/facet_value_f64_codec.rs | 50 ------- src/heed_codec/facet/facet_value_i64_codec.rs | 28 ---- src/heed_codec/facet/mod.rs | 8 +- src/search.rs | 115 +++++++++++++--- src/subcommand/infos.rs | 103 +++++++-------- src/update/index_documents/facet_level.rs | 125 ++++++++++++++++++ src/update/index_documents/mod.rs | 41 +++++- src/update/index_documents/store.rs | 6 +- src/update/settings.rs | 6 +- 11 files changed, 423 insertions(+), 164 deletions(-) create mode 100644 src/heed_codec/facet/facet_level_value_f64_codec.rs create mode 100644 src/heed_codec/facet/facet_level_value_i64_codec.rs delete mode 100644 src/heed_codec/facet/facet_value_f64_codec.rs delete mode 100644 src/heed_codec/facet/facet_value_i64_codec.rs create mode 100644 src/update/index_documents/facet_level.rs diff --git a/src/heed_codec/facet/facet_level_value_f64_codec.rs b/src/heed_codec/facet/facet_level_value_f64_codec.rs new file mode 100644 index 000000000..55fd4b132 --- /dev/null +++ b/src/heed_codec/facet/facet_level_value_f64_codec.rs @@ -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 { + 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> { + 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)); + } +} diff --git a/src/heed_codec/facet/facet_level_value_i64_codec.rs b/src/heed_codec/facet/facet_level_value_i64_codec.rs new file mode 100644 index 000000000..7cf9a714b --- /dev/null +++ b/src/heed_codec/facet/facet_level_value_i64_codec.rs @@ -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 { + 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> { + 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)) + } +} diff --git a/src/heed_codec/facet/facet_value_f64_codec.rs b/src/heed_codec/facet/facet_value_f64_codec.rs deleted file mode 100644 index 228514de5..000000000 --- a/src/heed_codec/facet/facet_value_f64_codec.rs +++ /dev/null @@ -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 { - 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> { - 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)); - } -} diff --git a/src/heed_codec/facet/facet_value_i64_codec.rs b/src/heed_codec/facet/facet_value_i64_codec.rs deleted file mode 100644 index f99b8a3ea..000000000 --- a/src/heed_codec/facet/facet_value_i64_codec.rs +++ /dev/null @@ -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 { - 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> { - 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)) - } -} diff --git a/src/heed_codec/facet/mod.rs b/src/heed_codec/facet/mod.rs index abe2c1d8a..ef97e6add 100644 --- a/src/heed_codec/facet/mod.rs +++ b/src/heed_codec/facet/mod.rs @@ -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; diff --git a/src/search.rs b/src/search.rs index 17f25edfc..3338d5222 100644 --- a/src/search.rs +++ b/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,29 +239,105 @@ 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))), - }; - let mut candidates = RoaringBitmap::new(); + fn explore_facet_levels( + rtxn: &heed::RoTxn, + db: &heed::Database, + field_id: u8, + level: u8, + left: Bound, + right: Bound, + candidates: &mut RoaringBitmap, + ) -> anyhow::Result<()> + { + let mut left_found = left; + let mut right_found = right; - let db = self.index.facet_field_id_value_docids; - let db = db.remap_types::(); - for result in db.range(self.rtxn, &range)? { - let ((_fid, _value), docids) = result?; - candidates.union_with(&docids); + 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) + }; + + 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(()) } - Some(candidates) + // 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::(); + + 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::(); + 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::().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, }; diff --git a/src/subcommand/infos.rs b/src/subcommand/infos.rs index 6a69d1bad..74b1505f9 100644 --- a/src/subcommand/infos.rs +++ b/src/subcommand/infos.rs @@ -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, + field_id: u8, + facet_type: crate::facet::FacetType, +) -> heed::Result> + '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::() + .map(|r| r.map(|((_, key), value)| (key.to_string(), value))); + Ok(Box::new(iter) as Box>) + }, + FacetType::Float => { + let iter = iter.remap_key_type::() + .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::() + .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::() - .map(|r| r.map(|((_, k), v)| (k.to_string(), v))); - Box::new(iter) as Box> - }, - FacetType::Float => { - let iter = iter.remap_types::() - .map(|r| r.map(|((_, k), v)| (k.to_string(), v))); - Box::new(iter) - }, - FacetType::Integer => { - let iter = iter.remap_types::() - .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::(); + 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 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> - }, - 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) diff --git a/src/update/index_documents/facet_level.rs b/src/update/index_documents/facet_level.rs new file mode 100644 index 000000000..0a87e21d6 --- /dev/null +++ b/src/update/index_documents/facet_level.rs @@ -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, + 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::() + .delete_range(wtxn, &range) + .map(drop) +} + +pub fn compute_facet_levels( + rtxn: &heed::RoTxn, + db: heed::Database, + compression_type: CompressionType, + compression_level: Option, + shrink_size: Option, + field_id: u8, + facet_type: FacetType, +) -> anyhow::Result> +{ + let last_level_size = 5; + let number_of_levels = 5; + let first_level_size = db.prefix_iter(rtxn, &[field_id])? + .remap_types::() + .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::(); + 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, + 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 +{ + // 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) +} diff --git a/src/update/index_documents/mod.rs b/src/update/index_documents/mod.rs index fe51c6b2b..f6587c3a8 100644 --- a/src/update/index_documents/mod.rs +++ b/src/update/index_documents/mod.rs @@ -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()); diff --git a/src/update/index_documents/store.rs b/src/update/index_documents/store.rs index 110ab7f25..289704b1a 100644 --- a/src/update/index_documents/store.rs +++ b/src/update/index_documents/store.rs @@ -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) diff --git a/src/update/settings.rs b/src/update/settings.rs index 141abcf00..cddd68ca3 100644 --- a/src/update/settings.rs +++ b/src/update/settings.rs @@ -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); }