refactor and apply to filter as id iterator

This commit is contained in:
Luis Cossío
2026-06-17 18:20:49 -04:00
parent af867b8b63
commit 47207f6b25
9 changed files with 219 additions and 70 deletions
@@ -70,9 +70,12 @@ where
&self,
values: impl Iterator<Item = FacetValue>,
hw_counter: &HardwareCounterCell,
f: impl FnMut(FacetValue, &mut dyn Iterator<Item = PointOffsetType>) -> OperationResult<()>,
mut f: impl FnMut(FacetValue, &mut dyn Iterator<Item = PointOffsetType>) -> 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::<FacetValueRef>::into(key).to_owned(), ids)
})
}
fn for_each_count_per_value(
@@ -153,9 +156,12 @@ where
&self,
values: impl Iterator<Item = FacetValue>,
hw_counter: &HardwareCounterCell,
f: impl FnMut(FacetValue, &mut dyn Iterator<Item = PointOffsetType>) -> OperationResult<()>,
mut f: impl FnMut(FacetValue, &mut dyn Iterator<Item = PointOffsetType>) -> 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::<FacetValueRef>::into(key).to_owned(), ids)
})
}
fn for_each_count_per_value(
@@ -16,9 +16,9 @@ pub trait MapIndexKey: Key + StoredValue + Eq + Display + Debug {
fn to_owned(&self) -> <Self as MapIndexKey>::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<<Self as MapIndexKey>::Owned>;
fn gridstore_block_size() -> usize {
size_of::<<Self as MapIndexKey>::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<<Self as MapIndexKey>::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<<Self as MapIndexKey>::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<<Self as MapIndexKey>::Owned> {
match value {
FacetValue::Uuid(uuid) => Some(uuid),
FacetValue::Keyword(_) | FacetValue::Int(_) | FacetValue::Bool(_) => None,
@@ -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<N: MapIndexKey + Key + ?Sized, S: UniversalRead> MapIndexRead<N> 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<V: Borrow<N>>(
&self,
values: impl Iterator<Item = FacetValue>,
values: impl Iterator<Item = V>,
hw_counter: &HardwareCounterCell,
mut f: impl FnMut(FacetValue, &mut dyn Iterator<Item = PointOffsetType>) -> OperationResult<()>,
mut f: impl FnMut(&N, &mut dyn Iterator<Item = PointOffsetType>) -> 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<FacetValue> = 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<V> = 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<N: MapIndexKey + Key + ?Sized, S: UniversalRead> MapIndexRead<N> 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<N> + 'a>(
&'a self,
values: impl Iterator<Item = V> + 'a,
hw_counter: &'a HardwareCounterCell,
) -> OperationResult<IdIter<'a>>
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)
}
@@ -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<IntPayloadType>>(
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,
@@ -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<str>>(
index: &'a T,
condition: &'a FieldCondition,
hw_counter: &'a HardwareCounterCell,
) -> OperationResult<Option<Box<dyn Iterator<Item = PointOffsetType> + 'a>>> {
let result: Option<Box<dyn Iterator<Item = PointOffsetType> + 'a>> = match &condition.r#match {
) -> OperationResult<Option<IdIter<'a>>> {
let result: Option<IdIter<'a>> = 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<str>>(
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()))
@@ -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<UuidIntType>>(
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() {
@@ -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<V: Borrow<N>>(
&self,
values: impl Iterator<Item = FacetValue>,
values: impl Iterator<Item = V>,
hw_counter: &HardwareCounterCell,
f: impl FnMut(FacetValue, &mut dyn Iterator<Item = PointOffsetType>) -> OperationResult<()>,
f: impl FnMut(&N, &mut dyn Iterator<Item = PointOffsetType>) -> 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<N> + 'a>(
&'a self,
values: impl Iterator<Item = V> + 'a,
hw_counter: &'a HardwareCounterCell,
) -> OperationResult<IdIter<'a>>
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(),
@@ -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<N: MapIndexKey + ?Sized> {
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<V: Borrow<N>>(
&self,
values: impl Iterator<Item = FacetValue>,
values: impl Iterator<Item = V>,
hw_counter: &HardwareCounterCell,
mut f: impl FnMut(FacetValue, &mut dyn Iterator<Item = PointOffsetType>) -> OperationResult<()>,
mut f: impl FnMut(&N, &mut dyn Iterator<Item = PointOffsetType>) -> 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<N> + 'a>(
&'a self,
values: impl Iterator<Item = V> + 'a,
hw_counter: &'a HardwareCounterCell,
) -> OperationResult<IdIter<'a>>
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<V: Borrow<N>>(
&self,
values: impl Iterator<Item = FacetValue>,
values: impl Iterator<Item = V>,
hw_counter: &HardwareCounterCell,
f: impl FnMut(FacetValue, &mut dyn Iterator<Item = PointOffsetType>) -> OperationResult<()>,
f: impl FnMut(&N, &mut dyn Iterator<Item = PointOffsetType>) -> 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<N> + 'a>(
&'a self,
values: impl Iterator<Item = V> + 'a,
hw_counter: &'a HardwareCounterCell,
) -> OperationResult<IdIter<'a>>
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(),
@@ -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<IntPayloadType>,
keys: &[IntPayloadType],
) -> BTreeMap<IntPayloadType, Vec<PointOffsetType>> {
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<PointOffsetType> = 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<IntPayloadType, Vec<PointOffsetType>> = 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::<IntPayloadType>(&data, temp_dir.path(), index_type, |v| (*v).into());
let index = load_map_index::<IntPayloadType>(&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::<IntPayloadType>(&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::<IntPayloadType>::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<IntPayloadType, Vec<PointOffsetType>> = [
(1, vec![3]),
(2, vec![1]),
(3, vec![1]),
(4, vec![1]),
(5, vec![3]),
]
.into_iter()
.collect();
assert_eq!(result, expected);
}