use std::borrow::Cow;
use std::path::Path;
use crate::common::counter::hardware_counter::HardwareCounterCell;
use crate::common::types::PointOffsetType;
use crate::common::universal_io::Result;
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::{DimId, DimOffset, Weight};
use crate::sparse::index::compressed_posting_list::{
CompressedPostingBuilder, CompressedPostingList, CompressedPostingListIterator,
};
use crate::sparse::index::posting_list_common::PostingListIter as _;
#[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,
}
impl<W: Weight> InvertedIndex for InvertedIndexCompressedImmutableRam<W> {
type Iter<'a> = CompressedPostingListIterator<'a, W>;
type Version = <InvertedIndexCompressedMmap<W> as InvertedIndex>::Version;
fn is_on_disk(&self) -> bool {
false
}
fn open(path: &Path) -> Result<Self> {
let mmap_inverted_index = InvertedIndexCompressedMmap::load(path)?;
let mut inverted_index = InvertedIndexCompressedImmutableRam {
postings: Vec::with_capacity(mmap_inverted_index.file_header.posting_count),
vector_count: mmap_inverted_index.file_header.vector_count,
total_sparse_size: mmap_inverted_index.total_sparse_vectors_size(),
};
let hw_counter = HardwareCounterCell::disposable();
for i in 0..mmap_inverted_index.file_header.posting_count as DimId {
let posting_list = mmap_inverted_index.get(i, &hw_counter)?;
inverted_index.postings.push(posting_list.to_owned());
}
mmap_inverted_index.clear_cache()?;
Ok(inverted_index)
}
fn save(&self, path: &Path) -> std::io::Result<()> {
InvertedIndexCompressedMmap::convert_and_save(self, path)?;
Ok(())
}
fn get<'a>(
&'a self,
id: DimOffset,
hw_counter: &'a HardwareCounterCell, ) -> Result<Self::Iter<'a>> {
let posting = self
.postings
.get(id as usize)
.ok_or(out_of_bounds(id, self.len()))?;
Ok(posting.iter(hw_counter))
}
fn len(&self) -> usize {
self.postings.len()
}
fn posting_list_len(&self, id: DimOffset, hw_counter: &HardwareCounterCell) -> Result<usize> {
Ok(self.get(id, hw_counter)?.len_to_end())
}
fn files(path: &Path) -> Vec<std::path::PathBuf> {
InvertedIndexCompressedMmap::<W>::files(path)
}
fn immutable_files(path: &Path) -> Vec<std::path::PathBuf> {
InvertedIndexCompressedMmap::<W>::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 from_ram_index<P: AsRef<Path>>(
ram_index: Cow<InvertedIndexRam>,
_path: P,
) -> std::io::Result<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,
})
}
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)
}
}
#[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(
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(tmp_dir_path.path()).unwrap();
assert_eq!(inverted_index_immutable_ram, loaded_inverted_index);
}
}