Skip to content

Commit 8095c51

Browse files
xzfcxzfc
committed
Integrate hnsw_with_vectors
1 parent acac786 commit 8095c51

6 files changed

Lines changed: 200 additions & 84 deletions

File tree

‎lib/common/common/src/flags.rs‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,16 @@ pub struct FeatureFlags {
5959
/// Use appendable quantization in appendable plain segments.
6060
// TODO(1.16.0): enable by default
6161
pub appendable_quantization: bool,
62+
63+
/// Allow writing the HNSW graph in the `CompressedWithVectors` format.
64+
pub hnsw_with_vectors: bool,
65+
66+
/// Forcefully enable or forcefully disable the `CompressedWithVectors`
67+
/// format, based on the value of [`Self::hnsw_with_vectors`].
68+
///
69+
/// When set, the graph will be converted on startup to the corresponding
70+
/// format (`Compressed` or `CompressedWithVectors`).
71+
pub hnsw_format_force: bool,
6272
}
6373

6474
impl Default for FeatureFlags {
@@ -74,6 +84,8 @@ impl Default for FeatureFlags {
7484
migrate_rocksdb_payload_storage: false,
7585
migrate_rocksdb_payload_indices: false,
7686
appendable_quantization: false,
87+
hnsw_with_vectors: false,
88+
hnsw_format_force: false,
7789
}
7890
}
7991
}
@@ -99,6 +111,8 @@ pub fn init_feature_flags(mut flags: FeatureFlags) {
99111
migrate_rocksdb_payload_storage,
100112
migrate_rocksdb_payload_indices,
101113
appendable_quantization,
114+
hnsw_with_vectors,
115+
hnsw_format_force,
102116
} = &mut flags;
103117

104118
// If all is set, explicitly set all feature flags
@@ -112,6 +126,8 @@ pub fn init_feature_flags(mut flags: FeatureFlags) {
112126
*migrate_rocksdb_payload_storage = true;
113127
*migrate_rocksdb_payload_indices = true;
114128
*appendable_quantization = true;
129+
*hnsw_with_vectors = true;
130+
*hnsw_format_force = true;
115131
}
116132

117133
let res = FEATURE_FLAGS.set(flags);

