[AI] implement ReadOnlyNumericIndex::preopen (#9743)

[AI] implement `ReadOnlyNumericIndex::preopen`

lib/segment/src/index/field_index/field_index_base/read_only/lifecycle.rs

@@ -88,16 +88,51 @@ impl<S: UniversalReadExt> ReadOnlyFieldIndex<S> {
                 }
             },
             PayloadIndexType::IntIndex => match mode {
-                ReadMode::Appendable => false,
-                ReadMode::Immutable { is_on_disk: _ } => false,
+                ReadMode::Appendable => {
+                    ReadOnlyNumericIndex::<IntPayloadType, IntPayloadType, S>::preopen_appendable(
+                        fs,
+                        numeric_dir(dir, field),
+                    )?
+                }
+                ReadMode::Immutable { is_on_disk } => {
+                    ReadOnlyNumericIndex::<IntPayloadType, IntPayloadType, S>::preopen_immutable(
+                        fs,
+                        &numeric_dir(dir, field),
+                        is_on_disk,
+                    )?
+                }
             },
             PayloadIndexType::DatetimeIndex => match mode {
-                ReadMode::Appendable => false,
-                ReadMode::Immutable { is_on_disk: _ } => false,
+                ReadMode::Appendable => ReadOnlyNumericIndex::<
+                    IntPayloadType,
+                    DateTimePayloadType,
+                    S,
+                >::preopen_appendable(
+                    fs, numeric_dir(dir, field)
+                )?,
+                ReadMode::Immutable { is_on_disk } => ReadOnlyNumericIndex::<
+                    IntPayloadType,
+                    DateTimePayloadType,
+                    S,
+                >::preopen_immutable(
+                    fs, &numeric_dir(dir, field), is_on_disk
+                )?,
             },
             PayloadIndexType::FloatIndex => match mode {
-                ReadMode::Appendable => false,
-                ReadMode::Immutable { is_on_disk: _ } => false,
+                ReadMode::Appendable => ReadOnlyNumericIndex::<
+                    FloatPayloadType,
+                    FloatPayloadType,
+                    S,
+                >::preopen_appendable(
+                    fs, numeric_dir(dir, field)
+                )?,
+                ReadMode::Immutable { is_on_disk } => ReadOnlyNumericIndex::<
+                    FloatPayloadType,
+                    FloatPayloadType,
+                    S,
+                >::preopen_immutable(
+                    fs, &numeric_dir(dir, field), is_on_disk
+                )?,
             },
             // Geo reuses the writable selector's `map_dir` (`-map` suffix).
             PayloadIndexType::GeoIndex => match mode {

lib/segment/src/index/field_index/histogram.rs

@@ -6,7 +6,7 @@ use std::path::{Path, PathBuf};
 
 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 common::universal_io::{CachedReadFs, UniversalReadFs, read_bin_via, read_json_via};
 use itertools::Itertools;
 use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
@@ -52,7 +52,14 @@ impl<T: Numericable + Serialize + DeserializeOwned> Histogram<T> {
         }
     }
 
