diff --git a/Cargo.lock b/Cargo.lock index 88f717749a..167ab474d9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4405,6 +4405,16 @@ version = "1.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc59d1bcc64fc5d021d67521f818db868368028108d37f0e98d74e33f68297b5" +[[package]] +name = "posting_list" +version = "0.0.0" +dependencies = [ + "bitpacking", + "common", + "rand 0.9.1", + "tempfile", +] + [[package]] name = "pprof" version = "0.14.0" diff --git a/Cargo.toml b/Cargo.toml index 03c3fe9928..4b66c0cbc8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -171,6 +171,7 @@ private_intra_doc_links = "allow" ahash = { version = "0.8.11", features = ["serde"] } atomicwrites = "0.4.4" bincode = "1.3.3" # no upgrade because 2.0.x is much slower https://github.com/qdrant/qdrant/pull/6134 +bitpacking = "0.9.2" bytemuck = { version = "1.23.0", features = ["extern_crate_alloc", "must_cast", "transparentwrapper_extra"] } bytes = "1.10.1" chrono = { version = "0.4.41", features = ["serde"] } @@ -263,6 +264,7 @@ members = [ "lib/collection", "lib/common/*", "lib/macros", + "lib/posting_list", "lib/segment", "lib/sparse", "lib/storage", diff --git a/lib/posting_list/Cargo.toml b/lib/posting_list/Cargo.toml new file mode 100644 index 0000000000..a2c1f14d43 --- /dev/null +++ b/lib/posting_list/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "posting_list" +version = "0.0.0" +authors = ["Qdrant Team "] +license = "Apache-2.0" +edition = "2024" +publish = false + +[lints] +workspace = true + +[dependencies] +common = { path = "../common/common" } +bitpacking = { workspace = true } + +[dev-dependencies] +rand = { workspace = true } +tempfile = { workspace = true } diff --git a/lib/posting_list/src/builder.rs b/lib/posting_list/src/builder.rs new file mode 100644 index 0000000000..b3e66a8417 --- /dev/null +++ b/lib/posting_list/src/builder.rs @@ -0,0 +1,148 @@ +use std::marker::PhantomData; + +use bitpacking::BitPacker; +use common::types::PointOffsetType; + +use crate::posting_list::{PostingChunk, PostingElement, PostingList}; +use crate::value_handler::{SizedHandler, UnsizedHandler, ValueHandler}; +use crate::{ + BitPackerImpl, CHUNK_LEN, SizedValue, UnsizedValue, VarPostingList, WeightsPostingList, +}; + +pub struct PostingBuilder { + elements: Vec>, +} + +impl Default for PostingBuilder { + fn default() -> Self { + Self { + elements: Vec::new(), + } + } +} + +impl PostingBuilder { + pub fn new() -> Self { + Self::default() + } + + pub fn add(&mut self, id: PointOffsetType, value: V) { + self.elements.push(PostingElement { id, value }); + } + + /// Unified implementation that works for both fixed-size and variable-size values + /// + /// This method uses the `ValueHandler::process_values` trait function to abstract the + /// differences between the two implementations, allowing us to share the common logic. + pub(crate) fn build_generic(mut self) -> PostingList + where + H: ValueHandler, + { + self.elements.sort_unstable_by_key(|e| e.id); + + let num_elements = self.elements.len(); + + // extract ids and values into separate lists + let (ids, values): (Vec<_>, Vec<_>) = + self.elements.into_iter().map(|e| (e.id, e.value)).unzip(); + + // process values + let (sized_values, var_size_data) = H::process_values(values); + + let bitpacker = BitPackerImpl::new(); + let mut chunks = Vec::with_capacity(ids.len() / CHUNK_LEN); + let mut id_data_size = 0; + + // process full chunks + let ids_chunks_iter = ids.chunks_exact(CHUNK_LEN); + let values_chunks_iter = sized_values.chunks_exact(CHUNK_LEN); + let remainder_ids = ids_chunks_iter.remainder(); + let remainder_values = values_chunks_iter.remainder(); + + 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_size = BitPackerImpl::compressed_block_size(chunk_bits); + + chunks.push(PostingChunk { + initial_id: initial, + offset: u32::try_from(id_data_size) + .expect("id_data_size should fit in u32, (smaller than 4GB)"), + sized_values: chunk_values + .try_into() + .expect("should be a valid chunk size"), + }); + id_data_size += chunk_size; + } + + // 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 }); + } + + // compress id_data + let mut id_data = vec![0u8; id_data_size]; + for (chunk_index, chunk_ids) in ids.chunks_exact(CHUNK_LEN).enumerate() { + 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), + chunk_ids, + &mut id_data[chunk.offset as usize..chunk.offset as usize + compressed_size], + chunk_bits as u8, + ); + } + + let last_id = ids.last().copied(); + + PostingList { + id_data, + var_size_data, + chunks, + remainders, + last_id, + _phantom: PhantomData, + } + } +} + +impl PostingBuilder<()> { + /// Add an id without a value. + pub fn add_id(&mut self, id: PointOffsetType) { + self.add(id, ()); + } + + /// Build a posting list with just the compressed ids. + pub fn build(self) -> PostingList> { + self.build_generic::>() + } +} + +impl PostingBuilder { + /// Build a posting list with fixed-sized values to store them directly in the PostingChunk + pub fn build_sized(self) -> WeightsPostingList { + self.build_generic::>() + } +} + +// Variable-sized value implementation. +impl PostingBuilder { + /// Build a posting list with variable-sized values to store them in the `var_size_data` field. + /// + /// For variable-size values, we store offsets along with each id to point into a flattened array + /// where the var-sized data lives. + pub fn build_unsized(self) -> VarPostingList { + self.build_generic::>() + } +} + +impl From> for PostingList +where + H: ValueHandler, +{ + fn from(value: PostingBuilder) -> Self { + value.build_generic::() + } +} diff --git a/lib/posting_list/src/iterator.rs b/lib/posting_list/src/iterator.rs new file mode 100644 index 0000000000..c4969b4d72 --- /dev/null +++ b/lib/posting_list/src/iterator.rs @@ -0,0 +1,51 @@ +use std::iter::FusedIterator; + +use common::types::PointOffsetType; + +use crate::PostingElement; +use crate::value_handler::ValueHandler; +use crate::visitor::PostingVisitor; + +pub struct PostingIterator<'a, H: ValueHandler> { + visitor: PostingVisitor<'a, H>, + current_id: Option, + offset: usize, +} + +impl<'a, H: ValueHandler> PostingIterator<'a, H> { + pub fn new(visitor: PostingVisitor<'a, H>) -> Self { + Self { + visitor, + current_id: None, + offset: 0, + } + } +} + +impl Iterator for PostingIterator<'_, H> { + type Item = PostingElement; + + fn next(&mut self) -> Option { + self.visitor.get_by_offset(self.offset).inspect(|elem| { + self.current_id = Some(elem.id); + self.offset += 1; + }) + } + + fn size_hint(&self) -> (usize, Option) { + let remaining_len = self.len(); + (remaining_len, Some(remaining_len)) + } + + fn count(self) -> usize { + self.size_hint().0 + } +} + +impl ExactSizeIterator for PostingIterator<'_, H> { + fn len(&self) -> usize { + self.visitor.list.len().saturating_sub(self.offset) + } +} + +impl FusedIterator for PostingIterator<'_, H> {} diff --git a/lib/posting_list/src/lib.rs b/lib/posting_list/src/lib.rs new file mode 100644 index 0000000000..04ae86792c --- /dev/null +++ b/lib/posting_list/src/lib.rs @@ -0,0 +1,53 @@ +mod builder; +mod iterator; +mod posting_list; +#[cfg(test)] +mod tests; +mod value_handler; +mod view; +mod visitor; + +use bitpacking::BitPacker; + +type BitPackerImpl = bitpacking::BitPacker4x; + +/// How many elements are packed in a single chunk. +const CHUNK_LEN: usize = 128; +const _: () = assert!(128 == BitPackerImpl::BLOCK_LEN); + +pub trait SizedValue: Sized + Copy + std::fmt::Debug {} + +impl SizedValue for () {} +impl SizedValue for u32 {} +impl SizedValue for u64 {} + +pub trait UnsizedValue: std::fmt::Debug { + fn write_len(&self) -> usize; + + fn write_to(&self, dst: &mut [u8]); + + fn from_bytes(data: &[u8]) -> Self; +} + +/// Posting list of ids, where ids are compressed. +pub type IdsPostingList = PostingList>; + +/// Posting list of ids + small fixed-sized values, where ids are compressed. +pub type WeightsPostingList = PostingList>; + +/// Posting list of ids + variably-sized values, where ids are compressed. +pub type VarPostingList = PostingList>; + +/// Non-owning posting list of ids, where ids are compressed. +pub type IdsPostingListView<'a> = PostingListView<'a, SizedHandler<()>>; + +/// Non-owning posting list of ids + small fixed-sized values, where ids are compressed. +pub type WeightsPostingListView<'a, W> = PostingListView<'a, SizedHandler>; + +/// Non-owning posting list of ids + variably-sized values, where ids are compressed. +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; diff --git a/lib/posting_list/src/posting_list.rs b/lib/posting_list/src/posting_list.rs new file mode 100644 index 0000000000..04d05c3da0 --- /dev/null +++ b/lib/posting_list/src/posting_list.rs @@ -0,0 +1,86 @@ +use std::marker::PhantomData; + +use common::types::PointOffsetType; + +use crate::CHUNK_LEN; +use crate::value_handler::ValueHandler; +use crate::view::PostingListView; +use crate::visitor::PostingVisitor; + +/// Generic compressed posting list. +/// +/// - `PostingList>` when there are no values (unit type `()`), there are just compressed ids + remainders +/// - `PostingList>` when there are `SizedValue` values, each id includes one value stored within the +/// 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 +pub struct PostingList { + pub(crate) id_data: Vec, + pub(crate) chunks: 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, +} + +#[derive(Debug, Clone)] +#[repr(C)] +pub struct PostingChunk { + /// Initial data point id. Used for decompression. + pub initial_id: PointOffsetType, + + /// An offset within id_data + pub offset: u32, + + /// Sized values for the chunk. + pub sized_values: [S; CHUNK_LEN], +} + +impl PostingChunk { + /// Get byte size of the compressed ids chunk. + pub(crate) fn get_compressed_size( + chunks: &[PostingChunk], + 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 + } else { + // Last chunk + data.len() - chunks[chunk_index].offset as usize + } + } +} + +impl PostingList { + pub fn view(&self) -> PostingListView { + let PostingList { + id_data, + chunks, + remainders, + var_size_data, + last_id, + _phantom, + } = self; + + PostingListView { + id_data, + chunks, + var_size_data, + remainders, + last_id: *last_id, + _phantom: PhantomData, + } + } + + pub fn visitor(&self) -> PostingVisitor<'_, H> { + let view = self.view(); + PostingVisitor::new(view) + } +} diff --git a/lib/posting_list/src/tests.rs b/lib/posting_list/src/tests.rs new file mode 100644 index 0000000000..d976dd6fba --- /dev/null +++ b/lib/posting_list/src/tests.rs @@ -0,0 +1,158 @@ +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}; + +// Simple struct that implements VarSizedValue for testing +#[derive(Debug, Clone, PartialEq)] +struct TestString(String); + +impl UnsizedValue for TestString { + fn write_len(&self) -> usize { + self.0.len() + } + + fn write_to(&self, dst: &mut [u8]) { + dst.copy_from_slice(self.0.as_bytes()); + } + + fn from_bytes(data: &[u8]) -> Self { + let s = String::from_utf8(data.to_vec()).expect("Failed to convert bytes to string"); + TestString(s) + } +} + +#[test] +fn test_just_ids_against_vec() { + check_various_lengths(|len| { + let posting_list = check_against_sorted_vec(|_rng, _id| (), |builder| builder.build(), len); + + // validate that chunks' sized values are empty + if let Some(first_chunk) = posting_list.chunks.first() { + let chunks_size = size_of_val(first_chunk); + let expected_chunk_size = size_of::() * 2; + assert_eq!(chunks_size, expected_chunk_size); + } + + // validate var_sized_data is empty + assert_eq!(posting_list.var_size_data.len(), 0); + }) +} + +#[test] +fn test_var_sized_against_vec() { + let alphanumeric = Alphanumeric; + check_various_lengths(|len| { + check_against_sorted_vec( + |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, + ); + }) +} + +#[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, + ); + }); +} + +fn generate_data( + amount: u32, + rng: &mut R, + gen_value: impl Fn(&mut R, u32) -> T, +) -> Vec<(u32, T)> { + let gen_id = |rng: &mut R| rng.random_range(0..amount); + + (0..amount) + .map(|_| { + let id = gen_id(rng); + (id, gen_value(rng, id)) + }) + .collect() +} + +fn check_various_lengths(check: impl Fn(u32)) { + let lengths = [ + 0, + 1, + 2, + 9, + 10, + CHUNK_LEN - 1, + CHUNK_LEN, + CHUNK_LEN + 1, + 100 * CHUNK_LEN, + 500 * CHUNK_LEN + 1, + 500 * CHUNK_LEN - 1, + 500 * CHUNK_LEN + CHUNK_LEN / 2, + ]; + for len in lengths { + check(len as u32); + } +} + +fn check_against_sorted_vec(gen_value: G, build: B, postings_count: u32) -> PostingList +where + G: Fn(&mut StdRng, PointOffsetType) -> H::Value, + H: ValueHandler, + B: FnOnce(PostingBuilder) -> PostingList, + H::Value: Clone + PartialEq, +{ + let rng = &mut StdRng::seed_from_u64(42); + let test_data = generate_data(postings_count, rng, gen_value); + + // Build our reference model + let mut model = test_data.clone(); + model.sort_unstable_by_key(|(id, _)| *id); + + // Create the posting list builder and add elements + let mut builder = PostingBuilder::new(); + for (id, value) in test_data { + builder.add(id, value); + } + + // Build the actual posting list + let posting_list = build(builder); + + // Access the posting list + let mut visitor = posting_list.visitor(); + + // Validate len() + assert_eq!(visitor.len(), model.len()); + + // Iterate through the elements in reference_model and check they can be found + for (offset, (expected_id, expected_value)) in model.iter().enumerate() { + let Some(elem) = visitor.get_by_offset(offset) else { + panic!("Element not found at offset {offset}"); + }; + + assert_eq!(elem.id, *expected_id); + assert_eq!(elem.value, *expected_value); + + // also check that contains function works + assert!(visitor.contains(*expected_id)); + } + + // Bounds check + assert!(visitor.get_by_offset(postings_count as usize).is_none()); + let out_of_range = (postings_count.next_multiple_of(CHUNK_LEN as u32)) as usize; + assert!(visitor.get_by_offset(out_of_range).is_none()); + + // There is no such id + assert!(!visitor.contains(postings_count)); + + posting_list +} diff --git a/lib/posting_list/src/value_handler.rs b/lib/posting_list/src/value_handler.rs new file mode 100644 index 0000000000..31a5dcb009 --- /dev/null +++ b/lib/posting_list/src/value_handler.rs @@ -0,0 +1,102 @@ +use std::marker::PhantomData; + +use crate::{SizedValue, UnsizedValue}; + +/// Trait to abstract the handling of values in PostingList +/// +/// This trait handles the differences between fixed-size and variable-size value +/// implementations, allowing us to have a unified implementation of `from_builder`. +/// +/// - For fixed-size values, the associated type [`ValueHandler::Sized`] is the same as the generic type V +/// - 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; + /// The value to store within each chunk, or alongside each id. + type Sized: std::fmt::Debug + std::marker::Sized + Copy; + + /// Process values before storage and return the necessary var_sized_data + /// + /// - For fixed-size values, this returns the values themselves and an empty var_sized_data. + /// - For variable-size values, this returns offsets and the flattened serialized data. + fn process_values(values: Vec) -> (Vec, Vec); + + /// Retrieve a value. + /// + /// - For sized values it returns the first argument. + /// - For variable-size values it returns the value between the two sized values in var_data. + fn get_value(sized_value: Self::Sized, next_sized_value: N, var_data: &[u8]) -> Self::Value + where + N: Fn() -> Option; +} + +/// Fixed-size value handler +pub struct SizedHandler(PhantomData); + +impl ValueHandler for SizedHandler { + type Value = V; + type Sized = V; + + fn process_values(values: Vec) -> (Vec, Vec) { + (values, Vec::new()) + } + + fn get_value(sized_value: V, _next_sized_value: N, _var_data: &[u8]) -> V + where + N: Fn() -> Option, + { + sized_value + } +} + +/// Var-size value handler +pub struct UnsizedHandler(PhantomData); + +impl ValueHandler for UnsizedHandler { + type Value = V; + type Sized = u32; + + fn process_values(values: Vec) -> (Vec, Vec) { + let mut offsets = Vec::with_capacity(values.len()); + let mut current_offset = 0u32; + + for value in &values { + offsets.push(current_offset); + let value_len = u32::try_from(value.write_len()) + .expect("Value larger than 4GB, use u64 offsets instead"); + // prepare next starting offset + current_offset = current_offset + .checked_add(value_len) + .expect("Size of all values exceeds 4GB"); + } + + let last_offset = offsets.last(); + let ranges = offsets + .windows(2) + .map(|w| w[0] as usize..w[1] as usize) + // the last one is not included in windows, but goes until the end + .chain( + last_offset + .iter() + .map(|&last| *last as usize..current_offset as usize), + ); + + let mut var_sized_data = vec![0; current_offset as usize]; + for (value, range) in values.iter().zip(ranges) { + value.write_to(&mut var_sized_data[range]); + } + + (offsets, var_sized_data) + } + + fn get_value(sized_value: Self::Sized, next_sized_value: N, var_data: &[u8]) -> Self::Value + where + N: Fn() -> Option, + { + let range = match next_sized_value() { + Some(next_value) => sized_value as usize..next_value as usize, + None => sized_value as usize..var_data.len(), + }; + V::from_bytes(&var_data[range]) + } +} diff --git a/lib/posting_list/src/view.rs b/lib/posting_list/src/view.rs new file mode 100644 index 0000000000..4eb599de03 --- /dev/null +++ b/lib/posting_list/src/view.rs @@ -0,0 +1,208 @@ +use std::marker::PhantomData; + +use bitpacking::BitPacker; +use common::types::PointOffsetType; + +use crate::value_handler::{SizedHandler, ValueHandler}; +use crate::visitor::PostingVisitor; +use crate::{ + BitPackerImpl, CHUNK_LEN, IdsPostingListView, PostingChunk, PostingElement, SizedValue, +}; + +/// A non-owning view of [`PostingList`]. +#[derive(Debug, Clone)] +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) last_id: Option, + pub(crate) _phantom: PhantomData, +} + +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, +} + +impl<'a> IdsPostingListView<'a> { + pub fn from_ids_components( + id_data: &'a [u8], + chunks: &'a [PostingChunk<()>], + remainders: &'a [PostingElement<()>], + last_id: Option, + ) -> Self { + Self { + id_data, + chunks, + var_size_data: &[], + remainders, + last_id, + _phantom: PhantomData, + } + } +} + +impl<'a, V: SizedValue> PostingListView<'a, SizedHandler> { + pub fn from_weighted_ids_components( + id_data: &'a [u8], + chunks: &'a [PostingChunk], + remainders: &'a [PostingElement], + last_id: Option, + ) -> Self { + Self { + id_data, + chunks, + var_size_data: &[], + remainders, + last_id, + _phantom: PhantomData, + } + } +} + +impl<'a, H: ValueHandler> PostingListView<'a, H> { + pub fn visitor(self) -> PostingVisitor<'a, H> { + PostingVisitor::new(self) + } + + pub fn components(&self) -> PostingListComponents { + let Self { + id_data, + chunks, + var_size_data, + remainders, + last_id, + _phantom, + } = self; + + PostingListComponents { + id_data, + chunks, + var_size_data, + remainders, + last_id: *last_id, + } + } + + pub fn from_components( + id_data: &'a [u8], + chunks: &'a [PostingChunk], + var_size_data: &'a [u8], + remainders: &'a [PostingElement], + last_id: Option, + ) -> Self { + Self { + id_data, + chunks, + var_size_data, + remainders, + last_id, + _phantom: PhantomData, + } + } + + pub(crate) fn decompress_chunk( + &self, + chunk_index: usize, + decompressed_chunk: &mut [PointOffsetType; CHUNK_LEN], + ) { + let chunk = &self.chunks[chunk_index]; + 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], + decompressed_chunk, + chunk_bits as u8, + ); + } + pub(crate) fn sized_values_unchecked(&self, chunk_idx: usize) -> &[H::Sized] { + &self.chunks[chunk_idx].sized_values + } + + pub(crate) fn sized_values(&self, chunk_idx: usize) -> Option<&[H::Sized; CHUNK_LEN]> { + 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; + }; + + let Some(initial_id) = self + .chunks + .first() + .map(|chunk| chunk.initial_id) + .or_else(|| self.remainders.first().map(|elem| elem.id)) + else { + return false; + }; + + id >= initial_id && id <= last_id + } + + /// Find the chunk that may contain the id. + /// It doesn't guarantee that the chunk contains the id, but if it is in the posting list, then it must be in the chunk. + /// + /// Assumes the id is in the posting list range. + pub fn find_chunk(&self, id: PointOffsetType, start_chunk: Option) -> Option { + 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) { + return None; + } + + let start_chunk = start_chunk.unwrap_or(0); + let chunks_slice = &chunks[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!(self.last_id.is_some_and(|last_id| id <= last_id)); + + match chunks_slice.binary_search_by(|chunk| chunk.initial_id.cmp(&id)) { + // id is the initial value of the chunk with index idx + Ok(idx) => Some(start_chunk + idx), + + // id is not the initial_id of any chunk + Err(insert_idx) if insert_idx > 0 => { + // this is the index of the chunk that could contain id + let idx = insert_idx - 1; + + // id could be within this chunk + Some(start_chunk + idx) + } + Err(_) => None, + } + } + + pub(crate) fn search_in_remainders(&self, id: PointOffsetType) -> Option { + self.remainders + .binary_search_by(|elem| elem.id.cmp(&id)) + .ok() + } + + /// The total number of elements in the posting list. + pub fn len(&self) -> usize { + self.chunks.len() * CHUNK_LEN + self.remainders.len() + } + + /// Checks if there are no elements in the posting list. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} diff --git a/lib/posting_list/src/visitor.rs b/lib/posting_list/src/visitor.rs new file mode 100644 index 0000000000..9624144dfb --- /dev/null +++ b/lib/posting_list/src/visitor.rs @@ -0,0 +1,118 @@ +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; + +/// A visitor for a posting list which caches the latest decompressed chunk of ids. +pub struct PostingVisitor<'a, H: ValueHandler> { + pub(crate) list: PostingListView<'a, H>, + + /// Index of the decompressed chunk. + /// It is used to shorten the search range of chunk index for the next value. + decompressed_chunk_idx: Option, + + /// Lazy decompressed chunk of ids. Never access this directly, prefer [`Self::decompressed_chunk`] function + decompressed_chunk: [PointOffsetType; CHUNK_LEN], +} + +impl<'a, H: ValueHandler> PostingVisitor<'a, H> { + pub(crate) fn new(view: PostingListView<'a, H>) -> Self { + Self { + list: view, + decompressed_chunk_idx: None, + decompressed_chunk: [0; CHUNK_LEN], + } + } + + pub fn len(&self) -> usize { + self.list.len() + } + + /// Returns the decompressed slice of ids for a chunk. + /// + /// Assumes the chunk_idx is valid. + fn decompressed_chunk(&mut self, chunk_idx: usize) -> &[PointOffsetType; CHUNK_LEN] { + if self.decompressed_chunk_idx != Some(chunk_idx) { + self.list + .decompress_chunk(chunk_idx, &mut self.decompressed_chunk); + self.decompressed_chunk_idx = Some(chunk_idx); + } + + &self.decompressed_chunk + } + + pub fn contains(&mut self, id: PointOffsetType) -> bool { + if !self.list.is_in_range(id) { + return false; + } + + // 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, None); + if let Some(chunk_index) = chunk_index { + if self.list.chunks[chunk_index].initial_id == id { + return true; + } + + self.decompressed_chunk(chunk_index) + .binary_search(&id) + .is_ok() + } else { + self.list.search_in_remainders(id).is_some() + } + } + + pub(crate) fn get_by_offset(&mut self, offset: usize) -> Option> { + let chunk_idx = offset / CHUNK_LEN; + let local_offset = offset % CHUNK_LEN; + + // bound check + if offset >= self.list.len() { + return None; + } + + // get from chunk + if chunk_idx < self.list.chunks.len() { + let id = self.decompressed_chunk(chunk_idx)[local_offset]; + let chunk_sized_values = self.list.sized_values_unchecked(chunk_idx); + let sized_value = chunk_sized_values[local_offset]; + let next_sized_value = || { + chunk_sized_values + .get(local_offset + 1) + .copied() + // or check first of the next chunk + .or_else(|| { + self.list + .sized_values(chunk_idx + 1) + .map(|sized_values| sized_values[0]) + }) + // or, if it is the last one, check first from remainders + .or_else(|| self.list.remainders.first().map(|e| e.value)) + }; + + let value = H::get_value(sized_value, next_sized_value, self.list.var_size_data); + + return Some(PostingElement { id, value }); + } + + // else, get from remainder + self.list.remainders.get(local_offset).map(|e| { + let id = e.id; + 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 } + }) + } +} + +impl<'a, H: ValueHandler> IntoIterator for PostingVisitor<'a, H> { + type Item = PostingElement; + type IntoIter = PostingIterator<'a, H>; + + fn into_iter(self) -> Self::IntoIter { + PostingIterator::new(self) + } +}