From 7a05502e0cbab5ad3cbc2ef241c1955ba85bb4c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20Coss=C3=ADo?= Date: Tue, 27 May 2025 18:53:41 -0400 Subject: [PATCH] Full text index with generic posting list (#6565) * integrate generic posting list in full-text index * improve intersection fn improve iterator intersection * make old module just for test * recover old implementation of MmapPostings for tests * add compatibility test * fix compression * clippy * nits * check if first id is already greater protects against the case where the iterator's current element hasn't been extracted yet * update current element when there is no greater or equal * reuse visitor in test * optimize alignment * bound checks and refactor is_in_range * review * edit comment --------- Co-authored-by: generall --- Cargo.lock | 2 + lib/posting_list/Cargo.toml | 1 + lib/posting_list/src/builder.rs | 22 +- lib/posting_list/src/iterator.rs | 68 ++++- lib/posting_list/src/lib.rs | 11 +- lib/posting_list/src/posting_list.rs | 72 +++++- lib/posting_list/src/tests.rs | 42 ++- lib/posting_list/src/value_handler.rs | 6 +- lib/posting_list/src/view.rs | 98 ++++--- lib/posting_list/src/visitor.rs | 87 ++++++- lib/segment/Cargo.toml | 1 + .../compressed_chunks_reader.rs | 11 - .../immutable_inverted_index.rs | 29 ++- .../full_text_index/inverted_index.rs | 16 +- .../mmap_inverted_index/mmap_postings.rs | 95 ++++--- .../mmap_inverted_index/mod.rs | 13 +- .../mmap_inverted_index/old_mmap_postings.rs | 239 ++++++++++++++++++ .../mmap_inverted_index/tests.rs | 61 +++++ .../index/field_index/full_text_index/mod.rs | 1 + .../full_text_index/posting_list.rs | 5 - .../full_text_index/postings_iterator.rs | 45 ++-- 21 files changed, 736 insertions(+), 189 deletions(-) create mode 100644 lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/old_mmap_postings.rs create mode 100644 lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 118a7e86c3..a326e87c90 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4413,6 +4413,7 @@ dependencies = [ "common", "rand 0.9.1", "tempfile", + "zerocopy 0.8.25", ] [[package]] @@ -5853,6 +5854,7 @@ dependencies = [ "num-traits", "ordered-float 5.0.0", "parking_lot", + "posting_list", "pprof", "procfs", "proptest", diff --git a/lib/posting_list/Cargo.toml b/lib/posting_list/Cargo.toml index a2c1f14d43..aef908e31d 100644 --- a/lib/posting_list/Cargo.toml +++ b/lib/posting_list/Cargo.toml @@ -12,6 +12,7 @@ workspace = true [dependencies] common = { path = "../common/common" } bitpacking = { workspace = true } +zerocopy = { workspace = true } [dev-dependencies] rand = { workspace = true } diff --git a/lib/posting_list/src/builder.rs b/lib/posting_list/src/builder.rs index b3e66a8417..93ce882787 100644 --- a/lib/posting_list/src/builder.rs +++ b/lib/posting_list/src/builder.rs @@ -2,8 +2,9 @@ use std::marker::PhantomData; use bitpacking::BitPacker; use common::types::PointOffsetType; +use zerocopy::little_endian::U32; -use crate::posting_list::{PostingChunk, PostingElement, PostingList}; +use crate::posting_list::{PostingChunk, PostingElement, PostingList, RemainderPosting}; use crate::value_handler::{SizedHandler, UnsizedHandler, ValueHandler}; use crate::{ BitPackerImpl, CHUNK_LEN, SizedValue, UnsizedValue, VarPostingList, WeightsPostingList, @@ -61,13 +62,14 @@ impl PostingBuilder { for (chunk_ids, chunk_values) in ids_chunks_iter.zip(values_chunks_iter) { let initial = chunk_ids[0]; - let chunk_bits = bitpacker.num_bits_strictly_sorted(initial.checked_sub(1), chunk_ids); + let chunk_bits = bitpacker.num_bits_sorted(initial, chunk_ids); let chunk_size = BitPackerImpl::compressed_block_size(chunk_bits); chunks.push(PostingChunk { - initial_id: initial, + initial_id: U32::from(initial), offset: u32::try_from(id_data_size) - .expect("id_data_size should fit in u32, (smaller than 4GB)"), + .expect("id_data_size should fit in u32, (smaller than 4GB)") + .into(), sized_values: chunk_values .try_into() .expect("should be a valid chunk size"), @@ -78,7 +80,10 @@ impl PostingBuilder { // now process remainders let mut remainders = Vec::with_capacity(num_elements % CHUNK_LEN); for (&id, &value) in remainder_ids.iter().zip(remainder_values) { - remainders.push(PostingElement { id, value }); + remainders.push(RemainderPosting { + id: U32::from(id), + value, + }); } // compress id_data @@ -87,10 +92,11 @@ impl PostingBuilder { let chunk = &chunks[chunk_index]; let compressed_size = PostingChunk::get_compressed_size(&chunks, &id_data, chunk_index); let chunk_bits = compressed_size * u8::BITS as usize / CHUNK_LEN; - bitpacker.compress_strictly_sorted( - chunk.initial_id.checked_sub(1), + bitpacker.compress_sorted( + chunk.initial_id.get(), chunk_ids, - &mut id_data[chunk.offset as usize..chunk.offset as usize + compressed_size], + &mut id_data + [chunk.offset.get() as usize..chunk.offset.get() as usize + compressed_size], chunk_bits as u8, ); } diff --git a/lib/posting_list/src/iterator.rs b/lib/posting_list/src/iterator.rs index c4969b4d72..702faf0b89 100644 --- a/lib/posting_list/src/iterator.rs +++ b/lib/posting_list/src/iterator.rs @@ -8,28 +8,75 @@ use crate::visitor::PostingVisitor; pub struct PostingIterator<'a, H: ValueHandler> { visitor: PostingVisitor<'a, H>, - current_id: Option, + current_elem: Option>, offset: usize, } -impl<'a, H: ValueHandler> PostingIterator<'a, H> { +impl<'a, H: ValueHandler> PostingIterator<'a, H> +where + H::Value: Clone, +{ pub fn new(visitor: PostingVisitor<'a, H>) -> Self { Self { visitor, - current_id: None, + current_elem: None, offset: 0, } } + + /// Advances the iterator until the current element id is greater than or equal to the given id. + /// + /// Returns `Some(PostingElement)` on the first element that is greater than or equal to the given id. It can be possible that this id is + /// the head of the iterator, so it does not need to be advanced. + /// + /// `None` means the iterator is exhausted. + pub fn advance_until_greater_or_equal( + &mut self, + target_id: PointOffsetType, + ) -> Option> { + if let Some(current) = &self.current_elem { + if current.id >= target_id { + return Some(current.clone()); + } + } + + if self.offset >= self.visitor.len() { + return None; + } + + let Some(offset) = self + .visitor + .search_greater_or_equal(target_id, Some(self.offset)) + else { + self.current_elem = None; + self.offset = self.visitor.len(); + return None; + }; + + debug_assert!(offset >= self.offset); + let greater_or_equal = self.visitor.get_by_offset(offset); + + self.current_elem = greater_or_equal.clone(); + self.offset = offset; + + greater_or_equal + } } -impl Iterator for PostingIterator<'_, H> { +impl Iterator for PostingIterator<'_, H> +where + H::Value: Clone, +{ type Item = PostingElement; fn next(&mut self) -> Option { - self.visitor.get_by_offset(self.offset).inspect(|elem| { - self.current_id = Some(elem.id); + let next_opt = self.visitor.get_by_offset(self.offset).inspect(|_| { self.offset += 1; - }) + }); + + self.current_elem = next_opt.clone(); + + next_opt } fn size_hint(&self) -> (usize, Option) { @@ -42,10 +89,13 @@ impl Iterator for PostingIterator<'_, H> { } } -impl ExactSizeIterator for PostingIterator<'_, H> { +impl ExactSizeIterator for PostingIterator<'_, H> +where + H::Value: Clone, +{ fn len(&self) -> usize { self.visitor.list.len().saturating_sub(self.offset) } } -impl FusedIterator for PostingIterator<'_, H> {} +impl FusedIterator for PostingIterator<'_, H> where H::Value: Clone {} diff --git a/lib/posting_list/src/lib.rs b/lib/posting_list/src/lib.rs index 04ae86792c..4590509e46 100644 --- a/lib/posting_list/src/lib.rs +++ b/lib/posting_list/src/lib.rs @@ -15,13 +15,13 @@ type BitPackerImpl = bitpacking::BitPacker4x; const CHUNK_LEN: usize = 128; const _: () = assert!(128 == BitPackerImpl::BLOCK_LEN); -pub trait SizedValue: Sized + Copy + std::fmt::Debug {} +pub trait SizedValue: Sized + Copy {} impl SizedValue for () {} impl SizedValue for u32 {} impl SizedValue for u64 {} -pub trait UnsizedValue: std::fmt::Debug { +pub trait UnsizedValue { fn write_len(&self) -> usize; fn write_to(&self, dst: &mut [u8]); @@ -48,6 +48,7 @@ pub type WeightsPostingListView<'a, W> = PostingListView<'a, SizedHandler>; pub type VarPostingListView<'a, V> = PostingListView<'a, UnsizedHandler>; pub use builder::PostingBuilder; -pub use posting_list::{PostingChunk, PostingElement, PostingList}; -use value_handler::{SizedHandler, UnsizedHandler}; -pub use view::PostingListView; +pub use posting_list::{PostingChunk, PostingElement, PostingList, RemainderPosting}; +pub use value_handler::{SizedHandler, UnsizedHandler, ValueHandler}; +pub use view::{PostingListComponents, PostingListView}; +pub use visitor::PostingVisitor; diff --git a/lib/posting_list/src/posting_list.rs b/lib/posting_list/src/posting_list.rs index 04d05c3da0..6d9c801727 100644 --- a/lib/posting_list/src/posting_list.rs +++ b/lib/posting_list/src/posting_list.rs @@ -1,11 +1,14 @@ use std::marker::PhantomData; use common::types::PointOffsetType; +use zerocopy::little_endian::U32; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; -use crate::CHUNK_LEN; +use crate::iterator::PostingIterator; use crate::value_handler::ValueHandler; use crate::view::PostingListView; use crate::visitor::PostingVisitor; +use crate::{CHUNK_LEN, PostingBuilder}; /// Generic compressed posting list. /// @@ -14,29 +17,45 @@ use crate::visitor::PostingVisitor; /// fixed-sized chunks /// - `PostingList>` when there are `VarSizedValue` values, each id includes one value in the chunk, /// which points to the actual value in the var_size_data +#[derive(Debug, Clone)] pub struct PostingList { pub(crate) id_data: Vec, pub(crate) chunks: Vec>, - pub(crate) remainders: Vec>, + pub(crate) remainders: Vec>, pub(crate) var_size_data: Vec, pub(crate) last_id: Option, pub(crate) _phantom: PhantomData, } -#[derive(Clone, Debug)] -pub struct PostingElement { - pub(crate) id: PointOffsetType, - pub(crate) value: S, +/// A single element in the posting list, which contains an id and a value. +/// +/// Stores a remainder of the posting list. The difference with [`PostingElement`] is +/// so that this is zerocopy friendly +#[derive(Clone, Debug, FromBytes, Immutable, IntoBytes, KnownLayout)] +#[repr(C)] // Required for IntoBytes to work correctly +pub struct RemainderPosting { + /// U32 is required for pinning endianness of the id + pub id: U32, + pub value: S, } -#[derive(Debug, Clone)] +/// A single element in the posting list, which contains an id and a value. +/// +/// Output-facing structure. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct PostingElement { + pub id: PointOffsetType, + pub value: V, +} + +#[derive(Debug, Clone, FromBytes, Immutable, IntoBytes, KnownLayout)] #[repr(C)] pub struct PostingChunk { /// Initial data point id. Used for decompression. - pub initial_id: PointOffsetType, + pub initial_id: U32, /// An offset within id_data - pub offset: u32, + pub offset: U32, /// Sized values for the chunk. pub sized_values: [S; CHUNK_LEN], @@ -46,14 +65,15 @@ impl PostingChunk { /// Get byte size of the compressed ids chunk. pub(crate) fn get_compressed_size( chunks: &[PostingChunk], - data: &[u8], + ids_data: &[u8], chunk_index: usize, ) -> usize { if chunk_index + 1 < chunks.len() { - chunks[chunk_index + 1].offset as usize - chunks[chunk_index].offset as usize + chunks[chunk_index + 1].offset.get() as usize + - chunks[chunk_index].offset.get() as usize } else { // Last chunk - data.len() - chunks[chunk_index].offset as usize + ids_data.len() - chunks[chunk_index].offset.get() as usize } } } @@ -83,4 +103,32 @@ impl PostingList { let view = self.view(); PostingVisitor::new(view) } + + pub fn iter(&self) -> PostingIterator<'_, H> + where + H::Value: Clone, + { + self.visitor().into_iter() + } + + pub fn len(&self) -> usize { + self.chunks.len() * CHUNK_LEN + self.remainders.len() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +impl FromIterator<(PointOffsetType, H::Value)> for PostingList +where + H: ValueHandler, +{ + fn from_iter>(iter: T) -> Self { + let mut builder = PostingBuilder::new(); + for (id, value) in iter { + builder.add(id, value); + } + builder.build_generic::() + } } diff --git a/lib/posting_list/src/tests.rs b/lib/posting_list/src/tests.rs index d976dd6fba..1959f6b7ba 100644 --- a/lib/posting_list/src/tests.rs +++ b/lib/posting_list/src/tests.rs @@ -1,10 +1,12 @@ +use std::collections::HashMap; + use common::types::PointOffsetType; use rand::distr::{Alphanumeric, SampleString}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; use crate::value_handler::ValueHandler; -use crate::{CHUNK_LEN, PostingBuilder, PostingList, UnsizedValue}; +use crate::{CHUNK_LEN, PostingBuilder, PostingList, SizedHandler, UnsizedHandler, UnsizedValue}; // Simple struct that implements VarSizedValue for testing #[derive(Debug, Clone, PartialEq)] @@ -28,7 +30,7 @@ impl UnsizedValue for TestString { #[test] fn test_just_ids_against_vec() { check_various_lengths(|len| { - let posting_list = check_against_sorted_vec(|_rng, _id| (), |builder| builder.build(), len); + let posting_list = check_against_sorted_vec::<_, SizedHandler<()>>(|_rng, _id| (), len); // validate that chunks' sized values are empty if let Some(first_chunk) = posting_list.chunks.first() { @@ -46,13 +48,12 @@ fn test_just_ids_against_vec() { fn test_var_sized_against_vec() { let alphanumeric = Alphanumeric; check_various_lengths(|len| { - check_against_sorted_vec( + check_against_sorted_vec::<_, UnsizedHandler>( |rng, id| { let len = rng.random_range(1..=20); let s = alphanumeric.sample_string(rng, len); TestString(format!("item_{id} {s}")) }, - |builder| builder.build_unsized(), len, ); }) @@ -61,11 +62,7 @@ fn test_var_sized_against_vec() { #[test] fn test_fixed_sized_against_vec() { check_various_lengths(|len| { - check_against_sorted_vec( - |_rng, id| u64::from(id) * 100, - |builder| builder.build_sized(), - len, - ); + check_against_sorted_vec::<_, SizedHandler>(|_rng, id| u64::from(id) * 100, len); }); } @@ -94,6 +91,8 @@ fn check_various_lengths(check: impl Fn(u32)) { CHUNK_LEN - 1, CHUNK_LEN, CHUNK_LEN + 1, + CHUNK_LEN + 2, + 2 * CHUNK_LEN + 10, 100 * CHUNK_LEN, 500 * CHUNK_LEN + 1, 500 * CHUNK_LEN - 1, @@ -104,12 +103,11 @@ fn check_various_lengths(check: impl Fn(u32)) { } } -fn check_against_sorted_vec(gen_value: G, build: B, postings_count: u32) -> PostingList +fn check_against_sorted_vec(gen_value: G, postings_count: u32) -> PostingList where G: Fn(&mut StdRng, PointOffsetType) -> H::Value, H: ValueHandler, - B: FnOnce(PostingBuilder) -> PostingList, - H::Value: Clone + PartialEq, + H::Value: Clone + PartialEq + std::fmt::Debug, { let rng = &mut StdRng::seed_from_u64(42); let test_data = generate_data(postings_count, rng, gen_value); @@ -125,10 +123,11 @@ where } // Build the actual posting list - let posting_list = build(builder); + let posting_list = builder.build_generic(); // Access the posting list let mut visitor = posting_list.visitor(); + let mut intersection_iter = posting_list.iter(); // Validate len() assert_eq!(visitor.len(), model.len()); @@ -144,6 +143,12 @@ where // also check that contains function works assert!(visitor.contains(*expected_id)); + + // also check that the intersection is full + let intersection = intersection_iter + .advance_until_greater_or_equal(*expected_id) + .unwrap(); + assert_eq!(intersection.id, *expected_id); } // Bounds check @@ -154,5 +159,16 @@ where // There is no such id assert!(!visitor.contains(postings_count)); + // intersect against all sequential ids in the posting range, model is a hashmap in this case + let model = model.into_iter().collect::>(); + let mut intersection_iter = posting_list.iter(); + for seq_id in 0..postings_count { + let model_contains = model.contains_key(&seq_id); + let iter_contains = intersection_iter + .advance_until_greater_or_equal(seq_id) + .is_some_and(|elem| elem.id == seq_id); + assert_eq!(model_contains, iter_contains, "Mismatch at seq_id {seq_id}"); + } + posting_list } diff --git a/lib/posting_list/src/value_handler.rs b/lib/posting_list/src/value_handler.rs index 31a5dcb009..36414ffb4c 100644 --- a/lib/posting_list/src/value_handler.rs +++ b/lib/posting_list/src/value_handler.rs @@ -11,9 +11,9 @@ use crate::{SizedValue, UnsizedValue}; /// - For variable-size values, [`ValueHandler::Sized`] is an offset into the var_sized_data pub trait ValueHandler { /// The type of value in each PostingElement. - type Value: std::fmt::Debug; + type Value; /// The value to store within each chunk, or alongside each id. - type Sized: std::fmt::Debug + std::marker::Sized + Copy; + type Sized: std::marker::Sized + Copy; /// Process values before storage and return the necessary var_sized_data /// @@ -31,6 +31,7 @@ pub trait ValueHandler { } /// Fixed-size value handler +#[derive(Debug, Clone, Copy)] pub struct SizedHandler(PhantomData); impl ValueHandler for SizedHandler { @@ -50,6 +51,7 @@ impl ValueHandler for SizedHandler { } /// Var-size value handler +#[derive(Debug, Clone, Copy)] pub struct UnsizedHandler(PhantomData); impl ValueHandler for UnsizedHandler { diff --git a/lib/posting_list/src/view.rs b/lib/posting_list/src/view.rs index 4eb599de03..254716e8a7 100644 --- a/lib/posting_list/src/view.rs +++ b/lib/posting_list/src/view.rs @@ -1,21 +1,23 @@ use std::marker::PhantomData; +use std::ops::RangeInclusive; use bitpacking::BitPacker; use common::types::PointOffsetType; +use zerocopy::little_endian::U32; +use crate::iterator::PostingIterator; +use crate::posting_list::RemainderPosting; use crate::value_handler::{SizedHandler, ValueHandler}; use crate::visitor::PostingVisitor; -use crate::{ - BitPackerImpl, CHUNK_LEN, IdsPostingListView, PostingChunk, PostingElement, SizedValue, -}; +use crate::{BitPackerImpl, CHUNK_LEN, IdsPostingListView, PostingChunk, PostingList, SizedValue}; /// A non-owning view of [`PostingList`]. -#[derive(Debug, Clone)] +#[derive(Debug)] pub struct PostingListView<'a, H: ValueHandler> { pub(crate) id_data: &'a [u8], pub(crate) chunks: &'a [PostingChunk], pub(crate) var_size_data: &'a [u8], - pub(crate) remainders: &'a [PostingElement], + pub(crate) remainders: &'a [RemainderPosting], pub(crate) last_id: Option, pub(crate) _phantom: PhantomData, } @@ -24,15 +26,15 @@ pub struct PostingListComponents<'a, S> { pub id_data: &'a [u8], pub chunks: &'a [PostingChunk], pub var_size_data: &'a [u8], - pub remainders: &'a [PostingElement], - pub last_id: Option, + pub remainders: &'a [RemainderPosting], + pub last_id: Option, } impl<'a> IdsPostingListView<'a> { pub fn from_ids_components( id_data: &'a [u8], chunks: &'a [PostingChunk<()>], - remainders: &'a [PostingElement<()>], + remainders: &'a [RemainderPosting<()>], last_id: Option, ) -> Self { Self { @@ -50,7 +52,7 @@ impl<'a, V: SizedValue> PostingListView<'a, SizedHandler> { pub fn from_weighted_ids_components( id_data: &'a [u8], chunks: &'a [PostingChunk], - remainders: &'a [PostingElement], + remainders: &'a [RemainderPosting], last_id: Option, ) -> Self { Self { @@ -64,11 +66,36 @@ impl<'a, V: SizedValue> PostingListView<'a, SizedHandler> { } } +impl<'a, H: ValueHandler> IntoIterator for PostingListView<'a, H> +where + H::Value: Clone, +{ + type Item = as Iterator>::Item; + type IntoIter = PostingIterator<'a, H>; + + fn into_iter(self) -> Self::IntoIter { + self.visitor().into_iter() + } +} + impl<'a, H: ValueHandler> PostingListView<'a, H> { pub fn visitor(self) -> PostingVisitor<'a, H> { PostingVisitor::new(self) } + // not implemented as ToOwned trait because it requires PostingList's Borrow to return + // a &PostingListView, which is not possible because it's a non-owning view + pub fn to_owned(self) -> PostingList { + PostingList { + id_data: self.id_data.to_vec(), + chunks: self.chunks.to_vec(), + var_size_data: self.var_size_data.to_vec(), + remainders: self.remainders.to_vec(), + last_id: self.last_id, + _phantom: PhantomData, + } + } + pub fn components(&self) -> PostingListComponents { let Self { id_data, @@ -84,7 +111,7 @@ impl<'a, H: ValueHandler> PostingListView<'a, H> { chunks, var_size_data, remainders, - last_id: *last_id, + last_id: last_id.map(U32::from), } } @@ -92,7 +119,7 @@ impl<'a, H: ValueHandler> PostingListView<'a, H> { id_data: &'a [u8], chunks: &'a [PostingChunk], var_size_data: &'a [u8], - remainders: &'a [PostingElement], + remainders: &'a [RemainderPosting], last_id: Option, ) -> Self { Self { @@ -114,9 +141,13 @@ impl<'a, H: ValueHandler> PostingListView<'a, H> { let compressed_size = PostingChunk::get_compressed_size(self.chunks, self.id_data, chunk_index); let chunk_bits = compressed_size * u8::BITS as usize / CHUNK_LEN; - BitPackerImpl::new().decompress_strictly_sorted( - chunk.initial_id.checked_sub(1), - &self.id_data[chunk.offset as usize..chunk.offset as usize + compressed_size], + + let start_offset = chunk.offset.get() as usize; + let end_offset = start_offset + compressed_size; + + BitPackerImpl::new().decompress_sorted( + chunk.initial_id.get(), + &self.id_data[start_offset..end_offset], decompressed_chunk, chunk_bits as u8, ); @@ -129,21 +160,17 @@ impl<'a, H: ValueHandler> PostingListView<'a, H> { self.chunks.get(chunk_idx).map(|chunk| &chunk.sized_values) } - pub(crate) fn is_in_range(&self, id: PointOffsetType) -> bool { - let Some(last_id) = self.last_id else { - return false; - }; + pub(crate) fn ids_range(&self, start_chunk: usize) -> Option> { + // if there is no last id, it means the posting list is empty + let last_id = self.last_id?; - let Some(initial_id) = self + let initial_id = self .chunks - .first() - .map(|chunk| chunk.initial_id) - .or_else(|| self.remainders.first().map(|elem| elem.id)) - else { - return false; - }; + .get(start_chunk) + .map(|chunk| chunk.initial_id.get()) + .or_else(|| self.remainders.first().map(|elem| elem.id.get()))?; - id >= initial_id && id <= last_id + Some(initial_id..=last_id) } /// Find the chunk that may contain the id. @@ -154,27 +181,23 @@ impl<'a, H: ValueHandler> PostingListView<'a, H> { let remainders = self.remainders; let chunks = self.chunks; - if chunks.is_empty() { - return None; - } - - // check if id is in the remainders list - if remainders.first().is_some_and(|elem| id >= elem.id) { + // check if id might be in the remainders list + if remainders.first().is_some_and(|elem| id >= elem.id.get()) { return None; } let start_chunk = start_chunk.unwrap_or(0); - let chunks_slice = &chunks[start_chunk..]; + let chunks_slice = chunks.get(start_chunk..)?; if chunks_slice.is_empty() { return None; } // No need to check if id is under range of posting list, // this function assumes it is within the range - debug_assert!(id >= chunks_slice[0].initial_id); + debug_assert!(id >= chunks_slice[0].initial_id.get()); debug_assert!(self.last_id.is_some_and(|last_id| id <= last_id)); - match chunks_slice.binary_search_by(|chunk| chunk.initial_id.cmp(&id)) { + match chunks_slice.binary_search_by(|chunk| chunk.initial_id.get().cmp(&id)) { // id is the initial value of the chunk with index idx Ok(idx) => Some(start_chunk + idx), @@ -190,10 +213,9 @@ impl<'a, H: ValueHandler> PostingListView<'a, H> { } } - pub(crate) fn search_in_remainders(&self, id: PointOffsetType) -> Option { + pub(crate) fn search_in_remainders(&self, id: PointOffsetType) -> Result { self.remainders - .binary_search_by(|elem| elem.id.cmp(&id)) - .ok() + .binary_search_by(|elem| elem.id.get().cmp(&id)) } /// The total number of elements in the posting list. diff --git a/lib/posting_list/src/visitor.rs b/lib/posting_list/src/visitor.rs index 9624144dfb..4723b31064 100644 --- a/lib/posting_list/src/visitor.rs +++ b/lib/posting_list/src/visitor.rs @@ -1,10 +1,9 @@ use common::types::PointOffsetType; -use crate::CHUNK_LEN; use crate::iterator::PostingIterator; -use crate::posting_list::PostingElement; use crate::value_handler::ValueHandler; use crate::view::PostingListView; +use crate::{CHUNK_LEN, PostingElement}; /// A visitor for a posting list which caches the latest decompressed chunk of ids. pub struct PostingVisitor<'a, H: ValueHandler> { @@ -31,6 +30,10 @@ impl<'a, H: ValueHandler> PostingVisitor<'a, H> { self.list.len() } + pub fn is_empty(&self) -> bool { + self.list.is_empty() + } + /// Returns the decompressed slice of ids for a chunk. /// /// Assumes the chunk_idx is valid. @@ -44,8 +47,74 @@ impl<'a, H: ValueHandler> PostingVisitor<'a, H> { &self.decompressed_chunk } + /// Returns the first offset whose element id is greater or equal to the given id. + /// + /// Returns `None` if there is no such element in the posting list + pub(crate) fn search_greater_or_equal( + &mut self, + id: PointOffsetType, + start_from_offset: Option, + ) -> Option { + let start_chunk = start_from_offset + .map(|offset| offset / CHUNK_LEN) + .unwrap_or(0); + + let ids_range = self.list.ids_range(start_chunk)?; + // check if the first in the chunk is already greater or equal to the target id + if ids_range.start() >= &id { + return Some(start_chunk * CHUNK_LEN); + } + // check if the target id is already greater than the last id + if ids_range.end() < &id { + return None; + } + + // Find the chunk that may contain the id and check if the id is in the chunk + let chunk_index = self.list.find_chunk(id, Some(start_chunk)); + + if let Some(chunk_index) = chunk_index { + let local_offset = match self.decompressed_chunk(chunk_index).binary_search(&id) { + Ok(found_local_offset) => found_local_offset, + Err(closest_local_offset) => { + // If the target id is bigger than all the values here, and smaller than the first id + // in the next chunk or remainders, then that next id is the closest greater id + if closest_local_offset >= CHUNK_LEN { + let next_offset = (chunk_index + 1) * CHUNK_LEN; + + let next_offset_exists = next_offset < self.len(); + + return next_offset_exists.then_some(next_offset); + } + + closest_local_offset + } + }; + + return Some(local_offset + (chunk_index * CHUNK_LEN)); + } + + // Check in remainders + let remainder_offset = match self.list.search_in_remainders(id) { + Ok(found_remainder_offset) => found_remainder_offset, + Err(closest_remainder_offset) => { + if closest_remainder_offset >= self.list.remainders.len() { + // There is no greater or equal id in the posting list + return None; + } + + closest_remainder_offset + } + }; + + Some(remainder_offset + self.list.chunks.len() * CHUNK_LEN) + } + pub fn contains(&mut self, id: PointOffsetType) -> bool { - if !self.list.is_in_range(id) { + if self + .list + .ids_range(0) + .is_none_or(|range| !range.contains(&id)) + { return false; } @@ -60,7 +129,7 @@ impl<'a, H: ValueHandler> PostingVisitor<'a, H> { .binary_search(&id) .is_ok() } else { - self.list.search_in_remainders(id).is_some() + self.list.search_in_remainders(id).is_ok() } } @@ -103,12 +172,18 @@ impl<'a, H: ValueHandler> PostingVisitor<'a, H> { let next_sized_value = || self.list.remainders.get(local_offset + 1).map(|r| r.value); let value = H::get_value(e.value, next_sized_value, self.list.var_size_data); - PostingElement { id, value } + PostingElement { + id: id.get(), + value, + } }) } } -impl<'a, H: ValueHandler> IntoIterator for PostingVisitor<'a, H> { +impl<'a, H: ValueHandler> IntoIterator for PostingVisitor<'a, H> +where + H::Value: Clone, +{ type Item = PostingElement; type IntoIter = PostingIterator<'a, H>; diff --git a/lib/segment/Cargo.toml b/lib/segment/Cargo.toml index d3660ff9e9..6e494c501e 100644 --- a/lib/segment/Cargo.toml +++ b/lib/segment/Cargo.toml @@ -117,6 +117,7 @@ common = { path = "../common/common" } io = { path = "../common/io" } macros = { path = "../macros" } memory = { path = "../common/memory" } +posting_list = { path = "../posting_list" } quantization = { path = "../quantization" } sparse = { path = "../sparse" } gpu = { path = "../gpu" } diff --git a/lib/segment/src/index/field_index/full_text_index/compressed_posting/compressed_chunks_reader.rs b/lib/segment/src/index/field_index/full_text_index/compressed_posting/compressed_chunks_reader.rs index 2f1c644446..25e76e209b 100644 --- a/lib/segment/src/index/field_index/full_text_index/compressed_posting/compressed_chunks_reader.rs +++ b/lib/segment/src/index/field_index/full_text_index/compressed_posting/compressed_chunks_reader.rs @@ -121,21 +121,10 @@ impl<'a> ChunkReader<'a> { ChunkReaderIter::new(self) } - pub fn to_vec(&self) -> Vec { - let postings: Vec = self.iter().collect(); - debug_assert!(postings.is_sorted()); - debug_assert_eq!(postings.len(), self.len()); - postings - } - pub fn len(&self) -> usize { self.chunks.len() * BitPackerImpl::BLOCK_LEN + self.remainder_postings.len() } - pub fn is_empty(&self) -> bool { - self.chunks.is_empty() && self.remainder_postings.is_empty() - } - pub fn chunks_len(&self) -> usize { self.chunks.len() } diff --git a/lib/segment/src/index/field_index/full_text_index/immutable_inverted_index.rs b/lib/segment/src/index/field_index/full_text_index/immutable_inverted_index.rs index 660354c7fc..910e3e0101 100644 --- a/lib/segment/src/index/field_index/full_text_index/immutable_inverted_index.rs +++ b/lib/segment/src/index/field_index/full_text_index/immutable_inverted_index.rs @@ -2,11 +2,11 @@ use std::collections::HashMap; use common::counter::hardware_counter::HardwareCounterCell; use common::types::PointOffsetType; +use posting_list::{IdsPostingList, IdsPostingListView, PostingBuilder}; use super::inverted_index::InvertedIndex; use super::mmap_inverted_index::MmapInvertedIndex; use crate::common::operation_error::{OperationError, OperationResult}; -use crate::index::field_index::full_text_index::compressed_posting::compressed_posting_list::CompressedPostingList; use crate::index::field_index::full_text_index::inverted_index::{ParsedQuery, TokenId}; use crate::index::field_index::full_text_index::mutable_inverted_index::MutableInvertedIndex; use crate::index::field_index::full_text_index::postings_iterator::intersect_compressed_postings_iterator; @@ -14,7 +14,7 @@ use crate::index::field_index::full_text_index::postings_iterator::intersect_com #[cfg_attr(test, derive(Clone))] #[derive(Default, Debug)] pub struct ImmutableInvertedIndex { - pub(in crate::index::field_index::full_text_index) postings: Vec, + pub(in crate::index::field_index::full_text_index) postings: Vec, pub(in crate::index::field_index::full_text_index) vocab: HashMap, pub(in crate::index::field_index::full_text_index) point_to_tokens_count: Vec>, pub(in crate::index::field_index::full_text_index) points_count: usize, @@ -66,7 +66,7 @@ impl InvertedIndex for ImmutableInvertedIndex { .iter() // We can safely pass hw_counter here because it's not measured. // Due to lifetime issues, we can't return a disposable counter. - .map(|posting| posting.reader()) + .map(|posting| posting.view()) .collect(); // in case of immutable index, deleted documents are still in the postings @@ -105,8 +105,8 @@ impl InvertedIndex for ImmutableInvertedIndex { // Check that all tokens are in document parsed_query.tokens.iter().all(|token_id| { - let postings = &self.postings[*token_id as usize]; - postings.reader().contains(point_id) + let posting_list = &self.postings[*token_id as usize]; + posting_list.visitor().contains(point_id) }) } @@ -159,12 +159,19 @@ impl From for ImmutableInvertedIndex { }) .collect(); - let postings: Vec = postings - .into_iter() - .map(|posting| CompressedPostingList::new(&posting.into_vec())) - .collect(); vocab.shrink_to_fit(); + let postings: Vec = postings + .into_iter() + .map(|posting| { + let mut builder = PostingBuilder::new(); + for id in posting.iter() { + builder.add_id(id); + } + builder.build() + }) + .collect(); + ImmutableInvertedIndex { postings, vocab, @@ -207,9 +214,9 @@ impl From<&MmapInvertedIndex> for ImmutableInvertedIndex { }) .collect(); - let postings: Vec = postings + let postings: Vec = postings .into_iter() - .map(|postings| CompressedPostingList::new(&postings.to_vec())) + .map(IdsPostingListView::to_owned) .collect(); vocab.shrink_to_fit(); diff --git a/lib/segment/src/index/field_index/full_text_index/inverted_index.rs b/lib/segment/src/index/field_index/full_text_index/inverted_index.rs index 409825dd33..755b228976 100644 --- a/lib/segment/src/index/field_index/full_text_index/inverted_index.rs +++ b/lib/segment/src/index/field_index/full_text_index/inverted_index.rs @@ -282,12 +282,14 @@ mod tests { let orig_posting = mutable.postings.get(*orig_token as usize).cloned().unwrap(); + let mut posting_visitor = new_posting.visitor(); let new_contains_orig = orig_posting .iter() - .all(|point_id| new_posting.reader().contains(point_id)); + .all(|point_id| posting_visitor.contains(point_id)); let orig_contains_new = new_posting .iter() + .map(|elem| elem.id) .all(|point_id| orig_posting.contains(point_id)); new_contains_orig && orig_contains_new @@ -326,20 +328,20 @@ mod tests { // Check same postings for (token_id, posting) in immutable.postings.iter().enumerate() { - let mutable_ids = posting.iter().collect::>(); + let mutable_elems = posting.iter().collect::>(); // Check mutable vs mmap - let mmap_ids = mmap + let mmap_elems = mmap .postings .get(token_id as u32, &hw_counter) .unwrap() - .iter() + .into_iter() .collect(); - assert_eq!(mutable_ids, mmap_ids); + assert_eq!(mutable_elems, mmap_elems); // Check mutable vs immutable mmap - let imm_mmap_ids = imm_mmap.postings[token_id].iter().collect(); - assert_eq!(mutable_ids, imm_mmap_ids); + let imm_mmap_elems = imm_mmap.postings[token_id].iter().collect(); + assert_eq!(mutable_elems, imm_mmap_elems); } for (point_id, count) in immutable.point_to_tokens_count.iter().enumerate() { diff --git a/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/mmap_postings.rs b/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/mmap_postings.rs index 484cc0c898..bda04fd6b2 100644 --- a/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/mmap_postings.rs +++ b/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/mmap_postings.rs @@ -9,11 +9,11 @@ use common::zeros::WriteZerosExt; use memmap2::Mmap; use memory::madvise::{Advice, AdviceSetting, Madviseable}; use memory::mmap_ops::open_read_mmap; +use posting_list::{ + IdsPostingList, IdsPostingListView, PostingChunk, PostingListComponents, RemainderPosting, +}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; -use crate::index::field_index::full_text_index::compressed_posting::compressed_chunks_reader::ChunkReader; -use crate::index::field_index::full_text_index::compressed_posting::compressed_common::CompressedPostingChunksIndex; -use crate::index::field_index::full_text_index::compressed_posting::compressed_posting_list::CompressedPostingList; use crate::index::field_index::full_text_index::inverted_index::TokenId; const ALIGNMENT: usize = 4; @@ -50,8 +50,8 @@ impl PostingListHeader { fn posting_size(&self) -> usize { self.data_bytes_count as usize + self.alignment_bytes_count as usize - + self.remainder_count as usize * size_of::() - + self.chunks_count as usize * size_of::() + + self.remainder_count as usize * size_of::>() + + self.chunks_count as usize * size_of::>() + size_of::() // last_doc_id } } @@ -89,16 +89,16 @@ impl MmapPostings { /// /// ```ignore /// last_doc_id: &'a PointOffsetType, - /// chunks_index: &'a [CompressedPostingChunksIndex], + /// chunks_index: &'a [PostingChunk<()>], /// data: &'a [u8], /// _alignment: &'a [u8], // 0-3 extra bytes to align the data /// remainder_postings: &'a [PointOffsetType], /// ``` - fn get_reader<'a>( + fn get_view<'a>( &'a self, header: &PostingListHeader, hw_counter: ConditionedCounter<'a>, - ) -> Option> { + ) -> Option> { let counter = hw_counter.payload_index_io_read_counter(); let bytes = self.mmap.get(header.offset as usize..)?; @@ -106,24 +106,24 @@ impl MmapPostings { let (last_doc_id, bytes) = PointOffsetType::read_from_prefix(bytes).ok()?; - counter.incr_delta(size_of::()); - let (chunks, bytes) = <[CompressedPostingChunksIndex]>::ref_from_prefix_with_elems( - bytes, - header.chunks_count as usize, - ) - .ok()?; + counter.incr_delta(size_of::>()); + let (chunks, bytes) = + <[PostingChunk<()>]>::ref_from_prefix_with_elems(bytes, header.chunks_count as usize) + .ok()?; let (data, bytes) = bytes.split_at(header.data_bytes_count as usize); let bytes = bytes.get(header.alignment_bytes_count as usize..)?; - let (remainder_postings, _) = - <[u32]>::ref_from_prefix_with_elems(bytes, header.remainder_count as usize).ok()?; + let (remainder_postings, _) = <[RemainderPosting<()>]>::ref_from_prefix_with_elems( + bytes, + header.remainder_count as usize, + ) + .ok()?; - Some(ChunkReader::new( - last_doc_id, - chunks, + Some(IdsPostingListView::from_ids_components( data, + chunks, remainder_postings, - hw_counter, + Some(last_doc_id), )) } @@ -131,7 +131,7 @@ impl MmapPostings { &'a self, token_id: TokenId, hw_counter: &'a HardwareCounterCell, - ) -> Option> { + ) -> Option> { let hw_counter = ConditionedCounter::new(self.on_disk, hw_counter); hw_counter @@ -140,12 +140,12 @@ impl MmapPostings { let header = self.get_header(token_id)?; - self.get_reader(header, hw_counter) + self.get_view(header, hw_counter) } /// Given a vector of compressed posting lists, this function writes them to the `path` file. /// The format of the file is compatible with the `MmapPostings` structure. - pub fn create(path: PathBuf, compressed_postings: &[CompressedPostingList]) -> io::Result<()> { + pub fn create(path: PathBuf, compressed_postings: &[IdsPostingList]) -> io::Result<()> { // Create a new empty file, where we will write the compressed posting lists and the header let file = tempfile::Builder::new() .prefix(path.file_name().ok_or(io::ErrorKind::InvalidInput)?) @@ -165,17 +165,24 @@ impl MmapPostings { let mut posting_offset = size_of::() + postings_lists_headers_size; for compressed_posting in compressed_postings { - let (data, chunks, remainder_postings) = compressed_posting.internal_structs(); + let view = compressed_posting.view(); + let PostingListComponents { + id_data, + chunks, + var_size_data: _, // not used with just ids postings + remainders, + last_id: _, // not used for the header + } = view.components(); - let data_len = data.len(); - let alignment_len = ALIGNMENT - data_len % ALIGNMENT; + let data_len = id_data.len(); + let alignment_len = data_len.next_multiple_of(ALIGNMENT) - data_len; let posting_list_header = PostingListHeader { offset: posting_offset as u64, chunks_count: chunks.len() as u32, - data_bytes_count: data.len() as u32, + data_bytes_count: data_len as u32, alignment_bytes_count: alignment_len as u8, - remainder_count: remainder_postings.len() as u8, + remainder_count: remainders.len() as u8, _reserved: [0; 6], }; @@ -186,25 +193,35 @@ impl MmapPostings { } for compressed_posting in compressed_postings { - let (data, chunks, remainder_postings) = compressed_posting.internal_structs(); + let view = compressed_posting.view(); + let PostingListComponents { + id_data, + chunks, + var_size_data: _, // not used with just ids postings + remainders, + last_id, + } = view.components(); - let last_doc_id = compressed_posting.last_doc_id(); - - bufw.write_all(last_doc_id.as_bytes())?; + bufw.write_all( + last_id + .expect("posting must have at least one element") + .as_bytes(), + )?; for chunk in chunks { bufw.write_all(chunk.as_bytes())?; } - bufw.write_all(data)?; + bufw.write_all(id_data)?; // Example: // For data size = 5, alignment = 3 as (5 + 3 = 8) - // alignment = 4 - 5 % 4 = 3 - bufw.write_zeros(ALIGNMENT - data.len() % ALIGNMENT)?; + // alignment = 8 - 5 = 3 + let data_len = id_data.len(); + bufw.write_zeros(data_len.next_multiple_of(ALIGNMENT) - data_len)?; - for posting in remainder_postings { - bufw.write_all(posting.as_bytes())?; + for element in remainders { + bufw.write_all(element.as_bytes())?; } } @@ -243,11 +260,11 @@ impl MmapPostings { self.mmap.populate(); } - /// Iterate over posting lists, returning chunk reader for each + /// Iterate over posting lists, returning a view for each pub fn iter_postings<'a>( &'a self, hw_counter: &'a HardwareCounterCell, - ) -> impl Iterator>> { + ) -> impl Iterator>> { (0..self.header.posting_count as u32).map(|posting_idx| self.get(posting_idx, hw_counter)) } } diff --git a/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/mod.rs b/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/mod.rs index a69609a32e..87be0d44d3 100644 --- a/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/mod.rs +++ b/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/mod.rs @@ -10,8 +10,8 @@ use memory::madvise::AdviceSetting; use memory::mmap_ops; use memory::mmap_type::{MmapBitSlice, MmapSlice}; use mmap_postings::MmapPostings; +use posting_list::IdsPostingListView; -use super::compressed_posting::compressed_chunks_reader::ChunkReader; use super::inverted_index::{InvertedIndex, ParsedQuery}; use super::postings_iterator::intersect_compressed_postings_iterator; use crate::common::mmap_bitslice_buffered_update_wrapper::MmapBitSliceBufferedUpdateWrapper; @@ -21,6 +21,12 @@ use crate::index::field_index::full_text_index::inverted_index::TokenId; mod mmap_postings; +/// Old implementation of mmap postings, used to test backwards compatibility temporarily +#[cfg(test)] +mod old_mmap_postings; +#[cfg(test)] +mod tests; + const POSTINGS_FILE: &str = "postings.dat"; const VOCAB_FILE: &str = "vocab.dat"; const POINT_TO_TOKENS_COUNT_FILE: &str = "point_to_tokens_count.dat"; @@ -123,12 +129,12 @@ impl MmapInvertedIndex { self.vocab.iter().map(|(k, v)| (k, v.first().unwrap())) } - /// Iterate over posting lists, returning chunk reader for each + /// Iterate over posting lists, returning a view for each #[inline] pub(super) fn iter_postings<'a>( &'a self, hw_counter: &'a HardwareCounterCell, - ) -> impl Iterator>> { + ) -> impl Iterator>> { self.postings.iter_postings(hw_counter) } @@ -273,6 +279,7 @@ impl InvertedIndex for MmapInvertedIndex { .get(*query_token, hw_counter) // unwrap safety: all tokens exist in the vocabulary, otherwise there'd be no query tokens .unwrap() + .visitor() .contains(point_id) }) } diff --git a/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/old_mmap_postings.rs b/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/old_mmap_postings.rs new file mode 100644 index 0000000000..f5e897c27f --- /dev/null +++ b/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/old_mmap_postings.rs @@ -0,0 +1,239 @@ +use std::io; +use std::io::Write; +use std::path::PathBuf; + +use common::counter::conditioned_counter::ConditionedCounter; +use common::counter::hardware_counter::HardwareCounterCell; +use common::types::PointOffsetType; +use common::zeros::WriteZerosExt; +use memmap2::Mmap; +use memory::madvise::{Advice, AdviceSetting}; +use memory::mmap_ops::open_read_mmap; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; + +use crate::index::field_index::full_text_index::compressed_posting::compressed_chunks_reader::ChunkReader; +use crate::index::field_index::full_text_index::compressed_posting::compressed_common::CompressedPostingChunksIndex; +use crate::index::field_index::full_text_index::compressed_posting::compressed_posting_list::CompressedPostingList; +use crate::index::field_index::full_text_index::inverted_index::TokenId; + +const ALIGNMENT: usize = 4; + +#[derive(Debug, Default, Clone, FromBytes, Immutable, IntoBytes, KnownLayout)] +#[repr(C)] +struct PostingsHeader { + /// Number of posting lists. One posting list per term + pub posting_count: usize, + _reserved: [u8; 32], +} + +/// This data structure should contain all the necessary information to +/// construct `CompressedMmapPostingList` from the mmap file. +#[derive(Debug, Default, Clone, FromBytes, Immutable, IntoBytes, KnownLayout)] +#[repr(C)] +struct PostingListHeader { + /// Offset in bytes from the start of the mmap file + /// where the posting list data starts + offset: u64, + /// Amount of chunks in compressed posting list + chunks_count: u32, + /// Length in bytes for the compressed postings data + data_bytes_count: u32, + /// Length in bytes for the alignment bytes + alignment_bytes_count: u8, + /// Length in bytes for the remainder postings + remainder_count: u8, + _reserved: [u8; 6], +} + +impl PostingListHeader { + /// Size of the posting list this header represents + fn posting_size(&self) -> usize { + self.data_bytes_count as usize + + self.alignment_bytes_count as usize + + self.remainder_count as usize * size_of::() + + self.chunks_count as usize * size_of::() + + size_of::() // last_doc_id + } +} + +/// MmapPostings Structure on disk: +/// +/// +/// `| PostingsHeader | +/// [ PostingListHeader, PostingListHeader, ... ] | +/// [ CompressedMmapPostingList, CompressedMmapPostingList, ... ] |` +pub struct MmapPostings { + _path: PathBuf, + mmap: Mmap, + header: PostingsHeader, + on_disk: bool, +} + +impl MmapPostings { + fn get_header(&self, token_id: TokenId) -> Option<&PostingListHeader> { + if self.header.posting_count <= token_id as usize { + return None; + } + + let header_offset = + size_of::() + token_id as usize * size_of::(); + + PostingListHeader::ref_from_prefix(self.mmap.get(header_offset..)?) + .ok() + .map(|(header, _)| header) + } + + /// Create ChunkReader from the given header + /// + /// Assume the following layout: + /// + /// ```ignore + /// last_doc_id: &'a PointOffsetType, + /// chunks_index: &'a [CompressedPostingChunksIndex], + /// data: &'a [u8], + /// _alignment: &'a [u8], // 0-3 extra bytes to align the data + /// remainder_postings: &'a [PointOffsetType], + /// ``` + fn get_reader<'a>( + &'a self, + header: &PostingListHeader, + hw_counter: ConditionedCounter<'a>, + ) -> Option> { + let counter = hw_counter.payload_index_io_read_counter(); + + let bytes = self.mmap.get(header.offset as usize..)?; + counter.incr_delta(size_of::()); + + let (last_doc_id, bytes) = PointOffsetType::read_from_prefix(bytes).ok()?; + + counter.incr_delta(size_of::()); + let (chunks, bytes) = <[CompressedPostingChunksIndex]>::ref_from_prefix_with_elems( + bytes, + header.chunks_count as usize, + ) + .ok()?; + let (data, bytes) = bytes.split_at(header.data_bytes_count as usize); + let bytes = bytes.get(header.alignment_bytes_count as usize..)?; + + let (remainder_postings, _) = + <[u32]>::ref_from_prefix_with_elems(bytes, header.remainder_count as usize).ok()?; + + Some(ChunkReader::new( + last_doc_id, + chunks, + data, + remainder_postings, + hw_counter, + )) + } + + pub fn get<'a>( + &'a self, + token_id: TokenId, + hw_counter: &'a HardwareCounterCell, + ) -> Option> { + let hw_counter = ConditionedCounter::new(self.on_disk, hw_counter); + + hw_counter + .payload_index_io_read_counter() + .incr_delta(size_of::()); + + let header = self.get_header(token_id)?; + + self.get_reader(header, hw_counter) + } + + /// Given a vector of compressed posting lists, this function writes them to the `path` file. + /// The format of the file is compatible with the `MmapPostings` structure. + pub fn create(path: PathBuf, compressed_postings: &[CompressedPostingList]) -> io::Result<()> { + // Create a new empty file, where we will write the compressed posting lists and the header + let file = tempfile::Builder::new() + .prefix(path.file_name().ok_or(io::ErrorKind::InvalidInput)?) + .tempfile_in(path.parent().ok_or(io::ErrorKind::InvalidInput)?)?; + let mut bufw = io::BufWriter::new(&file); + + let postings_header = PostingsHeader { + posting_count: compressed_postings.len(), + _reserved: [0; 32], + }; + + // Write the header to the buffer + bufw.write_all(postings_header.as_bytes())?; + + let postings_lists_headers_size = + compressed_postings.len() * size_of::(); + let mut posting_offset = size_of::() + postings_lists_headers_size; + + for compressed_posting in compressed_postings { + let (data, chunks, remainder_postings) = compressed_posting.internal_structs(); + + let data_len = data.len(); + let alignment_len = ALIGNMENT - data_len % ALIGNMENT; + + let posting_list_header = PostingListHeader { + offset: posting_offset as u64, + chunks_count: chunks.len() as u32, + data_bytes_count: data.len() as u32, + alignment_bytes_count: alignment_len as u8, + remainder_count: remainder_postings.len() as u8, + _reserved: [0; 6], + }; + + // Write the posting list header to the buffer + bufw.write_all(posting_list_header.as_bytes())?; + + posting_offset += posting_list_header.posting_size(); + } + + for compressed_posting in compressed_postings { + let (data, chunks, remainder_postings) = compressed_posting.internal_structs(); + + let last_doc_id = compressed_posting.last_doc_id(); + + bufw.write_all(last_doc_id.as_bytes())?; + + for chunk in chunks { + bufw.write_all(chunk.as_bytes())?; + } + + bufw.write_all(data)?; + + // Example: + // For data size = 5, alignment = 3 as (5 + 3 = 8) + // alignment = 4 - 5 % 4 = 3 + bufw.write_zeros(ALIGNMENT - data.len() % ALIGNMENT)?; + + for posting in remainder_postings { + bufw.write_all(posting.as_bytes())?; + } + } + + // Explicitly flush write buffer so we can catch IO errors + bufw.flush()?; + drop(bufw); + + file.as_file().sync_all()?; + file.persist(path)?; + + Ok(()) + } + + pub fn open(path: impl Into, populate: bool) -> io::Result { + let path = path.into(); + let mmap = open_read_mmap(&path, AdviceSetting::Advice(Advice::Normal), populate)?; + + let (header, _) = PostingsHeader::read_from_prefix(&mmap).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid header deserialization in {}", path.display()), + ) + })?; + + Ok(Self { + _path: path, + mmap, + header, + on_disk: !populate, + }) + } +} diff --git a/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/tests.rs b/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/tests.rs new file mode 100644 index 0000000000..d7020fde54 --- /dev/null +++ b/lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/tests.rs @@ -0,0 +1,61 @@ +use common::counter::hardware_counter::HardwareCounterCell; +use common::types::PointOffsetType; +use rand::Rng; + +use super::{mmap_postings, old_mmap_postings}; +use crate::index::field_index::full_text_index::compressed_posting::compressed_posting_list::CompressedPostingList; + +fn generate_ids(rng: &mut impl Rng, amount: usize) -> Vec { + let distr = rand::distr::Uniform::new(0, amount as u32).unwrap(); + rng.sample_iter(distr).take(amount).collect() +} + +#[test] +fn test_mmap_posting_lists_compatibility() { + let rng = &mut rand::rng(); + let lengths = [138, 14, 1889, 128, 129, 127]; + + // postings in a vec of vecs + let postings = lengths + .into_iter() + .map(|len| generate_ids(rng, len)) + .collect::>(); + + // old compressed postings implementation + let old_compressed_postings = postings + .iter() + .map(|ids| CompressedPostingList::new(ids)) + .collect::>(); + + let dir = tempfile::tempdir().unwrap(); + let postings_path = dir.path().join("postings.dat"); + + // Create mmap postings file + old_mmap_postings::MmapPostings::create(postings_path.clone(), &old_compressed_postings) + .unwrap(); + + // open with old impl + let old_postings = old_mmap_postings::MmapPostings::open(postings_path.clone(), true).unwrap(); + + // open with new impl + let new_postings = mmap_postings::MmapPostings::open(postings_path.clone(), true).unwrap(); + + let hw_counter = HardwareCounterCell::disposable(); + + for token_id in 0..postings.len() as u32 { + let old = old_postings.get(token_id, &hw_counter).unwrap(); + let new = new_postings.get(token_id, &hw_counter).unwrap(); + let model = &postings[token_id as usize]; + + // check all impls iterate the same ids + for (offset, id_old, elem_new, &id_model) in + itertools::multizip((0u32.., old.iter(), new.into_iter(), model.iter())) + { + assert!( + id_model == id_old && id_model == elem_new.id, + "Mismatch at token_id {token_id}: offset: {offset}, old: {id_old}, new: {}, model: {id_model}", + elem_new.id + ); + } + } +} diff --git a/lib/segment/src/index/field_index/full_text_index/mod.rs b/lib/segment/src/index/field_index/full_text_index/mod.rs index 99eaa8cdc1..7da06431fc 100644 --- a/lib/segment/src/index/field_index/full_text_index/mod.rs +++ b/lib/segment/src/index/field_index/full_text_index/mod.rs @@ -8,6 +8,7 @@ mod postings_iterator; pub mod text_index; mod tokenizers; +#[cfg(test)] mod compressed_posting; mod immutable_inverted_index; mod mutable_inverted_index; diff --git a/lib/segment/src/index/field_index/full_text_index/posting_list.rs b/lib/segment/src/index/field_index/full_text_index/posting_list.rs index cc497bb027..b7fd1cffbb 100644 --- a/lib/segment/src/index/field_index/full_text_index/posting_list.rs +++ b/lib/segment/src/index/field_index/full_text_index/posting_list.rs @@ -44,9 +44,4 @@ impl PostingList { pub fn iter(&self) -> impl Iterator + '_ { self.list.iter().copied() } - - #[inline] - pub fn into_vec(self) -> Vec { - self.list - } } diff --git a/lib/segment/src/index/field_index/full_text_index/postings_iterator.rs b/lib/segment/src/index/field_index/full_text_index/postings_iterator.rs index 3b64d1b02f..16b17402ee 100644 --- a/lib/segment/src/index/field_index/full_text_index/postings_iterator.rs +++ b/lib/segment/src/index/field_index/full_text_index/postings_iterator.rs @@ -1,9 +1,7 @@ use common::types::PointOffsetType; +use posting_list::{IdsPostingListView, PostingListView}; use super::posting_list::PostingList; -use crate::index::field_index::full_text_index::compressed_posting::compressed_chunks_reader::ChunkReader; -use crate::index::field_index::full_text_index::compressed_posting::compressed_posting_iterator::CompressedPostingIterator; -use crate::index::field_index::full_text_index::compressed_posting::compressed_posting_visitor::CompressedPostingVisitor; pub fn intersect_postings_iterator<'a>( mut postings: Vec<&'a PostingList>, @@ -24,7 +22,7 @@ pub fn intersect_postings_iterator<'a>( } pub fn intersect_compressed_postings_iterator<'a>( - mut postings: Vec>, + mut postings: Vec>, filter: impl Fn(PointOffsetType) -> bool + 'a, ) -> Box + 'a> { let smallest_posting_idx = postings @@ -34,20 +32,26 @@ pub fn intersect_compressed_postings_iterator<'a>( .map(|(idx, _posting)| idx) .unwrap(); let smallest_posting = postings.remove(smallest_posting_idx); - let smallest_posting_iterator = - CompressedPostingIterator::new(CompressedPostingVisitor::new(smallest_posting)); + let smallest_posting_iterator = smallest_posting.into_iter(); - let mut posting_visitors = postings + let mut posting_iterators = postings .into_iter() - .map(CompressedPostingVisitor::new) + .map(PostingListView::into_iter) .collect::>(); let and_iter = smallest_posting_iterator - .filter(move |doc_id| filter(*doc_id)) - .filter(move |doc_id| { - posting_visitors - .iter_mut() - .all(|posting_visitor| posting_visitor.contains_next_and_advance(*doc_id)) + .map(|elem| elem.id) + .filter(move |id| { + filter(*id) + && posting_iterators.iter_mut().all(|posting_iterator| { + // Custom "contains" check, which leverages the fact that smallest posting is sorted, + // so the next id that must be in all postings is strictly greater than the previous one. + // + // This means that the other iterators can remember the last id they returned to avoid extra work + posting_iterator + .advance_until_greater_or_equal(*id) + .is_some_and(|elem| elem.id == *id) + }) }); Box::new(and_iter) @@ -56,8 +60,9 @@ pub fn intersect_compressed_postings_iterator<'a>( #[cfg(test)] mod tests { + use posting_list::IdsPostingList; + use super::*; - use crate::index::field_index::full_text_index::compressed_posting::compressed_posting_list::CompressedPostingList; #[test] fn test_postings_iterator() { @@ -86,13 +91,13 @@ mod tests { assert_eq!(res, vec![2, 5]); - let p1_compressed = CompressedPostingList::new(&p1.into_vec()); - let p2_compressed = CompressedPostingList::new(&p2.into_vec()); - let p3_compressed = CompressedPostingList::new(&p3.into_vec()); + let p1_compressed: IdsPostingList = p1.iter().map(|id| (id, ())).collect(); + let p2_compressed: IdsPostingList = p2.iter().map(|id| (id, ())).collect(); + let p3_compressed: IdsPostingList = p3.iter().map(|id| (id, ())).collect(); let compressed_posting_reades = vec![ - p1_compressed.reader(), - p2_compressed.reader(), - p3_compressed.reader(), + p1_compressed.view(), + p2_compressed.view(), + p3_compressed.view(), ]; let merged = intersect_compressed_postings_iterator(compressed_posting_reades, |_| true);