Integrate hnsw_with_vectors

xzfc · xzfc · commit 85b456e3e28b · 2025-09-08T23:28:59.000Z
diff --git a/docs/redoc/master/openapi.json b/docs/redoc/master/openapi.json
@@ -11753,6 +11753,16 @@
             "description": "Use appendable quantization in appendable plain segments.",
             "default": false,
             "type": "boolean"
+          },
+          "hnsw_with_vectors": {
+            "description": "Allow writing the HNSW graph in the `CompressedWithVectors` format.",
+            "default": false,
+            "type": "boolean"
+          },
+          "hnsw_format_force": {
+            "description": "Forcefully enable or forcefully disable the `CompressedWithVectors` format, based on the value of [`Self::hnsw_with_vectors`].\n\nWhen set, the graph will be converted on startup to the corresponding format (`Compressed` or `CompressedWithVectors`).",
+            "default": false,
+            "type": "boolean"
           }
         }
       },
diff --git a/lib/common/common/src/flags.rs b/lib/common/common/src/flags.rs
@@ -59,6 +59,16 @@ pub struct FeatureFlags {
     /// Use appendable quantization in appendable plain segments.
     // TODO(1.16.0): enable by default
     pub appendable_quantization: bool,
+
+    /// Allow writing the HNSW graph in the `CompressedWithVectors` format.
+    pub hnsw_with_vectors: bool,
+
+    /// Forcefully enable or forcefully disable the `CompressedWithVectors`
+    /// format, based on the value of [`Self::hnsw_with_vectors`].
+    ///
+    /// When set, the graph will be converted on startup to the corresponding
+    /// format (`Compressed` or `CompressedWithVectors`).
+    pub hnsw_format_force: bool,
 }
 
 impl Default for FeatureFlags {
@@ -74,6 +84,8 @@ impl Default for FeatureFlags {
             migrate_rocksdb_payload_storage: false,
             migrate_rocksdb_payload_indices: false,
             appendable_quantization: false,
+            hnsw_with_vectors: false,
+            hnsw_format_force: false,
         }
     }
 }
@@ -99,6 +111,8 @@ pub fn init_feature_flags(mut flags: FeatureFlags) {
         migrate_rocksdb_payload_storage,
         migrate_rocksdb_payload_indices,
         appendable_quantization,
+        hnsw_with_vectors,
+        hnsw_format_force,
     } = &mut flags;
 
     // If all is set, explicitly set all feature flags
@@ -112,6 +126,8 @@ pub fn init_feature_flags(mut flags: FeatureFlags) {
         *migrate_rocksdb_payload_storage = true;
         *migrate_rocksdb_payload_indices = true;
         *appendable_quantization = true;
+        *hnsw_with_vectors = true;
+        *hnsw_format_force = true;
     }
 
     let res = FEATURE_FLAGS.set(flags);
diff --git a/lib/segment/benches/fixture.rs b/lib/segment/benches/fixture.rs
@@ -53,7 +53,7 @@ where
         let updated_ago = updated_ago(&graph_layers_path).unwrap_or_else(|_| "???".to_string());
         eprintln!("Loading cached links (built {updated_ago} ago) from {graph_layers_path:?}.");
         eprintln!("Delete the directory above if code related to HNSW graph building is changed");
