use std::cmp::max;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use crate::common::counter::hardware_counter::HardwareCounterCell;
use crate::common::fs::atomic_save_json;
use crate::common::mmap::AdviceSetting;
use crate::common::universal_io::{
OpenOptions, Populate, StoredStruct, UniversalKind, UniversalReadFileOps, UniversalWrite,
};
use num_traits::AsPrimitive;
use super::chunks::{create_chunk, read_chunks};
use super::config::{ChunkedVectorsConfig, Status};
use super::read::ChunkedVectorsRead;
use crate::segment::common::Flusher;
use crate::segment::common::operation_error::{OperationError, OperationResult};
use crate::segment::vector_storage::VectorOffsetType;
use crate::segment::vector_storage::common::CHUNK_SIZE;
#[derive(Debug)]
pub struct ChunkedVectors<T, S>
where
T: bytemuck::Pod + Send,
S: UniversalWrite + Send + 'static,
{
inner: ChunkedVectorsRead<T, S>,
status: StoredStruct<S, Status>,
fs: S::Fs,
}
impl<T, S> Deref for ChunkedVectors<T, S>
where
T: bytemuck::Pod + Send,
S: UniversalWrite + Send + 'static,
{
type Target = ChunkedVectorsRead<T, S>;
fn deref(&self) -> &Self::Target {
&self.inner
}
}
impl<T, S> ChunkedVectors<T, S>
where
T: bytemuck::Pod + Send,
S: UniversalWrite + Send + 'static,
{
pub fn storage_kind() -> UniversalKind {
S::kind()
}
pub fn ensure_status_file(fs: &S::Fs, directory: &Path) -> OperationResult<PathBuf> {
let status_file = ChunkedVectorsRead::<T, S>::status_file(directory);
if !fs.exists(&status_file)? {
{
let length = std::mem::size_of::<Status>();
crate::common::mmap::create_and_ensure_length(&status_file, length)?;
}
}
Ok(status_file)
}
fn ensure_config(
fs: &S::Fs,
directory: &Path,
dim: usize,
populate: Option<bool>,
) -> OperationResult<ChunkedVectorsConfig> {
let config_file = ChunkedVectorsRead::<T, S>::config_file(directory);
match ChunkedVectorsRead::<T, S>::load_config(fs, &config_file) {
Ok(Some(config)) => {
if config.dim == dim {
Ok(config)
} else {
Err(OperationError::service_error(format!(
"Wrong configuration in {}: expected {}, found {dim}",
config_file.display(),
config.dim,
)))
}
}
Ok(None) => Self::create_config(&config_file, dim, populate),
Err(e) => {
log::error!("Failed to deserialize config file {:?}: {e}", &config_file);
Self::create_config(&config_file, dim, populate)
}
}
}
fn create_config(
config_file: &Path,
dim: usize,
populate: Option<bool>,
) -> OperationResult<ChunkedVectorsConfig> {
if dim == 0 {
return Err(OperationError::service_error(
"The vector's dimension cannot be 0",
));
}
let chunk_size_bytes = CHUNK_SIZE;
let vector_size_bytes = dim * std::mem::size_of::<T>();
let chunk_size_vectors = chunk_size_bytes / vector_size_bytes;
let corrected_chunk_size_bytes = chunk_size_vectors * vector_size_bytes;
let config = ChunkedVectorsConfig {
chunk_size_bytes: corrected_chunk_size_bytes,
chunk_size_vectors,
dim,
populate,
};
atomic_save_json(config_file, &config)?;
Ok(config)
}
pub fn open(
fs: S::Fs,
directory: &Path,
dim: usize,
advice: AdviceSetting,
populate: Option<bool>,
) -> OperationResult<Self> {
fs_err::create_dir_all(directory)?;
let status_path = Self::ensure_status_file(&fs, directory)?;
let status: StoredStruct<S, Status> = StoredStruct::open(
&fs,
status_path,
OpenOptions {
writeable: true,
need_sequential: false,
populate: populate.map(Populate::from).unwrap_or_default(),
advice: AdviceSetting::Global,
},
Default::default(),
)?;
let config = Self::ensure_config(&fs, directory, dim, populate)?;
let chunks = read_chunks(&fs, directory, advice, populate.unwrap_or_default(), true)?;
let inner = ChunkedVectorsRead {
config,
len: status.len,
chunks,
directory: directory.to_owned(),
};
Ok(Self { inner, status, fs })
}
fn add_chunk(&mut self) -> OperationResult<()> {
let chunk = create_chunk(
&self.fs,
&self.inner.directory,
self.inner.chunks.len(),
self.inner.config.chunk_size_bytes,
)?;
self.inner.chunks.push(chunk);
Ok(())
}
pub fn insert(
&mut self,
key: VectorOffsetType,
vector: &[T],
hw_counter: &HardwareCounterCell,
) -> OperationResult<()> {
self.insert_many(key, vector, 1, hw_counter)
}
#[inline]
pub fn insert_many(
&mut self,
start_key: VectorOffsetType,
vectors: &[T],
count: usize,
hw_counter: &HardwareCounterCell,
) -> OperationResult<()> {
assert_eq!(
vectors.len(),
count * self.inner.config.dim,
"Vector size mismatch"
);
let start_key = start_key.as_();
let chunk_idx = self.inner.get_chunk_index(start_key);
let chunk_offset = self.inner.get_chunk_offset(start_key);
if chunk_offset + vectors.len()
> self.inner.config.dim * self.inner.config.chunk_size_vectors
{
return Err(OperationError::service_error(format!(
"Vectors do not fit in the chunk. Chunk idx {chunk_idx}, chunk offset {chunk_offset}, vectors count {count}",
)));
}
while chunk_idx >= self.inner.chunks.len() {
self.add_chunk()?;
}
let chunk = &mut self.inner.chunks[chunk_idx];
chunk.write((chunk_offset * size_of::<T>()) as u64, vectors)?;
hw_counter
.vector_io_write_counter()
.incr_delta(size_of_val(vectors));
let new_len = max(self.status.len, start_key + count);
if new_len > self.status.len {
self.status.len = new_len;
self.inner.len = new_len;
}
Ok(())
}
pub fn push(
&mut self,
vector: &[T],
hw_counter: &HardwareCounterCell,
) -> OperationResult<VectorOffsetType> {
let new_id = self.status.len;
self.insert(new_id, vector, hw_counter)?;
Ok(new_id)
}
pub fn flusher(&self) -> Flusher {
Box::new({
let status_flusher = self.status.flusher();
let chunks_flushers: Vec<_> = self
.inner
.chunks
.iter()
.map(|chunk| chunk.flusher())
.collect();
move || {
for flusher in chunks_flushers {
flusher()?;
}
status_flusher()?;
Ok(())
}
})
}
}