[AI] implement ReadOnlyGeoIndex::preopen (#9744)

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

@@ -136,8 +136,12 @@ impl<S: UniversalReadExt> ReadOnlyFieldIndex<S> {
             },
             // Geo reuses the writable selector's `map_dir` (`-map` suffix).
             PayloadIndexType::GeoIndex => match mode {
-                ReadMode::Appendable => false,
-                ReadMode::Immutable { is_on_disk: _ } => false,
+                ReadMode::Appendable => {
+                    ReadOnlyGeoIndex::<S>::preopen_appendable(fs, map_dir(dir, field))?
+                }
+                ReadMode::Immutable { is_on_disk } => {
+                    ReadOnlyGeoIndex::<S>::preopen_immutable(fs, &map_dir(dir, field), is_on_disk)?
+                }
             },
             PayloadIndexType::FullTextIndex => match mode {
                 ReadMode::Appendable => false,

lib/segment/src/index/field_index/geo_index/mutable_geo_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::GridstoreReader;
 
 use super::super::inner::InMemoryGeoIndex;
@@ -10,6 +10,19 @@ use crate::common::operation_error::{OperationError, OperationResult};
 use crate::types::{GeoPoint, RawGeoPoint};
 
 impl<S: UniversalRead> ReadOnlyAppendableGeoIndex<S> {
+    /// Schedule background prefetch of the Gridstore files [`open`](Self::open)
+    /// will read.
+    ///
+    /// Returns whether the on-disk directory exists.
+    pub fn preopen(fs: &impl CachedReadFs<File = S>, dir: PathBuf) -> OperationResult<bool> {
+        // Gridstore reader
+        Ok(
+            GridstoreReader::<Vec<RawGeoPoint>, S>::preopen(fs, dir, Populate::PreferBackground)
+                .ok_not_found()?
+                .is_some(),
+        )
+    }
+
     /// Open the appendable (Gridstore) geo index read-only, threading every
     /// file open through the filesystem handle `fs`.
     ///

lib/segment/src/index/field_index/geo_index/on_disk_geo_index/lifecycle.rs

@@ -14,8 +14,8 @@ use common::mmap::{AdviceSetting, MmapSlice, create_and_ensure_length};
 use common::stored_bitslice::{MmapBitSlice, StoredBitSlice};
 use common::types::PointOffsetType;
 use common::universal_io::{
-    MmapFile, MmapFs, OkNotFound, OpenOptions, Populate, ReadRange, TypedStorage, UniversalRead,
-    UniversalReadFs, read_json_via,
+    CachedReadFs, MmapFile, MmapFs, OkNotFound, OpenOptions, Populate, ReadRange, TypedStorage,
+    UniversalRead, UniversalReadFs, read_json_via,
 };
 use fs_err as fs;
 use memmap2::MmapMut;
@@ -160,6 +160,54 @@ impl<S: UniversalRead> OnDiskGeoIndex<S> {
         })
     }
 
+    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` when the segment is not in the
+    /// on-disk format.
+    pub fn preopen(
+        fs: &impl CachedReadFs<File = S>,
+        path: &Path,
+        populate: Populate,
+    ) -> OperationResult<bool> {
+        // Stats
+        let stats_path = path.join(STATS_PATH);
+        if fs
+            .schedule_prefetch(&stats_path, None, None)
+            .ok_not_found()?
+            .is_none()
+        {
+            // If stats file doesn't exist, assume the index doesn't exist on disk
+            return Ok(false);
+        }
+
+        // Geohash counts, points map, and point-id list
+        let options = Self::open_options(populate);
+        fs.schedule_prefetch(&path.join(COUNTS_PER_HASH), Some(options), None)?;
+        fs.schedule_prefetch(&path.join(POINTS_MAP), Some(options), None)?;
+        fs.schedule_prefetch(&path.join(POINTS_MAP_IDS), Some(options), None)?;
+
+        // Point to values
+        OnDiskPointToValues::<GeoPoint, S>::preopen(fs, path, populate)?;
+
+        // Deleted bitslice
+        fs.schedule_prefetch(
+            &path.join(DELETED_PATH),
+            Some(Self::open_options(Populate::PreferBackground)),
+            None,
+        )?;
+
+        Ok(true)
+    }
+
     pub fn open(
         fs: &impl UniversalReadFs<File = S>,
         path: &Path,
@@ -178,12 +226,7 @@ impl<S: UniversalRead> OnDiskGeoIndex<S> {
             return Ok(None);
         };
 
-        let open_options = OpenOptions {
-            writeable: false,
-            need_sequential: false,
-            populate,
-            advice: AdviceSetting::Global,
-        };
+        let open_options = Self::open_options(populate);
 
         let counts_per_hash =
             TypedStorage::new(fs.open(&counts_per_hash_path, open_options, Default::default())?);
@@ -198,12 +241,7 @@ impl<S: UniversalRead> OnDiskGeoIndex<S> {
         let deleted_payload_mmap = StoredBitSlice::<S>::open(
             fs,
             &deleted_path,
-            OpenOptions {
-                writeable: false,
-                need_sequential: false,
-                populate,
-                advice: AdviceSetting::Global,
-            },
+            Self::open_options(populate),
             Default::default(),
         )?;
         let deleted_payloads_bitslice = deleted_payload_mmap.read_all()?;

lib/segment/src/index/field_index/geo_index/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 super::super::mutable_geo_index::read_only::ReadOnlyAppendableGeoIndex;
 use super::super::on_disk_geo_index::OnDiskGeoIndex;
@@ -10,6 +10,35 @@ use crate::common::operation_error::OperationResult;
 use crate::index::field_index::geo_index::immutable_geo_index::ImmutableGeoIndex;
 
 impl<S: UniversalRead> ReadOnlyGeoIndex<S> {
+    /// 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> {
+        ReadOnlyAppendableGeoIndex::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,
+        };
+
+        OnDiskGeoIndex::preopen(fs, path, populate)
+    }
+
     /// Read-only mirror of [`GeoIndex::new_mutable`][1]: open the
     /// appendable (Gridstore-backed) geo index read-only, threading every
     /// file open through the filesystem handle `fs`.