diff --git a/lib/segment/src/index/field_index/map_index/facet_index_impl.rs b/lib/segment/src/index/field_index/map_index/facet_index_impl.rs index 2a417932b4..e93ce9db29 100644 --- a/lib/segment/src/index/field_index/map_index/facet_index_impl.rs +++ b/lib/segment/src/index/field_index/map_index/facet_index_impl.rs @@ -70,9 +70,12 @@ where &self, values: impl Iterator, hw_counter: &HardwareCounterCell, - f: impl FnMut(FacetValue, &mut dyn Iterator) -> OperationResult<()>, + mut f: impl FnMut(FacetValue, &mut dyn Iterator) -> OperationResult<()>, ) -> OperationResult<()> { - MapIndexRead::for_values_map(self, values, hw_counter, f) + let keys = values.filter_map(N::from_facet_value); + MapIndexRead::for_values_map(self, keys, hw_counter, |key, ids| { + f(Into::::into(key).to_owned(), ids) + }) } fn for_each_count_per_value( @@ -153,9 +156,12 @@ where &self, values: impl Iterator, hw_counter: &HardwareCounterCell, - f: impl FnMut(FacetValue, &mut dyn Iterator) -> OperationResult<()>, + mut f: impl FnMut(FacetValue, &mut dyn Iterator) -> OperationResult<()>, ) -> OperationResult<()> { - MapIndexRead::for_values_map(self, values, hw_counter, f) + let keys = values.filter_map(N::from_facet_value); + MapIndexRead::for_values_map(self, keys, hw_counter, |key, ids| { + f(Into::::into(key).to_owned(), ids) + }) } fn for_each_count_per_value( diff --git a/lib/segment/src/index/field_index/map_index/key.rs b/lib/segment/src/index/field_index/map_index/key.rs index cb122f848a..9c8fdc8772 100644 --- a/lib/segment/src/index/field_index/map_index/key.rs +++ b/lib/segment/src/index/field_index/map_index/key.rs @@ -16,9 +16,9 @@ pub trait MapIndexKey: Key + StoredValue + Eq + Display + Debug { fn to_owned(&self) -> ::Owned; - /// Borrow this key type out of a [`FacetValue`], or `None` if the variant - /// doesn't match (e.g. a keyword value against an integer index). - fn from_facet_value(value: &FacetValue) -> Option<&Self>; + /// Convert a [`FacetValue`] into this key's owned type, or `None` if the + /// variant doesn't match (e.g. a keyword value against an integer index). + fn from_facet_value(value: FacetValue) -> Option<::Owned>; fn gridstore_block_size() -> usize { size_of::<::Owned>() @@ -38,9 +38,9 @@ impl MapIndexKey for str { EcoString::from(self) } - fn from_facet_value(value: &FacetValue) -> Option<&Self> { + fn from_facet_value(value: FacetValue) -> Option<::Owned> { match value { - FacetValue::Keyword(keyword) => Some(keyword.as_str()), + FacetValue::Keyword(keyword) => Some(EcoString::from(keyword)), FacetValue::Uuid(_) | FacetValue::Int(_) | FacetValue::Bool(_) => None, } } @@ -65,7 +65,7 @@ impl MapIndexKey for IntPayloadType { *self } - fn from_facet_value(value: &FacetValue) -> Option<&Self> { + fn from_facet_value(value: FacetValue) -> Option<::Owned> { match value { FacetValue::Int(int) => Some(int), FacetValue::Keyword(_) | FacetValue::Uuid(_) | FacetValue::Bool(_) => None, @@ -80,7 +80,7 @@ impl MapIndexKey for UuidIntType { *self } - fn from_facet_value(value: &FacetValue) -> Option<&Self> { + fn from_facet_value(value: FacetValue) -> Option<::Owned> { match value { FacetValue::Uuid(uuid) => Some(uuid), FacetValue::Keyword(_) | FacetValue::Int(_) | FacetValue::Bool(_) => None, diff --git a/lib/segment/src/index/field_index/map_index/on_disk_map_index/read_ops.rs b/lib/segment/src/index/field_index/map_index/on_disk_map_index/read_ops.rs index 81f61058ce..6cc9dbbf36 100644 --- a/lib/segment/src/index/field_index/map_index/on_disk_map_index/read_ops.rs +++ b/lib/segment/src/index/field_index/map_index/on_disk_map_index/read_ops.rs @@ -1,4 +1,4 @@ -use std::borrow::Cow; +use std::borrow::{Borrow, Cow}; use std::iter; use common::bitvec::BitSliceExt; @@ -9,12 +9,12 @@ use common::persisted_hashmap::{Key, READ_ENTRY_OVERHEAD}; use common::types::PointOffsetType; use common::universal_io::UniversalRead; use itertools::Itertools; +use roaring::RoaringBitmap; use super::super::read_ops::MapIndexRead; use super::super::{IdIter, MapIndexKey}; use super::OnDiskMapIndex; use crate::common::operation_error::OperationResult; -use crate::data_types::facets::FacetValue; use crate::index::field_index::on_disk_point_to_values::ValuesIter; use crate::index::payload_config::StorageType; @@ -158,27 +158,26 @@ impl MapIndexRead for OnDisk } /// Batched override of [`MapIndexRead::for_values_map`]. - fn for_values_map( + /// + /// `f` may be called in any order, since batched reads can complete out of + /// order. + fn for_values_map>( &self, - values: impl Iterator, + values: impl Iterator, hw_counter: &HardwareCounterCell, - mut f: impl FnMut(FacetValue, &mut dyn Iterator) -> OperationResult<()>, + mut f: impl FnMut(&N, &mut dyn Iterator) -> OperationResult<()>, ) -> OperationResult<()> { let hw_counter = ConditionedCounter::always(hw_counter); - // Materialize the values into a stable buffer so the keys we borrow from - // them stay valid for the whole batched read, which may reorder requests. - let values: Vec = values.collect(); - - // Build `(value index, key)` requests, skipping values whose variant - // doesn't match this index's key type (mirrors the default impl). - let requests = values.iter().enumerate().filter_map(|(value_idx, value)| { - N::from_facet_value(value).map(|key| (value_idx, key)) + let values: Vec = values.collect(); + let requests = values.iter().map(|value| { + let key: &N = value.borrow(); + (key, key) }); self.storage .value_to_points - .for_each_entry_in_iter(requests, |value_idx, point_ids| { + .for_each_entry_in_iter(requests, |key, point_ids| { // Mirror `get_iterator`'s IO accounting. let io_read = match point_ids { Some(ids) => size_of_val(ids) + READ_ENTRY_OVERHEAD, @@ -196,10 +195,30 @@ impl MapIndexRead for OnDisk .unwrap_or(false) }); - f(values[value_idx].clone(), &mut ids) + f(key, &mut ids) }) } + /// Batched override of [`MapIndexRead::iter_for_values`]. + /// + /// Resolves every value's posting in a single batched read and collects the + /// union into a [`RoaringBitmap`] + fn iter_for_values<'a, V: Borrow + 'a>( + &'a self, + values: impl Iterator + 'a, + hw_counter: &'a HardwareCounterCell, + ) -> OperationResult> + where + N: 'a, + { + let mut ids = RoaringBitmap::new(); + self.for_values_map(values, hw_counter, |_value, posting| { + ids.extend(&mut *posting); + Ok(()) + })?; + Ok(Box::new(ids.into_iter())) + } + fn for_each_value(&self, f: impl FnMut(&N) -> OperationResult<()>) -> OperationResult<()> { self.storage.value_to_points.for_each_key(f) } 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 index b9deef857f..760339f9de 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 @@ -5,7 +5,6 @@ use common::counter::hardware_counter::HardwareCounterCell; use common::types::PointOffsetType; use common::universal_io::UniversalRead; use gridstore::Blob; -use itertools::Itertools; use super::super::MapIndex; use super::super::key::MapIndexKey; @@ -148,12 +147,9 @@ fn filter_impl<'a, T: MapIndexRead>( None } } - AnyVariants::Integers(integers) => Some(Box::new( - integers - .iter() - .flat_map(move |integer| index.get_iterator(integer, hw_counter)) - .unique(), - )), + AnyVariants::Integers(integers) => { + Some(index.iter_for_values(integers.iter(), hw_counter)?) + } }, Some(Match::Except(MatchExcept { except })) => match except { AnyVariants::Strings(_) => None, 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 index d4cb1ccb58..d1b3be6be7 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,6 @@ use common::counter::hardware_counter::HardwareCounterCell; use common::types::PointOffsetType; use common::universal_io::UniversalRead; use gridstore::Blob; -use itertools::Itertools; use super::super::MapIndex; use super::super::key::MapIndexKey; @@ -14,6 +13,7 @@ use super::super::read_only::ReadOnlyMapIndex; use super::super::read_ops::MapIndexRead; use crate::common::Flusher; use crate::common::operation_error::OperationResult; +use crate::index::field_index::map_index::IdIter; use crate::index::field_index::{ CardinalityEstimation, PayloadBlockCondition, PayloadFieldIndex, PayloadFieldIndexRead, PrimaryCondition, @@ -133,8 +133,8 @@ fn filter_impl<'a, T: MapIndexRead>( index: &'a T, condition: &'a FieldCondition, hw_counter: &'a HardwareCounterCell, -) -> OperationResult + 'a>>> { - let result: Option + 'a>> = match &condition.r#match { +) -> OperationResult>> { + let result: Option> = match &condition.r#match { Some(Match::Value(MatchValue { value })) => match value { ValueVariants::String(keyword) => { Some(Box::new(index.get_iterator(keyword.as_str(), hw_counter))) @@ -143,12 +143,9 @@ fn filter_impl<'a, T: MapIndexRead>( ValueVariants::Bool(_) => None, }, Some(Match::Any(MatchAny { any: any_variant })) => match any_variant { - AnyVariants::Strings(keywords) => Some(Box::new( - keywords - .iter() - .flat_map(move |keyword| index.get_iterator(keyword.as_str(), hw_counter)) - .unique(), - )), + AnyVariants::Strings(keywords) => { + Some(index.iter_for_values(keywords.iter().map(AsRef::as_ref), hw_counter)?) + } AnyVariants::Integers(integers) => { if integers.is_empty() { Some(Box::new(iter::empty())) 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 index 1bbef101f5..bd94cc31eb 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 @@ -8,7 +8,6 @@ use common::types::PointOffsetType; use common::universal_io::UniversalRead; use gridstore::Blob; use indexmap::IndexSet; -use itertools::Itertools; use uuid::Uuid; use super::super::MapIndex; @@ -157,12 +156,7 @@ fn filter_impl<'a, T: MapIndexRead>( return Ok(None); }; - Some(Box::new( - uuids - .into_iter() - .flat_map(move |uuid| index.get_iterator(&uuid, hw_counter)) - .unique(), - )) + Some(index.iter_for_values(uuids.into_iter(), hw_counter)?) } AnyVariants::Integers(integers) => { if integers.is_empty() { diff --git a/lib/segment/src/index/field_index/map_index/read_only/read_ops.rs b/lib/segment/src/index/field_index/map_index/read_only/read_ops.rs index b4169b54ab..f6d4a44301 100644 --- a/lib/segment/src/index/field_index/map_index/read_only/read_ops.rs +++ b/lib/segment/src/index/field_index/map_index/read_only/read_ops.rs @@ -1,4 +1,4 @@ -use std::borrow::Cow; +use std::borrow::{Borrow, Cow}; use common::counter::hardware_counter::HardwareCounterCell; use common::persisted_hashmap::Key; @@ -10,7 +10,6 @@ use super::super::read_ops::MapIndexRead; use super::super::{IdIter, MapIndexKey}; use super::ReadOnlyMapIndex; use crate::common::operation_error::OperationResult; -use crate::data_types::facets::FacetValue; use crate::index::payload_config::StorageType; /// Dispatcher impl: forwards every [`MapIndexRead`] method to the active @@ -140,11 +139,11 @@ where } // Dispatch instead of using default impl, for on-disk impl to use batched reads - fn for_values_map( + fn for_values_map>( &self, - values: impl Iterator, + values: impl Iterator, hw_counter: &HardwareCounterCell, - f: impl FnMut(FacetValue, &mut dyn Iterator) -> OperationResult<()>, + f: impl FnMut(&N, &mut dyn Iterator) -> OperationResult<()>, ) -> OperationResult<()> { match self { ReadOnlyMapIndex::Appendable(index) => index.for_values_map(values, hw_counter, f), @@ -153,6 +152,22 @@ where } } + // Dispatch instead of using default impl, for on-disk impl to use batched reads + fn iter_for_values<'a, V: Borrow + 'a>( + &'a self, + values: impl Iterator + 'a, + hw_counter: &'a HardwareCounterCell, + ) -> OperationResult> + where + N: 'a, + { + match self { + ReadOnlyMapIndex::Appendable(index) => index.iter_for_values(values, hw_counter), + ReadOnlyMapIndex::Immutable(index) => index.iter_for_values(values, hw_counter), + ReadOnlyMapIndex::OnDisk(index) => index.iter_for_values(values, hw_counter), + } + } + fn storage_type(&self) -> StorageType { match self { ReadOnlyMapIndex::Appendable(index) => index.storage_type(), 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 index 31abc669f9..77333c060f 100644 --- a/lib/segment/src/index/field_index/map_index/read_ops.rs +++ b/lib/segment/src/index/field_index/map_index/read_ops.rs @@ -5,11 +5,11 @@ use common::counter::hardware_counter::HardwareCounterCell; use common::types::PointOffsetType; use gridstore::Blob; use indexmap::IndexSet; +use itertools::Itertools; use super::key::MapIndexKey; use super::{IdIter, MapIndex}; use crate::common::operation_error::OperationResult; -use crate::data_types::facets::FacetValue; use crate::index::field_index::CardinalityEstimation; use crate::index::field_index::stat_tools::number_of_selected_points; use crate::index::payload_config::{IndexMutability, StorageType}; @@ -103,26 +103,37 @@ pub trait MapIndexRead { self.values_count(idx).unwrap_or(0) == 0 } - /// Backs [`FacetIndex::for_values_map`]: yield each value's posting, - /// skipping values whose variant doesn't match this index's key type. - /// - /// [`FacetIndex::for_values_map`]: crate::index::field_index::FacetIndex::for_values_map - fn for_values_map( + /// Use each value's posting of point ids, invoking `f` once per value. + fn for_values_map>( &self, - values: impl Iterator, + values: impl Iterator, hw_counter: &HardwareCounterCell, - mut f: impl FnMut(FacetValue, &mut dyn Iterator) -> OperationResult<()>, + mut f: impl FnMut(&N, &mut dyn Iterator) -> OperationResult<()>, ) -> OperationResult<()> { for value in values { - let Some(key) = N::from_facet_value(&value) else { - continue; - }; - let mut ids = self.get_iterator(key, hw_counter); + let value = value.borrow(); + let mut ids = self.get_iterator(value, hw_counter); f(value, &mut ids)?; } Ok(()) } + /// Iterator over deduplicated points matching **any** of the given `values`. + fn iter_for_values<'a, V: Borrow + 'a>( + &'a self, + values: impl Iterator + 'a, + hw_counter: &'a HardwareCounterCell, + ) -> OperationResult> + where + N: 'a, + { + Ok(Box::new( + values + .flat_map(move |value| self.get_iterator(value.borrow(), hw_counter)) + .unique(), + )) + } + fn match_cardinality( &self, value: &N, @@ -337,11 +348,11 @@ where } // Dispatch instead of using default impl, for on-disk impl to use batched reads - fn for_values_map( + fn for_values_map>( &self, - values: impl Iterator, + values: impl Iterator, hw_counter: &HardwareCounterCell, - f: impl FnMut(FacetValue, &mut dyn Iterator) -> OperationResult<()>, + f: impl FnMut(&N, &mut dyn Iterator) -> OperationResult<()>, ) -> OperationResult<()> { match self { MapIndex::Mutable(index) => index.for_values_map(values, hw_counter, f), @@ -350,6 +361,22 @@ where } } + // Dispatch instead of using default impl, for on-disk impl to use batched reads + fn iter_for_values<'a, V: Borrow + 'a>( + &'a self, + values: impl Iterator + 'a, + hw_counter: &'a HardwareCounterCell, + ) -> OperationResult> + where + N: 'a, + { + match self { + MapIndex::Mutable(index) => index.iter_for_values(values, hw_counter), + MapIndex::Immutable(index) => index.iter_for_values(values, hw_counter), + MapIndex::OnDisk(index) => index.iter_for_values(values, hw_counter), + } + } + fn storage_type(&self) -> StorageType { match self { MapIndex::Mutable(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 index 50dc947996..d5460a583d 100644 --- a/lib/segment/src/index/field_index/map_index/tests.rs +++ b/lib/segment/src/index/field_index/map_index/tests.rs @@ -1,5 +1,5 @@ use std::borrow::Borrow; -use std::collections::HashSet; +use std::collections::{BTreeMap, HashSet}; use std::hint::black_box; use std::path::Path; @@ -458,3 +458,98 @@ fn test_map_index_reload_short_deleted_bitslice(#[case] index_type: IndexType) { hits.sort(); assert_eq!(hits, vec![3]); } + +/// Run `for_values_map` for the given keys and collect `value -> sorted ids`. +/// Batched reads (the on-disk variant) may invoke the callback out of order, so +/// we sort each posting and key by a `BTreeMap`. +fn collect_for_values_map( + index: &MapIndex, + keys: &[IntPayloadType], +) -> BTreeMap> { + let hw_counter = HardwareCounterCell::new(); + let mut out = BTreeMap::new(); + MapIndexRead::for_values_map(index, keys.iter(), &hw_counter, |key, ids| { + let mut ids: Vec = ids.collect(); + ids.sort_unstable(); + out.insert(*key, ids); + Ok(()) + }) + .unwrap(); + out +} + +/// The on-disk variant overrides `for_values_map` with a batched read; assert it +/// agrees with the in-memory variants (which use the default per-key impl) and +/// with the directly-computed expectation. +#[test] +fn test_for_values_map_congruence() { + // Points 0..5; value 3 spans three points, value 6 only one. + #[rustfmt::skip] + let data = vec![ + vec![1, 2, 3], + vec![2, 3, 4], + vec![3], + vec![6], + vec![1, 5], + ]; + // Present values (multi- and single-point) plus an absent one (99), which + // must still be reported with an empty posting. + let query = [1, 2, 3, 4, 5, 99]; + + let expected: BTreeMap> = query + .iter() + .map(|&q| { + let ids = data + .iter() + .enumerate() + .filter(|(_, vals)| vals.contains(&q)) + .map(|(i, _)| i as PointOffsetType) + .collect(); + (q, ids) + }) + .collect(); + + for index_type in [ + IndexType::MutableGridstore, + IndexType::Mmap, + IndexType::RamMmap, + ] { + 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 result = collect_for_values_map(&index, &query); + assert_eq!(result, expected, "mismatch for {index_type:?}"); + } +} + +/// The batched on-disk `for_values_map` must drop deleted points from postings, +/// mirroring `get_iterator`. +#[test] +fn test_for_values_map_on_disk_deleted() { + let data = vec![vec![1, 2, 3], vec![2, 3, 4], vec![3], vec![1, 5]]; + + let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); + save_map_index::(&data, temp_dir.path(), IndexType::Mmap, |v| (*v).into()); + + // Load the on-disk variant with points 0 and 2 marked deleted. + let deleted = deleted_with(&[0, 2]); + let index = MapIndex::::new_immutable(temp_dir.path(), true, &deleted) + .unwrap() + .unwrap(); + assert!(matches!(index, MapIndex::OnDisk(_))); + + let result = collect_for_values_map(&index, &[1, 2, 3, 4, 5]); + + // Postings with deleted points {0, 2} removed. + let expected: BTreeMap> = [ + (1, vec![3]), + (2, vec![1]), + (3, vec![1]), + (4, vec![1]), + (5, vec![3]), + ] + .into_iter() + .collect(); + assert_eq!(result, expected); +}