-    pub fn load_universal<Fs: UniversalReadFs>(fs: &Fs, path: &Path) -> OperationResult<Self> {
+    /// Schedule background prefetch of the two files [`open`](Self::open) reads.
+    pub fn preopen(fs: &impl CachedReadFs, path: &Path) -> OperationResult<()> {
+        fs.schedule_prefetch(&path.join(CONFIG_PATH), None, None)?;
+        fs.schedule_prefetch(&path.join(BORDERS_PATH), None, None)?;
+        Ok(())
+    }
+
+    pub fn open<Fs: UniversalReadFs>(fs: &Fs, path: &Path) -> OperationResult<Self> {
         let config_path = path.join(CONFIG_PATH);
         let borders_path = path.join(BORDERS_PATH);
 

lib/segment/src/index/field_index/numeric_index/mutable_numeric_index/read_only/lifecycle.rs

@@ -1,7 +1,7 @@
 use std::path::PathBuf;
 
 use common::counter::hardware_counter::HardwareCounterCell;
-use common::universal_io::{OkNotFound, Populate, UniversalRead, UniversalReadFs};
+use common::universal_io::{CachedReadFs, OkNotFound, Populate, UniversalRead, UniversalReadFs};
 use gridstore::{Blob, GridstoreReader};
 
 use super::super::InMemoryNumericIndex;
@@ -15,6 +15,18 @@ impl<T: Encodable + Numericable + Send + Sync + Default, S: UniversalRead>
 where
     Vec<T>: Blob,
 {
+    /// Schedule background prefetch of the Gridstore files [`open`](Self::open)
+    /// will read. Returns whether the on-disk directory exists (`false` = the
+    /// index is not in the appendable format).
+    pub fn preopen(fs: &impl CachedReadFs<File = S>, dir: PathBuf) -> OperationResult<bool> {
+        // Gridstore reader
+        Ok(
+            GridstoreReader::<Vec<T>, S>::preopen(fs, dir, Populate::PreferBackground)
+                .ok_not_found()?
+                .is_some(),
+        )
+    }
+
     /// Open the appendable (Gridstore) numeric index read-only, threading every
     /// file open through the filesystem handle `fs`.
     ///

lib/segment/src/index/field_index/numeric_index/on_disk_numeric_index/lifecycle.rs

@@ -8,8 +8,8 @@ use common::mmap::{AdviceSetting, MmapSlice, create_and_ensure_length};
 use common::stored_bitslice::{MmapBitSlice, StoredBitSlice};
 use common::types::PointOffsetType;
 use common::universal_io::{
-    MmapFs, OkNotFound, OpenOptions, Populate, TypedStorage, UniversalRead, UniversalReadFs,
-    read_json_via,
+    CachedReadFs, MmapFs, OkNotFound, OpenOptions, Populate, TypedStorage, UniversalRead,
+    UniversalReadFs, read_json_via,
 };
 use fs_err as fs;
 use memmap2::MmapMut;
@@ -114,6 +114,56 @@ where
         })
     }
 
+    fn open_options(populate: Populate) -> OpenOptions {
+        OpenOptions {
+            writeable: false,
+            need_sequential: false,
+            populate,
+            advice: AdviceSetting::Global,
+        }
+    }
+
+    /// Schedule background prefetch of every file [`open`](Self::open) will read.
+    ///
+    /// Returns `false` (nothing scheduled) when the segment is not in the
+    /// on-disk format.
+    pub fn preopen(
+        fs: &impl CachedReadFs<File = S>,
+        path: &Path,
+        populate: Populate,
+    ) -> OperationResult<bool> {
+        // Config
+        let config_path = path.join(CONFIG_PATH);
+        if fs
+            .schedule_prefetch(&config_path, None, None)
+            .ok_not_found()?
+            .is_none()
+        {
+            // If config doesn't exist, assume the index doesn't exist on disk
+            return Ok(false);
+        }
+
+        // Histogram
+        Histogram::<T>::preopen(fs, path)?;
+
+        // Value pairs
+        let pairs_path = path.join(PAIRS_PATH);
+        fs.schedule_prefetch(&pairs_path, Some(Self::open_options(populate)), None)?;
+
+        // Point to values
+        OnDiskPointToValues::<T, S>::preopen(fs, path, populate)?;
+
+        // Deleted bitslice
+        let deleted_path = path.join(DELETED_PATH);
+        fs.schedule_prefetch(
+            &deleted_path,
+            Some(Self::open_options(Populate::PreferBackground)),
+            None,
+        )?;
+
+        Ok(true)
+    }
+
     /// Open and load mmap numeric index from the given path
     pub fn open(
         fs: &impl UniversalReadFs<File = S>,
@@ -132,15 +182,13 @@ where
             return Ok(None);
         };
 
-        let histogram = Histogram::<T>::load_universal(fs, path)?;
+        let histogram = Histogram::<T>::open(fs, path)?;
 
-        let pairs_options = OpenOptions {
-            writeable: false,
-            need_sequential: false,
-            populate,
-            advice: AdviceSetting::Global,
-        };
-        let pairs = TypedStorage::new(fs.open(&pairs_path, pairs_options, Default::default())?);
+        let pairs = TypedStorage::new(fs.open(
+            &pairs_path,
+            Self::open_options(populate),
+            Default::default(),
+        )?);
 
         let point_to_values = OnDiskPointToValues::open(fs, path, populate)?;
         let mut deleted = deleted_points.to_owned();
@@ -148,12 +196,7 @@ where
         let deleted_payload_mmap = StoredBitSlice::<S>::open(
             fs,
             &deleted_path,
-            OpenOptions {
-                writeable: false,
-                need_sequential: false,
-                populate: Populate::Auto,
-                advice: AdviceSetting::Global,
-            },
+            Self::open_options(Populate::Auto),
             Default::default(),
         )?;
         let deleted_payloads_bitslice = deleted_payload_mmap.read_all()?;

