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

[UIO] migrate SliceBufferedUpdateWrapper (#8518)

Browse Source 
* migrate SliceBufferedUpdateWrapper

* fix compressed versions creation

* @xzfc's review improvements

Co-authored-by: xzfc <5121426+xzfc@users.noreply.github.com>

* fix rebase

* Drop usage of `TypedStorage`

---------

Co-authored-by: xzfc <5121426+xzfc@users.noreply.github.com>
Co-authored-by: xzfc <xzfcpw@gmail.com>
This commit is contained in:
Luis Cossío
2026-04-08 13:09:55 -04:00
committed by GitHub
parent c9c4227691
commit b1e9c9a9df
4 changed files with 83 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
lib/segment/src/common/mmap_slice_buffered_update_wrapper.rs → lib/common/common/src/universal_io/wrappers/buffered_update.rs
View File

@@ -1,41 +1,40 @@
use std::sync::Arc;
use ahash::AHashMap;
use common::is_alive_lock::IsAliveLock;
use common::mmap::MmapSlice;
use common::types::PointOffsetType;
use parking_lot::{Mutex, RwLock};
use crate::common::Flusher;
use crate::is_alive_lock::IsAliveLock;
use crate::types::PointOffsetType;
use crate::universal_io::{Flusher, UniversalIoError, UniversalWrite};
/// A wrapper around `MmapSlice` that delays writing changes to the underlying file until they get
/// flushed manually.
/// This expects the underlying MmapSlice not to grow in size.
/// A wrapper around [`UniversalWrite`] that delays writing changes to the
/// underlying file until they get flushed manually.
/// This expects the underlying storage not to grow in size.
///
/// WARN: this structure is expected to be write-only.
#[derive(Debug)]
pub struct MmapSliceBufferedUpdateWrapper<T>
pub struct SliceBufferedUpdateWrapper<S: UniversalWrite<T>, T: Copy>
where
T: 'static,
{
mmap_slice: Arc<RwLock<MmapSlice<T>>>,
len: usize,
slice: Arc<RwLock<S>>,
len: u64,
pending_updates: Arc<Mutex<AHashMap<PointOffsetType, T>>>,
is_alive_lock: IsAliveLock,
}
impl<T> MmapSliceBufferedUpdateWrapper<T>
impl<S: UniversalWrite<T>, T: Copy> SliceBufferedUpdateWrapper<S, T>
where
T: 'static,
{
pub fn new(mmap_slice: MmapSlice<T>) -> Self {
let len = mmap_slice.len();
Self {
mmap_slice: Arc::new(RwLock::new(mmap_slice)),
pub fn new(slice_storage: S) -> Result<Self, UniversalIoError> {
let len = slice_storage.len()?;
Ok(Self {
slice: Arc::new(RwLock::new(slice_storage)),
len,
pending_updates: Arc::new(Mutex::new(AHashMap::new())),
is_alive_lock: IsAliveLock::new(),
}
})
}
/// Sets the item at `index` to `value` buffered.
@@ -44,7 +43,7 @@ where
/// Panics if the index is out of bounds.
pub fn set(&self, index: PointOffsetType, value: T) {
assert!(
(index as usize) < self.len,
u64::from(index) < self.len,
"index {index} out of range: {}",
self.len
);
@@ -52,9 +51,10 @@ where
}
}
impl<T> MmapSliceBufferedUpdateWrapper<T>
impl<S, T> SliceBufferedUpdateWrapper<S, T>
where
T: 'static + Sync + Send + Clone + PartialEq,
S: UniversalWrite<T> + Send + Sync + 'static,
T: Sync + Send + Copy + Clone + PartialEq + 'static,
{
pub fn flusher(&self) -> Flusher {
let updates = {
@@ -66,7 +66,7 @@ where
};
let pending_updates_weak = Arc::downgrade(&self.pending_updates);
let slice = Arc::downgrade(&self.mmap_slice);
let slice = Arc::downgrade(&self.slice);
let is_alive_handle = self.is_alive_lock.handle();
Box::new(move || {
let (Some(is_alive_guard), Some(pending_updates_arc), Some(slice)) = (
@@ -74,17 +74,32 @@ where
pending_updates_weak.upgrade(),
slice.upgrade(),
) else {
log::debug!(
"Aborted flushing on a dropped MmapSliceBufferedUpdateWrapper instance"
);
log::debug!("Aborted flushing on a dropped SliceBufferedUpdateWrapper instance");
return Ok(());
};
let mut mmap_slice_write = slice.write();
for (&index, value) in &updates {
mmap_slice_write[index as usize] = value.clone();
}
mmap_slice_write.flusher()()?;
let mut slice_guard = slice.write();
// Coalesce contiguous updates into the same write
let mut items: Vec<(PointOffsetType, T)> =
updates.iter().map(|(k, v)| (*k, *v)).collect();
items.sort_by_key(|(index, _)| *index);
let all_values: Vec<T> = items.iter().map(|(_, value)| *value).collect();
// Batch writes
let mut remaining_values: &[T] = &all_values[..];
let it = items.chunk_by(|(a, _), (b, _)| *a + 1 == *b).map(|chunk| {
let elem_start: PointOffsetType = chunk[0].0;
let byte_start = u64::from(elem_start) * size_of::<T>() as u64;
let chunk_values = remaining_values
.split_off(..chunk.len())
.expect("`chunk.len()` is sourced from the same slice as `remaining_values`");
(byte_start, chunk_values)
});
slice_guard.write_batch(it)?;
slice_guard.flusher()()?;
// Keep the guard till here to prevent concurrent drop/flushes
// We don't touch files from here on and can drop the alive guard

2
lib/common/common/src/universal_io/wrappers/mod.rs
View File

@@ -1,5 +1,7 @@
mod buffered_update;
mod read_only;
mod typed;
pub use buffered_update::SliceBufferedUpdateWrapper;
pub use read_only::ReadOnly;
pub use typed::TypedStorage;

1
lib/segment/src/common/mod.rs
View File

@@ -3,7 +3,6 @@ pub mod error_logging;
pub mod flags;
pub mod macros;
pub mod mmap_bitslice_buffered_update_wrapper;
pub mod mmap_slice_buffered_update_wrapper;
pub mod operation_error;
pub mod operation_time_statistics;
pub mod reciprocal_rank_fusion;

62
lib/segment/src/id_tracker/immutable_id_tracker.rs
View File

@@ -4,15 +4,16 @@ use std::path::{Path, PathBuf};
use byteorder::{ReadBytesExt, WriteBytesExt};
use common::bitvec::{BitSlice, BitSliceExt as _, BitVec};
use common::mmap::{AdviceSetting, MmapSlice, create_and_ensure_length, open_write_mmap};
use common::mmap::create_and_ensure_length;
use common::types::PointOffsetType;
use common::universal_io::OpenOptions;
use common::universal_io::{
MmapFile, OpenOptions, SliceBufferedUpdateWrapper, TypedStorage, UniversalRead, UniversalWrite,
};
use fs_err::File;
use uuid::Uuid;
use crate::common::Flusher;
use crate::common::mmap_bitslice_buffered_update_wrapper::MmapBitSliceBufferedUpdateWrapper;
use crate::common::mmap_slice_buffered_update_wrapper::MmapSliceBufferedUpdateWrapper;
use crate::common::operation_error::{OperationError, OperationResult};
use crate::common::stored_bitslice::MmapBitSlice;
use crate::id_tracker::compressed::compressed_point_mappings::CompressedPointMappings;
@@ -59,7 +60,7 @@ pub struct ImmutableIdTracker {
deleted_wrapper: MmapBitSliceBufferedUpdateWrapper,
internal_to_version: CompressedVersions,
internal_to_version_wrapper: MmapSliceBufferedUpdateWrapper<SeqNumberType>,
internal_to_version_wrapper: SliceBufferedUpdateWrapper<MmapFile, SeqNumberType>,
mappings: CompressedPointMappings,
}
@@ -257,16 +258,23 @@ impl ImmutableIdTracker {
let deleted_wrapper = MmapBitSliceBufferedUpdateWrapper::new(deleted_storage);
let internal_to_version_map = open_write_mmap(
&Self::version_mapping_file_path(segment_path),
AdviceSetting::Global,
true,
let internal_to_version_file = TypedStorage::<MmapFile, SeqNumberType>::open(
Self::version_mapping_file_path(segment_path),
OpenOptions {
writeable: true,
need_sequential: false,
disk_parallel: None,
populate: Some(true),
advice: None,
prevent_caching: None,
},
)?;
let internal_to_version_mapslice: MmapSlice<SeqNumberType> =
unsafe { MmapSlice::try_from(internal_to_version_map)? };
let internal_to_version = CompressedVersions::from_slice(&internal_to_version_mapslice);
let internal_to_version_slice = internal_to_version_file.read_whole()?;
let internal_to_version = CompressedVersions::from_slice(&internal_to_version_slice);
let internal_to_version_wrapper =
MmapSliceBufferedUpdateWrapper::new(internal_to_version_mapslice);
SliceBufferedUpdateWrapper::new(internal_to_version_file.inner)?;
let reader = BufReader::new(File::open(Self::mappings_file_path(segment_path))?);
let mappings = Self::load_mapping(reader, Some(deleted_bitvec))?;
@@ -326,22 +334,27 @@ impl ImmutableIdTracker {
let version_size = mmap_size::<SeqNumberType>(min_size);
create_and_ensure_length(&version_filepath, version_size)?;
}
let mut internal_to_version_wrapper = unsafe {
MmapSlice::try_from(open_write_mmap(
&version_filepath,
AdviceSetting::Global,
false,
)?)?
};
internal_to_version_wrapper[..internal_to_version.len()]
.copy_from_slice(internal_to_version);
let internal_to_version = CompressedVersions::from_slice(&internal_to_version_wrapper);
let mut internal_to_version_file = TypedStorage::<MmapFile, SeqNumberType>::open(
&version_filepath,
OpenOptions {
writeable: true,
need_sequential: false,
disk_parallel: None,
populate: Some(false),
advice: None,
prevent_caching: None,
},
)?;
internal_to_version_file.write(0, internal_to_version)?;
let internal_to_version =
CompressedVersions::from_slice(&internal_to_version_file.read_whole()?);
debug_assert_eq!(internal_to_version.len(), mappings.total_point_count());
let internal_to_version_wrapper =
MmapSliceBufferedUpdateWrapper::new(internal_to_version_wrapper);
SliceBufferedUpdateWrapper::new(internal_to_version_file.inner)?;
// Write mappings to disk.
let file = File::create(Self::mappings_file_path(path))?;
@@ -457,7 +470,8 @@ impl IdTracker for ImmutableIdTracker {
/// Creates a flusher function, that writes the points versions to disk.
fn versions_flusher(&self) -> Flusher {
self.internal_to_version_wrapper.flusher()
let flusher = self.internal_to_version_wrapper.flusher();
Box::new(move || flusher().map_err(OperationError::from))
}
fn total_point_count(&self) -> usize {