upstream-archive/qdrant
1
0
Fork 0
mirror of https://github.com/qdrant/qdrant.git synced 2026-08-05 09:31:05 -05:00
Code Issues Packages Projects Releases Wiki Activity

Full text index with generic posting list (#6565)

Browse Source 
* 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 <andrey@vasnetsov.com>
This commit is contained in:
Luis Cossío
2025-05-27 18:53:41 -04:00
committed by GitHub
parent 1dfebb5ea0
commit 7a05502e0c
21 changed files with 736 additions and 189 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
Cargo.lock generated
View File

@@ -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",

1
lib/posting_list/Cargo.toml
View File

@@ -12,6 +12,7 @@ workspace = true
[dependencies]
common = { path = "../common/common" }
bitpacking = { workspace = true }
zerocopy = { workspace = true }
[dev-dependencies]
rand = { workspace = true }

22
lib/posting_list/src/builder.rs
View File

@@ -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<V> PostingBuilder<V> {
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<V> PostingBuilder<V> {
// 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<V> PostingBuilder<V> {
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,
);
}

68
lib/posting_list/src/iterator.rs
View File

@@ -8,28 +8,75 @@ use crate::visitor::PostingVisitor;
pub struct PostingIterator<'a, H: ValueHandler> {
visitor: PostingVisitor<'a, H>,
current_id: Option<PointOffsetType>,
current_elem: Option<PostingElement<H::Value>>,
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<PostingElement<H::Value>> {
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<H: ValueHandler> Iterator for PostingIterator<'_, H> {
impl<H: ValueHandler> Iterator for PostingIterator<'_, H>
where
H::Value: Clone,
{
type Item = PostingElement<H::Value>;
fn next(&mut self) -> Option<Self::Item> {
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<usize>) {
@@ -42,10 +89,13 @@ impl<H: ValueHandler> Iterator for PostingIterator<'_, H> {
}
}
impl<H: ValueHandler> ExactSizeIterator for PostingIterator<'_, H> {
impl<H: ValueHandler> ExactSizeIterator for PostingIterator<'_, H>
where
H::Value: Clone,
{
fn len(&self) -> usize {
self.visitor.list.len().saturating_sub(self.offset)
}
}
impl<H: ValueHandler> FusedIterator for PostingIterator<'_, H> {}
impl<H: ValueHandler> FusedIterator for PostingIterator<'_, H> where H::Value: Clone {}

11
lib/posting_list/src/lib.rs
View File

@@ -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<W>>;
pub type VarPostingListView<'a, V> = PostingListView<'a, UnsizedHandler<V>>;
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;

72
lib/posting_list/src/posting_list.rs
View File

@@ -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<VarSized<V>>` 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<H: ValueHandler> {
pub(crate) id_data: Vec<u8>,
pub(crate) chunks: Vec<PostingChunk<H::Sized>>,
pub(crate) remainders: Vec<PostingElement<H::Sized>>,
pub(crate) remainders: Vec<RemainderPosting<H::Sized>>,
pub(crate) var_size_data: Vec<u8>,
pub(crate) last_id: Option<PointOffsetType>,
pub(crate) _phantom: PhantomData<H>,
}
#[derive(Clone, Debug)]
pub struct PostingElement<S> {
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<S: Sized> {
/// 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<V> {
pub id: PointOffsetType,
pub value: V,
}
#[derive(Debug, Clone, FromBytes, Immutable, IntoBytes, KnownLayout)]
#[repr(C)]
pub struct PostingChunk<S: Sized> {
/// 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<S: Sized> PostingChunk<S> {
/// Get byte size of the compressed ids chunk.
pub(crate) fn get_compressed_size(
chunks: &[PostingChunk<S>],
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<H: ValueHandler> PostingList<H> {
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<H> FromIterator<(PointOffsetType, H::Value)> for PostingList<H>
where
H: ValueHandler,
{
fn from_iter<T: IntoIterator<Item = (PointOffsetType, H::Value)>>(iter: T) -> Self {
let mut builder = PostingBuilder::new();
for (id, value) in iter {
builder.add(id, value);
}
builder.build_generic::<H>()
}
}

42
lib/posting_list/src/tests.rs
View File

@@ -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<TestString>>(
|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<u64>>(|_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<G, H, B>(gen_value: G, build: B, postings_count: u32) -> PostingList<H>
fn check_against_sorted_vec<G, H>(gen_value: G, postings_count: u32) -> PostingList<H>
where
G: Fn(&mut StdRng, PointOffsetType) -> H::Value,
H: ValueHandler,
B: FnOnce(PostingBuilder<H::Value>) -> PostingList<H>,
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::<HashMap<_, _>>();
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
}

6
lib/posting_list/src/value_handler.rs
View File

@@ -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<V>(PhantomData<V>);
impl<V: SizedValue + Copy> ValueHandler for SizedHandler<V> {
@@ -50,6 +51,7 @@ impl<V: SizedValue + Copy> ValueHandler for SizedHandler<V> {
}
/// Var-size value handler
#[derive(Debug, Clone, Copy)]
pub struct UnsizedHandler<V>(PhantomData<V>);
impl<V: UnsizedValue> ValueHandler for UnsizedHandler<V> {

98
lib/posting_list/src/view.rs
View File

@@ -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<H::Sized>],
pub(crate) var_size_data: &'a [u8],
pub(crate) remainders: &'a [PostingElement<H::Sized>],
pub(crate) remainders: &'a [RemainderPosting<H::Sized>],
pub(crate) last_id: Option<PointOffsetType>,
pub(crate) _phantom: PhantomData<H>,
}
@@ -24,15 +26,15 @@ pub struct PostingListComponents<'a, S> {
pub id_data: &'a [u8],
pub chunks: &'a [PostingChunk<S>],
pub var_size_data: &'a [u8],
pub remainders: &'a [PostingElement<S>],
pub last_id: Option<PointOffsetType>,
pub remainders: &'a [RemainderPosting<S>],
pub last_id: Option<U32>,
}
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<PointOffsetType>,
) -> Self {
Self {
@@ -50,7 +52,7 @@ impl<'a, V: SizedValue> PostingListView<'a, SizedHandler<V>> {
pub fn from_weighted_ids_components(
id_data: &'a [u8],
chunks: &'a [PostingChunk<V>],
remainders: &'a [PostingElement<V>],
remainders: &'a [RemainderPosting<V>],
last_id: Option<PointOffsetType>,
) -> Self {
Self {
@@ -64,11 +66,36 @@ impl<'a, V: SizedValue> PostingListView<'a, SizedHandler<V>> {
}
}
impl<'a, H: ValueHandler> IntoIterator for PostingListView<'a, H>
where
H::Value: Clone,
{
type Item = <PostingIterator<'a, H> 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<H> {
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<H::Sized> {
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<H::Sized>],
var_size_data: &'a [u8],
remainders: &'a [PostingElement<H::Sized>],
remainders: &'a [RemainderPosting<H::Sized>],
last_id: Option<PointOffsetType>,
) -> 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<RangeInclusive<u32>> {
// 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<usize> {
pub(crate) fn search_in_remainders(&self, id: PointOffsetType) -> Result<usize, usize> {
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.

87
lib/posting_list/src/visitor.rs
View File

@@ -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<usize>,
) -> Option<usize> {
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<H::Value>;
type IntoIter = PostingIterator<'a, H>;

1
lib/segment/Cargo.toml
View File

@@ -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" }

11
lib/segment/src/index/field_index/full_text_index/compressed_posting/compressed_chunks_reader.rs
View File

@@ -121,21 +121,10 @@ impl<'a> ChunkReader<'a> {
ChunkReaderIter::new(self)
}
pub fn to_vec(&self) -> Vec<PointOffsetType> {
let postings: Vec<PointOffsetType> = 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()
}

29
lib/segment/src/index/field_index/full_text_index/immutable_inverted_index.rs
View File

@@ -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<CompressedPostingList>,
pub(in crate::index::field_index::full_text_index) postings: Vec<IdsPostingList>,
pub(in crate::index::field_index::full_text_index) vocab: HashMap<String, TokenId>,
pub(in crate::index::field_index::full_text_index) point_to_tokens_count: Vec<Option<usize>>,
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<MutableInvertedIndex> for ImmutableInvertedIndex {
})
.collect();
let postings: Vec<CompressedPostingList> = postings
.into_iter()
.map(|posting| CompressedPostingList::new(&posting.into_vec()))
.collect();
vocab.shrink_to_fit();
let postings: Vec<IdsPostingList> = 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<CompressedPostingList> = postings
let postings: Vec<IdsPostingList> = postings
.into_iter()
.map(|postings| CompressedPostingList::new(&postings.to_vec()))
.map(IdsPostingListView::to_owned)
.collect();
vocab.shrink_to_fit();

16
lib/segment/src/index/field_index/full_text_index/inverted_index.rs
View File

@@ -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::<HashSet<_>>();
let mutable_elems = posting.iter().collect::<HashSet<_>>();
// 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() {

95
lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/mmap_postings.rs
View File

@@ -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::<PointOffsetType>()
+ self.chunks_count as usize * size_of::<CompressedPostingChunksIndex>()
+ self.remainder_count as usize * size_of::<RemainderPosting<()>>()
+ self.chunks_count as usize * size_of::<PostingChunk<()>>()
+ size_of::<PointOffsetType>() // 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<ChunkReader<'a>> {
) -> Option<IdsPostingListView<'a>> {
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::<CompressedPostingChunksIndex>());
let (chunks, bytes) = <[CompressedPostingChunksIndex]>::ref_from_prefix_with_elems(
bytes,
header.chunks_count as usize,
)
.ok()?;
counter.incr_delta(size_of::<PostingChunk<()>>());
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<ChunkReader<'a>> {
) -> Option<IdsPostingListView<'a>> {
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::<PostingsHeader>() + 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<Item = Option<ChunkReader<'a>>> {
) -> impl Iterator<Item = Option<IdsPostingListView<'a>>> {
(0..self.header.posting_count as u32).map(|posting_idx| self.get(posting_idx, hw_counter))
}
}

13
lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/mod.rs
View File

@@ -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<Item = Option<ChunkReader<'a>>> {
) -> impl Iterator<Item = Option<IdsPostingListView<'a>>> {
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)
})
}

239
lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/old_mmap_postings.rs Normal file
View File

@@ -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::<PointOffsetType>()
+ self.chunks_count as usize * size_of::<CompressedPostingChunksIndex>()
+ size_of::<PointOffsetType>() // 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::<PostingsHeader>() + token_id as usize * size_of::<PostingListHeader>();
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<ChunkReader<'a>> {
let counter = hw_counter.payload_index_io_read_counter();
let bytes = self.mmap.get(header.offset as usize..)?;
counter.incr_delta(size_of::<PointOffsetType>());
let (last_doc_id, bytes) = PointOffsetType::read_from_prefix(bytes).ok()?;
counter.incr_delta(size_of::<CompressedPostingChunksIndex>());
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<ChunkReader<'a>> {
let hw_counter = ConditionedCounter::new(self.on_disk, hw_counter);
hw_counter
.payload_index_io_read_counter()
.incr_delta(size_of::<PostingListHeader>());
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::<PostingListHeader>();
let mut posting_offset = size_of::<PostingsHeader>() + 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<PathBuf>, populate: bool) -> io::Result<Self> {
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,
})
}
}

61
lib/segment/src/index/field_index/full_text_index/mmap_inverted_index/tests.rs Normal file
View File

@@ -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<PointOffsetType> {
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::<Vec<_>>();
// old compressed postings implementation
let old_compressed_postings = postings
.iter()
.map(|ids| CompressedPostingList::new(ids))
.collect::<Vec<_>>();
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
);
}
}
}

1
lib/segment/src/index/field_index/full_text_index/mod.rs
View File

@@ -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;

5
lib/segment/src/index/field_index/full_text_index/posting_list.rs
View File

@@ -44,9 +44,4 @@ impl PostingList {
pub fn iter(&self) -> impl Iterator<Item = PointOffsetType> + '_ {
self.list.iter().copied()
}
#[inline]
pub fn into_vec(self) -> Vec<PointOffsetType> {
self.list
}
}

45
lib/segment/src/index/field_index/full_text_index/postings_iterator.rs
View File

@@ -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<ChunkReader<'a>>,
mut postings: Vec<IdsPostingListView<'a>>,
filter: impl Fn(PointOffsetType) -> bool + 'a,
) -> Box<dyn Iterator<Item = PointOffsetType> + '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::<Vec<_>>();
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);