lib/segment/src/index/field_index/numeric_index/read_only/lifecycle.rs

@@ -2,7 +2,7 @@ use std::marker::PhantomData;
 use std::path::{Path, PathBuf};
 
 use common::bitvec::BitSlice;
-use common::universal_io::{UniversalRead, UniversalReadFs};
+use common::universal_io::{CachedReadFs, UniversalRead, UniversalReadFs};
 use gridstore::Blob;
 
 use super::super::Encodable;
@@ -18,6 +18,25 @@ impl<T: Encodable + Numericable + StoredValue + Send + Sync + Default, P, S: Uni
 where
     Vec<T>: Blob,
 {
+    /// Schedule background prefetch for the appendable (Gridstore) format,
+    /// forwarding to [`ReadOnlyNumericIndexInner::preopen_appendable`].
+    pub fn preopen_appendable(
+        fs: &impl CachedReadFs<File = S>,
+        dir: PathBuf,
+    ) -> OperationResult<bool> {
+        ReadOnlyNumericIndexInner::preopen_appendable(fs, dir)
+    }
+
+    /// Schedule background prefetch for the immutable (mmap) format, forwarding
+    /// to [`ReadOnlyNumericIndexInner::preopen_immutable`].
+    pub fn preopen_immutable(
+        fs: &impl CachedReadFs<File = S>,
+        path: &Path,
+        is_on_disk: bool,
+    ) -> OperationResult<bool> {
+        ReadOnlyNumericIndexInner::preopen_immutable(fs, path, is_on_disk)
+    }
+
     /// Read-only mirror of [`NumericIndex::new_gridstore`][1]: forwards to
     /// [`ReadOnlyNumericIndexInner::open_appendable`] and wraps the inner with
     /// the typed payload-value phantom `P`.

lib/segment/src/index/field_index/numeric_index/storage/read_only/lifecycle.rs

@@ -1,7 +1,7 @@
 use std::path::{Path, PathBuf};
 
 use common::bitvec::BitSlice;
-use common::universal_io::{Populate, UniversalRead, UniversalReadFs};
+use common::universal_io::{CachedReadFs, Populate, UniversalRead, UniversalReadFs};
 use gridstore::Blob;
 
 use super::super::super::Encodable;
@@ -19,6 +19,35 @@ impl<T: Encodable + Numericable + StoredValue + Send + Sync + Default, S: Univer
 where
     Vec<T>: Blob,
 {
+    /// Schedule background prefetch for the appendable (Gridstore) format.
+    ///
+    /// Returns `false` when nothing was scheduled (directory absent).
+    pub fn preopen_appendable(
+        fs: &impl CachedReadFs<File = S>,
+        dir: PathBuf,
+    ) -> OperationResult<bool> {
+        ReadOnlyAppendableNumericIndex::preopen(fs, dir)
+    }
+
+    /// Schedule background prefetch for the immutable (mmap) format.
+    ///
+    /// Returns `false` when the on-disk index doesn't exist.
+    pub fn preopen_immutable(
+        fs: &impl CachedReadFs<File = S>,
+        path: &Path,
+        is_on_disk: bool,
+    ) -> OperationResult<bool> {
+        let effective_is_on_disk =
+            is_on_disk || common::low_memory::low_memory_mode().prefer_disk();
+
+        let populate = match effective_is_on_disk {
+            true => Populate::No,
+            false => Populate::PreferBackground,
+        };
+
+        OnDiskNumericIndex::<T, S>::preopen(fs, path, populate)
+    }
+
     /// Read-only mirror of [`NumericIndexInner::new_gridstore`][1]: open the
     /// appendable (Gridstore-backed) numeric index read-only, threading every
     /// file open through the filesystem handle `fs`.

lib/segment/src/index/field_index/tests/histogram_tests.rs

@@ -280,6 +280,6 @@ fn test_save_load_histogram() {
         .unwrap();
     histogram.save(dir.path()).unwrap();
 
-    let loaded_histogram = Histogram::<f64>::load_universal(&MmapFs, dir.path()).unwrap();
+    let loaded_histogram = Histogram::<f64>::open(&MmapFs, dir.path()).unwrap();
     assert_eq!(histogram, loaded_histogram);
 }