use std::cmp::Ordering;
use std::fmt::Debug;
use std::mem::size_of;
use bitpacking::BitPacker as _;
use crate::common::counter::hardware_counter::HardwareCounterCell;
use crate::common::counter::iterator_hw_measurement::HwMeasurementIteratorExt;
use crate::common::types::PointOffsetType;
#[cfg(debug_assertions)]
use itertools::Itertools as _;
use zerocopy::{FromBytes, Immutable, KnownLayout};
use super::posting_list_common::{
GenericPostingElement, PostingElement, PostingElementEx, PostingListIter,
};
use crate::sparse::common::types::{DimWeight, Weight};
use crate::sparse::index::inverted_index::inverted_index_compressed_mmap::PostingListFileHeader;
type BitPackerImpl = bitpacking::BitPacker4x;
pub(super) const CHUNK_SIZE: usize = BitPackerImpl::BLOCK_LEN;
#[derive(Default, Debug, Clone, PartialEq)]
pub struct CompressedPostingList<W: Weight> {
id_data: Vec<u8>,
chunks: Vec<CompressedPostingChunk<W>>,
remainders: Vec<GenericPostingElement<W>>,
last_id: Option<PointOffsetType>,
quantization_params: W::QuantizationParams,
}
#[derive(Debug, Clone)]
pub struct CompressedPostingListView<'a, W: Weight> {
id_data: &'a [u8],
chunks: &'a [CompressedPostingChunk<W>],
remainders: &'a [GenericPostingElement<W>],
last_id: Option<PointOffsetType>,
multiplier: W::QuantizationParams,
hw_counter: &'a HardwareCounterCell,
}
#[derive(Debug, Clone, PartialEq, FromBytes, Immutable, KnownLayout)]
#[repr(C)]
pub struct CompressedPostingChunk<W: Weight> {
initial: PointOffsetType,
offset: u32,
weights: [W; CHUNK_SIZE],
}
impl<W: Weight> CompressedPostingList<W> {
pub(super) fn view<'a>(
&'a self,
hw_counter: &'a HardwareCounterCell,
) -> CompressedPostingListView<'a, W> {
CompressedPostingListView {
id_data: &self.id_data,
chunks: &self.chunks,
remainders: &self.remainders,
last_id: self.last_id,
multiplier: self.quantization_params,
hw_counter,
}
}
pub fn iter<'a>(
&'a self,
hw_counter: &'a HardwareCounterCell,
) -> CompressedPostingListIterator<'a, W> {
self.view(hw_counter).iter()
}
#[cfg(test)]
pub fn from(records: Vec<(PointOffsetType, DimWeight)>) -> CompressedPostingList<W> {
let mut posting_list = CompressedPostingBuilder::new();
for (id, weight) in records {
posting_list.add(id, weight);
}
posting_list.build()
}
}
pub struct CompressedPostingListStoreSize {
pub total: usize,
pub id_data_bytes: usize,
pub chunks_count: usize,
}
impl CompressedPostingListStoreSize {
fn new<W: Weight>(id_data_bytes: usize, chunks_count: usize, remainders_count: usize) -> Self {
CompressedPostingListStoreSize {
total: id_data_bytes
+ chunks_count * size_of::<CompressedPostingChunk<W>>()
+ remainders_count * size_of::<GenericPostingElement<W>>(),
id_data_bytes,
chunks_count,
}
}
}
enum IdChunkPosition {
Before,
Chunk(usize),
After,
}
impl<'a, W: Weight> CompressedPostingListView<'a, W> {
pub(super) fn new(
header: PostingListFileHeader<W>,
data: &'a [u8],
hw_counter: &'a HardwareCounterCell,
) -> Option<Self> {
let ids_len = header.ids_len as usize;
let chunks_bytes = header.chunks_count as usize * size_of::<CompressedPostingChunk<W>>();
let id_data = data.get(..ids_len)?;
let chunks = <[CompressedPostingChunk<W>]>::ref_from_bytes(
data.get(ids_len..ids_len + chunks_bytes)?,
)
.ok()?;
let remainders =
<[GenericPostingElement<W>]>::ref_from_bytes(data.get(ids_len + chunks_bytes..)?)
.ok()?;
Some(CompressedPostingListView {
id_data,
chunks,
remainders,
last_id: header.last_id.checked_sub(1),
multiplier: header.quantization_params,
hw_counter,
})
}
pub(super) fn parts(
&self,
) -> (
&'a [u8],
&'a [CompressedPostingChunk<W>],
&'a [GenericPostingElement<W>],
) {
(self.id_data, self.chunks, self.remainders)
}
pub fn last_id(&self) -> Option<PointOffsetType> {
self.last_id
}
pub fn multiplier(&self) -> W::QuantizationParams {
self.multiplier
}
pub(super) fn store_size(&self) -> CompressedPostingListStoreSize {
CompressedPostingListStoreSize::new::<W>(
self.id_data.len(),
self.chunks.len(),
self.remainders.len(),
)
}
pub fn to_owned(&self) -> CompressedPostingList<W> {
CompressedPostingList {
id_data: self.id_data.to_vec(),
chunks: self.chunks.to_vec(),
remainders: self.remainders.to_vec(),
last_id: self.last_id,
quantization_params: self.multiplier,
}
}
pub fn len(&self) -> usize {
self.chunks.len() * CHUNK_SIZE + self.remainders.len()
}
pub fn is_empty(&self) -> bool {
self.chunks.is_empty() && self.remainders.is_empty()
}
fn decompress_chunk(
&self,
chunk_index: usize,
decompressed_chunk: &mut [PointOffsetType; CHUNK_SIZE],
) {
let chunk = &self.chunks[chunk_index];
let chunk_size = Self::get_chunk_size(self.chunks, self.id_data, chunk_index);
self.hw_counter.vector_io_read().incr_delta(chunk_size);
let chunk_bits = chunk_size * u8::BITS as usize / CHUNK_SIZE;
BitPackerImpl::new().decompress_strictly_sorted(
chunk.initial.checked_sub(1),
&self.id_data[chunk.offset as usize..chunk.offset as usize + chunk_size],
decompressed_chunk,
chunk_bits as u8,
);
}
fn chunk_id_by_position(&self, position: usize) -> Option<usize> {
let chunk_index = position / CHUNK_SIZE;
if chunk_index < self.chunks.len() {
Some(chunk_index)
} else {
None
}
}
fn chunk_id_by_id(&self, id: PointOffsetType) -> IdChunkPosition {
let chunk_index_result = self.chunks.binary_search_by(|c| c.initial.cmp(&id));
match chunk_index_result {
Ok(chunk_id) => {
IdChunkPosition::Chunk(chunk_id)
}
Err(position) => {
if position == self.chunks.len() {
if let Some(first_remainder) = self.remainders.first() {
if id < first_remainder.record_id {
if position > 0 {
IdChunkPosition::Chunk(position - 1)
} else {
IdChunkPosition::Before
}
} else {
IdChunkPosition::After
}
} else {
IdChunkPosition::Chunk(position - 1)
}
} else if position == 0 {
IdChunkPosition::Before
} else {
IdChunkPosition::Chunk(position - 1)
}
}
}
}
fn get_chunk_size(
chunks: &[CompressedPostingChunk<W>],
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 {
data.len() - chunks[chunk_index].offset as usize
}
}
#[inline]
fn get_remainder_id(&self, index: usize) -> Option<&GenericPostingElement<W>> {
self.hw_counter
.vector_io_read()
.incr_delta(size_of::<GenericPostingElement<W>>());
self.remainders.get(index)
}
#[inline]
fn iter_remainder_from(
&self,
index: usize,
) -> impl Iterator<Item = &'_ GenericPostingElement<W>> + '_ {
self.remainders[index..].iter().measure_hw_with_cell(
self.hw_counter,
size_of::<GenericPostingElement<W>>(),
|hw_counter| hw_counter.vector_io_read(),
)
}
#[inline]
fn remainder_len(&self) -> usize {
self.remainders.len()
}
#[inline]
fn chunks_len(&self) -> usize {
self.chunks.len()
}
#[inline]
fn get_weight(&self, pos: usize) -> W {
self.hw_counter.vector_io_read().incr_delta(size_of::<W>());
let chunk = &self.chunks[pos / CHUNK_SIZE];
chunk.weights[pos % CHUNK_SIZE]
}
#[inline]
fn weights_range(&self, pos: usize, count: usize) -> &[W] {
debug_assert!(count <= CHUNK_SIZE);
self.hw_counter
.vector_io_read()
.incr_delta(size_of::<W>() * count);
let chunk = &self.chunks[pos / CHUNK_SIZE];
let start = pos % CHUNK_SIZE;
chunk.weights[start..start + count].as_ref()
}
pub fn iter(&self) -> CompressedPostingListIterator<'a, W> {
CompressedPostingListIterator::new(self)
}
}
pub struct CompressedPostingBuilder {
elements: Vec<PostingElement>,
}
impl CompressedPostingBuilder {
#[allow(clippy::new_without_default)]
pub fn new() -> Self {
CompressedPostingBuilder {
elements: Vec::new(),
}
}
pub fn add(&mut self, record_id: PointOffsetType, weight: DimWeight) {
self.elements.push(PostingElement { record_id, weight });
}
pub fn build<W: Weight>(mut self) -> CompressedPostingList<W> {
self.elements.sort_unstable_by_key(|e| e.record_id);
let quantization_params =
W::quantization_params_for(self.elements.iter().map(|e| e.weight));
#[cfg(debug_assertions)]
if let Some(e) = self.elements.iter().duplicates_by(|e| e.record_id).next() {
panic!("Duplicate id {} in posting list", e.record_id);
}
let mut this_chunk = Vec::with_capacity(CHUNK_SIZE);
let bitpacker = BitPackerImpl::new();
let mut chunks = Vec::with_capacity(self.elements.len() / CHUNK_SIZE);
let mut data_size = 0;
let mut remainders = Vec::with_capacity(self.elements.len() % CHUNK_SIZE);
for chunk in self.elements.chunks(CHUNK_SIZE) {
if chunk.len() == CHUNK_SIZE {
this_chunk.clear();
this_chunk.extend(chunk.iter().map(|e| e.record_id));
let initial = this_chunk[0];
let chunk_bits =
bitpacker.num_bits_strictly_sorted(initial.checked_sub(1), &this_chunk);
let chunk_size = BitPackerImpl::compressed_block_size(chunk_bits);
chunks.push(CompressedPostingChunk {
initial,
offset: data_size as u32,
weights: chunk
.iter()
.map(|e| Weight::from_f32(quantization_params, e.weight))
.collect::<Vec<_>>()
.try_into()
.expect("Invalid chunk size"),
});
data_size += chunk_size;
} else {
for e in chunk {
remainders.push(GenericPostingElement {
record_id: e.record_id,
weight: Weight::from_f32(quantization_params, e.weight),
});
}
}
}
let mut id_data = vec![0u8; data_size];
for (chunk_index, chunk_data) in
self.elements.as_chunks::<CHUNK_SIZE>().0.iter().enumerate()
{
this_chunk.clear();
this_chunk.extend(chunk_data.iter().map(|e| e.record_id));
let chunk = &chunks[chunk_index];
let chunk_size =
CompressedPostingListView::get_chunk_size(&chunks, &id_data, chunk_index);
let chunk_bits = chunk_size * u8::BITS as usize / CHUNK_SIZE;
bitpacker.compress_strictly_sorted(
chunk.initial.checked_sub(1),
&this_chunk,
&mut id_data[chunk.offset as usize..chunk.offset as usize + chunk_size],
chunk_bits as u8,
);
}
CompressedPostingList {
id_data,
chunks,
remainders,
last_id: self.elements.last().map(|e| e.record_id),
quantization_params,
}
}
}
#[derive(Clone)]
pub struct CompressedPostingListIterator<'a, W: Weight> {
list: CompressedPostingListView<'a, W>,
unpacked: bool,
decompressed_chunk: [PointOffsetType; CHUNK_SIZE],
pos: (usize, Option<PointOffsetType>),
}
impl<'a, W: Weight> CompressedPostingListIterator<'a, W> {
#[inline]
fn new(list: &CompressedPostingListView<'a, W>) -> Self {
Self {
list: list.clone(),
unpacked: false,
decompressed_chunk: [0; CHUNK_SIZE],
pos: (0, None),
}
}
#[inline]
fn next_from(&mut self, peek: PostingElementEx) -> PostingElement {
if self.pos.0 / CHUNK_SIZE < self.list.chunks.len() {
self.pos = (self.pos.0 + 1, None);
if self.pos.0.is_multiple_of(CHUNK_SIZE) {
self.unpacked = false;
}
} else {
self.pos = (self.pos.0 + 1, None);
}
peek.into()
}
#[inline]
fn next(&mut self) -> Option<PostingElement> {
let result = self.peek()?;
Some(self.next_from(result))
}
}
impl<W: Weight> PostingListIter for CompressedPostingListIterator<'_, W> {
#[inline]
fn peek(&mut self) -> Option<PostingElementEx> {
let pos = self.pos.0;
if pos / CHUNK_SIZE < self.list.chunks_len() {
if !self.unpacked {
self.list
.decompress_chunk(pos / CHUNK_SIZE, &mut self.decompressed_chunk);
self.unpacked = true;
}
return Some(PostingElementEx {
record_id: self.decompressed_chunk[pos % CHUNK_SIZE],
weight: self.list.get_weight(pos).to_f32(self.list.multiplier),
max_next_weight: Default::default(),
});
}
self.list
.get_remainder_id(pos - self.list.chunks_len() * CHUNK_SIZE)
.map(|e| PostingElementEx {
record_id: e.record_id,
weight: e.weight.to_f32(self.list.multiplier),
max_next_weight: Default::default(),
})
}
#[inline]
fn last_id(&self) -> Option<PointOffsetType> {
self.list.last_id
}
fn element_size(&self) -> usize {
size_of::<W>()
}
fn skip_to(&mut self, record_id: PointOffsetType) -> Option<PostingElementEx> {
if let Some(current_record_id) = self.pos.1.as_ref()
&& record_id < *current_record_id
{
return None;
}
let current_chunk_id_opt = self.list.chunk_id_by_position(self.pos.0);
let required_chunk_id = self.list.chunk_id_by_id(record_id);
match (required_chunk_id, current_chunk_id_opt) {
(IdChunkPosition::Chunk(chunk_id), Some(current_chunk_id)) => {
match chunk_id.cmp(¤t_chunk_id) {
Ordering::Less => {
return None;
}
Ordering::Equal => {
let min_pos = chunk_id * CHUNK_SIZE;
self.pos = (std::cmp::max(self.pos.0, min_pos), None);
}
Ordering::Greater => {
self.pos = (chunk_id * CHUNK_SIZE, None);
self.unpacked = false;
}
}
}
(IdChunkPosition::Chunk(_), None) => {
return None;
}
(IdChunkPosition::Before, _) => {
return None;
}
(IdChunkPosition::After, _) => {
let min_pos = self.list.chunks_len() * CHUNK_SIZE;
self.pos = (std::cmp::max(self.pos.0, min_pos), None);
self.unpacked = false;
}
};
while let Some(current_element) = self.peek() {
self.pos = (self.pos.0, Some(current_element.record_id));
match current_element.record_id.cmp(&record_id) {
Ordering::Equal => return Some(current_element),
Ordering::Greater => return None,
Ordering::Less => {
self.next_from(current_element);
}
}
}
None
}
#[inline]
fn skip_to_end(&mut self) {
self.pos = (
self.list.chunks_len() * CHUNK_SIZE + self.list.remainder_len(),
None,
);
}
#[inline]
fn len_to_end(&self) -> usize {
self.list.len() - self.pos.0
}
#[inline]
fn current_index(&self) -> usize {
self.pos.0
}
#[inline]
fn for_each_till_id<Ctx: ?Sized>(
&mut self,
id: PointOffsetType,
ctx: &mut Ctx,
mut f: impl FnMut(&mut Ctx, PointOffsetType, DimWeight),
) {
let mut pos = self.pos.0;
let mut weights_buf = [0.0; CHUNK_SIZE];
let mut need_unpack = !self.unpacked;
while pos / CHUNK_SIZE < self.list.chunks_len() {
if need_unpack {
self.list
.decompress_chunk(pos / CHUNK_SIZE, &mut self.decompressed_chunk);
}
need_unpack = true;
let start = pos % CHUNK_SIZE;
let count = count_le_sorted(id, &self.decompressed_chunk[start..]);
let weights = self.list.weights_range(pos, count);
let weights =
W::into_f32_slice(self.list.multiplier, weights, &mut weights_buf[..count]);
for (idx, weight) in
std::iter::zip(&self.decompressed_chunk[start..start + count], weights)
{
f(ctx, *idx, *weight);
}
pos += count;
if start + count != CHUNK_SIZE {
self.unpacked = true;
self.pos = (pos, None);
return;
}
}
for e in self
.list
.iter_remainder_from(pos - self.list.chunks_len() * CHUNK_SIZE)
{
if e.record_id > id {
self.pos = (pos, None);
return;
}
f(ctx, e.record_id, e.weight.to_f32(self.list.multiplier));
pos += 1;
}
self.pos = (pos, None);
}
fn reliable_max_next_weight() -> bool {
false
}
fn into_std_iter(self) -> impl Iterator<Item = PostingElement> {
CompressedPostingListStdIterator(self)
}
}
#[derive(Clone)]
pub struct CompressedPostingListStdIterator<'a, W: Weight>(CompressedPostingListIterator<'a, W>);
impl<W: Weight> Iterator for CompressedPostingListStdIterator<'_, W> {
type Item = PostingElement;
fn next(&mut self) -> Option<Self::Item> {
self.0.next()
}
}
fn count_le_sorted<T: Copy + Eq + Ord>(val: T, data: &[T]) -> usize {
if data.last().is_none_or(|&x| x < val) {
return data.len();
}
data.binary_search(&val).map_or_else(|x| x, |x| x + 1)
}
#[cfg(test)]
mod tests {
use super::*;
const CASES: [usize; 6] = [0, 64, 128, 192, 256, 320];
fn mk_case(count: usize) -> Vec<(PointOffsetType, DimWeight)> {
(0..count)
.map(|i| (i as u32 + 10000, i as DimWeight))
.collect()
}
fn cases() -> Vec<Vec<(PointOffsetType, DimWeight)>> {
CASES.iter().copied().map(mk_case).collect()
}
#[test]
fn test_iter() {
for case in cases() {
let list = CompressedPostingList::<f32>::from(case.clone());
let hw_counter = HardwareCounterCell::new();
let mut iter = list.iter(&hw_counter);
let mut count = 0;
assert_eq!(iter.len_to_end(), case.len(), "len_to_end");
while let Some(e) = iter.next() {
assert_eq!(e.record_id, case[count].0);
assert_eq!(e.weight, case[count].1);
assert_eq!(iter.len_to_end(), case.len() - count - 1);
count += 1;
}
}
}
#[test]
#[allow(clippy::needless_range_loop)] fn test_try_till_id() {
let hw_counter = HardwareCounterCell::new();
for i in 0..CASES.len() {
for j in i..CASES.len() {
for k in j..CASES.len() {
eprintln!("\n\n\n{} {} {}", CASES[i], CASES[j], CASES[k]);
let case = mk_case(CASES[k]);
let pl = CompressedPostingList::<f32>::from(case.clone());
let mut iter = pl.iter(&hw_counter);
let mut data = Vec::new();
let mut counter = 0;
iter.for_each_till_id(
case.get(CASES[i]).map_or(PointOffsetType::MAX, |x| x.0) - 1,
&mut (),
|_, id, weight| {
eprintln!(" {id}");
data.push((id, weight));
counter += 1;
},
);
assert_eq!(data, &case[..CASES[i]]);
eprintln!(" ;");
let mut data = Vec::new();
let mut counter = 0;
iter.for_each_till_id(
case.get(CASES[j]).map_or(PointOffsetType::MAX, |x| x.0) - 1,
&mut (),
|_, id, weight| {
eprintln!(" {id}");
data.push((id, weight));
counter += 1;
},
);
assert_eq!(data, &case[CASES[i]..CASES[j]]);
}
}
}
}
#[test]
fn test_count_le_sorted() {
let data = [1, 2, 4, 5];
for val in 0..9 {
let pos = count_le_sorted(val, &data);
assert!(data[..pos].iter().all(|&x| x <= val));
assert!(data[pos..].iter().all(|&x| x > val));
}
}
}