use std::borrow::Cow;
use std::path::Path;
use blink_alloc::Blink;
use crate::common::counter::hardware_counter::HardwareCounterCell;
use crate::common::ext::VecExt;
use crate::common::types::PointOffsetType;
use crate::common::universal_io::{
MmapFs, UioResult, UniversalRead, UniversalReadFs, UniversalWrite, UserData,
};
use super::inverted_index_compressed_mmap::InvertedIndexCompressedMmap;
use super::inverted_index_ram::InvertedIndexRam;
use super::{InvertedIndex, out_of_bounds};
use crate::sparse::common::sparse_vector::RemappedSparseVector;
use crate::sparse::common::types::{DimOffset, Weight};
use crate::sparse::index::compressed_posting_list::{
CompressedPostingBuilder, CompressedPostingList, CompressedPostingListIterator,
CompressedPostingListView,
};
use crate::sparse::index::inverted_index::inverted_index_compressed_mmap::Version;
use crate::sparse::index::inverted_index::{InvertedIndexReadOnly, InvertedIndexReadWrite};
#[derive(Debug, Clone, PartialEq)]
pub struct InvertedIndexCompressedImmutableRam<W: Weight> {
pub(super) postings: Vec<CompressedPostingList<W>>,
pub(super) vector_count: usize,
pub(super) total_sparse_size: usize,
}
type Storage = crate::common::universal_io::MmapFile;
impl<W: Weight, S: UniversalRead + 'static> InvertedIndexReadOnly<S>
for InvertedIndexCompressedImmutableRam<W>
{
fn open_ro_impl<Fs: UniversalReadFs<File = S>>(fs: &Fs, path: &Path) -> UioResult<Self> {
let mmap_inverted_index = InvertedIndexCompressedMmap::<W, S>::open_ro(fs, path)?;
Self::from_mmap_index(mmap_inverted_index)
}
}
impl<W: Weight, S: UniversalWrite + 'static> InvertedIndexReadWrite<S>
for InvertedIndexCompressedImmutableRam<W>
{
fn open_rw_impl(fs: &<S as UniversalRead>::Fs, path: &Path) -> UioResult<Self> {
let mmap_inverted_index = InvertedIndexCompressedMmap::<W, S>::open_rw(fs, path)?;
Self::from_mmap_index(mmap_inverted_index)
}
fn from_ram_index_impl<P: AsRef<Path>>(
_fs: &S::Fs,
ram_index: Cow<InvertedIndexRam>,
_path: P,
) -> UioResult<Self> {
let mut postings = Vec::with_capacity(ram_index.postings.len());
for old_posting_list in &ram_index.postings {
let mut new_posting_list = CompressedPostingBuilder::new();
for elem in &old_posting_list.elements {
new_posting_list.add(elem.record_id, elem.weight);
}
postings.push(new_posting_list.build());
}
let hw_counter = HardwareCounterCell::disposable();
let total_sparse_size = postings
.iter()
.map(|p| p.view(&hw_counter).store_size().total)
.sum();
Ok(InvertedIndexCompressedImmutableRam {
postings,
vector_count: ram_index.vector_count,
total_sparse_size,
})
}
}
impl<W: Weight> InvertedIndex for InvertedIndexCompressedImmutableRam<W> {
type Iter<'a> = CompressedPostingListIterator<'a, W>;
type Version = Version;
fn is_on_disk(&self) -> bool {
false
}
fn save(&self, path: &Path) -> UioResult<()> {
InvertedIndexCompressedMmap::<W, Storage>::convert_and_save(&MmapFs, self, path)?;
Ok(())
}
fn get_batch<'a, U: UserData>(
&'a self,
ids: impl Iterator<Item = (U, DimOffset)>,
_arena: &'a Blink,
hw_counter: &'a HardwareCounterCell, mut callback: impl FnMut(U, Self::Iter<'a>) -> UioResult<()>,
) -> UioResult<()> {
for (user_data, id) in ids {
callback(user_data, self.get(id, hw_counter)?.iter())?;
}
Ok(())
}
fn len(&self) -> usize {
self.postings.len()
}
fn posting_list_len_batch<U: UserData>(
&self,
ids: impl Iterator<Item = (U, DimOffset)>,
hw_counter: &HardwareCounterCell,
mut callback: impl FnMut(U, usize) -> UioResult<()>,
) -> UioResult<()> {
for (user_data, id) in ids {
callback(user_data, self.get(id, hw_counter)?.len())?;
}
Ok(())
}
fn files(path: &Path) -> Vec<std::path::PathBuf> {
InvertedIndexCompressedMmap::<W, Storage>::files(path)
}
fn immutable_files(path: &Path) -> Vec<std::path::PathBuf> {
InvertedIndexCompressedMmap::<W, Storage>::immutable_files(path)
}
fn remove(&mut self, _id: PointOffsetType, _old_vector: RemappedSparseVector) {
panic!("Cannot remove from a read-only RAM inverted index")
}
fn upsert(
&mut self,
_id: PointOffsetType,
_vector: RemappedSparseVector,
_old_vector: Option<RemappedSparseVector>,
) {
panic!("Cannot upsert into a read-only RAM inverted index")
}
fn vector_count(&self) -> usize {
self.vector_count
}
fn total_sparse_vectors_size(&self) -> usize {
self.total_sparse_size
}
fn max_index(&self) -> Option<DimOffset> {
self.postings
.len()
.checked_sub(1)
.map(|len| len as DimOffset)
}
}
impl<W: Weight> InvertedIndexCompressedImmutableRam<W> {
fn from_mmap_index<S>(mmap_inverted_index: InvertedIndexCompressedMmap<W, S>) -> UioResult<Self>
where
S: UniversalRead + 'static,
{
let hw_counter = HardwareCounterCell::disposable();
let mut postings = vec![None; mmap_inverted_index.file_header.posting_count];
mmap_inverted_index.for_each_view(&hw_counter, |id, view| {
postings[id as usize] = Some(view.to_owned());
Ok(())
})?;
mmap_inverted_index.clear_cache()?;
Ok(InvertedIndexCompressedImmutableRam {
postings: postings.transform_in_place(Option::unwrap),
vector_count: mmap_inverted_index.file_header.vector_count,
total_sparse_size: mmap_inverted_index
.file_header
.total_sparse_size
.unwrap_or(0),
})
}
#[inline]
fn get<'a>(
&'a self,
id: DimOffset,
hw_counter: &'a HardwareCounterCell,
) -> UioResult<CompressedPostingListView<'a, W>> {
let Some(posting) = self.postings.get(id as usize) else {
return Err(out_of_bounds(id, self.len()));
};
Ok(posting.view(hw_counter))
}
}
#[cfg(test)]
mod tests {
use tempfile::Builder;
use super::*;
use crate::sparse::common::sparse_vector_fixture::random_sparse_vector;
use crate::sparse::common::types::QuantizedU8;
use crate::sparse::index::inverted_index::inverted_index_ram_builder::InvertedIndexBuilder;
#[test]
fn test_save_load_tiny() {
let mut builder = InvertedIndexBuilder::new();
builder.add(1, vec![(1, 10.0), (2, 10.0), (3, 10.0)].try_into().unwrap());
builder.add(2, vec![(1, 20.0), (2, 20.0), (3, 20.0)].try_into().unwrap());
builder.add(3, vec![(1, 30.0), (2, 30.0), (3, 30.0)].try_into().unwrap());
let inverted_index_ram = builder.build();
check_save_load::<f32>(&inverted_index_ram);
check_save_load::<half::f16>(&inverted_index_ram);
check_save_load::<u8>(&inverted_index_ram);
check_save_load::<QuantizedU8>(&inverted_index_ram);
}
#[test]
fn test_save_load_large() {
let mut rnd_gen = rand::rng();
let mut builder = InvertedIndexBuilder::new();
for i in 0..1024 {
builder.add(i, random_sparse_vector(&mut rnd_gen, 3).into_remapped());
}
let inverted_index_ram = builder.build();
check_save_load::<f32>(&inverted_index_ram);
check_save_load::<half::f16>(&inverted_index_ram);
check_save_load::<u8>(&inverted_index_ram);
check_save_load::<QuantizedU8>(&inverted_index_ram);
}
fn check_save_load<W: Weight>(inverted_index_ram: &InvertedIndexRam) {
let tmp_dir_path = Builder::new().prefix("test_index_dir").tempdir().unwrap();
let inverted_index_immutable_ram =
InvertedIndexCompressedImmutableRam::<W>::from_ram_index(
&MmapFs,
Cow::Borrowed(inverted_index_ram),
tmp_dir_path.path(),
)
.unwrap();
inverted_index_immutable_ram
.save(tmp_dir_path.path())
.unwrap();
let loaded_inverted_index =
InvertedIndexCompressedImmutableRam::<W>::open_ro(&MmapFs, tmp_dir_path.path())
.unwrap();
assert_eq!(inverted_index_immutable_ram, loaded_inverted_index);
}
}