From 12cc1e50e6f8ded2550a306c7c1f99ea4fa024da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20Coss=C3=ADo?= Date: Fri, 29 May 2026 16:12:17 -0400 Subject: [PATCH] [UIO] Universal load numeric index (#9234) * rename module, use read_via * propagate S to `UniversalNumericIndex` --- .../benches/numeric_index_check_values.rs | 13 +++++-- .../src/index/field_index/histogram.rs | 9 +++-- .../field_index/numeric_index/builders.rs | 4 +- .../immutable_numeric_index/lifecycle.rs | 2 +- .../immutable_numeric_index/mod.rs | 2 +- .../index/field_index/numeric_index/mod.rs | 2 +- .../mutable_numeric_index/lifecycle.rs | 2 +- .../numeric_index/storage/lifecycle.rs | 5 ++- .../field_index/numeric_index/storage/mod.rs | 2 +- .../numeric_index/storage/read_only/mod.rs | 2 +- .../lifecycle.rs | 37 ++++++++++++------- .../mod.rs | 0 .../read_ops.rs | 0 .../field_index/tests/histogram_tests.rs | 3 +- 14 files changed, 52 insertions(+), 31 deletions(-) rename lib/segment/src/index/field_index/numeric_index/{mmap_numeric_index => universal_numeric_index}/lifecycle.rs (89%) rename lib/segment/src/index/field_index/numeric_index/{mmap_numeric_index => universal_numeric_index}/mod.rs (100%) rename lib/segment/src/index/field_index/numeric_index/{mmap_numeric_index => universal_numeric_index}/read_ops.rs (100%) diff --git a/lib/segment/benches/numeric_index_check_values.rs b/lib/segment/benches/numeric_index_check_values.rs index e37d21df54..50c4268e2c 100644 --- a/lib/segment/benches/numeric_index_check_values.rs +++ b/lib/segment/benches/numeric_index_check_values.rs @@ -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(|| { diff --git a/lib/segment/src/index/field_index/histogram.rs b/lib/segment/src/index/field_index/histogram.rs index cffd575c34..e9a983cb06 100644 --- a/lib/segment/src/index/field_index/histogram.rs +++ b/lib/segment/src/index/field_index/histogram.rs @@ -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 Histogram { } } - pub fn load(path: &Path) -> OperationResult { + pub fn load_via(fs: &Fs, path: &Path) -> OperationResult { 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, Counts)> = read_bin(&borders_path)?; + let histogram_config: HistogramConfig = read_json_via(fs, &config_path)?; + let histogram_buckets: Vec<(Point, Counts)> = read_bin_via(fs, &borders_path)?; Ok(Self { max_bucket_size: histogram_config.max_bucket_size, diff --git a/lib/segment/src/index/field_index/numeric_index/builders.rs b/lib/segment/src/index/field_index/numeric_index/builders.rs index c058978bdf..80c2883e77 100644 --- a/lib/segment/src/index/field_index/numeric_index/builders.rs +++ b/lib/segment/src/index/field_index/numeric_index/builders.rs @@ -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 { let inner = UniversalNumericIndex::build( + &MmapFs, self.in_memory_index, &self.path, self.is_on_disk, diff --git a/lib/segment/src/index/field_index/numeric_index/immutable_numeric_index/lifecycle.rs b/lib/segment/src/index/field_index/numeric_index/immutable_numeric_index/lifecycle.rs index 298b10d157..f76124c7c7 100644 --- a/lib/segment/src/index/field_index/numeric_index/immutable_numeric_index/lifecycle.rs +++ b/lib/segment/src/index/field_index/numeric_index/immutable_numeric_index/lifecycle.rs @@ -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; diff --git a/lib/segment/src/index/field_index/numeric_index/immutable_numeric_index/mod.rs b/lib/segment/src/index/field_index/numeric_index/immutable_numeric_index/mod.rs index a5682ba307..024f6ccd35 100644 --- a/lib/segment/src/index/field_index/numeric_index/immutable_numeric_index/mod.rs +++ b/lib/segment/src/index/field_index/numeric_index/immutable_numeric_index/mod.rs @@ -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}; diff --git a/lib/segment/src/index/field_index/numeric_index/mod.rs b/lib/segment/src/index/field_index/numeric_index/mod.rs index 4730b4309d..3f38e60c2c 100644 --- a/lib/segment/src/index/field_index/numeric_index/mod.rs +++ b/lib/segment/src/index/field_index/numeric_index/mod.rs @@ -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; diff --git a/lib/segment/src/index/field_index/numeric_index/mutable_numeric_index/lifecycle.rs b/lib/segment/src/index/field_index/numeric_index/mutable_numeric_index/lifecycle.rs index a07632710a..d4de8212c0 100644 --- a/lib/segment/src/index/field_index/numeric_index/mutable_numeric_index/lifecycle.rs +++ b/lib/segment/src/index/field_index/numeric_index/mutable_numeric_index/lifecycle.rs @@ -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}; diff --git a/lib/segment/src/index/field_index/numeric_index/storage/lifecycle.rs b/lib/segment/src/index/field_index/numeric_index/storage/lifecycle.rs index 276853222d..881736afa5 100644 --- a/lib/segment/src/index/field_index/numeric_index/storage/lifecycle.rs +++ b/lib/segment/src/index/field_index/numeric_index/storage/lifecycle.rs @@ -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); diff --git a/lib/segment/src/index/field_index/numeric_index/storage/mod.rs b/lib/segment/src/index/field_index/numeric_index/storage/mod.rs index 5cda920e24..5be5512c35 100644 --- a/lib/segment/src/index/field_index/numeric_index/storage/mod.rs +++ b/lib/segment/src/index/field_index/numeric_index/storage/mod.rs @@ -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; diff --git a/lib/segment/src/index/field_index/numeric_index/storage/read_only/mod.rs b/lib/segment/src/index/field_index/numeric_index/storage/read_only/mod.rs index 9eac6be29e..ca7ead78eb 100644 --- a/lib/segment/src/index/field_index/numeric_index/storage/read_only/mod.rs +++ b/lib/segment/src/index/field_index/numeric_index/storage/read_only/mod.rs @@ -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; diff --git a/lib/segment/src/index/field_index/numeric_index/mmap_numeric_index/lifecycle.rs b/lib/segment/src/index/field_index/numeric_index/universal_numeric_index/lifecycle.rs similarity index 89% rename from lib/segment/src/index/field_index/numeric_index/mmap_numeric_index/lifecycle.rs rename to lib/segment/src/index/field_index/numeric_index/universal_numeric_index/lifecycle.rs index f3ec8f2fd3..a3bd02f187 100644 --- a/lib/segment/src/index/field_index/numeric_index/mmap_numeric_index/lifecycle.rs +++ b/lib/segment/src/index/field_index/numeric_index/universal_numeric_index/lifecycle.rs @@ -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 UniversalNumericIndex { +impl UniversalNumericIndex +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, path: &Path, is_on_disk: bool, @@ -48,8 +56,8 @@ impl Univers in_memory_index.histogram.save(path)?; - StoredPointToValues::::from_iter( - &MmapFs, + StoredPointToValues::::from_iter( + fs, path, in_memory_index .point_to_values @@ -101,13 +109,14 @@ impl 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 Univers return Ok(None); } - let histogram = Histogram::::load(path)?; - let config: UniversalNumericIndexConfig = read_json(&config_path)?; + let histogram = Histogram::::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 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::::open( + fs, &deleted_path, OpenOptions { writeable: true, @@ -145,7 +154,7 @@ impl Univers populate: Populate::Auto, advice: AdviceSetting::Global, }, - (), + Default::default(), )?; let deleted_payloads_bitslice = deleted_payload_mmap.read_all()?; diff --git a/lib/segment/src/index/field_index/numeric_index/mmap_numeric_index/mod.rs b/lib/segment/src/index/field_index/numeric_index/universal_numeric_index/mod.rs similarity index 100% rename from lib/segment/src/index/field_index/numeric_index/mmap_numeric_index/mod.rs rename to lib/segment/src/index/field_index/numeric_index/universal_numeric_index/mod.rs diff --git a/lib/segment/src/index/field_index/numeric_index/mmap_numeric_index/read_ops.rs b/lib/segment/src/index/field_index/numeric_index/universal_numeric_index/read_ops.rs similarity index 100% rename from lib/segment/src/index/field_index/numeric_index/mmap_numeric_index/read_ops.rs rename to lib/segment/src/index/field_index/numeric_index/universal_numeric_index/read_ops.rs diff --git a/lib/segment/src/index/field_index/tests/histogram_tests.rs b/lib/segment/src/index/field_index/tests/histogram_tests.rs index 52195b404b..b0b144e8b3 100644 --- a/lib/segment/src/index/field_index/tests/histogram_tests.rs +++ b/lib/segment/src/index/field_index/tests/histogram_tests.rs @@ -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::::load(dir.path()).unwrap(); + let loaded_histogram = Histogram::::load_via(&MmapFs, dir.path()).unwrap(); assert_eq!(histogram, loaded_histogram); }