‎lib/segment/benches/fixture.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ where
5353
let updated_ago = updated_ago(&graph_layers_path).unwrap_or_else(|_| "???".to_string());
5454
eprintln!("Loading cached links (built {updated_ago} ago) from {graph_layers_path:?}.");
5555
eprintln!("Delete the directory above if code related to HNSW graph building is changed");
56-
GraphLayers::load(&path, false, false).unwrap()
56+
GraphLayers::load(&path, false, None).unwrap()
5757
} else {
5858
let mut graph_layers_builder =
5959
GraphLayersBuilder::new(num_vectors, HnswM::new2(m), ef_construct, 10, use_heuristic);

‎lib/segment/src/index/hnsw_index/graph_layers.rs‎

Lines changed: 108 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use common::fixed_length_priority_queue::FixedLengthPriorityQueue;
77
use common::types::{PointOffsetType, ScoredPointOffset};
88
use io::file_operations::{atomic_save, read_bin};
99
use itertools::Itertools;
10+
use memory::madvise::Advice;
1011
use serde::{Deserialize, Serialize};
1112

1213
use super::HnswM;
@@ -500,72 +501,97 @@ impl GraphLayers {
500501
}
501502

502503
impl GraphLayers {
503-
pub fn load(dir: &Path, on_disk: bool, compress: bool) -> OperationResult<Self> {
504+
pub fn load(
505+
dir: &Path,
506+
on_disk: bool,
507+
force_format: Option<&GraphLinksFormatParam>,
508+
) -> OperationResult<Self> {
504509
let graph_data: GraphLayerData = read_bin(&GraphLayers::get_path(dir))?;
505-
506-
if compress {
507-
Self::convert_to_compressed(dir, HnswM::new(graph_data.m, graph_data.m0))?;
508-
}
509-
510+
let hnsw_m = HnswM::new(graph_data.m, graph_data.m0);
510511
Ok(Self {
511-
hnsw_m: HnswM::new(graph_data.m, graph_data.m0),
512-
links: Self::load_links(dir, on_disk)?,
512+
hnsw_m,
513+
links: Self::load_links(dir, on_disk, force_format, hnsw_m)?,
513514
entry_points: graph_data.entry_points.into_owned(),
514515
visited_pool: VisitedPool::new(),
515516
})
516517
}
517518

518-
fn load_links(dir: &Path, on_disk: bool) -> OperationResult<GraphLinks> {
519-
for format in [
519+
fn load_links(
520+
dir: &Path,
521+
on_disk: bool,
522+
force_format: Option<&GraphLinksFormatParam>,
523+
hnsw_m: HnswM,
524+
) -> OperationResult<GraphLinks> {
525+
let formats = [
520526
GraphLinksFormat::CompressedWithVectors,
521527
GraphLinksFormat::Compressed,
522528
GraphLinksFormat::Plain,
523-
] {
524-
let path = GraphLayers::get_links_path(dir, format);
525-
if path.exists() {
526-
return GraphLinks::load_from_file(&path, on_disk, format);
529+
];
530+
531+
match force_format {
532+
// Normal flow - load the first available format, in order of preference.
533+
None => {
534+
for format in formats {
535+
let path = GraphLayers::get_links_path(dir, format);
536+
if path.exists() {
537+
return GraphLinks::load_from_file(&path, on_disk, format, Advice::Random);
538+
}
539+
}
527540
}
528-
}
529-
Err(OperationError::service_error("No links file found"))
530-
}
531-
532-
/// Convert the "plain" format into the "compressed" format.
533-
/// Note: conversion into the "compressed with vectors" format is not
534-
/// supported at the moment, though it is possible to implement.
535-
/// As far as [`super::hnsw::LINK_COMPRESSION_CONVERT_EXISTING`] is false,
536-
/// this code is not used in production.
537-
fn convert_to_compressed(dir: &Path, hnsw_m: HnswM) -> OperationResult<()> {
538-
let plain_path = Self::get_links_path(dir, GraphLinksFormat::Plain);
539-
let compressed_path = Self::get_links_path(dir, GraphLinksFormat::Compressed);
540-
let compressed_with_vectors_path =
541-
Self::get_links_path(dir, GraphLinksFormat::CompressedWithVectors);
542-
543-
if compressed_path.exists() || compressed_with_vectors_path.exists() {
544-
return Ok(());
545-
}
541+
// Forced format (tests/benchmarking only) - convert if necessary.
542+
Some(force_format) => {
543+
// Happy path - the file already exists
544+
let path = GraphLayers::get_links_path(dir, force_format.as_format());
545+
if path.exists() {
546+
return GraphLinks::load_from_file(
547+
&path,
548+
on_disk,
549+
force_format.as_format(),
550+
Advice::Random,
551+
);
552+
}
546553

547-
let start = std::time::Instant::now();
548-
549-
let links = GraphLinks::load_from_file(&plain_path, true, GraphLinksFormat::Plain)?;
550-
let original_size = plain_path.metadata()?.len();
551-
atomic_save(&compressed_path, |writer| {
552-
let edges = links.to_edges();
553-
serialize_graph_links(edges, GraphLinksFormatParam::Compressed, hnsw_m, writer)
554-
})?;
555-
let new_size = compressed_path.metadata()?.len();
556-
557-
// Remove the original file
558-
std::fs::remove_file(plain_path)?;
559-
560-
log::debug!(
561-
"Compressed HNSW graph links in {:.1?}: {:.1}MB -> {:.1}MB ({:.1}%)",
562-
start.elapsed(),
563-
original_size as f64 / 1024.0 / 1024.0,
564-
new_size as f64 / 1024.0 / 1024.0,
565-
new_size as f64 / original_size as f64 * 100.0,
566-
);
554+
// Unhappy path - convert from another format.
555+
for format in formats {
556+
let original_path = GraphLayers::get_links_path(dir, format);
557+
if format == force_format.as_format() || !original_path.exists() {
558+
continue;
559+
}
567560

568-
Ok(())
561+
let start = std::time::Instant::now();
562+
let edges = GraphLinks::load_from_file(
563+
&original_path,
564+
on_disk,
565+
format,
566+
Advice::Sequential,
567+
)?
568+
.to_edges();
569+
let original_size = original_path.metadata()?.len();
570+
atomic_save(&path, |writer| {
571+
serialize_graph_links(edges, *force_format, hnsw_m, writer)
572+
})?;
573+
let new_size = path.metadata()?.len();
574+
575+
// NOTE: The original file is not removed.
576+
577+
log::debug!(
578+
"Converted HNSW graph links in {:.1?}: {:.1}MB -> {:.1}MB ({:.1}%)",
579+
start.elapsed(),
580+
original_size as f64 / 1024.0 / 1024.0,
581+
new_size as f64 / 1024.0 / 1024.0,
582+
new_size as f64 / original_size as f64 * 100.0,
583+
);
584+
585+
return GraphLinks::load_from_file(
586+
&path,
587+
on_disk,
588+
force_format.as_format(),
589+
Advice::Random,
590+
);
591+
}
592+
}
593+
}
594+
Err(OperationError::service_error("No links file found"))
569595
}
570596

571597
#[cfg(feature = "testing")]
@@ -684,12 +710,10 @@ mod tests {
684710
}
685711

686712
#[rstest]
687-
#[case::uncompressed((GraphLinksFormat::Plain, false))]
688-
#[case::converted((GraphLinksFormat::Plain, true))]
689-
#[case::compressed((GraphLinksFormat::Compressed, false))]
690-
#[case::recompressed((GraphLinksFormat::Compressed, true))]
691-
#[case::compressed_with_vectors((GraphLinksFormat::CompressedWithVectors, false))]
692-
fn test_save_and_load(#[case] (initial_format, compress): (GraphLinksFormat, bool)) {
713+
#[case::uncompressed(GraphLinksFormat::Plain)]
714+
#[case::compressed(GraphLinksFormat::Compressed)]
715+
#[case::compressed_with_vectors(GraphLinksFormat::CompressedWithVectors)]
716+
fn test_save_and_load(#[case] initial_format: GraphLinksFormat) {
693717
let distance = Distance::Cosine;
694718
let num_vectors = 100;
695719
let dim = 8;
@@ -706,30 +730,45 @@ mod tests {
706730
M,
707731
dim,
708732
false,
709-
initial_format.is_with_vectors(),
733+
true,
710734
distance,
711735
&mut rng,
712736
);
737+
let graph_links_vectors = vector_holder.graph_links_vectors();
713738
let graph1 = graph_layers_builder
714739
.into_graph_layers(
715740
dir.path(),
716-
initial_format.with_param_for_tests(vector_holder.graph_links_vectors().as_ref()),
741+
initial_format.with_param_for_tests(graph_links_vectors.as_ref()),
717742
true,
718743
)
719744
.unwrap();
720745
assert_eq!(graph1.links.format(), initial_format);
721746
let res1 = search_in_graph(&query, top, &vector_holder, &graph1);
722747
drop(graph1);
723748

724-
let graph2 = GraphLayers::load(dir.path(), false, compress).unwrap();
725-
if compress {
726-
assert_eq!(graph2.links.format(), GraphLinksFormat::Compressed);
727-
} else {
728-
assert_eq!(graph2.links.format(), initial_format);
729-
}
730-
let res2 = search_in_graph(&query, top, &vector_holder, &graph2);
749+
for force_format in [
750+
None,
751+
Some(GraphLinksFormat::Plain),
752+
Some(GraphLinksFormat::Compressed),
753+
Some(GraphLinksFormat::CompressedWithVectors),
754+
] {
755+
eprintln!("force_format = {force_format:?}");
756+
let graph2 = GraphLayers::load(
757+
dir.path(),
758+
false,
759+
force_format
760+
.map(|fmt| fmt.with_param(graph_links_vectors.as_ref()))
761+
.as_ref(),
762+
)
763+
.unwrap();
764+
assert_eq!(
765+
graph2.links.format(),
766+
force_format.unwrap_or(initial_format)
767+
);
768+
let res2 = search_in_graph(&query, top, &vector_holder, &graph2);
731769

732-
assert_eq!(res1, res2)
770+
assert_eq!(res1, res2)
771+
}
733772
}
734773

735774
#[rstest]

‎lib/segment/src/index/hnsw_index/graph_layers_builder.rs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use bitvec::prelude::BitVec;
88
use common::ext::BitSliceExt;
99
use common::types::{PointOffsetType, ScoredPointOffset};
1010
use io::file_operations::{atomic_save, atomic_save_bin};
11+
use memory::madvise::Advice;
1112
use parking_lot::{Mutex, MutexGuard, RwLock};
1213
use rand::Rng;
1314
use rand::distr::Uniform;
@@ -197,7 +198,12 @@ impl GraphLayersBuilder {
197198
atomic_save(&links_path, |writer| {
198199
serialize_graph_links(edges, format_param, self.hnsw_m, writer)
199200
})?;
200-
links = GraphLinks::load_from_file(&links_path, true, format_param.as_format())?;
201+
links = GraphLinks::load_from_file(
202+
&links_path,
203+
true,
204+
format_param.as_format(),
205+
Advice::Random,
206+
)?;
201207
} else {
202208
// Since we'll keep it in the RAM anyway, we can afford to build in the RAM too.
203209
links = GraphLinks::new_from_edges(edges, format_param, self.hnsw_m)?;

‎lib/segment/src/index/hnsw_index/graph_links.rs‎

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,15 @@ use std::io::Cursor;
33
use std::path::Path;
44
use std::sync::Arc;
55

6+
use common::flags::FeatureFlags;
67
use common::types::PointOffsetType;
78
use memmap2::Mmap;
89
use memory::madvise::{Advice, AdviceSetting, Madviseable};
910
use memory::mmap_ops::open_read_mmap;
1011

1112
use crate::common::operation_error::{OperationError, OperationResult};
1213
use crate::index::hnsw_index::HnswM;
14+
use crate::vector_storage::common::PAGE_SIZE_BYTES;
1315
use crate::vector_storage::quantized::quantized_vectors::QuantizedVectors;
1416
use crate::vector_storage::{Sequential, VectorStorageEnum};
1517

@@ -182,7 +184,32 @@ impl GraphLinksFormat {
182184
}
183185
}
184186

185-
impl GraphLinksFormatParam<'_> {
187+
impl<'a> GraphLinksFormatParam<'a> {
188+
/// Decide which format to use.
189+
pub fn recommended<Q: GraphLinksVectors>(
190+
feature_flags: &FeatureFlags,
191+
hnsw_m0: usize,
192+
vectors: Option<&'a Q>,
193+
) -> Self {
194+
let is_small_enough = |layout: &GraphLinksVectorsLayout| {
195+
let mut node_size_bytes = 0;
196+
node_size_bytes += layout.base.size();
197+
node_size_bytes += hnsw_m0 * 3; // Rough estimate: 3 bytes per compressed id
198+
node_size_bytes += hnsw_m0 * layout.link.size();
199+
node_size_bytes <= 2 * PAGE_SIZE_BYTES
200+
};
201+
202+
if feature_flags.hnsw_with_vectors
203+
&& let Some(vectors) = vectors
204+
&& let Ok(layout) = vectors.vectors_layout()
205+
&& (is_small_enough(&layout) || feature_flags.hnsw_format_force)
206+
{
207+
GraphLinksFormatParam::CompressedWithVectors(vectors)
208+
} else {
209+
GraphLinksFormatParam::Compressed
210+
}
211+
}
212+
186213
pub fn as_format(&self) -> GraphLinksFormat {
187214
match self {
188215
GraphLinksFormatParam::Plain => GraphLinksFormat::Plain,
@@ -224,9 +251,10 @@ impl GraphLinks {
224251
path: &Path,
225252
on_disk: bool,
226253
format: GraphLinksFormat,
254+
advice: Advice,
227255
) -> OperationResult<Self> {
228256
let populate = !on_disk;
229-
let mmap = open_read_mmap(path, AdviceSetting::Advice(Advice::Random), populate)?;
257+
let mmap = open_read_mmap(path, AdviceSetting::Advice(advice), populate)?;
230258
Self::try_new(GraphLinksEnum::Mmap(Arc::new(mmap)), |x| {
231259
GraphLinksView::load(x.as_bytes(), format)
232260
})
@@ -340,6 +368,7 @@ pub(super) fn normalize_links(m: usize, mut links: Vec<PointOffsetType>) -> Vec<
340368
#[cfg(test)]
341369
mod tests {
342370
use io::file_operations::atomic_save;
371+
use memory::madvise::Advice;
343372
use rand::Rng;
344373
use rstest::rstest;
345374
use tempfile::Builder;
@@ -481,7 +510,8 @@ mod tests {
481510
})
482511
.unwrap();
483512

484-
let cmp_links = GraphLinks::load_from_file(&links_file, true, format).unwrap();
513+
let cmp_links =
514+
GraphLinks::load_from_file(&links_file, true, format, Advice::Sequential).unwrap();
485515
check_links(links, &cmp_links, &vectors);
486516
}
487517

0 commit comments

Comments
 (0)