From 67295ef87f4d93ea4612393bd41787d9ee714c85 Mon Sep 17 00:00:00 2001 From: Andrey Vasnetsov Date: Tue, 12 May 2026 16:21:37 +0200 Subject: [PATCH] refactor(map_index): split mod.rs into read_ops, lifecycle, tests (#9015) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(map_index): split mod.rs into read_ops, lifecycle, tests mod.rs is now ~40 lines containing only the MapIndex enum, type aliases, and module declarations. - read_ops.rs: read-only inherent methods (get_values, get_iterator, for_each_*, except_cardinality, except_set, telemetry, ram_usage, mutability/storage type) - lifecycle.rs: open/builder/flush/wipe/remove_point/files/populate/ clear_cache - tests.rs: all #[cfg(test)] tests Also folded payload_index_impl_{int,str,uuid}.rs into a dedicated payload_index_impl/ submodule. * refactor(map_index): split storage submodules into dedicated dirs (#9016) * refactor(map_index): split storage submodules into dedicated module dirs Turn each of the three storage implementation files into a directory module split into read_ops + lifecycle, mirroring the parent module layout introduced in #9015. - mutable_map_index/{mod,lifecycle,read_ops}.rs - immutable_map_index/{mod,lifecycle,read_ops}.rs - mmap_map_index/{mod,lifecycle,read_ops}.rs Each mod.rs holds only the struct definitions, internal Storage type, and config/constants. read_ops.rs holds the read-only query methods; lifecycle.rs holds open/build/flush/wipe/files/remove_point and internal mutation helpers. Pure refactor — no behavior change. * refactor(map_index): introduce MapIndexRead trait Define a unified read-only trait `MapIndexRead` describing the methods every storage variant exposes (check_values_any, get_values, get_iterator, for_each_*, storage_type, ram_usage_bytes, etc.). Each storage variant's read_ops.rs now contains a trait impl instead of inherent methods. Signatures are unified across variants: - `hw_counter` is accepted by every method that needs it for the mmap variant; mutable / immutable accept and ignore it. - `check_values_any` returns `bool` (mmap absorbs IO errors internally with the existing FIXME, matching the parent's prior `.unwrap_or`). - `for_each_count_per_value` takes `deferred_internal_id` uniformly; the immutable variant `debug_assert!`s it is `None`. `for_points_values` keeps its variant-specific callback signatures and stays as an inherent method — it's only used by FacetIndex with explicit pattern matching. Pure refactor — no behavior change. * refactor(mutable_map_index): drop single-variant Storage enum The `Storage` enum had only one variant (`Gridstore`), so every `match &self.storage { Storage::Gridstore(s) => ... }` was just unwrapping the same path. Replace the field with `Gridstore>` directly and inline every match. --- .../lifecycle.rs} | 196 +--- .../map_index/immutable_map_index/mod.rs | 41 + .../map_index/immutable_map_index/read_ops.rs | 180 ++++ .../index/field_index/map_index/lifecycle.rs | 133 +++ .../field_index/map_index/mmap_map_index.rs | 497 ---------- .../map_index/mmap_map_index/lifecycle.rs | 222 +++++ .../map_index/mmap_map_index/mod.rs | 64 ++ .../map_index/mmap_map_index/read_ops.rs | 253 +++++ .../src/index/field_index/map_index/mod.rs | 906 +----------------- .../map_index/mutable_map_index.rs | 360 ------- .../map_index/mutable_map_index/lifecycle.rs | 195 ++++ .../map_index/mutable_map_index/mod.rs | 21 + .../map_index/mutable_map_index/read_ops.rs | 150 +++ .../int.rs} | 2 +- .../map_index/payload_index_impl/mod.rs | 3 + .../str.rs} | 2 +- .../uuid.rs} | 2 +- .../index/field_index/map_index/read_ops.rs | 381 ++++++++ .../src/index/field_index/map_index/tests.rs | 452 +++++++++ 19 files changed, 2113 insertions(+), 1947 deletions(-) rename lib/segment/src/index/field_index/map_index/{immutable_map_index.rs => immutable_map_index/lifecycle.rs} (59%) create mode 100644 lib/segment/src/index/field_index/map_index/immutable_map_index/mod.rs create mode 100644 lib/segment/src/index/field_index/map_index/immutable_map_index/read_ops.rs create mode 100644 lib/segment/src/index/field_index/map_index/lifecycle.rs delete mode 100644 lib/segment/src/index/field_index/map_index/mmap_map_index.rs create mode 100644 lib/segment/src/index/field_index/map_index/mmap_map_index/lifecycle.rs create mode 100644 lib/segment/src/index/field_index/map_index/mmap_map_index/mod.rs create mode 100644 lib/segment/src/index/field_index/map_index/mmap_map_index/read_ops.rs delete mode 100644 lib/segment/src/index/field_index/map_index/mutable_map_index.rs create mode 100644 lib/segment/src/index/field_index/map_index/mutable_map_index/lifecycle.rs create mode 100644 lib/segment/src/index/field_index/map_index/mutable_map_index/mod.rs create mode 100644 lib/segment/src/index/field_index/map_index/mutable_map_index/read_ops.rs rename lib/segment/src/index/field_index/map_index/{payload_index_impl_int.rs => payload_index_impl/int.rs} (99%) create mode 100644 lib/segment/src/index/field_index/map_index/payload_index_impl/mod.rs rename lib/segment/src/index/field_index/map_index/{payload_index_impl_str.rs => payload_index_impl/str.rs} (99%) rename lib/segment/src/index/field_index/map_index/{payload_index_impl_uuid.rs => payload_index_impl/uuid.rs} (99%) create mode 100644 lib/segment/src/index/field_index/map_index/read_ops.rs create mode 100644 lib/segment/src/index/field_index/map_index/tests.rs diff --git a/lib/segment/src/index/field_index/map_index/immutable_map_index.rs b/lib/segment/src/index/field_index/map_index/immutable_map_index/lifecycle.rs similarity index 59% rename from lib/segment/src/index/field_index/map_index/immutable_map_index.rs rename to lib/segment/src/index/field_index/map_index/immutable_map_index/lifecycle.rs index bb7e058318..a714c0e527 100644 --- a/lib/segment/src/index/field_index/map_index/immutable_map_index.rs +++ b/lib/segment/src/index/field_index/map_index/immutable_map_index/lifecycle.rs @@ -1,56 +1,26 @@ -use std::borrow::{Borrow as _, Cow}; +use std::borrow::Borrow as _; use std::collections::HashMap; -use std::iter; -use std::ops::Range; use std::path::PathBuf; use bitvec::vec::BitVec; use common::counter::hardware_counter::HardwareCounterCell; -use common::persisted_hashmap::Key; use common::types::PointOffsetType; use gridstore::Blob; -use super::mmap_map_index::MmapMapIndex; -use super::{IdIter, MapIndexKey}; +use super::super::MapIndexKey; +use super::super::mmap_map_index::MmapMapIndex; +use super::super::read_ops::MapIndexRead; +use super::{ContainerSegment, ImmutableMapIndex, Storage}; use crate::common::Flusher; use crate::common::operation_error::OperationResult; use crate::index::field_index::immutable_point_to_values::ImmutablePointToValues; -use crate::index::payload_config::StorageType; - -pub struct ImmutableMapIndex { - value_to_points: HashMap<::Owned, ContainerSegment>, - /// Container holding a slice of point IDs per value. `value_to_point` holds the range per value. - /// Each slice MUST be sorted so that we can binary search over it. - value_to_points_container: Vec, - deleted_value_to_points_container: BitVec, - point_to_values: ImmutablePointToValues<::Owned>, - /// Amount of point which have at least one indexed payload value - indexed_points: usize, - values_count: usize, - // Backing storage, source of state, persists deletions - storage: Storage, - /// Snapshot of approximate RAM usage at construction time. - /// Not refreshed on `remove_point`. - cached_ram_usage_bytes: usize, -} - -enum Storage { - Mmap(Box>), -} - -pub(super) struct ContainerSegment { - /// Range in the container which holds point IDs for the value. - range: Range, - /// Number of available point IDs in the range, excludes number of deleted points. - count: u32, -} impl ImmutableMapIndex where Vec<::Owned>: Blob + Send + Sync, { /// Open and load immutable map index from mmap storage - pub(super) fn open_mmap(index: MmapMapIndex) -> OperationResult { + pub(in super::super) fn open_mmap(index: MmapMapIndex) -> OperationResult { let hw_counter = HardwareCounterCell::disposable(); // Internal operation let mut indexed_points = 0; @@ -258,7 +228,7 @@ where } #[inline] - pub(super) fn wipe(self) -> OperationResult<()> { + pub(in super::super) fn wipe(self) -> OperationResult<()> { match self.storage { Storage::Mmap(index) => index.wipe(), } @@ -275,169 +245,23 @@ where } #[inline] - pub(super) fn files(&self) -> Vec { + pub(in super::super) fn files(&self) -> Vec { match self.storage { Storage::Mmap(ref index) => index.files(), } } #[inline] - pub(super) fn immutable_files(&self) -> Vec { + pub(in super::super) fn immutable_files(&self) -> Vec { match &self.storage { Storage::Mmap(index) => index.immutable_files(), } } #[inline] - pub(super) fn flusher(&self) -> Flusher { + pub(in super::super) fn flusher(&self) -> Flusher { match self.storage { Storage::Mmap(ref index) => index.flusher(), } } - - pub fn check_values_any(&self, idx: PointOffsetType, check_fn: impl Fn(&N) -> bool) -> bool { - self.point_to_values - .check_values_any(idx, |v| check_fn(v.borrow())) - } - - pub fn get_values( - &self, - idx: PointOffsetType, - ) -> Option> + '_> { - Some( - self.point_to_values - .get_values(idx)? - .map(|v| Cow::Borrowed(v.borrow())), - ) - } - - pub fn values_count(&self, idx: PointOffsetType) -> Option { - Some(self.point_to_values.get_values(idx)?.count()) - } - - pub fn get_indexed_points(&self) -> usize { - self.indexed_points - } - - pub fn get_values_count(&self) -> usize { - self.values_count - } - - pub fn get_unique_values_count(&self) -> usize { - self.value_to_points.len() - } - - pub fn get_count_for_value(&self, value: &N) -> Option { - self.value_to_points - .get(value) - .map(|entry| entry.count as usize) - } - - pub fn for_points_values( - &self, - points: impl Iterator, - mut f: impl FnMut(PointOffsetType, &[::Owned]), - ) { - points.for_each(|idx| { - if let Some(values) = self.point_to_values.get_values_slice(idx) { - f(idx, values); - } - }); - } - - pub fn for_each_count_per_value( - &self, - mut f: impl FnMut(&N, usize) -> OperationResult<()>, - ) -> OperationResult<()> { - self.value_to_points - .iter() - .try_for_each(|(k, entry)| f(k.borrow(), entry.count as usize)) - } - - pub fn for_each_value_map( - &self, - mut f: impl FnMut(&N, &mut dyn Iterator) -> OperationResult<()>, - ) -> OperationResult<()> { - self.value_to_points - .iter() - .try_for_each(|(k, entry)| f(k.borrow(), &mut self.get_entry_iterator(entry))) - } - - pub fn get_iterator(&self, value: &N) -> IdIter<'_> { - if let Some(entry) = self.value_to_points.get(value) { - Box::new(self.get_entry_iterator(entry)) - } else { - Box::new(iter::empty::()) - } - } - - fn get_entry_iterator( - &self, - entry: &ContainerSegment, - ) -> impl Iterator { - let range = entry.range.start as usize..entry.range.end as usize; - - let deleted_flags = self - .deleted_value_to_points_container - .iter() - .by_vals() - .skip(range.start) - .chain(std::iter::repeat(false)); - - self.value_to_points_container[range] - .iter() - .zip(deleted_flags) - .filter(|(_, is_deleted)| !is_deleted) - .map(|(idx, _)| *idx) - } - - pub fn for_each_value( - &self, - mut f: impl FnMut(&N) -> OperationResult<()>, - ) -> OperationResult<()> { - self.value_to_points.keys().try_for_each(|v| f(v.borrow())) - } - - pub fn storage_type(&self) -> StorageType { - match &self.storage { - Storage::Mmap(index) => StorageType::Mmap { - is_on_disk: index.is_on_disk(), - }, - } - } - - /// Approximate RAM usage in bytes (cached at construction). - pub fn ram_usage_bytes(&self) -> usize { - self.cached_ram_usage_bytes - } - - fn compute_ram_usage_bytes(&self) -> usize { - let Self { - value_to_points, - value_to_points_container, - deleted_value_to_points_container, - point_to_values, - indexed_points: _, - values_count: _, - storage: _, - cached_ram_usage_bytes: _, - } = self; - - let hashmap_entry_overhead = size_of::() + size_of::(); - let vtp_base_bytes: usize = value_to_points.capacity() - * (size_of::<::Owned>() - + size_of::() - + hashmap_entry_overhead); - // Account for heap-allocated key data (e.g., long strings) - let vtp_heap_bytes: usize = value_to_points.keys().map(|k| N::owned_heap_bytes(k)).sum(); - let container_bytes = value_to_points_container.capacity() * size_of::(); - let deleted_bytes = deleted_value_to_points_container - .capacity() - .div_ceil(u8::BITS as usize); - vtp_base_bytes - + vtp_heap_bytes - + container_bytes - + deleted_bytes - + point_to_values.ram_usage_bytes() - } } diff --git a/lib/segment/src/index/field_index/map_index/immutable_map_index/mod.rs b/lib/segment/src/index/field_index/map_index/immutable_map_index/mod.rs new file mode 100644 index 0000000000..e003c7e0a2 --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/immutable_map_index/mod.rs @@ -0,0 +1,41 @@ +use std::collections::HashMap; +use std::ops::Range; + +use bitvec::vec::BitVec; +use common::persisted_hashmap::Key; +use common::types::PointOffsetType; + +use super::MapIndexKey; +use super::mmap_map_index::MmapMapIndex; +use crate::index::field_index::immutable_point_to_values::ImmutablePointToValues; + +mod lifecycle; +mod read_ops; + +pub struct ImmutableMapIndex { + pub(super) value_to_points: HashMap<::Owned, ContainerSegment>, + /// Container holding a slice of point IDs per value. `value_to_point` holds the range per value. + /// Each slice MUST be sorted so that we can binary search over it. + pub(super) value_to_points_container: Vec, + pub(super) deleted_value_to_points_container: BitVec, + pub(super) point_to_values: ImmutablePointToValues<::Owned>, + /// Amount of point which have at least one indexed payload value + pub(super) indexed_points: usize, + pub(super) values_count: usize, + // Backing storage, source of state, persists deletions + pub(super) storage: Storage, + /// Snapshot of approximate RAM usage at construction time. + /// Not refreshed on `remove_point`. + pub(super) cached_ram_usage_bytes: usize, +} + +pub(super) enum Storage { + Mmap(Box>), +} + +pub(super) struct ContainerSegment { + /// Range in the container which holds point IDs for the value. + range: Range, + /// Number of available point IDs in the range, excludes number of deleted points. + count: u32, +} diff --git a/lib/segment/src/index/field_index/map_index/immutable_map_index/read_ops.rs b/lib/segment/src/index/field_index/map_index/immutable_map_index/read_ops.rs new file mode 100644 index 0000000000..4f28b8cbcc --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/immutable_map_index/read_ops.rs @@ -0,0 +1,180 @@ +use std::borrow::{Borrow as _, Cow}; +use std::iter; + +use common::counter::hardware_counter::HardwareCounterCell; +use common::types::PointOffsetType; +use gridstore::Blob; + +use super::super::read_ops::MapIndexRead; +use super::super::{IdIter, MapIndexKey}; +use super::{ContainerSegment, ImmutableMapIndex, Storage}; +use crate::common::operation_error::OperationResult; +use crate::index::payload_config::StorageType; + +impl MapIndexRead for ImmutableMapIndex +where + Vec<::Owned>: Blob + Send + Sync, +{ + fn check_values_any( + &self, + idx: PointOffsetType, + _hw_counter: &HardwareCounterCell, + check_fn: impl Fn(&N) -> bool, + ) -> bool { + self.point_to_values + .check_values_any(idx, |v| check_fn(v.borrow())) + } + + fn get_values<'a>( + &'a self, + idx: PointOffsetType, + _hw_counter: &HardwareCounterCell, + ) -> Option> + 'a> + where + N: 'a, + { + Some( + self.point_to_values + .get_values(idx)? + .map(|v| Cow::Borrowed(v.borrow())), + ) + } + + fn values_count(&self, idx: PointOffsetType) -> Option { + Some(self.point_to_values.get_values(idx)?.count()) + } + + fn get_indexed_points(&self) -> usize { + self.indexed_points + } + + fn get_values_count(&self) -> usize { + self.values_count + } + + fn get_unique_values_count(&self) -> usize { + self.value_to_points.len() + } + + fn get_count_for_value(&self, value: &N, _hw_counter: &HardwareCounterCell) -> Option { + self.value_to_points + .get(value) + .map(|entry| entry.count as usize) + } + + fn get_iterator(&self, value: &N, _hw_counter: &HardwareCounterCell) -> IdIter<'_> { + if let Some(entry) = self.value_to_points.get(value) { + Box::new(self.get_entry_iterator(entry)) + } else { + Box::new(iter::empty::()) + } + } + + fn for_each_value(&self, mut f: impl FnMut(&N) -> OperationResult<()>) -> OperationResult<()> { + self.value_to_points.keys().try_for_each(|v| f(v.borrow())) + } + + fn for_each_count_per_value( + &self, + deferred_internal_id: Option, + mut f: impl FnMut(&N, usize) -> OperationResult<()>, + ) -> OperationResult<()> { + // Immutable indexes don't support deferred filtering; callers must + // pass `None`. See `MapIndex::for_each_count_per_value` for context. + debug_assert!(deferred_internal_id.is_none()); + let _ = deferred_internal_id; + self.value_to_points + .iter() + .try_for_each(|(k, entry)| f(k.borrow(), entry.count as usize)) + } + + fn for_each_value_map( + &self, + _hw_counter: &HardwareCounterCell, + mut f: impl FnMut(&N, &mut dyn Iterator) -> OperationResult<()>, + ) -> OperationResult<()> { + self.value_to_points + .iter() + .try_for_each(|(k, entry)| f(k.borrow(), &mut self.get_entry_iterator(entry))) + } + + fn storage_type(&self) -> StorageType { + match &self.storage { + Storage::Mmap(index) => StorageType::Mmap { + is_on_disk: index.is_on_disk(), + }, + } + } + + /// Approximate RAM usage in bytes (cached at construction). + fn ram_usage_bytes(&self) -> usize { + self.cached_ram_usage_bytes + } +} + +impl ImmutableMapIndex +where + Vec<::Owned>: Blob + Send + Sync, +{ + pub fn for_points_values( + &self, + points: impl Iterator, + mut f: impl FnMut(PointOffsetType, &[::Owned]), + ) { + points.for_each(|idx| { + if let Some(values) = self.point_to_values.get_values_slice(idx) { + f(idx, values); + } + }); + } + + fn get_entry_iterator( + &self, + entry: &ContainerSegment, + ) -> impl Iterator { + let range = entry.range.start as usize..entry.range.end as usize; + + let deleted_flags = self + .deleted_value_to_points_container + .iter() + .by_vals() + .skip(range.start) + .chain(std::iter::repeat(false)); + + self.value_to_points_container[range] + .iter() + .zip(deleted_flags) + .filter(|(_, is_deleted)| !is_deleted) + .map(|(idx, _)| *idx) + } + + pub(super) fn compute_ram_usage_bytes(&self) -> usize { + let Self { + value_to_points, + value_to_points_container, + deleted_value_to_points_container, + point_to_values, + indexed_points: _, + values_count: _, + storage: _, + cached_ram_usage_bytes: _, + } = self; + + let hashmap_entry_overhead = size_of::() + size_of::(); + let vtp_base_bytes: usize = value_to_points.capacity() + * (size_of::<::Owned>() + + size_of::() + + hashmap_entry_overhead); + // Account for heap-allocated key data (e.g., long strings) + let vtp_heap_bytes: usize = value_to_points.keys().map(|k| N::owned_heap_bytes(k)).sum(); + let container_bytes = value_to_points_container.capacity() * size_of::(); + let deleted_bytes = deleted_value_to_points_container + .capacity() + .div_ceil(u8::BITS as usize); + vtp_base_bytes + + vtp_heap_bytes + + container_bytes + + deleted_bytes + + point_to_values.ram_usage_bytes() + } +} diff --git a/lib/segment/src/index/field_index/map_index/lifecycle.rs b/lib/segment/src/index/field_index/map_index/lifecycle.rs new file mode 100644 index 0000000000..28b8bb13fe --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/lifecycle.rs @@ -0,0 +1,133 @@ +use std::path::{Path, PathBuf}; + +use common::bitvec::BitSlice; +use common::types::PointOffsetType; +use gridstore::Blob; + +use super::MapIndex; +use super::builders::MapIndexMmapBuilder; +use super::immutable_map_index::ImmutableMapIndex; +use super::key::MapIndexKey; +use super::mmap_map_index::MmapMapIndex; +use super::mutable_map_index::MutableMapIndex; +use crate::common::Flusher; +use crate::common::operation_error::OperationResult; + +impl MapIndex +where + Vec<::Owned>: Blob + Send + Sync, +{ + /// Load immutable mmap based index, either in RAM or on disk + pub fn new_mmap( + path: &Path, + is_on_disk: bool, + deleted_points: &BitSlice, + ) -> OperationResult> { + // Low-memory mode downgrades the in-RAM `Immutable` wrapper to the + // pure-mmap `Storage` variant at load time. Files are shared between + // variants; the persisted `is_on_disk` flag in `mmap_index` is + // untouched. + let effective_is_on_disk = + is_on_disk || common::low_memory::low_memory_mode().prefer_disk(); + + let Some(mmap_index) = MmapMapIndex::open(path, effective_is_on_disk, deleted_points)? + else { + return Ok(None); + }; + + let index = if effective_is_on_disk { + MapIndex::Mmap(Box::new(mmap_index)) + } else { + // Load into RAM, use mmap as backing storage + MapIndex::Immutable(ImmutableMapIndex::open_mmap(mmap_index)?) + }; + Ok(Some(index)) + } + + pub fn new_gridstore(dir: PathBuf, create_if_missing: bool) -> OperationResult> { + let index = MutableMapIndex::open_gridstore(dir, create_if_missing)?; + Ok(index.map(MapIndex::Mutable)) + } + + pub fn builder_mmap( + path: &Path, + is_on_disk: bool, + deleted_points: &BitSlice, + ) -> MapIndexMmapBuilder { + MapIndexMmapBuilder { + path: path.to_owned(), + point_to_values: Default::default(), + values_to_points: Default::default(), + is_on_disk, + deleted_points: deleted_points.to_owned(), + } + } + + pub fn builder_gridstore(dir: PathBuf) -> super::builders::MapIndexGridstoreBuilder { + super::builders::MapIndexGridstoreBuilder::new(dir) + } + + pub(crate) fn flusher(&self) -> Flusher { + match self { + MapIndex::Mutable(index) => index.flusher(), + MapIndex::Immutable(index) => index.flusher(), + MapIndex::Mmap(index) => index.flusher(), + } + } + + pub(crate) fn wipe(self) -> OperationResult<()> { + match self { + MapIndex::Mutable(index) => index.wipe(), + MapIndex::Immutable(index) => index.wipe(), + MapIndex::Mmap(index) => index.wipe(), + } + } + + pub(crate) fn remove_point(&mut self, id: PointOffsetType) -> OperationResult<()> { + match self { + MapIndex::Mutable(index) => index.remove_point(id), + MapIndex::Immutable(index) => index.remove_point(id), + MapIndex::Mmap(index) => { + index.remove_point(id); + Ok(()) + } + } + } + + pub(crate) fn files(&self) -> Vec { + match self { + MapIndex::Mutable(index) => index.files(), + MapIndex::Immutable(index) => index.files(), + MapIndex::Mmap(index) => index.files(), + } + } + + pub(crate) fn immutable_files(&self) -> Vec { + match self { + MapIndex::Mutable(_) => vec![], + MapIndex::Immutable(index) => index.immutable_files(), + MapIndex::Mmap(index) => index.immutable_files(), + } + } + + /// Populate all pages in the mmap. + /// Block until all pages are populated. + pub fn populate(&self) -> OperationResult<()> { + match self { + MapIndex::Mutable(_) => {} + MapIndex::Immutable(_) => {} + MapIndex::Mmap(index) => index.populate()?, + } + Ok(()) + } + + /// Drop disk cache. + pub fn clear_cache(&self) -> OperationResult<()> { + match self { + MapIndex::Mutable(index) => index.clear_cache()?, + MapIndex::Immutable(index) => index.clear_cache()?, + MapIndex::Mmap(index) => index.clear_cache()?, + } + Ok(()) + } +} diff --git a/lib/segment/src/index/field_index/map_index/mmap_map_index.rs b/lib/segment/src/index/field_index/map_index/mmap_map_index.rs deleted file mode 100644 index 33c438f5c2..0000000000 --- a/lib/segment/src/index/field_index/map_index/mmap_map_index.rs +++ /dev/null @@ -1,497 +0,0 @@ -use std::borrow::{Borrow, Cow}; -use std::iter; -use std::ops::BitOrAssign; -use std::path::{Path, PathBuf}; - -use ahash::HashMap; -use common::bitvec::{BitSlice, BitSliceExt, BitVec}; -use common::counter::conditioned_counter::ConditionedCounter; -use common::counter::hardware_counter::HardwareCounterCell; -use common::counter::iterator_hw_measurement::HwMeasurementIteratorExt; -use common::fs::{atomic_save_json, clear_disk_cache, read_json}; -use common::mmap::create_and_ensure_length; -use common::persisted_hashmap::{Key, READ_ENTRY_OVERHEAD, UniversalHashMap, serialize_hashmap}; -use common::stored_bitslice::MmapBitSlice; -use common::types::PointOffsetType; -use common::universal_io::{MmapFile, OpenOptions}; -use fs_err as fs; -use itertools::Itertools; -use serde::{Deserialize, Serialize}; - -use super::{IdIter, MapIndexKey}; -use crate::common::Flusher; -use crate::common::operation_error::{OperationError, OperationResult}; -use crate::index::field_index::stored_point_to_values::{StoredPointToValues, ValuesIter}; - -const DELETED_PATH: &str = "deleted.bin"; -const HASHMAP_PATH: &str = "values_to_points.bin"; -const CONFIG_PATH: &str = "mmap_field_index_config.json"; - -/// Mmap-backed immutable map index. -/// -/// On-disk state (`values_to_points.bin`, `deleted.bin`, `point_to_values.*`, -/// `mmap_field_index_config.json`) is written once during [`Self::build`] and -/// not mutated afterwards: `deleted.bin` records only the points whose payload -/// was empty at build time. -/// -/// Runtime deletions live in the in-memory `Storage::deleted` bitvec. They are -/// **not persisted** — [`Self::flusher`] is a no-op and [`Self::remove_point`] -/// only updates the in-memory bitvec. Callers must re-supply the authoritative -/// deletion set (typically `id_tracker.deleted_point_bitslice()`) via the -/// `deleted_points` argument to [`Self::open`] on reload. -pub struct MmapMapIndex { - path: PathBuf, - pub(super) storage: Storage, - deleted_count: usize, - total_key_value_pairs: usize, - is_on_disk: bool, -} - -pub(super) struct Storage { - pub(super) value_to_points: UniversalHashMap, - point_to_values: StoredPointToValues, - /// In-memory deletion bitmap. Reconstructed at load time as the union of - /// the build-time empty-payload bits read from `deleted.bin` and the - /// segment-level deleted bitslice supplied by the id-tracker. Not persisted. - pub(super) deleted: BitVec, -} - -impl Storage { - pub(crate) fn ram_usage_bytes(&self) -> usize { - let Self { - value_to_points: _, - point_to_values, - deleted, - } = self; - - // `value_to_points` is a mmap-backed hashmap with no in-memory state. - point_to_values.ram_usage_bytes() + deleted.capacity().div_ceil(u8::BITS as usize) - } -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct MmapMapIndexConfig { - total_key_value_pairs: usize, -} - -impl MmapMapIndex { - /// Open and load mmap map index from the given path - pub fn open( - path: &Path, - is_on_disk: bool, - deleted_points: &BitSlice, - ) -> OperationResult> { - let hashmap_path = path.join(HASHMAP_PATH); - let deleted_path = path.join(DELETED_PATH); - let config_path = path.join(CONFIG_PATH); - - // If config doesn't exist, assume the index doesn't exist on disk - if !config_path.is_file() { - return Ok(None); - } - - let config: MmapMapIndexConfig = read_json(&config_path)?; - - let do_populate = !is_on_disk; - - let value_to_points = UniversalHashMap::open( - &hashmap_path, - OpenOptions { - writeable: false, - populate: Some(do_populate), - ..OpenOptions::default() - }, - )?; - let point_to_values = StoredPointToValues::open(path, do_populate)?; - - let mut deleted = deleted_points.to_owned(); - - let deleted_payload_mmap = MmapBitSlice::open(&deleted_path, OpenOptions::default())?; - let deleted_payloads_bitslice = deleted_payload_mmap.read_all()?; - - // `deleted` length must match `point_to_values.len()` because it only - // tracks the index's contents. The id-tracker's deleted mask can be - // shorter or longer; if shorter, the missing entries default to live - // (the id-tracker is the source of truth for deletions, and a shorter - // mask just means it doesn't yet know about those higher offsets). - deleted.resize(point_to_values.len(), false); - deleted.bitor_assign(deleted_payloads_bitslice.as_ref()); - - let deleted_count = deleted.count_ones(); - - Ok(Some(Self { - path: path.to_path_buf(), - storage: Storage { - value_to_points, - point_to_values, - deleted, - }, - deleted_count, - total_key_value_pairs: config.total_key_value_pairs, - is_on_disk, - })) - } - - pub fn build( - path: &Path, - point_to_values: Vec::Owned>>, - values_to_points: HashMap<::Owned, Vec>, - is_on_disk: bool, - deleted_points: &BitSlice, - ) -> OperationResult { - fs::create_dir_all(path)?; - - let hashmap_path = path.join(HASHMAP_PATH); - let deleted_path = path.join(DELETED_PATH); - let config_path = path.join(CONFIG_PATH); - - atomic_save_json( - &config_path, - &MmapMapIndexConfig { - total_key_value_pairs: point_to_values.iter().map(|v| v.len()).sum(), - }, - )?; - - serialize_hashmap( - &hashmap_path, - values_to_points - .iter() - .map(|(value, ids)| (value.borrow(), ids.iter().copied())), - )?; - - StoredPointToValues::::from_iter( - path, - point_to_values.iter().enumerate().map(|(idx, values)| { - ( - idx as PointOffsetType, - values.iter().map(|value| value.borrow()), - ) - }), - )?; - - { - let deleted_flags_count = point_to_values.len(); - let _ = create_and_ensure_length( - &deleted_path, - deleted_flags_count - .div_ceil(u8::BITS as usize) - .next_multiple_of(size_of::()), - )?; - - let mut deleted = MmapBitSlice::open(&deleted_path, OpenOptions::default())?; - deleted.set_ascending_bits_batch( - point_to_values - .iter() - .enumerate() - .filter(|(_, values)| values.is_empty()) - .map(|(idx, _)| (idx as u64, true)), - )?; - deleted.flusher()()?; - } - - Self::open(path, is_on_disk, deleted_points)?.ok_or_else(|| { - OperationError::service_error("Failed to open MmapMapIndex after building it") - }) - } - - /// No-op flusher: the on-disk state is build-time only. See the type-level - /// docs on [`MmapMapIndex`] for the deletion durability contract. - pub fn flusher(&self) -> Flusher { - Box::new(|| Ok(())) - } - - pub fn wipe(self) -> OperationResult<()> { - let files = self.files(); - let path = self.path.clone(); - // drop mmap handles before deleting files - drop(self); - for file in files { - fs::remove_file(file)?; - } - let _ = fs::remove_dir(path); - Ok(()) - } - - pub fn files(&self) -> Vec { - let mut files = vec![ - self.path.join(HASHMAP_PATH), - self.path.join(DELETED_PATH), - self.path.join(CONFIG_PATH), - ]; - files.extend(self.storage.point_to_values.files()); - files - } - - pub fn immutable_files(&self) -> Vec { - let mut files = vec![ - self.path.join(HASHMAP_PATH), - self.path.join(DELETED_PATH), - self.path.join(CONFIG_PATH), - ]; - files.extend(self.storage.point_to_values.immutable_files()); - files - } - - /// Marks `idx` as deleted in the in-memory deletion bitvec. - /// - /// Not persisted: on reopen, deletions must be re-supplied via the - /// `deleted_points` argument to [`Self::open`]. - pub fn remove_point(&mut self, idx: PointOffsetType) { - let idx = idx as usize; - if idx < self.storage.deleted.len() && !self.storage.deleted.get_bit(idx).unwrap_or(true) { - self.storage.deleted.set(idx, true); - self.deleted_count += 1; - } - } - - pub fn check_values_any( - &self, - idx: PointOffsetType, - hw_counter: &HardwareCounterCell, - check_fn: impl Fn(&N) -> bool, - ) -> OperationResult { - let hw_counter = self.make_conditioned_counter(hw_counter); - - // Measure self.deleted access. - hw_counter - .payload_index_io_read_counter() - .incr_delta(size_of::()); - - let is_deleted = self - .storage - .deleted - .get_bit(idx as usize) - .is_some_and(|b| b); - - Ok(!is_deleted - && self - .storage - .point_to_values - .check_values_any(idx, |v| check_fn(v), &hw_counter)?) - } - - pub fn get_values( - &self, - idx: PointOffsetType, - hw_counter: &HardwareCounterCell, - ) -> Option> + '_>> { - let hw_counter = self.make_conditioned_counter(hw_counter); - - // We can account cost of reading `bool`, but it will likely be more expensive, than - // actually reading bool itself. - - if self.storage.deleted.get_bit(idx as usize) == Some(false) { - self.storage - .point_to_values - .values_iter(idx, hw_counter) - .ok()? - .map(|iter| Box::new(iter) as Box>>) - } else { - None - } - } - - pub fn for_points_values( - &self, - mut points: impl Iterator, - hw_counter: &HardwareCounterCell, - mut f: impl FnMut(PointOffsetType, ValuesIter<'_, N>), - ) -> OperationResult<()> { - let hw_counter = self.make_conditioned_counter(hw_counter); - - points.try_for_each(|idx| { - if self.storage.deleted.get_bit(idx as usize) != Some(false) { - return Ok(()); - } - if let Some(iter) = self.storage.point_to_values.values_iter(idx, hw_counter)? { - f(idx, iter); - } - Ok(()) - }) - } - - pub fn values_count(&self, idx: PointOffsetType) -> Option { - if self.storage.deleted.get_bit(idx as usize) == Some(false) { - self.storage.point_to_values.get_values_count(idx).ok()? - } else { - None - } - } - - pub fn get_indexed_points(&self) -> usize { - self.storage - .point_to_values - .len() - .saturating_sub(self.deleted_count) - } - - /// Returns the number of key-value pairs in the index. - /// Note that is doesn't count deleted pairs. - pub fn get_values_count(&self) -> usize { - self.total_key_value_pairs - } - - pub fn get_unique_values_count(&self) -> usize { - self.storage.value_to_points.keys_count() - } - - pub fn get_count_for_value( - &self, - value: &N, - hw_counter: &HardwareCounterCell, - ) -> Option { - let hw_counter = self.make_conditioned_counter(hw_counter); - - // Since `value_to_points.get` doesn't actually force read from disk for all values - // we need to only account for the overhead of hashmap lookup - hw_counter - .payload_index_io_read_counter() - .incr_delta(READ_ENTRY_OVERHEAD); - - match self - .storage - .value_to_points - .unbatched_get_values_count(value) - { - Ok(Some(count)) => Some(count), - Ok(None) => None, - Err(err) => { - debug_assert!( - false, - "Error while getting count for value {value:?}: {err:?}", - ); - log::error!("Error while getting count for value {value:?}: {err:?}"); - None - } - } - } - - pub fn get_iterator(&self, value: &N, hw_counter: &HardwareCounterCell) -> IdIter<'_> { - let hw_counter = self.make_conditioned_counter(hw_counter); - - match self.storage.value_to_points.unbatched_get(value) { - Ok(Some(values)) => { - // We're iterating over the whole (mmapped) slice - hw_counter - .payload_index_io_read_counter() - .incr_delta(size_of_val(values.as_slice()) + READ_ENTRY_OVERHEAD); - - Box::new( - values.into_iter().filter(|idx| { - !self.storage.deleted.get_bit(*idx as usize).unwrap_or(false) - }), - ) - } - Ok(None) => { - hw_counter - .payload_index_io_read_counter() - .incr_delta(READ_ENTRY_OVERHEAD); - - Box::new(iter::empty()) - } - Err(err) => { - debug_assert!( - false, - "Error while getting iterator for value {value:?}: {err:?}", - ); - log::error!("Error while getting iterator for value {value:?}: {err:?}"); - Box::new(iter::empty()) - } - } - } - - pub fn for_each_value(&self, f: impl FnMut(&N) -> OperationResult<()>) -> OperationResult<()> { - self.storage.value_to_points.for_each_key(f) - } - - pub fn for_each_count_per_value( - &self, - deferred_internal_id: Option, - mut f: impl FnMut(&N, usize) -> OperationResult<()>, - ) -> OperationResult<()> { - self.storage.value_to_points.for_each_entry(|k, v| { - let count = v - .iter() - .filter(|&&idx| { - !self.storage.deleted.get_bit(idx as usize).unwrap_or(true) - - // TODO(deferred): Maybe we can improve this filter and use take_while instead. For this we - // need to make sure that `v` is always sorted which we _can_ enforce when finalizing the index. - && deferred_internal_id.is_none_or(|deferred| idx < deferred) - }) - .unique() - .count(); - f(k, count) - }) - } - - pub fn for_each_value_map( - &self, - hw_counter: &HardwareCounterCell, - mut f: impl FnMut(&N, &mut dyn Iterator) -> OperationResult<()>, - ) -> OperationResult<()> { - let hw_counter = self.make_conditioned_counter(hw_counter); - let deleted = &self.storage.deleted; - - self.storage.value_to_points.for_each_entry(|k, v| { - hw_counter - .payload_index_io_read_counter() - .incr_delta(k.write_bytes()); - - let mut iter = v - .iter() - .copied() - .filter(|idx| !deleted.get_bit(*idx as usize).unwrap_or(true)) - .measure_hw_with_acc( - hw_counter.new_accumulator(), - size_of::(), - |i| i.payload_index_io_read_counter(), - ); - - f(k, &mut iter) - }) - } - - fn make_conditioned_counter<'a>( - &self, - hw_counter: &'a HardwareCounterCell, - ) -> ConditionedCounter<'a> { - ConditionedCounter::new(self.is_on_disk, hw_counter) - } - - pub fn is_on_disk(&self) -> bool { - self.is_on_disk - } - - /// Populate all pages in the mmap. - /// Block until all pages are populated. - pub fn populate(&self) -> OperationResult<()> { - self.storage.value_to_points.populate()?; - self.storage.point_to_values.populate()?; - Ok(()) - } - - /// Drop disk cache. - pub fn clear_cache(&self) -> OperationResult<()> { - let Self { - path, - storage, - deleted_count: _, - total_key_value_pairs: _, - is_on_disk: _, - } = self; - let Storage { - value_to_points, - point_to_values, - deleted: _, - } = storage; - value_to_points.clear_ram_cache()?; - clear_disk_cache(&path.join(DELETED_PATH))?; - point_to_values.clear_cache()?; - Ok(()) - } - - pub(crate) fn ram_usage_bytes(&self) -> usize { - self.storage.ram_usage_bytes() - } -} diff --git a/lib/segment/src/index/field_index/map_index/mmap_map_index/lifecycle.rs b/lib/segment/src/index/field_index/map_index/mmap_map_index/lifecycle.rs new file mode 100644 index 0000000000..a65272c06a --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/mmap_map_index/lifecycle.rs @@ -0,0 +1,222 @@ +use std::borrow::Borrow; +use std::ops::BitOrAssign; +use std::path::{Path, PathBuf}; + +use ahash::HashMap; +use common::bitvec::{BitSlice, BitSliceExt}; +use common::fs::{atomic_save_json, clear_disk_cache, read_json}; +use common::mmap::create_and_ensure_length; +use common::persisted_hashmap::{Key, UniversalHashMap, serialize_hashmap}; +use common::stored_bitslice::MmapBitSlice; +use common::types::PointOffsetType; +use common::universal_io::{MmapFile, OpenOptions}; +use fs_err as fs; + +use super::super::MapIndexKey; +use super::{CONFIG_PATH, DELETED_PATH, HASHMAP_PATH, MmapMapIndex, MmapMapIndexConfig, Storage}; +use crate::common::Flusher; +use crate::common::operation_error::{OperationError, OperationResult}; +use crate::index::field_index::stored_point_to_values::StoredPointToValues; + +impl MmapMapIndex { + /// Open and load mmap map index from the given path + pub fn open( + path: &Path, + is_on_disk: bool, + deleted_points: &BitSlice, + ) -> OperationResult> { + let hashmap_path = path.join(HASHMAP_PATH); + let deleted_path = path.join(DELETED_PATH); + let config_path = path.join(CONFIG_PATH); + + // If config doesn't exist, assume the index doesn't exist on disk + if !config_path.is_file() { + return Ok(None); + } + + let config: MmapMapIndexConfig = read_json(&config_path)?; + + let do_populate = !is_on_disk; + + let value_to_points = UniversalHashMap::open( + &hashmap_path, + OpenOptions { + writeable: false, + populate: Some(do_populate), + ..OpenOptions::default() + }, + )?; + let point_to_values = StoredPointToValues::open(path, do_populate)?; + + let mut deleted = deleted_points.to_owned(); + + let deleted_payload_mmap = MmapBitSlice::open(&deleted_path, OpenOptions::default())?; + let deleted_payloads_bitslice = deleted_payload_mmap.read_all()?; + + // `deleted` length must match `point_to_values.len()` because it only + // tracks the index's contents. The id-tracker's deleted mask can be + // shorter or longer; if shorter, the missing entries default to live + // (the id-tracker is the source of truth for deletions, and a shorter + // mask just means it doesn't yet know about those higher offsets). + deleted.resize(point_to_values.len(), false); + deleted.bitor_assign(deleted_payloads_bitslice.as_ref()); + + let deleted_count = deleted.count_ones(); + + Ok(Some(Self { + path: path.to_path_buf(), + storage: Storage { + value_to_points, + point_to_values, + deleted, + }, + deleted_count, + total_key_value_pairs: config.total_key_value_pairs, + is_on_disk, + })) + } + + pub fn build( + path: &Path, + point_to_values: Vec::Owned>>, + values_to_points: HashMap<::Owned, Vec>, + is_on_disk: bool, + deleted_points: &BitSlice, + ) -> OperationResult { + fs::create_dir_all(path)?; + + let hashmap_path = path.join(HASHMAP_PATH); + let deleted_path = path.join(DELETED_PATH); + let config_path = path.join(CONFIG_PATH); + + atomic_save_json( + &config_path, + &MmapMapIndexConfig { + total_key_value_pairs: point_to_values.iter().map(|v| v.len()).sum(), + }, + )?; + + serialize_hashmap( + &hashmap_path, + values_to_points + .iter() + .map(|(value, ids)| (value.borrow(), ids.iter().copied())), + )?; + + StoredPointToValues::::from_iter( + path, + point_to_values.iter().enumerate().map(|(idx, values)| { + ( + idx as PointOffsetType, + values.iter().map(|value| value.borrow()), + ) + }), + )?; + + { + let deleted_flags_count = point_to_values.len(); + let _ = create_and_ensure_length( + &deleted_path, + deleted_flags_count + .div_ceil(u8::BITS as usize) + .next_multiple_of(size_of::()), + )?; + + let mut deleted = MmapBitSlice::open(&deleted_path, OpenOptions::default())?; + deleted.set_ascending_bits_batch( + point_to_values + .iter() + .enumerate() + .filter(|(_, values)| values.is_empty()) + .map(|(idx, _)| (idx as u64, true)), + )?; + deleted.flusher()()?; + } + + Self::open(path, is_on_disk, deleted_points)?.ok_or_else(|| { + OperationError::service_error("Failed to open MmapMapIndex after building it") + }) + } + + /// No-op flusher: the on-disk state is build-time only. See the type-level + /// docs on [`MmapMapIndex`] for the deletion durability contract. + pub fn flusher(&self) -> Flusher { + Box::new(|| Ok(())) + } + + pub fn wipe(self) -> OperationResult<()> { + let files = self.files(); + let path = self.path.clone(); + // drop mmap handles before deleting files + drop(self); + for file in files { + fs::remove_file(file)?; + } + let _ = fs::remove_dir(path); + Ok(()) + } + + pub fn files(&self) -> Vec { + let mut files = vec![ + self.path.join(HASHMAP_PATH), + self.path.join(DELETED_PATH), + self.path.join(CONFIG_PATH), + ]; + files.extend(self.storage.point_to_values.files()); + files + } + + pub fn immutable_files(&self) -> Vec { + let mut files = vec![ + self.path.join(HASHMAP_PATH), + self.path.join(DELETED_PATH), + self.path.join(CONFIG_PATH), + ]; + files.extend(self.storage.point_to_values.immutable_files()); + files + } + + /// Marks `idx` as deleted in the in-memory deletion bitvec. + /// + /// Not persisted: on reopen, deletions must be re-supplied via the + /// `deleted_points` argument to [`Self::open`]. + pub fn remove_point(&mut self, idx: PointOffsetType) { + let idx = idx as usize; + if idx < self.storage.deleted.len() && !self.storage.deleted.get_bit(idx).unwrap_or(true) { + self.storage.deleted.set(idx, true); + self.deleted_count += 1; + } + } + + /// Populate all pages in the mmap. + /// Block until all pages are populated. + pub fn populate(&self) -> OperationResult<()> { + self.storage.value_to_points.populate()?; + self.storage.point_to_values.populate()?; + Ok(()) + } + + /// Drop disk cache. + pub fn clear_cache(&self) -> OperationResult<()> { + let Self { + path, + storage, + deleted_count: _, + total_key_value_pairs: _, + is_on_disk: _, + } = self; + let Storage { + value_to_points, + point_to_values, + deleted: _, + } = storage; + value_to_points.clear_ram_cache()?; + clear_disk_cache(&path.join(DELETED_PATH))?; + point_to_values.clear_cache()?; + Ok(()) + } + + pub(crate) fn ram_usage_bytes(&self) -> usize { + self.storage.ram_usage_bytes() + } +} diff --git a/lib/segment/src/index/field_index/map_index/mmap_map_index/mod.rs b/lib/segment/src/index/field_index/map_index/mmap_map_index/mod.rs new file mode 100644 index 0000000000..a5aebc9468 --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/mmap_map_index/mod.rs @@ -0,0 +1,64 @@ +use std::path::PathBuf; + +use common::bitvec::BitVec; +use common::persisted_hashmap::{Key, UniversalHashMap}; +use common::types::PointOffsetType; +use common::universal_io::MmapFile; +use serde::{Deserialize, Serialize}; + +use super::MapIndexKey; +use crate::index::field_index::stored_point_to_values::StoredPointToValues; + +mod lifecycle; +mod read_ops; + +pub(super) const DELETED_PATH: &str = "deleted.bin"; +pub(super) const HASHMAP_PATH: &str = "values_to_points.bin"; +pub(super) const CONFIG_PATH: &str = "mmap_field_index_config.json"; + +/// Mmap-backed immutable map index. +/// +/// On-disk state (`values_to_points.bin`, `deleted.bin`, `point_to_values.*`, +/// `mmap_field_index_config.json`) is written once during [`Self::build`] and +/// not mutated afterwards: `deleted.bin` records only the points whose payload +/// was empty at build time. +/// +/// Runtime deletions live in the in-memory `Storage::deleted` bitvec. They are +/// **not persisted** — [`Self::flusher`] is a no-op and [`Self::remove_point`] +/// only updates the in-memory bitvec. Callers must re-supply the authoritative +/// deletion set (typically `id_tracker.deleted_point_bitslice()`) via the +/// `deleted_points` argument to [`Self::open`] on reload. +pub struct MmapMapIndex { + pub(super) path: PathBuf, + pub(super) storage: Storage, + pub(super) deleted_count: usize, + pub(super) total_key_value_pairs: usize, + pub(super) is_on_disk: bool, +} + +pub(super) struct Storage { + pub(super) value_to_points: UniversalHashMap, + pub(super) point_to_values: StoredPointToValues, + /// In-memory deletion bitmap. Reconstructed at load time as the union of + /// the build-time empty-payload bits read from `deleted.bin` and the + /// segment-level deleted bitslice supplied by the id-tracker. Not persisted. + pub(super) deleted: BitVec, +} + +impl Storage { + pub(super) fn ram_usage_bytes(&self) -> usize { + let Self { + value_to_points: _, + point_to_values, + deleted, + } = self; + + // `value_to_points` is a mmap-backed hashmap with no in-memory state. + point_to_values.ram_usage_bytes() + deleted.capacity().div_ceil(u8::BITS as usize) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(super) struct MmapMapIndexConfig { + pub(super) total_key_value_pairs: usize, +} diff --git a/lib/segment/src/index/field_index/map_index/mmap_map_index/read_ops.rs b/lib/segment/src/index/field_index/map_index/mmap_map_index/read_ops.rs new file mode 100644 index 0000000000..4dea975f3f --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/mmap_map_index/read_ops.rs @@ -0,0 +1,253 @@ +use std::borrow::Cow; +use std::iter; + +use common::bitvec::BitSliceExt; +use common::counter::conditioned_counter::ConditionedCounter; +use common::counter::hardware_counter::HardwareCounterCell; +use common::counter::iterator_hw_measurement::HwMeasurementIteratorExt; +use common::persisted_hashmap::{Key, READ_ENTRY_OVERHEAD}; +use common::types::PointOffsetType; +use itertools::Itertools; + +use super::super::read_ops::MapIndexRead; +use super::super::{IdIter, MapIndexKey}; +use super::MmapMapIndex; +use crate::common::operation_error::OperationResult; +use crate::index::field_index::stored_point_to_values::ValuesIter; +use crate::index::payload_config::StorageType; + +impl MapIndexRead for MmapMapIndex { + fn check_values_any( + &self, + idx: PointOffsetType, + hw_counter: &HardwareCounterCell, + check_fn: impl Fn(&N) -> bool, + ) -> bool { + let hw_counter = self.make_conditioned_counter(hw_counter); + + // Measure self.deleted access. + hw_counter + .payload_index_io_read_counter() + .incr_delta(size_of::()); + + let is_deleted = self + .storage + .deleted + .get_bit(idx as usize) + .is_some_and(|b| b); + + if is_deleted { + return false; + } + + // FIXME: don't silently ignore errors. Log error? Update ConditionCheckerFn? + self.storage + .point_to_values + .check_values_any(idx, |v| check_fn(v), &hw_counter) + .unwrap_or(false) + } + + fn get_values<'a>( + &'a self, + idx: PointOffsetType, + hw_counter: &HardwareCounterCell, + ) -> Option> + 'a> + where + N: 'a, + { + let hw_counter = self.make_conditioned_counter(hw_counter); + + // We can account cost of reading `bool`, but it will likely be more expensive, than + // actually reading bool itself. + + if self.storage.deleted.get_bit(idx as usize) == Some(false) { + self.storage + .point_to_values + .values_iter(idx, hw_counter) + .ok()? + .map(|iter| Box::new(iter) as Box>>) + } else { + None + } + } + + fn values_count(&self, idx: PointOffsetType) -> Option { + if self.storage.deleted.get_bit(idx as usize) == Some(false) { + self.storage.point_to_values.get_values_count(idx).ok()? + } else { + None + } + } + + fn get_indexed_points(&self) -> usize { + self.storage + .point_to_values + .len() + .saturating_sub(self.deleted_count) + } + + /// Returns the number of key-value pairs in the index. + /// Note that is doesn't count deleted pairs. + fn get_values_count(&self) -> usize { + self.total_key_value_pairs + } + + fn get_unique_values_count(&self) -> usize { + self.storage.value_to_points.keys_count() + } + + fn get_count_for_value(&self, value: &N, hw_counter: &HardwareCounterCell) -> Option { + let hw_counter = self.make_conditioned_counter(hw_counter); + + // Since `value_to_points.get` doesn't actually force read from disk for all values + // we need to only account for the overhead of hashmap lookup + hw_counter + .payload_index_io_read_counter() + .incr_delta(READ_ENTRY_OVERHEAD); + + match self + .storage + .value_to_points + .unbatched_get_values_count(value) + { + Ok(Some(count)) => Some(count), + Ok(None) => None, + Err(err) => { + debug_assert!( + false, + "Error while getting count for value {value:?}: {err:?}", + ); + log::error!("Error while getting count for value {value:?}: {err:?}"); + None + } + } + } + + fn get_iterator(&self, value: &N, hw_counter: &HardwareCounterCell) -> IdIter<'_> { + let hw_counter = self.make_conditioned_counter(hw_counter); + + match self.storage.value_to_points.unbatched_get(value) { + Ok(Some(values)) => { + // We're iterating over the whole (mmapped) slice + hw_counter + .payload_index_io_read_counter() + .incr_delta(size_of_val(values.as_slice()) + READ_ENTRY_OVERHEAD); + + Box::new( + values.into_iter().filter(|idx| { + !self.storage.deleted.get_bit(*idx as usize).unwrap_or(false) + }), + ) + } + Ok(None) => { + hw_counter + .payload_index_io_read_counter() + .incr_delta(READ_ENTRY_OVERHEAD); + + Box::new(iter::empty()) + } + Err(err) => { + debug_assert!( + false, + "Error while getting iterator for value {value:?}: {err:?}", + ); + log::error!("Error while getting iterator for value {value:?}: {err:?}"); + Box::new(iter::empty()) + } + } + } + + fn for_each_value(&self, f: impl FnMut(&N) -> OperationResult<()>) -> OperationResult<()> { + self.storage.value_to_points.for_each_key(f) + } + + fn for_each_count_per_value( + &self, + deferred_internal_id: Option, + mut f: impl FnMut(&N, usize) -> OperationResult<()>, + ) -> OperationResult<()> { + self.storage.value_to_points.for_each_entry(|k, v| { + let count = v + .iter() + .filter(|&&idx| { + !self.storage.deleted.get_bit(idx as usize).unwrap_or(true) + + // TODO(deferred): Maybe we can improve this filter and use take_while instead. For this we + // need to make sure that `v` is always sorted which we _can_ enforce when finalizing the index. + && deferred_internal_id.is_none_or(|deferred| idx < deferred) + }) + .unique() + .count(); + f(k, count) + }) + } + + fn for_each_value_map( + &self, + hw_counter: &HardwareCounterCell, + mut f: impl FnMut(&N, &mut dyn Iterator) -> OperationResult<()>, + ) -> OperationResult<()> { + let hw_counter = self.make_conditioned_counter(hw_counter); + let deleted = &self.storage.deleted; + + self.storage.value_to_points.for_each_entry(|k, v| { + hw_counter + .payload_index_io_read_counter() + .incr_delta(k.write_bytes()); + + let mut iter = v + .iter() + .copied() + .filter(|idx| !deleted.get_bit(*idx as usize).unwrap_or(true)) + .measure_hw_with_acc( + hw_counter.new_accumulator(), + size_of::(), + |i| i.payload_index_io_read_counter(), + ); + + f(k, &mut iter) + }) + } + + fn storage_type(&self) -> StorageType { + StorageType::Mmap { + is_on_disk: self.is_on_disk, + } + } + + fn ram_usage_bytes(&self) -> usize { + self.storage.ram_usage_bytes() + } +} + +impl MmapMapIndex { + pub fn for_points_values( + &self, + mut points: impl Iterator, + hw_counter: &HardwareCounterCell, + mut f: impl FnMut(PointOffsetType, ValuesIter<'_, N>), + ) -> OperationResult<()> { + let hw_counter = self.make_conditioned_counter(hw_counter); + + points.try_for_each(|idx| { + if self.storage.deleted.get_bit(idx as usize) != Some(false) { + return Ok(()); + } + if let Some(iter) = self.storage.point_to_values.values_iter(idx, hw_counter)? { + f(idx, iter); + } + Ok(()) + }) + } + + pub(super) fn make_conditioned_counter<'a>( + &self, + hw_counter: &'a HardwareCounterCell, + ) -> ConditionedCounter<'a> { + ConditionedCounter::new(self.is_on_disk, hw_counter) + } + + pub fn is_on_disk(&self) -> bool { + self.is_on_disk + } +} diff --git a/lib/segment/src/index/field_index/map_index/mod.rs b/lib/segment/src/index/field_index/map_index/mod.rs index 02ba16f788..e7823cd62e 100644 --- a/lib/segment/src/index/field_index/map_index/mod.rs +++ b/lib/segment/src/index/field_index/map_index/mod.rs @@ -1,34 +1,23 @@ -use std::borrow::{Borrow, Cow}; -use std::hash::{BuildHasher, Hash}; -use std::path::{Path, PathBuf}; - -use common::bitvec::BitSlice; -use common::counter::hardware_counter::HardwareCounterCell; use common::types::PointOffsetType; use gridstore::Blob; -use indexmap::IndexSet; pub use self::builders::{MapIndexBuilder, MapIndexGridstoreBuilder, MapIndexMmapBuilder}; use self::immutable_map_index::ImmutableMapIndex; pub use self::key::MapIndexKey; use self::mmap_map_index::MmapMapIndex; use self::mutable_map_index::MutableMapIndex; -use crate::common::Flusher; -use crate::common::operation_error::OperationResult; -use crate::index::field_index::CardinalityEstimation; -use crate::index::field_index::stat_tools::number_of_selected_points; -use crate::index::payload_config::{IndexMutability, StorageType}; -use crate::telemetry::PayloadIndexTelemetry; mod builders; mod facet_index_impl; pub mod immutable_map_index; pub mod key; +mod lifecycle; pub mod mmap_map_index; pub mod mutable_map_index; -mod payload_index_impl_int; -mod payload_index_impl_str; -mod payload_index_impl_uuid; +mod payload_index_impl; +mod read_ops; +#[cfg(test)] +mod tests; mod value_indexer_impl; /// Block size in Gridstore for keyword map index. @@ -48,888 +37,3 @@ where Immutable(ImmutableMapIndex), Mmap(Box>), } - -impl MapIndex -where - Vec<::Owned>: Blob + Send + Sync, -{ - /// Load immutable mmap based index, either in RAM or on disk - pub fn new_mmap( - path: &Path, - is_on_disk: bool, - deleted_points: &BitSlice, - ) -> OperationResult> { - // Low-memory mode downgrades the in-RAM `Immutable` wrapper to the - // pure-mmap `Storage` variant at load time. Files are shared between - // variants; the persisted `is_on_disk` flag in `mmap_index` is - // untouched. - let effective_is_on_disk = - is_on_disk || common::low_memory::low_memory_mode().prefer_disk(); - - let Some(mmap_index) = MmapMapIndex::open(path, effective_is_on_disk, deleted_points)? - else { - return Ok(None); - }; - - let index = if effective_is_on_disk { - MapIndex::Mmap(Box::new(mmap_index)) - } else { - // Load into RAM, use mmap as backing storage - MapIndex::Immutable(ImmutableMapIndex::open_mmap(mmap_index)?) - }; - Ok(Some(index)) - } - - pub fn new_gridstore(dir: PathBuf, create_if_missing: bool) -> OperationResult> { - let index = MutableMapIndex::open_gridstore(dir, create_if_missing)?; - Ok(index.map(MapIndex::Mutable)) - } - - pub fn builder_mmap( - path: &Path, - is_on_disk: bool, - deleted_points: &BitSlice, - ) -> MapIndexMmapBuilder { - MapIndexMmapBuilder { - path: path.to_owned(), - point_to_values: Default::default(), - values_to_points: Default::default(), - is_on_disk, - deleted_points: deleted_points.to_owned(), - } - } - - pub fn builder_gridstore(dir: PathBuf) -> MapIndexGridstoreBuilder { - MapIndexGridstoreBuilder::new(dir) - } - - pub fn check_values_any( - &self, - idx: PointOffsetType, - hw_counter: &HardwareCounterCell, - check_fn: impl Fn(&N) -> bool, - ) -> bool { - match self { - MapIndex::Mutable(index) => index.check_values_any(idx, check_fn), - MapIndex::Immutable(index) => index.check_values_any(idx, check_fn), - // FIXME: don't silently ignore errors. Log error? Update ConditionCheckerFn? - MapIndex::Mmap(index) => index - .check_values_any(idx, hw_counter, check_fn) - .unwrap_or(false), - } - } - - pub fn get_values( - &self, - idx: PointOffsetType, - hw_counter: &HardwareCounterCell, - ) -> Option> + '_>> { - match self { - MapIndex::Mutable(index) => Some(Box::new(index.get_values(idx)?)), - MapIndex::Immutable(index) => Some(Box::new(index.get_values(idx)?)), - MapIndex::Mmap(index) => Some(Box::new(index.get_values(idx, hw_counter)?)), - } - } - - pub fn values_count(&self, idx: PointOffsetType) -> usize { - match self { - MapIndex::Mutable(index) => index.values_count(idx).unwrap_or_default(), - MapIndex::Immutable(index) => index.values_count(idx).unwrap_or_default(), - MapIndex::Mmap(index) => index.values_count(idx).unwrap_or_default(), - } - } - - pub(crate) fn get_indexed_points(&self) -> usize { - match self { - MapIndex::Mutable(index) => index.get_indexed_points(), - MapIndex::Immutable(index) => index.get_indexed_points(), - MapIndex::Mmap(index) => index.get_indexed_points(), - } - } - - fn get_values_count(&self) -> usize { - match self { - MapIndex::Mutable(index) => index.get_values_count(), - MapIndex::Immutable(index) => index.get_values_count(), - MapIndex::Mmap(index) => index.get_values_count(), - } - } - - pub fn get_unique_values_count(&self) -> usize { - match self { - MapIndex::Mutable(index) => index.get_unique_values_count(), - MapIndex::Immutable(index) => index.get_unique_values_count(), - MapIndex::Mmap(index) => index.get_unique_values_count(), - } - } - - pub(crate) fn get_count_for_value( - &self, - value: &N, - hw_counter: &HardwareCounterCell, - ) -> Option { - match self { - MapIndex::Mutable(index) => index.get_count_for_value(value), - MapIndex::Immutable(index) => index.get_count_for_value(value), - MapIndex::Mmap(index) => index.get_count_for_value(value, hw_counter), - } - } - - pub(crate) fn get_iterator(&self, value: &N, hw_counter: &HardwareCounterCell) -> IdIter<'_> { - match self { - MapIndex::Mutable(index) => index.get_iterator(value), - MapIndex::Immutable(index) => index.get_iterator(value), - MapIndex::Mmap(index) => index.get_iterator(value, hw_counter), - } - } - - pub fn for_each_value(&self, f: impl FnMut(&N) -> OperationResult<()>) -> OperationResult<()> { - match self { - MapIndex::Mutable(index) => index.for_each_value(f), - MapIndex::Immutable(index) => index.for_each_value(f), - MapIndex::Mmap(index) => index.for_each_value(f), - } - } - - pub fn for_each_count_per_value( - &self, - deferred_internal_id: Option, - f: impl FnMut(&N, usize) -> OperationResult<()>, - ) -> OperationResult<()> { - match self { - MapIndex::Mutable(index) => index.for_each_count_per_value(deferred_internal_id, f), - - // Two reasons we don't implement deferred filtering here: - // - We don't have both deferred points and an immutable index. - // - It is not trivial (nor performant) to implement correct filtering for this index variant as - // it doesn't work well in combination with the way it handles deletions. - MapIndex::Immutable(index) => { - debug_assert!(deferred_internal_id.is_none()); - index.for_each_count_per_value(f) - } - - MapIndex::Mmap(index) => index.for_each_count_per_value(deferred_internal_id, f), - } - } - - pub fn for_each_value_map( - &self, - hw_cell: &HardwareCounterCell, - f: impl FnMut(&N, &mut dyn Iterator) -> OperationResult<()>, - ) -> OperationResult<()> { - match self { - MapIndex::Mutable(index) => index.for_each_value_map(f), - MapIndex::Immutable(index) => index.for_each_value_map(f), - MapIndex::Mmap(index) => index.for_each_value_map(hw_cell, f), - } - } - - pub(crate) fn flusher(&self) -> Flusher { - match self { - MapIndex::Mutable(index) => index.flusher(), - MapIndex::Immutable(index) => index.flusher(), - MapIndex::Mmap(index) => index.flusher(), - } - } - - pub(crate) fn match_cardinality( - &self, - value: &N, - hw_counter: &HardwareCounterCell, - ) -> CardinalityEstimation { - let values_count = self.get_count_for_value(value, hw_counter).unwrap_or(0); - - CardinalityEstimation::exact(values_count) - } - - pub fn get_telemetry_data(&self) -> PayloadIndexTelemetry { - PayloadIndexTelemetry { - field_name: None, - points_count: self.get_indexed_points(), - points_values_count: self.get_values_count(), - histogram_bucket_size: None, - index_type: match self { - MapIndex::Mutable(_) => "mutable_map", - MapIndex::Immutable(_) => "immutable_map", - MapIndex::Mmap(_) => "mmap_map", - }, - } - } - - pub fn values_is_empty(&self, idx: PointOffsetType) -> bool { - self.values_count(idx) == 0 - } - - pub(crate) fn wipe(self) -> OperationResult<()> { - match self { - MapIndex::Mutable(index) => index.wipe(), - MapIndex::Immutable(index) => index.wipe(), - MapIndex::Mmap(index) => index.wipe(), - } - } - - pub(crate) fn remove_point(&mut self, id: PointOffsetType) -> OperationResult<()> { - match self { - MapIndex::Mutable(index) => index.remove_point(id), - MapIndex::Immutable(index) => index.remove_point(id), - MapIndex::Mmap(index) => { - index.remove_point(id); - Ok(()) - } - } - } - - pub(crate) fn files(&self) -> Vec { - match self { - MapIndex::Mutable(index) => index.files(), - MapIndex::Immutable(index) => index.files(), - MapIndex::Mmap(index) => index.files(), - } - } - - pub(crate) fn immutable_files(&self) -> Vec { - match self { - MapIndex::Mutable(_) => vec![], - MapIndex::Immutable(index) => index.immutable_files(), - MapIndex::Mmap(index) => index.immutable_files(), - } - } - - /// Estimates cardinality for `except` clause - /// - /// # Arguments - /// - /// * 'excluded' - values, which are not considered as matching - /// - /// # Returns - /// - /// * `CardinalityEstimation` - estimation of cardinality - pub(crate) fn except_cardinality<'a>( - &'a self, - excluded: impl Iterator, - hw_counter: &HardwareCounterCell, - ) -> CardinalityEstimation { - // Minimal case: we exclude as many points as possible. - // In this case, excluded points do not have any other values except excluded ones. - // So the first step - we estimate how many other points is needed to fit unused values. - - // Example: - // Values: 20, 20 - // Unique values: 5 - // Total points: 100 - // Total values: 110 - // total_excluded_value_count = 40 - // non_excluded_values_count = 110 - 40 = 70 - // max_values_per_point = 5 - 2 = 3 - // min_not_excluded_by_values = 70 / 3 = 24 - // min = max(24, 100 - 40) = 60 - // exp = ... - // max = min(20, 70) = 20 - - // Values: 60, 60 - // Unique values: 5 - // Total points: 100 - // Total values: 200 - // total_excluded_value_count = 120 - // non_excluded_values_count = 200 - 120 = 80 - // max_values_per_point = 5 - 2 = 3 - // min_not_excluded_by_values = 80 / 3 = 27 - // min = max(27, 100 - 120) = 27 - // exp = ... - // max = min(60, 80) = 60 - - // Values: 60, 60, 60 - // Unique values: 5 - // Total points: 100 - // Total values: 200 - // total_excluded_value_count = 180 - // non_excluded_values_count = 200 - 180 = 20 - // max_values_per_point = 5 - 3 = 2 - // min_not_excluded_by_values = 20 / 2 = 10 - // min = max(10, 100 - 180) = 10 - // exp = ... - // max = min(60, 20) = 20 - - let excluded_value_counts: Vec<_> = excluded - .map(|val| { - self.get_count_for_value(val.borrow(), hw_counter) - .unwrap_or(0) - }) - .collect(); - let total_excluded_value_count: usize = excluded_value_counts.iter().sum(); - - debug_assert!(total_excluded_value_count <= self.get_values_count()); - - let non_excluded_values_count = self - .get_values_count() - .saturating_sub(total_excluded_value_count); - let max_values_per_point = self - .get_unique_values_count() - .saturating_sub(excluded_value_counts.len()); - - if max_values_per_point == 0 { - debug_assert_eq!(non_excluded_values_count, 0); - return CardinalityEstimation::exact(0); - } - - let min_not_excluded_by_values = non_excluded_values_count.div_ceil(max_values_per_point); - - let min = min_not_excluded_by_values.max( - self.get_indexed_points() - .saturating_sub(total_excluded_value_count), - ); - - let max_excluded_value_count = excluded_value_counts.iter().max().copied().unwrap_or(0); - - let max = self - .get_indexed_points() - .saturating_sub(max_excluded_value_count) - .min(non_excluded_values_count); - - let exp = number_of_selected_points(self.get_indexed_points(), non_excluded_values_count) - .max(min) - .min(max); - - CardinalityEstimation { - primary_clauses: vec![], - min, - exp, - max, - } - } - - pub(crate) fn except_set<'a, K, A>( - &'a self, - excluded: &'a IndexSet, - hw_counter: &'a HardwareCounterCell, - ) -> OperationResult + 'a>> - where - A: BuildHasher, - K: Borrow + Hash + Eq, - { - let mut points = IndexSet::new(); - self.for_each_value(|key| { - if !excluded.contains(key.borrow()) { - self.get_iterator(key.borrow(), hw_counter).for_each(|p| { - points.insert(p); - }); - } - Ok(()) - })?; - Ok(Box::new(points.into_iter())) - } - - /// Approximate RAM usage in bytes for in-memory structures. - pub fn ram_usage_bytes(&self) -> usize { - match self { - MapIndex::Mutable(index) => index.ram_usage_bytes(), - MapIndex::Immutable(index) => index.ram_usage_bytes(), - MapIndex::Mmap(index) => index.ram_usage_bytes(), - } - } - - pub fn is_on_disk(&self) -> bool { - match self { - MapIndex::Mutable(_) => false, - MapIndex::Immutable(_) => false, - MapIndex::Mmap(index) => index.is_on_disk(), - } - } - - /// Populate all pages in the mmap. - /// Block until all pages are populated. - pub fn populate(&self) -> OperationResult<()> { - match self { - MapIndex::Mutable(_) => {} - MapIndex::Immutable(_) => {} - MapIndex::Mmap(index) => index.populate()?, - } - Ok(()) - } - - /// Drop disk cache. - pub fn clear_cache(&self) -> OperationResult<()> { - match self { - MapIndex::Mutable(index) => index.clear_cache()?, - MapIndex::Immutable(index) => index.clear_cache()?, - MapIndex::Mmap(index) => index.clear_cache()?, - } - Ok(()) - } - - pub fn get_mutability_type(&self) -> IndexMutability { - match self { - Self::Mutable(_) => IndexMutability::Mutable, - Self::Immutable(_) => IndexMutability::Immutable, - Self::Mmap(_) => IndexMutability::Immutable, - } - } - - pub fn get_storage_type(&self) -> StorageType { - match self { - Self::Mutable(index) => index.storage_type(), - Self::Immutable(index) => index.storage_type(), - Self::Mmap(index) => StorageType::Mmap { - is_on_disk: index.is_on_disk(), - }, - } - } -} - -#[cfg(test)] -mod tests { - use std::collections::HashSet; - use std::hint::black_box; - use std::path::Path; - - use common::bitvec::BitVec; - use common::counter::hardware_counter::HardwareCounterCell; - use common::types::PointOffsetType; - use ecow::EcoString; - use gridstore::Blob; - use rstest::rstest; - use serde_json::Value; - use tempfile::Builder; - - use super::*; - use crate::index::field_index::{ - CardinalityEstimation, FieldIndexBuilderTrait, PayloadFieldIndex, PayloadFieldIndexRead, - ValueIndexer, - }; - use crate::types::{IntPayloadType, PayloadKeyType, UuidIntType}; - - /// Generous default size for the deleted-points bitslice used in tests. - /// - /// Must be larger than the stored mmap deletion bitslice for any test in - /// this file (which is sized to the highest point id, rounded up to a - /// `usize` boundary). 4096 bits comfortably covers all current tests. - const TEST_DELETED_BITS: usize = 4096; - - /// All-zero deletion bitslice for tests that don't care about deletions. - fn empty_deleted() -> BitVec { - BitVec::repeat(false, TEST_DELETED_BITS) - } - - /// Deletion bitslice with specific points marked as deleted. - fn deleted_with(points: &[PointOffsetType]) -> BitVec { - let mut v = empty_deleted(); - for &p in points { - v.set(p as usize, true); - } - v - } - - #[derive(Clone, Copy, PartialEq, Debug)] - enum IndexType { - MutableGridstore, - Mmap, - RamMmap, - } - - fn save_map_index( - data: &[Vec<::Owned>], - path: &Path, - index_type: IndexType, - into_value: impl Fn(&::Owned) -> Value, - ) where - N: MapIndexKey + ?Sized, - Vec<::Owned>: Blob + Send + Sync, - MapIndex: PayloadFieldIndex + ValueIndexer, - as ValueIndexer>::ValueType: Into<::Owned>, - { - let hw_counter = HardwareCounterCell::new(); - - match index_type { - IndexType::MutableGridstore => { - let mut builder = MapIndex::::builder_gridstore(path.to_path_buf()); - builder.init().unwrap(); - for (idx, values) in data.iter().enumerate() { - let values: Vec = values.iter().map(&into_value).collect(); - let values: Vec<_> = values.iter().collect(); - builder - .add_point(idx as PointOffsetType, &values, &hw_counter) - .unwrap(); - } - builder.finalize().unwrap(); - } - IndexType::Mmap | IndexType::RamMmap => { - let mut builder = MapIndex::::builder_mmap(path, false, &empty_deleted()); - builder.init().unwrap(); - for (idx, values) in data.iter().enumerate() { - let values: Vec = values.iter().map(&into_value).collect(); - let values: Vec<_> = values.iter().collect(); - builder - .add_point(idx as PointOffsetType, &values, &hw_counter) - .unwrap(); - } - builder.finalize().unwrap(); - } - } - } - - fn load_map_index( - data: &[Vec<::Owned>], - path: &Path, - index_type: IndexType, - ) -> MapIndex - where - Vec<::Owned>: Blob + Send + Sync, - { - let index = match index_type { - IndexType::MutableGridstore => MapIndex::::new_gridstore(path.to_path_buf(), true) - .unwrap() - .unwrap(), - IndexType::Mmap => MapIndex::::new_mmap(path, true, &empty_deleted()) - .unwrap() - .unwrap(), - IndexType::RamMmap => MapIndex::::new_mmap(path, false, &empty_deleted()) - .unwrap() - .unwrap(), - }; - let hw_counter = HardwareCounterCell::new(); - for (idx, values) in data.iter().enumerate() { - let index_values: HashSet<::Owned> = index - .get_values(idx as PointOffsetType, &hw_counter) - .unwrap() - .map(|v| MapIndexKey::to_owned(v.as_ref())) - .collect(); - let index_values: HashSet<&N> = index_values.iter().map(|v| v.borrow()).collect(); - let check_values: HashSet<&N> = values.iter().map(|v| v.borrow()).collect(); - assert_eq!(index_values, check_values); - } - - index - } - - #[test] - fn test_uuid_payload_index() { - let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); - let mut builder = - MapIndex::::builder_mmap(temp_dir.path(), false, &empty_deleted()); - - builder.init().unwrap(); - - let hw_counter = HardwareCounterCell::new(); - - let uuid: Value = Value::String("baa56dfc-e746-4ec1-bf50-94822535a46c".to_string()); - - for idx in 0..100 { - builder - .add_point(idx as PointOffsetType, &[&uuid], &hw_counter) - .unwrap(); - } - - let index = builder.finalize().unwrap(); - - index - .for_each_payload_block(50, PayloadKeyType::new("test_uuid"), &mut |block| { - black_box(block); - Ok(()) - }) - .unwrap(); - } - - #[test] - fn test_index_non_ascending_insertion() { - let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); - let mut builder = - MapIndex::::builder_mmap(temp_dir.path(), false, &empty_deleted()); - builder.init().unwrap(); - - let data = [vec![1, 2, 3, 4, 5, 6], vec![25], vec![10, 11]]; - - let hw_counter = HardwareCounterCell::new(); - - for (idx, values) in data.iter().enumerate().rev() { - let values: Vec = values.iter().map(|i| (*i).into()).collect(); - let values: Vec<_> = values.iter().collect(); - builder - .add_point(idx as PointOffsetType, &values, &hw_counter) - .unwrap(); - } - - let index = builder.finalize().unwrap(); - let hw_counter = HardwareCounterCell::new(); - for (idx, values) in data.iter().enumerate().rev() { - let res: Vec<_> = index - .get_values(idx as u32, &hw_counter) - .unwrap() - .map(|i| *i as i32) - .collect(); - assert_eq!(res, *values); - } - } - - #[rstest] - #[case(IndexType::MutableGridstore)] - #[case(IndexType::Mmap)] - #[case(IndexType::RamMmap)] - fn test_int_disk_map_index(#[case] index_type: IndexType) { - let data = vec![ - vec![1, 2, 3, 4, 5, 6], - vec![1, 2, 3, 4, 5, 6], - vec![13, 14, 15, 16, 17, 18], - vec![19, 20, 21, 22, 23, 24], - vec![25], - ]; - - let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); - save_map_index::(&data, temp_dir.path(), index_type, |v| (*v).into()); - let index = load_map_index::(&data, temp_dir.path(), index_type); - - let hw_counter = HardwareCounterCell::new(); - - assert!( - !index - .except_cardinality(std::iter::empty(), &hw_counter) - .equals_min_exp_max(&CardinalityEstimation::exact(0)) - ); - } - - #[rstest] - #[case(IndexType::MutableGridstore)] - #[case(IndexType::Mmap)] - #[case(IndexType::RamMmap)] - fn test_string_disk_map_index(#[case] index_type: IndexType) { - let data = vec![ - vec![ - EcoString::from("AABB"), - EcoString::from("UUFF"), - EcoString::from("IIBB"), - ], - vec![ - EcoString::from("PPMM"), - EcoString::from("QQXX"), - EcoString::from("YYBB"), - ], - vec![ - EcoString::from("FFMM"), - EcoString::from("IICC"), - EcoString::from("IIBB"), - ], - vec![ - EcoString::from("AABB"), - EcoString::from("UUFF"), - EcoString::from("IIBB"), - ], - vec![EcoString::from("PPGG")], - ]; - - let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); - save_map_index::(&data, temp_dir.path(), index_type, |v| v.to_string().into()); - let index = load_map_index::(&data, temp_dir.path(), index_type); - - let hw_counter = HardwareCounterCell::new(); - - assert!( - !index - .except_cardinality(vec![].into_iter(), &hw_counter) - .equals_min_exp_max(&CardinalityEstimation::exact(0)) - ); - } - - #[rstest] - #[case(IndexType::MutableGridstore)] - #[case(IndexType::Mmap)] - #[case(IndexType::RamMmap)] - fn test_empty_index(#[case] index_type: IndexType) { - let data: Vec> = vec![]; - - let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); - save_map_index::(&data, temp_dir.path(), index_type, |v| v.to_string().into()); - let index = load_map_index::(&data, temp_dir.path(), index_type); - - let hw_counter = HardwareCounterCell::new(); - - assert!( - index - .except_cardinality(std::iter::empty(), &hw_counter) - .equals_min_exp_max(&CardinalityEstimation::exact(0)) - ); - } - - /// Test that `get_values` on an on-disk mmap index actually increments the hardware counter. - #[test] - fn test_mmap_get_values_hw_counter() { - let data = vec![vec![1i64, 2, 3], vec![4, 5], vec![6]]; - - let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); - save_map_index::(&data, temp_dir.path(), IndexType::Mmap, |v| (*v).into()); - let index = load_map_index::(&data, temp_dir.path(), IndexType::Mmap); - - let hw_counter = HardwareCounterCell::new(); - for idx in 0..data.len() { - let _values: Vec<_> = index - .get_values(idx as PointOffsetType, &hw_counter) - .unwrap() - .collect(); - } - - assert!( - hw_counter.payload_index_io_read_counter().get() > 0, - "Expected on-disk mmap get_values to track payload index IO reads, but counter was 0" - ); - - let temp_dir2 = Builder::new().prefix("store_dir").tempdir().unwrap(); - save_map_index::(&data, temp_dir2.path(), IndexType::RamMmap, |v| { - (*v).into() - }); - let index2 = load_map_index::(&data, temp_dir2.path(), IndexType::RamMmap); - - let hw_counter2 = HardwareCounterCell::new(); - for idx in 0..data.len() { - let _values: Vec<_> = index2 - .get_values(idx as PointOffsetType, &hw_counter2) - .unwrap() - .collect(); - } - - assert_eq!( - hw_counter2.payload_index_io_read_counter().get(), - 0, - "Expected RAM mmap get_values NOT to track IO reads, but counter was non-zero" - ); - } - - /// Reload contract: runtime deletions are not persisted by the mmap map - /// index. Callers must re-supply the deletion bitslice on reload. - /// - /// Test data is chosen so that every value retains at least one live - /// point after deletions — otherwise `ImmutableMapIndex::open_mmap` hits a - /// pre-existing debug-only assertion when a value's slice becomes empty. - #[rstest] - #[case(IndexType::MutableGridstore)] - #[case(IndexType::Mmap)] - #[case(IndexType::RamMmap)] - fn test_map_index_reload(#[case] index_type: IndexType) { - let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); - let data: Vec> = vec![ - vec![1, 2], // id 0 - vec![1], // id 1 - vec![2], // id 2 - vec![1, 3], // id 3 - vec![2, 3], // id 4 - vec![3], // id 5 - ]; - - { - save_map_index::(&data, temp_dir.path(), index_type, |v| (*v).into()); - let mut index = load_map_index::(&data, temp_dir.path(), index_type); - index.remove_point(1).unwrap(); - index.remove_point(2).unwrap(); - index.remove_point(5).unwrap(); - index.flusher()().unwrap(); - assert_eq!(index.get_indexed_points(), 3); - drop(index); - } - - let deleted = deleted_with(&[1, 2, 5]); - let new_index = match index_type { - IndexType::MutableGridstore => { - MapIndex::::new_gridstore(temp_dir.path().to_path_buf(), true) - .unwrap() - .unwrap() - } - IndexType::Mmap => { - MapIndex::::new_mmap(temp_dir.path(), true, &deleted) - .unwrap() - .unwrap() - } - IndexType::RamMmap => { - MapIndex::::new_mmap(temp_dir.path(), false, &deleted) - .unwrap() - .unwrap() - } - }; - - assert_eq!(new_index.get_indexed_points(), 3); - - let hw_counter = HardwareCounterCell::new(); - for id in [1u32, 2, 5] { - assert_eq!( - new_index.values_count(id), - 0, - "deleted point {id} should have no values after reload", - ); - } - for id in [0u32, 3, 4] { - assert!( - new_index.values_count(id) > 0, - "live point {id} should have values after reload", - ); - } - - let mut hits: Vec = new_index.get_iterator(&1, &hw_counter).collect(); - hits.sort(); - assert_eq!(hits, vec![0, 3]); - - let mut hits: Vec = new_index.get_iterator(&2, &hw_counter).collect(); - hits.sort(); - assert_eq!(hits, vec![0, 4]); - - let mut hits: Vec = new_index.get_iterator(&3, &hw_counter).collect(); - hits.sort(); - assert_eq!(hits, vec![3, 4]); - } - - /// Regression test: when reloading an mmap map index with a `deleted_points` - /// bitslice shorter than `point_to_values.len()`, missing entries must - /// default to live, not deleted. Empty-payload bits from the on-disk - /// `deleted.bin` and any deletions encoded inside the short bitslice must - /// still be honored. - #[rstest] - #[case(IndexType::Mmap)] - #[case(IndexType::RamMmap)] - fn test_map_index_reload_short_deleted_bitslice(#[case] index_type: IndexType) { - let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); - - let data: Vec> = vec![ - vec![1], // id 0 - vec![1, 2], // id 1 - vec![], // id 2 — empty payload - vec![2, 3], // id 3 - vec![3], // id 4 - ]; - - save_map_index::(&data, temp_dir.path(), index_type, |v| (*v).into()); - - let mut short_deleted = BitVec::repeat(false, 2); - short_deleted.set(1, true); - - let new_index = match index_type { - IndexType::Mmap => { - MapIndex::::new_mmap(temp_dir.path(), true, &short_deleted) - .unwrap() - .unwrap() - } - IndexType::RamMmap => { - MapIndex::::new_mmap(temp_dir.path(), false, &short_deleted) - .unwrap() - .unwrap() - } - IndexType::MutableGridstore => unreachable!(), - }; - - let hw_counter = HardwareCounterCell::new(); - - assert!(new_index.values_count(0) > 0, "id 0 should be live"); - assert_eq!(new_index.values_count(1), 0, "id 1 deleted via bitslice"); - assert_eq!( - new_index.values_count(2), - 0, - "id 2 deleted via build-time empty" - ); - assert!( - new_index.values_count(3) > 0, - "id 3 should be live (beyond bitslice)" - ); - assert!( - new_index.values_count(4) > 0, - "id 4 should be live (beyond bitslice)" - ); - - let mut hits: Vec = new_index.get_iterator(&2, &hw_counter).collect(); - hits.sort(); - assert_eq!(hits, vec![3]); - } -} diff --git a/lib/segment/src/index/field_index/map_index/mutable_map_index.rs b/lib/segment/src/index/field_index/map_index/mutable_map_index.rs deleted file mode 100644 index 13acce0f05..0000000000 --- a/lib/segment/src/index/field_index/map_index/mutable_map_index.rs +++ /dev/null @@ -1,360 +0,0 @@ -use std::borrow::{Borrow, Cow}; -use std::collections::HashMap; -use std::iter; -use std::path::PathBuf; - -use common::counter::hardware_counter::HardwareCounterCell; -use common::types::PointOffsetType; -use gridstore::config::StorageOptions; -use gridstore::error::GridstoreError; -use gridstore::{Blob, Gridstore}; -use roaring::RoaringBitmap; - -use super::{IdIter, MapIndexKey}; -use crate::common::Flusher; -use crate::common::operation_error::{OperationError, OperationResult}; -use crate::index::payload_config::StorageType; - -/// Default options for Gridstore storage -const fn default_gridstore_options(block_size: usize) -> StorageOptions { - StorageOptions { - // Size dependent on map value type - block_size_bytes: Some(block_size), - compression: Some(gridstore::config::Compression::None), - page_size_bytes: Some(block_size * 8192 * 32), // 4 to 8 MiB = block_size * region_blocks * regions, - region_size_blocks: None, - } -} - -pub struct MutableMapIndex -where - Vec<::Owned>: Blob + Send + Sync, -{ - pub(super) map: HashMap<::Owned, RoaringBitmap>, - pub(super) point_to_values: Vec::Owned>>, - /// Amount of point which have at least one indexed payload value - pub(super) indexed_points: usize, - pub(super) values_count: usize, - storage: Storage<::Owned>, -} - -enum Storage -where - Vec: Blob + Send + Sync, -{ - Gridstore(Gridstore>), -} - -impl MutableMapIndex -where - Vec<::Owned>: Blob + Send + Sync, -{ - /// Open and load mutable map index from Gridstore storage - /// - /// The `create_if_missing` parameter indicates whether to create a new Gridstore if it does - /// not exist. If false and files don't exist, the load function will indicate nothing could be - /// loaded. - pub fn open_gridstore(path: PathBuf, create_if_missing: bool) -> OperationResult> { - let store = if create_if_missing { - let options = default_gridstore_options(N::gridstore_block_size()); - Gridstore::open_or_create(path, options).map_err(|err| { - OperationError::service_error(format!( - "failed to open mutable map index on gridstore: {err}" - )) - })? - } else if path.exists() { - Gridstore::open(path).map_err(|err| { - OperationError::service_error(format!( - "failed to open mutable map index on gridstore: {err}" - )) - })? - } else { - // Files don't exist, cannot load - return Ok(None); - }; - - // Load in-memory index from Gridstore - let mut map = HashMap::<_, RoaringBitmap>::new(); - let mut point_to_values = Vec::new(); - let mut indexed_points = 0; - let mut values_count = 0; - - let hw_counter = HardwareCounterCell::disposable(); - let hw_counter_ref = hw_counter.ref_payload_index_io_write_counter(); - store - .iter::<_, GridstoreError>( - |idx, values: Vec<_>| { - for value in values { - if point_to_values.len() <= idx as usize { - point_to_values.resize_with(idx as usize + 1, Vec::new) - } - let point_values = &mut point_to_values[idx as usize]; - - if point_values.is_empty() { - indexed_points += 1; - } - values_count += 1; - - point_values.push(value.clone()); - map.entry(value).or_default().insert(idx); - } - - Ok(true) - }, - hw_counter_ref, - ) - // unwrap safety: never returns an error - .unwrap(); - - Ok(Some(Self { - map, - point_to_values, - indexed_points, - values_count, - storage: Storage::Gridstore(store), - })) - } - - pub fn add_many_to_map( - &mut self, - idx: PointOffsetType, - values: Vec, - hw_counter: &HardwareCounterCell, - ) -> OperationResult<()> - where - Q: Into<::Owned> + Clone, - { - if values.is_empty() { - return Ok(()); - } - - self.values_count += values.len(); - if self.point_to_values.len() <= idx as usize { - self.point_to_values.resize_with(idx as usize + 1, Vec::new) - } - - self.point_to_values[idx as usize] = Vec::with_capacity(values.len()); - - match &mut self.storage { - Storage::Gridstore(store) => { - let hw_counter_ref = hw_counter.ref_payload_index_io_write_counter(); - - for value in values.clone() { - let entry = self.map.entry(value.into()); - self.point_to_values[idx as usize].push(entry.key().clone()); - entry.or_default().insert(idx); - } - - let values = values.into_iter().map(Into::into).collect::>(); - store - .put_value(idx, &values, hw_counter_ref) - .map_err(|err| { - OperationError::service_error(format!( - "failed to put value in mutable map index gridstore: {err}" - )) - })?; - } - } - - self.indexed_points += 1; - Ok(()) - } - - pub fn remove_point(&mut self, idx: PointOffsetType) -> OperationResult<()> { - if self.point_to_values.len() <= idx as usize { - return Ok(()); - } - - let removed_values = std::mem::take(&mut self.point_to_values[idx as usize]); - - if !removed_values.is_empty() { - self.indexed_points -= 1; - } - self.values_count -= removed_values.len(); - - for value in &removed_values { - if let Some(vals) = self.map.get_mut(value.borrow()) { - vals.remove(idx); - } - } - - match &mut self.storage { - Storage::Gridstore(store) => { - store.delete_value(idx)?; - } - } - - Ok(()) - } - - #[inline] - pub(super) fn clear(&mut self) -> OperationResult<()> { - match &mut self.storage { - Storage::Gridstore(store) => store.clear().map_err(|err| { - OperationError::service_error(format!("Failed to clear mutable map index: {err}",)) - }), - } - } - - #[inline] - pub(super) fn wipe(self) -> OperationResult<()> { - match self.storage { - Storage::Gridstore(store) => store.wipe().map_err(|err| { - OperationError::service_error(format!("Failed to wipe mutable map index: {err}",)) - }), - } - } - - /// Clear cache - /// - /// Only clears cache of Gridstore storage if used. Does not clear in-memory representation of - /// index. - pub fn clear_cache(&self) -> OperationResult<()> { - match &self.storage { - Storage::Gridstore(index) => index.clear_cache().map_err(|err| { - OperationError::service_error(format!( - "Failed to clear mutable map index gridstore cache: {err}" - )) - }), - } - } - - #[inline] - pub(super) fn files(&self) -> Vec { - match &self.storage { - Storage::Gridstore(store) => store.files(), - } - } - - #[inline] - pub(super) fn flusher(&self) -> Flusher { - match &self.storage { - Storage::Gridstore(store) => { - let storage_flusher = store.flusher(); - Box::new(move || storage_flusher().map_err(OperationError::from)) - } - } - } - - pub fn check_values_any(&self, idx: PointOffsetType, check_fn: impl Fn(&N) -> bool) -> bool { - self.point_to_values - .get(idx as usize) - .map(|values| values.iter().any(|v| check_fn(v.borrow()))) - .unwrap_or(false) - } - - pub fn get_values( - &self, - idx: PointOffsetType, - ) -> Option> + '_> { - Some( - self.point_to_values - .get(idx as usize)? - .iter() - .map(|v| Cow::Borrowed(v.borrow())), - ) - } - - pub fn values_count(&self, idx: PointOffsetType) -> Option { - self.point_to_values.get(idx as usize).map(Vec::len) - } - - pub fn get_indexed_points(&self) -> usize { - self.indexed_points - } - - pub fn get_values_count(&self) -> usize { - self.values_count - } - - pub fn get_unique_values_count(&self) -> usize { - self.map.len() - } - - pub fn get_count_for_value(&self, value: &N) -> Option { - self.map.get(value).map(|p| p.len() as usize) - } - - pub fn for_points_values( - &self, - points: impl Iterator, - mut f: impl FnMut(PointOffsetType, &[::Owned]), - ) { - points.for_each(|idx| { - if let Some(values) = self.point_to_values.get(idx as usize) { - f(idx, values); - } - }); - } - - pub fn for_each_count_per_value( - &self, - deferred_internal_id: Option, - mut f: impl FnMut(&N, usize) -> OperationResult<()>, - ) -> OperationResult<()> { - self.map.iter().try_for_each(|(k, v)| { - let count = match deferred_internal_id { - Some(deferred_internal_id) => v.range_cardinality(..deferred_internal_id) as usize, - None => v.len() as usize, - }; - f(k.borrow(), count) - }) - } - - pub fn for_each_value_map( - &self, - mut f: impl FnMut(&N, &mut dyn Iterator) -> OperationResult<()>, - ) -> OperationResult<()> { - self.map - .iter() - .try_for_each(|(k, v)| f(k.borrow(), &mut v.iter())) - } - - pub fn get_iterator(&self, value: &N) -> IdIter<'_> { - self.map - .get(value) - .map(|ids| Box::new(ids.iter()) as IdIter) - .unwrap_or_else(|| Box::new(iter::empty::())) - } - - pub fn for_each_value( - &self, - mut f: impl FnMut(&N) -> OperationResult<()>, - ) -> OperationResult<()> { - self.map.keys().try_for_each(|v| f(v.borrow())) - } - - pub fn storage_type(&self) -> StorageType { - match &self.storage { - Storage::Gridstore(_) => StorageType::Gridstore, - } - } - - /// Approximate RAM usage in bytes for in-memory index structures. - pub fn ram_usage_bytes(&self) -> usize { - let Self { - map, - point_to_values, - indexed_points: _, - values_count: _, - storage: _, // disk-backed, accounted via files - } = self; - - let hashmap_entry_overhead = std::mem::size_of::() + std::mem::size_of::(); - let map_base_bytes = map.capacity() - * (std::mem::size_of::<::Owned>() - + std::mem::size_of::() - + hashmap_entry_overhead); - // Account for heap-allocated key data (e.g., long strings) - let map_key_heap_bytes: usize = map.keys().map(|k| N::owned_heap_bytes(k)).sum(); - let map_bitmap_bytes: usize = map.values().map(|bitmap| bitmap.serialized_size()).sum(); - let map_bytes = map_base_bytes + map_key_heap_bytes + map_bitmap_bytes; - let ptv_bytes: usize = point_to_values.capacity() - * std::mem::size_of::::Owned>>() - + point_to_values - .iter() - .map(|v| v.capacity() * std::mem::size_of::<::Owned>()) - .sum::(); - map_bytes + ptv_bytes - } -} diff --git a/lib/segment/src/index/field_index/map_index/mutable_map_index/lifecycle.rs b/lib/segment/src/index/field_index/map_index/mutable_map_index/lifecycle.rs new file mode 100644 index 0000000000..11276d2dcf --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/mutable_map_index/lifecycle.rs @@ -0,0 +1,195 @@ +use std::borrow::Borrow; +use std::collections::HashMap; +use std::path::PathBuf; + +use common::counter::hardware_counter::HardwareCounterCell; +use common::types::PointOffsetType; +use gridstore::config::StorageOptions; +use gridstore::error::GridstoreError; +use gridstore::{Blob, Gridstore}; +use roaring::RoaringBitmap; + +use super::super::MapIndexKey; +use super::MutableMapIndex; +use crate::common::Flusher; +use crate::common::operation_error::{OperationError, OperationResult}; + +/// Default options for Gridstore storage +const fn default_gridstore_options(block_size: usize) -> StorageOptions { + StorageOptions { + // Size dependent on map value type + block_size_bytes: Some(block_size), + compression: Some(gridstore::config::Compression::None), + page_size_bytes: Some(block_size * 8192 * 32), // 4 to 8 MiB = block_size * region_blocks * regions, + region_size_blocks: None, + } +} + +impl MutableMapIndex +where + Vec<::Owned>: Blob + Send + Sync, +{ + /// Open and load mutable map index from Gridstore storage + /// + /// The `create_if_missing` parameter indicates whether to create a new Gridstore if it does + /// not exist. If false and files don't exist, the load function will indicate nothing could be + /// loaded. + pub fn open_gridstore(path: PathBuf, create_if_missing: bool) -> OperationResult> { + let store = if create_if_missing { + let options = default_gridstore_options(N::gridstore_block_size()); + Gridstore::open_or_create(path, options).map_err(|err| { + OperationError::service_error(format!( + "failed to open mutable map index on gridstore: {err}" + )) + })? + } else if path.exists() { + Gridstore::open(path).map_err(|err| { + OperationError::service_error(format!( + "failed to open mutable map index on gridstore: {err}" + )) + })? + } else { + // Files don't exist, cannot load + return Ok(None); + }; + + // Load in-memory index from Gridstore + let mut map = HashMap::<_, RoaringBitmap>::new(); + let mut point_to_values = Vec::new(); + let mut indexed_points = 0; + let mut values_count = 0; + + let hw_counter = HardwareCounterCell::disposable(); + let hw_counter_ref = hw_counter.ref_payload_index_io_write_counter(); + store + .iter::<_, GridstoreError>( + |idx, values: Vec<_>| { + for value in values { + if point_to_values.len() <= idx as usize { + point_to_values.resize_with(idx as usize + 1, Vec::new) + } + let point_values = &mut point_to_values[idx as usize]; + + if point_values.is_empty() { + indexed_points += 1; + } + values_count += 1; + + point_values.push(value.clone()); + map.entry(value).or_default().insert(idx); + } + + Ok(true) + }, + hw_counter_ref, + ) + // unwrap safety: never returns an error + .unwrap(); + + Ok(Some(Self { + map, + point_to_values, + indexed_points, + values_count, + storage: store, + })) + } + + pub fn add_many_to_map( + &mut self, + idx: PointOffsetType, + values: Vec, + hw_counter: &HardwareCounterCell, + ) -> OperationResult<()> + where + Q: Into<::Owned> + Clone, + { + if values.is_empty() { + return Ok(()); + } + + self.values_count += values.len(); + if self.point_to_values.len() <= idx as usize { + self.point_to_values.resize_with(idx as usize + 1, Vec::new) + } + + self.point_to_values[idx as usize] = Vec::with_capacity(values.len()); + + let hw_counter_ref = hw_counter.ref_payload_index_io_write_counter(); + + for value in values.clone() { + let entry = self.map.entry(value.into()); + self.point_to_values[idx as usize].push(entry.key().clone()); + entry.or_default().insert(idx); + } + + let values = values.into_iter().map(Into::into).collect::>(); + self.storage + .put_value(idx, &values, hw_counter_ref) + .map_err(|err| { + OperationError::service_error(format!( + "failed to put value in mutable map index gridstore: {err}" + )) + })?; + + self.indexed_points += 1; + Ok(()) + } + + pub fn remove_point(&mut self, idx: PointOffsetType) -> OperationResult<()> { + if self.point_to_values.len() <= idx as usize { + return Ok(()); + } + + let removed_values = std::mem::take(&mut self.point_to_values[idx as usize]); + + if !removed_values.is_empty() { + self.indexed_points -= 1; + } + self.values_count -= removed_values.len(); + + for value in &removed_values { + if let Some(vals) = self.map.get_mut(value.borrow()) { + vals.remove(idx); + } + } + + self.storage.delete_value(idx)?; + + Ok(()) + } + + #[inline] + pub(in super::super) fn clear(&mut self) -> OperationResult<()> { + self.storage.clear().map_err(|err| { + OperationError::service_error(format!("Failed to clear mutable map index: {err}")) + }) + } + + #[inline] + pub(in super::super) fn wipe(self) -> OperationResult<()> { + self.storage.wipe().map_err(|err| { + OperationError::service_error(format!("Failed to wipe mutable map index: {err}")) + }) + } + + /// Clear gridstore disk cache. Does not affect the in-memory index. + pub fn clear_cache(&self) -> OperationResult<()> { + self.storage.clear_cache().map_err(|err| { + OperationError::service_error(format!( + "Failed to clear mutable map index gridstore cache: {err}" + )) + }) + } + + #[inline] + pub(in super::super) fn files(&self) -> Vec { + self.storage.files() + } + + #[inline] + pub(in super::super) fn flusher(&self) -> Flusher { + let storage_flusher = self.storage.flusher(); + Box::new(move || storage_flusher().map_err(OperationError::from)) + } +} diff --git a/lib/segment/src/index/field_index/map_index/mutable_map_index/mod.rs b/lib/segment/src/index/field_index/map_index/mutable_map_index/mod.rs new file mode 100644 index 0000000000..c219b104e3 --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/mutable_map_index/mod.rs @@ -0,0 +1,21 @@ +use std::collections::HashMap; + +use gridstore::{Blob, Gridstore}; +use roaring::RoaringBitmap; + +use super::MapIndexKey; + +mod lifecycle; +mod read_ops; + +pub struct MutableMapIndex +where + Vec<::Owned>: Blob + Send + Sync, +{ + pub(super) map: HashMap<::Owned, RoaringBitmap>, + pub(super) point_to_values: Vec::Owned>>, + /// Amount of point which have at least one indexed payload value + pub(super) indexed_points: usize, + pub(super) values_count: usize, + pub(super) storage: Gridstore::Owned>>, +} diff --git a/lib/segment/src/index/field_index/map_index/mutable_map_index/read_ops.rs b/lib/segment/src/index/field_index/map_index/mutable_map_index/read_ops.rs new file mode 100644 index 0000000000..95ee656498 --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/mutable_map_index/read_ops.rs @@ -0,0 +1,150 @@ +use std::borrow::{Borrow, Cow}; +use std::iter; + +use common::counter::hardware_counter::HardwareCounterCell; +use common::types::PointOffsetType; +use gridstore::Blob; +use roaring::RoaringBitmap; + +use super::super::read_ops::MapIndexRead; +use super::super::{IdIter, MapIndexKey}; +use super::MutableMapIndex; +use crate::common::operation_error::OperationResult; +use crate::index::payload_config::StorageType; + +impl MapIndexRead for MutableMapIndex +where + Vec<::Owned>: Blob + Send + Sync, +{ + fn check_values_any( + &self, + idx: PointOffsetType, + _hw_counter: &HardwareCounterCell, + check_fn: impl Fn(&N) -> bool, + ) -> bool { + self.point_to_values + .get(idx as usize) + .map(|values| values.iter().any(|v| check_fn(v.borrow()))) + .unwrap_or(false) + } + + fn get_values<'a>( + &'a self, + idx: PointOffsetType, + _hw_counter: &HardwareCounterCell, + ) -> Option> + 'a> + where + N: 'a, + { + Some( + self.point_to_values + .get(idx as usize)? + .iter() + .map(|v| Cow::Borrowed(v.borrow())), + ) + } + + fn values_count(&self, idx: PointOffsetType) -> Option { + self.point_to_values.get(idx as usize).map(Vec::len) + } + + fn get_indexed_points(&self) -> usize { + self.indexed_points + } + + fn get_values_count(&self) -> usize { + self.values_count + } + + fn get_unique_values_count(&self) -> usize { + self.map.len() + } + + fn get_count_for_value(&self, value: &N, _hw_counter: &HardwareCounterCell) -> Option { + self.map.get(value).map(|p| p.len() as usize) + } + + fn get_iterator(&self, value: &N, _hw_counter: &HardwareCounterCell) -> IdIter<'_> { + self.map + .get(value) + .map(|ids| Box::new(ids.iter()) as IdIter) + .unwrap_or_else(|| Box::new(iter::empty::())) + } + + fn for_each_value(&self, mut f: impl FnMut(&N) -> OperationResult<()>) -> OperationResult<()> { + self.map.keys().try_for_each(|v| f(v.borrow())) + } + + fn for_each_count_per_value( + &self, + deferred_internal_id: Option, + mut f: impl FnMut(&N, usize) -> OperationResult<()>, + ) -> OperationResult<()> { + self.map.iter().try_for_each(|(k, v)| { + let count = match deferred_internal_id { + Some(deferred_internal_id) => v.range_cardinality(..deferred_internal_id) as usize, + None => v.len() as usize, + }; + f(k.borrow(), count) + }) + } + + fn for_each_value_map( + &self, + _hw_counter: &HardwareCounterCell, + mut f: impl FnMut(&N, &mut dyn Iterator) -> OperationResult<()>, + ) -> OperationResult<()> { + self.map + .iter() + .try_for_each(|(k, v)| f(k.borrow(), &mut v.iter())) + } + + fn storage_type(&self) -> StorageType { + StorageType::Gridstore + } + + /// Approximate RAM usage in bytes for in-memory index structures. + fn ram_usage_bytes(&self) -> usize { + let Self { + map, + point_to_values, + indexed_points: _, + values_count: _, + storage: _, // disk-backed, accounted via files + } = self; + + let hashmap_entry_overhead = std::mem::size_of::() + std::mem::size_of::(); + let map_base_bytes = map.capacity() + * (std::mem::size_of::<::Owned>() + + std::mem::size_of::() + + hashmap_entry_overhead); + // Account for heap-allocated key data (e.g., long strings) + let map_key_heap_bytes: usize = map.keys().map(|k| N::owned_heap_bytes(k)).sum(); + let map_bitmap_bytes: usize = map.values().map(|bitmap| bitmap.serialized_size()).sum(); + let map_bytes = map_base_bytes + map_key_heap_bytes + map_bitmap_bytes; + let ptv_bytes: usize = point_to_values.capacity() + * std::mem::size_of::::Owned>>() + + point_to_values + .iter() + .map(|v| v.capacity() * std::mem::size_of::<::Owned>()) + .sum::(); + map_bytes + ptv_bytes + } +} + +impl MutableMapIndex +where + Vec<::Owned>: Blob + Send + Sync, +{ + pub fn for_points_values( + &self, + points: impl Iterator, + mut f: impl FnMut(PointOffsetType, &[::Owned]), + ) { + points.for_each(|idx| { + if let Some(values) = self.point_to_values.get(idx as usize) { + f(idx, values); + } + }); + } +} diff --git a/lib/segment/src/index/field_index/map_index/payload_index_impl_int.rs b/lib/segment/src/index/field_index/map_index/payload_index_impl/int.rs similarity index 99% rename from lib/segment/src/index/field_index/map_index/payload_index_impl_int.rs rename to lib/segment/src/index/field_index/map_index/payload_index_impl/int.rs index f3157cf247..21472137fc 100644 --- a/lib/segment/src/index/field_index/map_index/payload_index_impl_int.rs +++ b/lib/segment/src/index/field_index/map_index/payload_index_impl/int.rs @@ -6,7 +6,7 @@ use common::counter::hardware_counter::HardwareCounterCell; use common::types::PointOffsetType; use itertools::Itertools; -use super::MapIndex; +use super::super::MapIndex; use crate::common::Flusher; use crate::common::operation_error::OperationResult; use crate::index::field_index::{ diff --git a/lib/segment/src/index/field_index/map_index/payload_index_impl/mod.rs b/lib/segment/src/index/field_index/map_index/payload_index_impl/mod.rs new file mode 100644 index 0000000000..01aef11d8a --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/payload_index_impl/mod.rs @@ -0,0 +1,3 @@ +mod int; +mod str; +mod uuid; diff --git a/lib/segment/src/index/field_index/map_index/payload_index_impl_str.rs b/lib/segment/src/index/field_index/map_index/payload_index_impl/str.rs similarity index 99% rename from lib/segment/src/index/field_index/map_index/payload_index_impl_str.rs rename to lib/segment/src/index/field_index/map_index/payload_index_impl/str.rs index ebfadcb7d6..7aa6ce1cd3 100644 --- a/lib/segment/src/index/field_index/map_index/payload_index_impl_str.rs +++ b/lib/segment/src/index/field_index/map_index/payload_index_impl/str.rs @@ -6,7 +6,7 @@ use common::counter::hardware_counter::HardwareCounterCell; use common::types::PointOffsetType; use itertools::Itertools; -use super::MapIndex; +use super::super::MapIndex; use crate::common::Flusher; use crate::common::operation_error::OperationResult; use crate::index::field_index::{ diff --git a/lib/segment/src/index/field_index/map_index/payload_index_impl_uuid.rs b/lib/segment/src/index/field_index/map_index/payload_index_impl/uuid.rs similarity index 99% rename from lib/segment/src/index/field_index/map_index/payload_index_impl_uuid.rs rename to lib/segment/src/index/field_index/map_index/payload_index_impl/uuid.rs index ff88d9733c..09159206d4 100644 --- a/lib/segment/src/index/field_index/map_index/payload_index_impl_uuid.rs +++ b/lib/segment/src/index/field_index/map_index/payload_index_impl/uuid.rs @@ -9,7 +9,7 @@ use indexmap::IndexSet; use itertools::Itertools; use uuid::Uuid; -use super::MapIndex; +use super::super::MapIndex; use crate::common::Flusher; use crate::common::operation_error::OperationResult; use crate::index::field_index::{ diff --git a/lib/segment/src/index/field_index/map_index/read_ops.rs b/lib/segment/src/index/field_index/map_index/read_ops.rs new file mode 100644 index 0000000000..b83957080d --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/read_ops.rs @@ -0,0 +1,381 @@ +use std::borrow::{Borrow, Cow}; +use std::hash::{BuildHasher, Hash}; + +use common::counter::hardware_counter::HardwareCounterCell; +use common::types::PointOffsetType; +use gridstore::Blob; +use indexmap::IndexSet; + +use super::key::MapIndexKey; +use super::{IdIter, MapIndex}; +use crate::common::operation_error::OperationResult; +use crate::index::field_index::CardinalityEstimation; +use crate::index::field_index::stat_tools::number_of_selected_points; +use crate::index::payload_config::{IndexMutability, StorageType}; +use crate::telemetry::PayloadIndexTelemetry; + +/// Read-only operations supported by every map-index storage variant +/// ([`super::mutable_map_index::MutableMapIndex`], +/// [`super::immutable_map_index::ImmutableMapIndex`], +/// [`super::mmap_map_index::MmapMapIndex`]). +/// +/// Signatures are unified across variants so the enum-level dispatcher in +/// [`MapIndex`] can call them generically. Variants that don't need +/// `hw_counter` (`Mutable` / `Immutable`) accept and ignore it; the mmap +/// variant uses it to track payload-index IO. +pub(super) trait MapIndexRead { + fn check_values_any( + &self, + idx: PointOffsetType, + hw_counter: &HardwareCounterCell, + check_fn: impl Fn(&N) -> bool, + ) -> bool; + + fn get_values<'a>( + &'a self, + idx: PointOffsetType, + hw_counter: &HardwareCounterCell, + ) -> Option> + 'a> + where + N: 'a; + + fn values_count(&self, idx: PointOffsetType) -> Option; + + fn get_indexed_points(&self) -> usize; + + fn get_values_count(&self) -> usize; + + fn get_unique_values_count(&self) -> usize; + + fn get_count_for_value(&self, value: &N, hw_counter: &HardwareCounterCell) -> Option; + + fn get_iterator(&self, value: &N, hw_counter: &HardwareCounterCell) -> IdIter<'_>; + + fn for_each_value(&self, f: impl FnMut(&N) -> OperationResult<()>) -> OperationResult<()>; + + /// Iterate `(value, count)` pairs. + /// + /// `deferred_internal_id` (mutable / mmap only) restricts the count to + /// point IDs strictly less than the given value. The immutable variant + /// does not support deferred filtering and asserts the argument is `None`. + fn for_each_count_per_value( + &self, + deferred_internal_id: Option, + f: impl FnMut(&N, usize) -> OperationResult<()>, + ) -> OperationResult<()>; + + fn for_each_value_map( + &self, + hw_counter: &HardwareCounterCell, + f: impl FnMut(&N, &mut dyn Iterator) -> OperationResult<()>, + ) -> OperationResult<()>; + + fn storage_type(&self) -> StorageType; + + fn ram_usage_bytes(&self) -> usize; +} + +impl MapIndex +where + Vec<::Owned>: Blob + Send + Sync, +{ + pub fn check_values_any( + &self, + idx: PointOffsetType, + hw_counter: &HardwareCounterCell, + check_fn: impl Fn(&N) -> bool, + ) -> bool { + match self { + MapIndex::Mutable(index) => index.check_values_any(idx, hw_counter, check_fn), + MapIndex::Immutable(index) => index.check_values_any(idx, hw_counter, check_fn), + MapIndex::Mmap(index) => index.check_values_any(idx, hw_counter, check_fn), + } + } + + pub fn get_values( + &self, + idx: PointOffsetType, + hw_counter: &HardwareCounterCell, + ) -> Option> + '_>> { + match self { + MapIndex::Mutable(index) => Some(Box::new(index.get_values(idx, hw_counter)?)), + MapIndex::Immutable(index) => Some(Box::new(index.get_values(idx, hw_counter)?)), + MapIndex::Mmap(index) => Some(Box::new(index.get_values(idx, hw_counter)?)), + } + } + + pub fn values_count(&self, idx: PointOffsetType) -> usize { + match self { + MapIndex::Mutable(index) => index.values_count(idx).unwrap_or_default(), + MapIndex::Immutable(index) => index.values_count(idx).unwrap_or_default(), + MapIndex::Mmap(index) => index.values_count(idx).unwrap_or_default(), + } + } + + pub(crate) fn get_indexed_points(&self) -> usize { + match self { + MapIndex::Mutable(index) => index.get_indexed_points(), + MapIndex::Immutable(index) => index.get_indexed_points(), + MapIndex::Mmap(index) => index.get_indexed_points(), + } + } + + fn get_values_count(&self) -> usize { + match self { + MapIndex::Mutable(index) => index.get_values_count(), + MapIndex::Immutable(index) => index.get_values_count(), + MapIndex::Mmap(index) => index.get_values_count(), + } + } + + pub fn get_unique_values_count(&self) -> usize { + match self { + MapIndex::Mutable(index) => index.get_unique_values_count(), + MapIndex::Immutable(index) => index.get_unique_values_count(), + MapIndex::Mmap(index) => index.get_unique_values_count(), + } + } + + pub(crate) fn get_count_for_value( + &self, + value: &N, + hw_counter: &HardwareCounterCell, + ) -> Option { + match self { + MapIndex::Mutable(index) => index.get_count_for_value(value, hw_counter), + MapIndex::Immutable(index) => index.get_count_for_value(value, hw_counter), + MapIndex::Mmap(index) => index.get_count_for_value(value, hw_counter), + } + } + + pub(crate) fn get_iterator(&self, value: &N, hw_counter: &HardwareCounterCell) -> IdIter<'_> { + match self { + MapIndex::Mutable(index) => index.get_iterator(value, hw_counter), + MapIndex::Immutable(index) => index.get_iterator(value, hw_counter), + MapIndex::Mmap(index) => index.get_iterator(value, hw_counter), + } + } + + pub fn for_each_value(&self, f: impl FnMut(&N) -> OperationResult<()>) -> OperationResult<()> { + match self { + MapIndex::Mutable(index) => index.for_each_value(f), + MapIndex::Immutable(index) => index.for_each_value(f), + MapIndex::Mmap(index) => index.for_each_value(f), + } + } + + pub fn for_each_count_per_value( + &self, + deferred_internal_id: Option, + f: impl FnMut(&N, usize) -> OperationResult<()>, + ) -> OperationResult<()> { + // The immutable variant does not support deferred filtering — it + // asserts the argument is `None`. Two reasons we don't implement it: + // - We don't have both deferred points and an immutable index. + // - It is not trivial (nor performant) to implement correct filtering + // for this index variant as it doesn't work well in combination + // with the way it handles deletions. + match self { + MapIndex::Mutable(index) => index.for_each_count_per_value(deferred_internal_id, f), + MapIndex::Immutable(index) => index.for_each_count_per_value(deferred_internal_id, f), + MapIndex::Mmap(index) => index.for_each_count_per_value(deferred_internal_id, f), + } + } + + pub fn for_each_value_map( + &self, + hw_cell: &HardwareCounterCell, + f: impl FnMut(&N, &mut dyn Iterator) -> OperationResult<()>, + ) -> OperationResult<()> { + match self { + MapIndex::Mutable(index) => index.for_each_value_map(hw_cell, f), + MapIndex::Immutable(index) => index.for_each_value_map(hw_cell, f), + MapIndex::Mmap(index) => index.for_each_value_map(hw_cell, f), + } + } + + pub(crate) fn match_cardinality( + &self, + value: &N, + hw_counter: &HardwareCounterCell, + ) -> CardinalityEstimation { + let values_count = self.get_count_for_value(value, hw_counter).unwrap_or(0); + + CardinalityEstimation::exact(values_count) + } + + pub fn get_telemetry_data(&self) -> PayloadIndexTelemetry { + PayloadIndexTelemetry { + field_name: None, + points_count: self.get_indexed_points(), + points_values_count: self.get_values_count(), + histogram_bucket_size: None, + index_type: match self { + MapIndex::Mutable(_) => "mutable_map", + MapIndex::Immutable(_) => "immutable_map", + MapIndex::Mmap(_) => "mmap_map", + }, + } + } + + pub fn values_is_empty(&self, idx: PointOffsetType) -> bool { + self.values_count(idx) == 0 + } + + /// Estimates cardinality for `except` clause + /// + /// # Arguments + /// + /// * 'excluded' - values, which are not considered as matching + /// + /// # Returns + /// + /// * `CardinalityEstimation` - estimation of cardinality + pub(crate) fn except_cardinality<'a>( + &'a self, + excluded: impl Iterator, + hw_counter: &HardwareCounterCell, + ) -> CardinalityEstimation { + // Minimal case: we exclude as many points as possible. + // In this case, excluded points do not have any other values except excluded ones. + // So the first step - we estimate how many other points is needed to fit unused values. + + // Example: + // Values: 20, 20 + // Unique values: 5 + // Total points: 100 + // Total values: 110 + // total_excluded_value_count = 40 + // non_excluded_values_count = 110 - 40 = 70 + // max_values_per_point = 5 - 2 = 3 + // min_not_excluded_by_values = 70 / 3 = 24 + // min = max(24, 100 - 40) = 60 + // exp = ... + // max = min(20, 70) = 20 + + // Values: 60, 60 + // Unique values: 5 + // Total points: 100 + // Total values: 200 + // total_excluded_value_count = 120 + // non_excluded_values_count = 200 - 120 = 80 + // max_values_per_point = 5 - 2 = 3 + // min_not_excluded_by_values = 80 / 3 = 27 + // min = max(27, 100 - 120) = 27 + // exp = ... + // max = min(60, 80) = 60 + + // Values: 60, 60, 60 + // Unique values: 5 + // Total points: 100 + // Total values: 200 + // total_excluded_value_count = 180 + // non_excluded_values_count = 200 - 180 = 20 + // max_values_per_point = 5 - 3 = 2 + // min_not_excluded_by_values = 20 / 2 = 10 + // min = max(10, 100 - 180) = 10 + // exp = ... + // max = min(60, 20) = 20 + + let excluded_value_counts: Vec<_> = excluded + .map(|val| { + self.get_count_for_value(val.borrow(), hw_counter) + .unwrap_or(0) + }) + .collect(); + let total_excluded_value_count: usize = excluded_value_counts.iter().sum(); + + debug_assert!(total_excluded_value_count <= self.get_values_count()); + + let non_excluded_values_count = self + .get_values_count() + .saturating_sub(total_excluded_value_count); + let max_values_per_point = self + .get_unique_values_count() + .saturating_sub(excluded_value_counts.len()); + + if max_values_per_point == 0 { + debug_assert_eq!(non_excluded_values_count, 0); + return CardinalityEstimation::exact(0); + } + + let min_not_excluded_by_values = non_excluded_values_count.div_ceil(max_values_per_point); + + let min = min_not_excluded_by_values.max( + self.get_indexed_points() + .saturating_sub(total_excluded_value_count), + ); + + let max_excluded_value_count = excluded_value_counts.iter().max().copied().unwrap_or(0); + + let max = self + .get_indexed_points() + .saturating_sub(max_excluded_value_count) + .min(non_excluded_values_count); + + let exp = number_of_selected_points(self.get_indexed_points(), non_excluded_values_count) + .max(min) + .min(max); + + CardinalityEstimation { + primary_clauses: vec![], + min, + exp, + max, + } + } + + pub(crate) fn except_set<'a, K, A>( + &'a self, + excluded: &'a IndexSet, + hw_counter: &'a HardwareCounterCell, + ) -> OperationResult + 'a>> + where + A: BuildHasher, + K: Borrow + Hash + Eq, + { + let mut points = IndexSet::new(); + self.for_each_value(|key| { + if !excluded.contains(key.borrow()) { + self.get_iterator(key.borrow(), hw_counter).for_each(|p| { + points.insert(p); + }); + } + Ok(()) + })?; + Ok(Box::new(points.into_iter())) + } + + /// Approximate RAM usage in bytes for in-memory structures. + pub fn ram_usage_bytes(&self) -> usize { + match self { + MapIndex::Mutable(index) => index.ram_usage_bytes(), + MapIndex::Immutable(index) => index.ram_usage_bytes(), + MapIndex::Mmap(index) => index.ram_usage_bytes(), + } + } + + pub fn is_on_disk(&self) -> bool { + match self { + MapIndex::Mutable(_) => false, + MapIndex::Immutable(_) => false, + MapIndex::Mmap(index) => index.is_on_disk(), + } + } + + pub fn get_mutability_type(&self) -> IndexMutability { + match self { + Self::Mutable(_) => IndexMutability::Mutable, + Self::Immutable(_) => IndexMutability::Immutable, + Self::Mmap(_) => IndexMutability::Immutable, + } + } + + pub fn get_storage_type(&self) -> StorageType { + match self { + Self::Mutable(index) => index.storage_type(), + Self::Immutable(index) => index.storage_type(), + Self::Mmap(index) => index.storage_type(), + } + } +} diff --git a/lib/segment/src/index/field_index/map_index/tests.rs b/lib/segment/src/index/field_index/map_index/tests.rs new file mode 100644 index 0000000000..63a8040ce6 --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/tests.rs @@ -0,0 +1,452 @@ +use std::borrow::Borrow; +use std::collections::HashSet; +use std::hint::black_box; +use std::path::Path; + +use common::bitvec::BitVec; +use common::counter::hardware_counter::HardwareCounterCell; +use common::types::PointOffsetType; +use ecow::EcoString; +use gridstore::Blob; +use rstest::rstest; +use serde_json::Value; +use tempfile::Builder; + +use super::MapIndex; +use super::key::MapIndexKey; +use crate::index::field_index::{ + CardinalityEstimation, FieldIndexBuilderTrait, PayloadFieldIndex, PayloadFieldIndexRead, + ValueIndexer, +}; +use crate::types::{IntPayloadType, PayloadKeyType, UuidIntType}; + +/// Generous default size for the deleted-points bitslice used in tests. +/// +/// Must be larger than the stored mmap deletion bitslice for any test in +/// this file (which is sized to the highest point id, rounded up to a +/// `usize` boundary). 4096 bits comfortably covers all current tests. +const TEST_DELETED_BITS: usize = 4096; + +/// All-zero deletion bitslice for tests that don't care about deletions. +fn empty_deleted() -> BitVec { + BitVec::repeat(false, TEST_DELETED_BITS) +} + +/// Deletion bitslice with specific points marked as deleted. +fn deleted_with(points: &[PointOffsetType]) -> BitVec { + let mut v = empty_deleted(); + for &p in points { + v.set(p as usize, true); + } + v +} + +#[derive(Clone, Copy, PartialEq, Debug)] +enum IndexType { + MutableGridstore, + Mmap, + RamMmap, +} + +fn save_map_index( + data: &[Vec<::Owned>], + path: &Path, + index_type: IndexType, + into_value: impl Fn(&::Owned) -> Value, +) where + N: MapIndexKey + ?Sized, + Vec<::Owned>: Blob + Send + Sync, + MapIndex: PayloadFieldIndex + ValueIndexer, + as ValueIndexer>::ValueType: Into<::Owned>, +{ + let hw_counter = HardwareCounterCell::new(); + + match index_type { + IndexType::MutableGridstore => { + let mut builder = MapIndex::::builder_gridstore(path.to_path_buf()); + builder.init().unwrap(); + for (idx, values) in data.iter().enumerate() { + let values: Vec = values.iter().map(&into_value).collect(); + let values: Vec<_> = values.iter().collect(); + builder + .add_point(idx as PointOffsetType, &values, &hw_counter) + .unwrap(); + } + builder.finalize().unwrap(); + } + IndexType::Mmap | IndexType::RamMmap => { + let mut builder = MapIndex::::builder_mmap(path, false, &empty_deleted()); + builder.init().unwrap(); + for (idx, values) in data.iter().enumerate() { + let values: Vec = values.iter().map(&into_value).collect(); + let values: Vec<_> = values.iter().collect(); + builder + .add_point(idx as PointOffsetType, &values, &hw_counter) + .unwrap(); + } + builder.finalize().unwrap(); + } + } +} + +fn load_map_index( + data: &[Vec<::Owned>], + path: &Path, + index_type: IndexType, +) -> MapIndex +where + Vec<::Owned>: Blob + Send + Sync, +{ + let index = match index_type { + IndexType::MutableGridstore => MapIndex::::new_gridstore(path.to_path_buf(), true) + .unwrap() + .unwrap(), + IndexType::Mmap => MapIndex::::new_mmap(path, true, &empty_deleted()) + .unwrap() + .unwrap(), + IndexType::RamMmap => MapIndex::::new_mmap(path, false, &empty_deleted()) + .unwrap() + .unwrap(), + }; + let hw_counter = HardwareCounterCell::new(); + for (idx, values) in data.iter().enumerate() { + let index_values: HashSet<::Owned> = index + .get_values(idx as PointOffsetType, &hw_counter) + .unwrap() + .map(|v| MapIndexKey::to_owned(v.as_ref())) + .collect(); + let index_values: HashSet<&N> = index_values.iter().map(|v| v.borrow()).collect(); + let check_values: HashSet<&N> = values.iter().map(|v| v.borrow()).collect(); + assert_eq!(index_values, check_values); + } + + index +} + +#[test] +fn test_uuid_payload_index() { + let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); + let mut builder = + MapIndex::::builder_mmap(temp_dir.path(), false, &empty_deleted()); + + builder.init().unwrap(); + + let hw_counter = HardwareCounterCell::new(); + + let uuid: Value = Value::String("baa56dfc-e746-4ec1-bf50-94822535a46c".to_string()); + + for idx in 0..100 { + builder + .add_point(idx as PointOffsetType, &[&uuid], &hw_counter) + .unwrap(); + } + + let index = builder.finalize().unwrap(); + + index + .for_each_payload_block(50, PayloadKeyType::new("test_uuid"), &mut |block| { + black_box(block); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn test_index_non_ascending_insertion() { + let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); + let mut builder = + MapIndex::::builder_mmap(temp_dir.path(), false, &empty_deleted()); + builder.init().unwrap(); + + let data = [vec![1, 2, 3, 4, 5, 6], vec![25], vec![10, 11]]; + + let hw_counter = HardwareCounterCell::new(); + + for (idx, values) in data.iter().enumerate().rev() { + let values: Vec = values.iter().map(|i| (*i).into()).collect(); + let values: Vec<_> = values.iter().collect(); + builder + .add_point(idx as PointOffsetType, &values, &hw_counter) + .unwrap(); + } + + let index = builder.finalize().unwrap(); + let hw_counter = HardwareCounterCell::new(); + for (idx, values) in data.iter().enumerate().rev() { + let res: Vec<_> = index + .get_values(idx as u32, &hw_counter) + .unwrap() + .map(|i| *i as i32) + .collect(); + assert_eq!(res, *values); + } +} + +#[rstest] +#[case(IndexType::MutableGridstore)] +#[case(IndexType::Mmap)] +#[case(IndexType::RamMmap)] +fn test_int_disk_map_index(#[case] index_type: IndexType) { + let data = vec![ + vec![1, 2, 3, 4, 5, 6], + vec![1, 2, 3, 4, 5, 6], + vec![13, 14, 15, 16, 17, 18], + vec![19, 20, 21, 22, 23, 24], + vec![25], + ]; + + let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); + save_map_index::(&data, temp_dir.path(), index_type, |v| (*v).into()); + let index = load_map_index::(&data, temp_dir.path(), index_type); + + let hw_counter = HardwareCounterCell::new(); + + assert!( + !index + .except_cardinality(std::iter::empty(), &hw_counter) + .equals_min_exp_max(&CardinalityEstimation::exact(0)) + ); +} + +#[rstest] +#[case(IndexType::MutableGridstore)] +#[case(IndexType::Mmap)] +#[case(IndexType::RamMmap)] +fn test_string_disk_map_index(#[case] index_type: IndexType) { + let data = vec![ + vec![ + EcoString::from("AABB"), + EcoString::from("UUFF"), + EcoString::from("IIBB"), + ], + vec![ + EcoString::from("PPMM"), + EcoString::from("QQXX"), + EcoString::from("YYBB"), + ], + vec![ + EcoString::from("FFMM"), + EcoString::from("IICC"), + EcoString::from("IIBB"), + ], + vec![ + EcoString::from("AABB"), + EcoString::from("UUFF"), + EcoString::from("IIBB"), + ], + vec![EcoString::from("PPGG")], + ]; + + let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); + save_map_index::(&data, temp_dir.path(), index_type, |v| v.to_string().into()); + let index = load_map_index::(&data, temp_dir.path(), index_type); + + let hw_counter = HardwareCounterCell::new(); + + assert!( + !index + .except_cardinality(vec![].into_iter(), &hw_counter) + .equals_min_exp_max(&CardinalityEstimation::exact(0)) + ); +} + +#[rstest] +#[case(IndexType::MutableGridstore)] +#[case(IndexType::Mmap)] +#[case(IndexType::RamMmap)] +fn test_empty_index(#[case] index_type: IndexType) { + let data: Vec> = vec![]; + + let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); + save_map_index::(&data, temp_dir.path(), index_type, |v| v.to_string().into()); + let index = load_map_index::(&data, temp_dir.path(), index_type); + + let hw_counter = HardwareCounterCell::new(); + + assert!( + index + .except_cardinality(std::iter::empty(), &hw_counter) + .equals_min_exp_max(&CardinalityEstimation::exact(0)) + ); +} + +/// Test that `get_values` on an on-disk mmap index actually increments the hardware counter. +#[test] +fn test_mmap_get_values_hw_counter() { + let data = vec![vec![1i64, 2, 3], vec![4, 5], vec![6]]; + + let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); + save_map_index::(&data, temp_dir.path(), IndexType::Mmap, |v| (*v).into()); + let index = load_map_index::(&data, temp_dir.path(), IndexType::Mmap); + + let hw_counter = HardwareCounterCell::new(); + for idx in 0..data.len() { + let _values: Vec<_> = index + .get_values(idx as PointOffsetType, &hw_counter) + .unwrap() + .collect(); + } + + assert!( + hw_counter.payload_index_io_read_counter().get() > 0, + "Expected on-disk mmap get_values to track payload index IO reads, but counter was 0" + ); + + let temp_dir2 = Builder::new().prefix("store_dir").tempdir().unwrap(); + save_map_index::(&data, temp_dir2.path(), IndexType::RamMmap, |v| (*v).into()); + let index2 = load_map_index::(&data, temp_dir2.path(), IndexType::RamMmap); + + let hw_counter2 = HardwareCounterCell::new(); + for idx in 0..data.len() { + let _values: Vec<_> = index2 + .get_values(idx as PointOffsetType, &hw_counter2) + .unwrap() + .collect(); + } + + assert_eq!( + hw_counter2.payload_index_io_read_counter().get(), + 0, + "Expected RAM mmap get_values NOT to track IO reads, but counter was non-zero" + ); +} + +/// Reload contract: runtime deletions are not persisted by the mmap map +/// index. Callers must re-supply the deletion bitslice on reload. +/// +/// Test data is chosen so that every value retains at least one live +/// point after deletions — otherwise `ImmutableMapIndex::open_mmap` hits a +/// pre-existing debug-only assertion when a value's slice becomes empty. +#[rstest] +#[case(IndexType::MutableGridstore)] +#[case(IndexType::Mmap)] +#[case(IndexType::RamMmap)] +fn test_map_index_reload(#[case] index_type: IndexType) { + let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); + let data: Vec> = vec![ + vec![1, 2], // id 0 + vec![1], // id 1 + vec![2], // id 2 + vec![1, 3], // id 3 + vec![2, 3], // id 4 + vec![3], // id 5 + ]; + + { + save_map_index::(&data, temp_dir.path(), index_type, |v| (*v).into()); + let mut index = load_map_index::(&data, temp_dir.path(), index_type); + index.remove_point(1).unwrap(); + index.remove_point(2).unwrap(); + index.remove_point(5).unwrap(); + index.flusher()().unwrap(); + assert_eq!(index.get_indexed_points(), 3); + drop(index); + } + + let deleted = deleted_with(&[1, 2, 5]); + let new_index = match index_type { + IndexType::MutableGridstore => { + MapIndex::::new_gridstore(temp_dir.path().to_path_buf(), true) + .unwrap() + .unwrap() + } + IndexType::Mmap => MapIndex::::new_mmap(temp_dir.path(), true, &deleted) + .unwrap() + .unwrap(), + IndexType::RamMmap => { + MapIndex::::new_mmap(temp_dir.path(), false, &deleted) + .unwrap() + .unwrap() + } + }; + + assert_eq!(new_index.get_indexed_points(), 3); + + let hw_counter = HardwareCounterCell::new(); + for id in [1u32, 2, 5] { + assert_eq!( + new_index.values_count(id), + 0, + "deleted point {id} should have no values after reload", + ); + } + for id in [0u32, 3, 4] { + assert!( + new_index.values_count(id) > 0, + "live point {id} should have values after reload", + ); + } + + let mut hits: Vec = new_index.get_iterator(&1, &hw_counter).collect(); + hits.sort(); + assert_eq!(hits, vec![0, 3]); + + let mut hits: Vec = new_index.get_iterator(&2, &hw_counter).collect(); + hits.sort(); + assert_eq!(hits, vec![0, 4]); + + let mut hits: Vec = new_index.get_iterator(&3, &hw_counter).collect(); + hits.sort(); + assert_eq!(hits, vec![3, 4]); +} + +/// Regression test: when reloading an mmap map index with a `deleted_points` +/// bitslice shorter than `point_to_values.len()`, missing entries must +/// default to live, not deleted. Empty-payload bits from the on-disk +/// `deleted.bin` and any deletions encoded inside the short bitslice must +/// still be honored. +#[rstest] +#[case(IndexType::Mmap)] +#[case(IndexType::RamMmap)] +fn test_map_index_reload_short_deleted_bitslice(#[case] index_type: IndexType) { + let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); + + let data: Vec> = vec![ + vec![1], // id 0 + vec![1, 2], // id 1 + vec![], // id 2 — empty payload + vec![2, 3], // id 3 + vec![3], // id 4 + ]; + + save_map_index::(&data, temp_dir.path(), index_type, |v| (*v).into()); + + let mut short_deleted = BitVec::repeat(false, 2); + short_deleted.set(1, true); + + let new_index = match index_type { + IndexType::Mmap => { + MapIndex::::new_mmap(temp_dir.path(), true, &short_deleted) + .unwrap() + .unwrap() + } + IndexType::RamMmap => { + MapIndex::::new_mmap(temp_dir.path(), false, &short_deleted) + .unwrap() + .unwrap() + } + IndexType::MutableGridstore => unreachable!(), + }; + + let hw_counter = HardwareCounterCell::new(); + + assert!(new_index.values_count(0) > 0, "id 0 should be live"); + assert_eq!(new_index.values_count(1), 0, "id 1 deleted via bitslice"); + assert_eq!( + new_index.values_count(2), + 0, + "id 2 deleted via build-time empty" + ); + assert!( + new_index.values_count(3) > 0, + "id 3 should be live (beyond bitslice)" + ); + assert!( + new_index.values_count(4) > 0, + "id 4 should be live (beyond bitslice)" + ); + + let mut hits: Vec = new_index.get_iterator(&2, &hw_counter).collect(); + hits.sort(); + assert_eq!(hits, vec![3]); +}