[UIO] Universal load numeric index (#9234)

* rename module, use read_via

* propagate S to `UniversalNumericIndex<T, S>`
This commit is contained in:
Luis Cossío
2026-05-29 16:12:17 -04:00
committed by timvisee
parent 779a66f783
commit 12cc1e50e6
14 changed files with 52 additions and 31 deletions

View File

@@ -1,13 +1,14 @@
use common::bitvec::BitVec;
use common::counter::hardware_counter::HardwareCounterCell;
use common::types::PointOffsetType;
use common::universal_io::{MmapFile, MmapFs};
use criterion::{Criterion, criterion_group, criterion_main};
use rand::prelude::StdRng;
use rand::{RngExt, SeedableRng};
use segment::common::operation_error::OperationResult;
use segment::index::field_index::numeric_index::NumericIndexRead;
use segment::index::field_index::numeric_index::mmap_numeric_index::UniversalNumericIndex;
use segment::index::field_index::numeric_index::mutable_numeric_index::InMemoryNumericIndex;
use segment::index::field_index::numeric_index::universal_numeric_index::UniversalNumericIndex;
use tempfile::Builder;
mod prof;
@@ -56,8 +57,14 @@ pub fn struct_numeric_check_values(c: &mut Criterion) {
})
});
let mmap_index =
UniversalNumericIndex::build(mutable_index, dir.path(), false, &deleted_points).unwrap();
let mmap_index = UniversalNumericIndex::<_, MmapFile>::build(
&MmapFs,
mutable_index,
dir.path(),
false,
&deleted_points,
)
.unwrap();
group.bench_function("mmap-numeric-index", |b| {
b.iter(|| {

View File

@@ -4,8 +4,9 @@ use std::fmt::Debug;
use std::ops::Bound;
use std::path::{Path, PathBuf};
use common::fs::{atomic_save_bin, atomic_save_json, read_bin, read_json};
use common::fs::{atomic_save_bin, atomic_save_json};
use common::types::PointOffsetType;
use common::universal_io::{UniversalReadFs, read_bin_via, read_json_via};
use itertools::Itertools;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
@@ -51,12 +52,12 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
}
}
pub fn load(path: &Path) -> OperationResult<Self> {
pub fn load_via<Fs: UniversalReadFs>(fs: &Fs, path: &Path) -> OperationResult<Self> {
let config_path = path.join(CONFIG_PATH);
let borders_path = path.join(BORDERS_PATH);
let histogram_config: HistogramConfig = read_json(&config_path)?;
let histogram_buckets: Vec<(Point<T>, Counts)> = read_bin(&borders_path)?;
let histogram_config: HistogramConfig = read_json_via(fs, &config_path)?;
let histogram_buckets: Vec<(Point<T>, Counts)> = read_bin_via(fs, &borders_path)?;
Ok(Self {
max_bucket_size: histogram_config.max_bucket_size,

View File

@@ -4,12 +4,13 @@ use std::path::PathBuf;
use common::bitvec::BitVec;
use common::counter::hardware_counter::HardwareCounterCell;
use common::types::PointOffsetType;
use common::universal_io::MmapFs;
use gridstore::Blob;
use serde_json::Value;
use super::mmap_numeric_index::UniversalNumericIndex;
use super::mutable_numeric_index::InMemoryNumericIndex;
use super::storage::NumericIndexInner;
use super::universal_numeric_index::UniversalNumericIndex;
use super::{Encodable, NumericIndex, NumericIndexIntoInnerValue};
use crate::common::operation_error::{OperationError, OperationResult};
use crate::index::field_index::numeric_point::Numericable;
@@ -123,6 +124,7 @@ where
fn finalize(self) -> OperationResult<Self::FieldIndexType> {
let inner = UniversalNumericIndex::build(
&MmapFs,
self.in_memory_index,
&self.path,
self.is_on_disk,

View File

@@ -5,8 +5,8 @@ use common::types::PointOffsetType;
use gridstore::Blob;
use super::super::Encodable;
use super::super::mmap_numeric_index::UniversalNumericIndex;
use super::super::mutable_numeric_index::InMemoryNumericIndex;
use super::super::universal_numeric_index::UniversalNumericIndex;
use super::{ImmutableNumericIndex, NumericKeySortedVec};
use crate::common::Flusher;
use crate::common::operation_error::OperationResult;

View File

@@ -3,7 +3,7 @@ use std::ops::Bound;
use common::bitvec::{BitSliceExt as _, BitVec};
use super::Encodable;
use super::mmap_numeric_index::UniversalNumericIndex;
use super::universal_numeric_index::UniversalNumericIndex;
use crate::index::field_index::histogram::Histogram;
use crate::index::field_index::immutable_point_to_values::ImmutablePointToValues;
use crate::index::field_index::numeric_point::{Numericable, Point};

View File

@@ -2,7 +2,6 @@ mod builders;
mod encodable;
pub mod immutable_numeric_index;
mod lifecycle;
pub mod mmap_numeric_index;
pub mod mutable_numeric_index;
mod numeric_field_index;
pub mod numeric_index_read;
@@ -10,6 +9,7 @@ mod query;
mod read_only;
mod read_ops;
mod storage;
pub mod universal_numeric_index;
mod value_indexer;
use std::marker::PhantomData;

View File

@@ -9,8 +9,8 @@ use gridstore::{Blob, Gridstore};
use super::super::Encodable;
use super::super::lifecycle::{HISTOGRAM_MAX_BUCKET_SIZE, HISTOGRAM_PRECISION};
use super::super::mmap_numeric_index::UniversalNumericIndex;
use super::super::numeric_index_read::NumericIndexRead;
use super::super::universal_numeric_index::UniversalNumericIndex;
use super::{InMemoryNumericIndex, MutableNumericIndex, default_gridstore_options};
use crate::common::Flusher;
use crate::common::operation_error::{OperationError, OperationResult};

View File

@@ -5,12 +5,13 @@ use std::path::{Path, PathBuf};
use common::bitvec::BitSlice;
use common::types::PointOffsetType;
use common::universal_io::MmapFs;
use gridstore::Blob;
use super::super::Encodable;
use super::super::immutable_numeric_index::ImmutableNumericIndex;
use super::super::mmap_numeric_index::UniversalNumericIndex;
use super::super::mutable_numeric_index::MutableNumericIndex;
use super::super::universal_numeric_index::UniversalNumericIndex;
use super::NumericIndexInner;
use crate::common::Flusher;
use crate::common::operation_error::OperationResult;
@@ -35,7 +36,7 @@ where
is_on_disk || common::low_memory::low_memory_mode().prefer_disk();
let Some(mmap_index) =
UniversalNumericIndex::open(path, effective_is_on_disk, deleted_points)?
UniversalNumericIndex::open(&MmapFs, path, effective_is_on_disk, deleted_points)?
else {
// Files don't exist, cannot load
return Ok(None);

View File

@@ -30,8 +30,8 @@ use gridstore::Blob;
use super::Encodable;
use super::immutable_numeric_index::ImmutableNumericIndex;
use super::mmap_numeric_index::UniversalNumericIndex;
use super::mutable_numeric_index::MutableNumericIndex;
use super::universal_numeric_index::UniversalNumericIndex;
use crate::index::field_index::numeric_point::Numericable;
use crate::index::field_index::stored_point_to_values::StoredValue;

View File

@@ -2,8 +2,8 @@ use common::universal_io::UniversalRead;
use gridstore::Blob;
use super::super::Encodable;
use super::super::mmap_numeric_index::UniversalNumericIndex;
use super::super::mutable_numeric_index::read_only::ReadOnlyAppendableNumericIndex;
use super::super::universal_numeric_index::UniversalNumericIndex;
use crate::index::field_index::numeric_point::Numericable;
use crate::index::field_index::stored_point_to_values::StoredValue;

View File

@@ -3,11 +3,13 @@ use std::ops::BitOrAssign;
use std::path::{Path, PathBuf};
use common::bitvec::{BitSlice, BitSliceExt};
use common::fs::{atomic_save_json, clear_disk_cache, read_json};
use common::fs::{atomic_save_json, clear_disk_cache};
use common::mmap::{AdviceSetting, MmapSlice, create_and_ensure_length};
use common::stored_bitslice::MmapBitSlice;
use common::stored_bitslice::{MmapBitSlice, StoredBitSlice};
use common::types::PointOffsetType;
use common::universal_io::{MmapFile, MmapFs, OpenOptions, Populate, TypedStorage};
use common::universal_io::{
MmapFs, OpenOptions, Populate, TypedStorage, UniversalRead, read_json_via,
};
use fs_err as fs;
use memmap2::MmapMut;
use serde::{Deserialize, Serialize};
@@ -26,8 +28,14 @@ struct UniversalNumericIndexConfig {
max_values_per_point: usize,
}
impl<T: Encodable + Numericable + Default + StoredValue + bytemuck::Pod> UniversalNumericIndex<T> {
impl<T, S> UniversalNumericIndex<T, S>
where
T: Encodable + Numericable + Default + StoredValue + bytemuck::Pod,
S: UniversalRead,
{
/// TODO: save using `S::Fs` too
pub fn build(
fs: &S::Fs,
in_memory_index: InMemoryNumericIndex<T>,
path: &Path,
is_on_disk: bool,
@@ -48,8 +56,8 @@ impl<T: Encodable + Numericable + Default + StoredValue + bytemuck::Pod> Univers
in_memory_index.histogram.save(path)?;
StoredPointToValues::<T, MmapFile>::from_iter(
&MmapFs,
StoredPointToValues::<T, S>::from_iter(
fs,
path,
in_memory_index
.point_to_values
@@ -101,13 +109,14 @@ impl<T: Encodable + Numericable + Default + StoredValue + bytemuck::Pod> Univers
deleted.flusher()()?;
}
Self::open(path, is_on_disk, deleted_points)?.ok_or_else(|| {
Self::open(fs, path, is_on_disk, deleted_points)?.ok_or_else(|| {
OperationError::service_error("Failed to open UniversalNumericIndex after building it")
})
}
/// Open and load mmap numeric index from the given path
pub fn open(
fs: &S::Fs,
path: &Path,
is_on_disk: bool,
deleted_points: &BitSlice,
@@ -121,8 +130,8 @@ impl<T: Encodable + Numericable + Default + StoredValue + bytemuck::Pod> Univers
return Ok(None);
}
let histogram = Histogram::<T>::load(path)?;
let config: UniversalNumericIndexConfig = read_json(&config_path)?;
let histogram = Histogram::<T>::load_via(fs, path)?;
let config: UniversalNumericIndexConfig = read_json_via(fs, &config_path)?;
let do_populate = !is_on_disk;
let pairs_options = OpenOptions {
@@ -131,13 +140,13 @@ impl<T: Encodable + Numericable + Default + StoredValue + bytemuck::Pod> Univers
populate: Populate::from(do_populate),
advice: AdviceSetting::Global,
};
let pairs = TypedStorage::open(&MmapFs, pairs_path, pairs_options, ())?;
let pairs = TypedStorage::open(fs, pairs_path, pairs_options, Default::default())?;
let point_to_values = StoredPointToValues::open(&MmapFs, path, do_populate)?;
let point_to_values = StoredPointToValues::open(fs, path, do_populate)?;
let mut deleted = deleted_points.to_owned();
let deleted_payload_mmap = MmapBitSlice::open(
&MmapFs,
let deleted_payload_mmap = StoredBitSlice::<S>::open(
fs,
&deleted_path,
OpenOptions {
writeable: true,
@@ -145,7 +154,7 @@ impl<T: Encodable + Numericable + Default + StoredValue + bytemuck::Pod> Univers
populate: Populate::Auto,
advice: AdviceSetting::Global,
},
(),
Default::default(),
)?;
let deleted_payloads_bitslice = deleted_payload_mmap.read_all()?;

View File

@@ -2,6 +2,7 @@ use std::cell::Cell;
use std::collections::BTreeSet;
use std::collections::Bound::{Excluded, Included, Unbounded};
use common::universal_io::MmapFs;
use itertools::Itertools;
use rand::prelude::StdRng;
use rand::{RngExt, SeedableRng};
@@ -278,6 +279,6 @@ fn test_save_load_histogram() {
.unwrap();
histogram.save(dir.path()).unwrap();
let loaded_histogram = Histogram::<f64>::load(dir.path()).unwrap();
let loaded_histogram = Histogram::<f64>::load_via(&MmapFs, dir.path()).unwrap();
assert_eq!(histogram, loaded_histogram);
}