-        GraphLayers::load(&path, false, false).unwrap()
+        GraphLayers::load(&path, false, None).unwrap()
     } else {
         let mut graph_layers_builder =
             GraphLayersBuilder::new(num_vectors, HnswM::new2(m), ef_construct, 10, use_heuristic);
diff --git a/lib/segment/src/index/hnsw_index/graph_layers.rs b/lib/segment/src/index/hnsw_index/graph_layers.rs
@@ -7,6 +7,7 @@ use common::fixed_length_priority_queue::FixedLengthPriorityQueue;
 use common::types::{PointOffsetType, ScoredPointOffset};
 use io::file_operations::{atomic_save, read_bin};
 use itertools::Itertools;
+use memory::madvise::Advice;
 use serde::{Deserialize, Serialize};
 
 use super::HnswM;
@@ -500,72 +501,99 @@ impl GraphLayers {
 }
 
 impl GraphLayers {
-    pub fn load(dir: &Path, on_disk: bool, compress: bool) -> OperationResult<Self> {
+    pub fn load(
+        dir: &Path,
+        on_disk: bool,
+        force_format: Option<&GraphLinksFormatParam>,
+    ) -> OperationResult<Self> {
         let graph_data: GraphLayerData = read_bin(&GraphLayers::get_path(dir))?;
-
-        if compress {
-            Self::convert_to_compressed(dir, HnswM::new(graph_data.m, graph_data.m0))?;
-        }
-
+        let hnsw_m = HnswM::new(graph_data.m, graph_data.m0);
         Ok(Self {
-            hnsw_m: HnswM::new(graph_data.m, graph_data.m0),
-            links: Self::load_links(dir, on_disk)?,
+            hnsw_m,
+            links: Self::load_links(dir, on_disk, force_format, hnsw_m)?,
             entry_points: graph_data.entry_points.into_owned(),
             visited_pool: VisitedPool::new(),
         })
     }
 
-    fn load_links(dir: &Path, on_disk: bool) -> OperationResult<GraphLinks> {
-        for format in [
+    fn load_links(
+        dir: &Path,
+        on_disk: bool,
+        force_format: Option<&GraphLinksFormatParam>,
+        hnsw_m: HnswM,
+    ) -> OperationResult<GraphLinks> {
+        let formats = [
             GraphLinksFormat::CompressedWithVectors,
             GraphLinksFormat::Compressed,
             GraphLinksFormat::Plain,
-        ] {
-            let path = GraphLayers::get_links_path(dir, format);
-            if path.exists() {
-                return GraphLinks::load_from_file(&path, on_disk, format);
+        ];
+
+        match force_format {
+            // Normal flow - load the first available format, in order of preference.
+            None => {
+                for format in formats {
+                    let path = GraphLayers::get_links_path(dir, format);
+                    if path.exists() {
+                        return GraphLinks::load_from_file(&path, on_disk, format, Advice::Random);
+                    }
+                }
             }
-        }
-        Err(OperationError::service_error("No links file found"))
-    }
-
-    /// Convert the "plain" format into the "compressed" format.
-    /// Note: conversion into the "compressed with vectors" format is not
-    /// supported at the moment, though it is possible to implement.
-    /// As far as [`super::hnsw::LINK_COMPRESSION_CONVERT_EXISTING`] is false,
-    /// this code is not used in production.
-    fn convert_to_compressed(dir: &Path, hnsw_m: HnswM) -> OperationResult<()> {
-        let plain_path = Self::get_links_path(dir, GraphLinksFormat::Plain);
-        let compressed_path = Self::get_links_path(dir, GraphLinksFormat::Compressed);
-        let compressed_with_vectors_path =
-            Self::get_links_path(dir, GraphLinksFormat::CompressedWithVectors);
-
-        if compressed_path.exists() || compressed_with_vectors_path.exists() {
-            return Ok(());
-        }
+            // Forced format (tests/benchmarking only) - convert if necessary.
+            Some(force_format) => {
+                // Happy path - the file already exists
+                let path = GraphLayers::get_links_path(dir, force_format.as_format());
+                if path.exists() {
+                    return GraphLinks::load_from_file(
+                        &path,
+                        on_disk,
+                        force_format.as_format(),
+                        Advice::Random,
+                    );
+                }
 
-        let start = std::time::Instant::now();
-
-        let links = GraphLinks::load_from_file(&plain_path, true, GraphLinksFormat::Plain)?;
-        let original_size = plain_path.metadata()?.len();
-        atomic_save(&compressed_path, |writer| {
-            let edges = links.to_edges();
-            serialize_graph_links(edges, GraphLinksFormatParam::Compressed, hnsw_m, writer)
-        })?;
-        let new_size = compressed_path.metadata()?.len();
-
-        // Remove the original file
-        std::fs::remove_file(plain_path)?;
-
-        log::debug!(
-            "Compressed HNSW graph links in {:.1?}: {:.1}MB -> {:.1}MB ({:.1}%)",
-            start.elapsed(),
-            original_size as f64 / 1024.0 / 1024.0,
-            new_size as f64 / 1024.0 / 1024.0,
-            new_size as f64 / original_size as f64 * 100.0,
-        );
+                // Unhappy path - convert from another format.
+                for format in formats {
+                    let original_path = GraphLayers::get_links_path(dir, format);
+                    if format == force_format.as_format() || !original_path.exists() {
+                        continue;
+                    }
 
-        Ok(())
+                    let start = std::time::Instant::now();
+                    let edges = GraphLinks::load_from_file(
+                        &original_path,
+                        on_disk,
+                        format,
+                        Advice::Sequential,
+                    )?
+                    .to_edges();
+                    let original_size = original_path.metadata()?.len();
+                    atomic_save(&path, |writer| {
+                        serialize_graph_links(edges, *force_format, hnsw_m, writer)
+                    })?;
+                    let new_size = path.metadata()?.len();
+
+                    // NOTE: The original file is not removed.
+
+                    log::info!(
+                        "Converted HNSW graph links in {:.1?}: {:.1}MB -> {:.1}MB ({:.1}%)",
+                        start.elapsed(),
+                        original_size as f64 / 1024.0 / 1024.0,
+                        new_size as f64 / 1024.0 / 1024.0,
+                        new_size as f64 / original_size as f64 * 100.0,
+                    );
+
+                    return GraphLinks::load_from_file(
+                        &path,
+                        on_disk,
+                        force_format.as_format(),
+                        Advice::Random,
+                    );
+                }
+            }
+        }
+        Err(OperationError::service_error(format!(
+            "No HNSW graph links file found in {dir:?}"
+        )))
     }
 
     #[cfg(feature = "testing")]
@@ -684,12 +712,10 @@ mod tests {
     }
 
     #[rstest]
-    #[case::uncompressed((GraphLinksFormat::Plain, false))]
-    #[case::converted((GraphLinksFormat::Plain, true))]
-    #[case::compressed((GraphLinksFormat::Compressed, false))]
-    #[case::recompressed((GraphLinksFormat::Compressed, true))]
-    #[case::compressed_with_vectors((GraphLinksFormat::CompressedWithVectors, false))]
-    fn test_save_and_load(#[case] (initial_format, compress): (GraphLinksFormat, bool)) {
+    #[case::uncompressed(GraphLinksFormat::Plain)]
+    #[case::compressed(GraphLinksFormat::Compressed)]
+    #[case::compressed_with_vectors(GraphLinksFormat::CompressedWithVectors)]
+    fn test_save_and_load(#[case] initial_format: GraphLinksFormat) {
         let distance = Distance::Cosine;
         let num_vectors = 100;
         let dim = 8;
@@ -706,30 +732,45 @@ mod tests {
             M,
             dim,
             false,
-            initial_format.is_with_vectors(),
+            true,
             distance,
             &mut rng,
         );
+        let graph_links_vectors = vector_holder.graph_links_vectors();
         let graph1 = graph_layers_builder
             .into_graph_layers(
                 dir.path(),
-                initial_format.with_param_for_tests(vector_holder.graph_links_vectors().as_ref()),
+                initial_format.with_param_for_tests(graph_links_vectors.as_ref()),
                 true,
             )
             .unwrap();
         assert_eq!(graph1.links.format(), initial_format);
         let res1 = search_in_graph(&query, top, &vector_holder, &graph1);
         drop(graph1);
 
-        let graph2 = GraphLayers::load(dir.path(), false, compress).unwrap();
-        if compress {
-            assert_eq!(graph2.links.format(), GraphLinksFormat::Compressed);
-        } else {
-            assert_eq!(graph2.links.format(), initial_format);
-        }
-        let res2 = search_in_graph(&query, top, &vector_holder, &graph2);
+        for force_format in [
+            None,
+            Some(GraphLinksFormat::Plain),
+            Some(GraphLinksFormat::Compressed),
+            Some(GraphLinksFormat::CompressedWithVectors),
+        ] {
+            eprintln!("force_format = {force_format:?}");
+            let graph2 = GraphLayers::load(
+                dir.path(),
+                false,
+                force_format
+                    .map(|fmt| fmt.with_param(graph_links_vectors.as_ref()))
+                    .as_ref(),
+            )
+            .unwrap();
+            assert_eq!(
+                graph2.links.format(),
+                force_format.unwrap_or(initial_format)
+            );
+            let res2 = search_in_graph(&query, top, &vector_holder, &graph2);
 
-        assert_eq!(res1, res2)
+            assert_eq!(res1, res2)
+        }
     }
 
     #[rstest]
diff --git a/lib/segment/src/index/hnsw_index/graph_layers_builder.rs b/lib/segment/src/index/hnsw_index/graph_layers_builder.rs
@@ -8,6 +8,7 @@ use bitvec::prelude::BitVec;
 use common::ext::BitSliceExt;
 use common::types::{PointOffsetType, ScoredPointOffset};
 use io::file_operations::{atomic_save, atomic_save_bin};
+use memory::madvise::Advice;
 use parking_lot::{Mutex, MutexGuard, RwLock};
 use rand::Rng;
 use rand::distr::Uniform;
@@ -197,7 +198,12 @@ impl GraphLayersBuilder {
             atomic_save(&links_path, |writer| {
                 serialize_graph_links(edges, format_param, self.hnsw_m, writer)
             })?;
-            links = GraphLinks::load_from_file(&links_path, true, format_param.as_format())?;
+            links = GraphLinks::load_from_file(
+                &links_path,
+                true,
+                format_param.as_format(),
+                Advice::Random,
+            )?;
         } else {
             // Since we'll keep it in the RAM anyway, we can afford to build in the RAM too.
             links = GraphLinks::new_from_edges(edges, format_param, self.hnsw_m)?;
diff --git a/lib/segment/src/index/hnsw_index/graph_links.rs b/lib/segment/src/index/hnsw_index/graph_links.rs
diff --git a/lib/segment/src/index/hnsw_index/hnsw.rs b/lib/segment/src/index/hnsw_index/hnsw.rs
