feat: support seismic

crates/base/src/index.rs

@@ -160,6 +160,18 @@ impl IndexOptions {
                     ));
                 }
             }
+            IndexingOptions::Seismic(_) => {
+                if !matches!(self.vector.d, DistanceKind::Dot) {
+                    return Err(ValidationError::new(
+                        "seismic is not support for distance that is not negative dot product",
+                    ));
+                }
+                if !matches!(self.vector.v, VectorKind::SVecf32) {
+                    return Err(ValidationError::new(
+                        "seismic is not support for vectors that are not sparse vectors",
+                    ));
+                }
+            }
         }
         Ok(())
     }
@@ -294,6 +306,7 @@ pub enum IndexingOptions {
     Hnsw(HnswIndexingOptions),
     InvertedIndex(InvertedIndexingOptions),
     Rabitq(RabitqIndexingOptions),
+    Seismic(SeismicIndexingOptions),
 }
 
 impl IndexingOptions {
@@ -321,6 +334,12 @@ impl IndexingOptions {
         };
         x
     }
+    pub fn unwrap_seismic(self) -> SeismicIndexingOptions {
+        let IndexingOptions::Seismic(x) = self else {
+            unreachable!()
+        };
+        x
+    }
 }
 
 impl Default for IndexingOptions {
@@ -337,6 +356,7 @@ impl Validate for IndexingOptions {
             Self::Hnsw(x) => x.validate(),
             Self::InvertedIndex(x) => x.validate(),
             Self::Rabitq(x) => x.validate(),
+            Self::Seismic(x) => x.validate(),
         }
     }
 }
@@ -462,6 +482,51 @@ impl Default for RabitqIndexingOptions {
     }
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
+#[validate(schema(function = "Self::validate_self"))]
+#[serde(deny_unknown_fields)]
+pub struct SeismicIndexingOptions {
+    #[serde(default = "SeismicIndexingOptions::default_n_postings")]
+    #[validate(range(min = 100, max = 100_000))]
+    pub n_postings: u32,
+    #[serde(default = "SeismicIndexingOptions::default_centroid_fraction")]
+    #[validate(range(min = 0.01, max = 1.))]
+    pub centroid_fraction: f32,
+    #[serde(default = "SeismicIndexingOptions::default_summary_energy")]
+    #[validate(range(min = 0., max = 1.))]
+    pub summary_energy: f32,
+}
+
+impl SeismicIndexingOptions {
+    fn default_n_postings() -> u32 {
+        4000
+    }
+    fn default_centroid_fraction() -> f32 {
+        0.1
+    }
+    fn default_summary_energy() -> f32 {
+        0.4
+    }
+    fn validate_self(&self) -> Result<(), ValidationError> {
+        if (self.n_postings as f32 * self.centroid_fraction) as u32 > 65535 {
+            return Err(ValidationError::new(
+                "centroids number cannot exceed 65535 in seismic indexing",
+            ));
+        }
+        Ok(())
+    }
+}
+
+impl Default for SeismicIndexingOptions {
+    fn default() -> Self {
+        Self {
+            n_postings: Self::default_n_postings(),
+            centroid_fraction: Self::default_centroid_fraction(),
+            summary_energy: Self::default_summary_energy(),
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(deny_unknown_fields)]
 #[serde(rename_all = "snake_case")]
@@ -583,6 +648,12 @@ pub struct SearchOptions {
     pub rabitq_fast_scan: bool,
     #[validate(range(min = 1, max = 65535))]
     pub diskann_ef_search: u32,
+    #[validate(range(min = 1, max = 65535))]
+    pub seismic_heap_size: u32,
+    #[validate(range(min = 1, max = 100_000))]
+    pub seismic_q_cut: u32,
+    #[validate(range(min = 0.01, max = 1.))]
+    pub seismic_heap_factor: f32,
 }
 
 #[derive(Debug, Serialize, Deserialize)]

crates/cli/src/args.rs

@@ -144,6 +144,9 @@ impl QueryArguments {
             ivf_pq_fast_scan: false,
             rabitq_fast_scan: true,
             rabitq_nprobe: self.probe,
+            seismic_heap_size: 100,
+            seismic_q_cut: 3,
+            seismic_heap_factor: 1.0,
         }
     }
 }

crates/common/Cargo.toml

@@ -8,6 +8,7 @@ base = { path = "../base" }
 
 log.workspace = true
 memmap2.workspace = true
+parking_lot.workspace = true
 rand.workspace = true
 rustix.workspace = true
 serde.workspace = true

crates/common/src/lib.rs

@@ -1,11 +1,15 @@
+#![cfg_attr(target_arch = "aarch64", feature(stdarch_aarch64_prefetch))]
+
 pub mod always_equal;
 pub mod clean;
 pub mod dir_ops;
 pub mod file_atomic;
 pub mod json;
 pub mod mmap_array;
+pub mod prefetch;
 pub mod rand;
 pub mod remap;
 pub mod sample;
 pub mod variants;
 pub mod vec2;
+pub mod visited;

crates/common/src/prefetch.rs

@@ -0,0 +1,26 @@
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+#[allow(non_snake_case)]
+#[inline]
+pub fn prefetch_read_NTA(ptr: *const i8) {
+    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+    {
+        #[cfg(target_arch = "x86")]
+        use core::arch::x86::{_mm_prefetch, _MM_HINT_NTA};
+
+        #[cfg(target_arch = "x86_64")]
+        use core::arch::x86_64::{_mm_prefetch, _MM_HINT_NTA};
+
+        unsafe {
+            _mm_prefetch(ptr, _MM_HINT_NTA);
+        }
+    }
+
+    #[cfg(target_arch = "aarch64")]
+    {
+        use core::arch::aarch64::{_prefetch, _PREFETCH_LOCALITY0, _PREFETCH_READ};
+
+        unsafe {
+            _prefetch(ptr, _PREFETCH_READ, _PREFETCH_LOCALITY0);
+        }
+    }
+}

crates/graph/Cargo.toml

@@ -4,7 +4,6 @@ version.workspace = true
 edition.workspace = true
 
 [dependencies]
-parking_lot.workspace = true
 rand.workspace = true
 
 base = { path = "../base" }

crates/graph/src/lib.rs

@@ -2,4 +2,3 @@
 
 pub mod prune;
 pub mod search;
-pub mod visited;

crates/graph/src/search.rs

@@ -1,9 +1,9 @@
-use crate::visited::VisitedGuard;
-use crate::visited::VisitedPool;
 use base::scalar::F32;
 use base::search::Element;
 use base::search::GraphReranker;
 use base::search::Payload;
+use common::visited::VisitedGuard;
+use common::visited::VisitedPool;
 use std::cmp::Reverse;
 use std::collections::BinaryHeap;
 

crates/hnsw/src/lib.rs

@@ -9,7 +9,7 @@ use base::vector::VectorBorrowed;
 use common::json::Json;
 use common::mmap_array::MmapArray;
 use common::remap::RemappedCollection;
-use graph::visited::VisitedPool;
+use common::visited::VisitedPool;
 use num_traits::Float;
 use parking_lot::RwLock;
 use quantization::operator::OperatorQuantization;

crates/index/Cargo.toml

@@ -31,6 +31,7 @@ hnsw = { path = "../hnsw" }
 inverted = { path = "../inverted" }
 ivf = { path = "../ivf" }
 rabitq = { path = "../rabitq" }
+seismic = { path = "../seismic" }
 
 [lints]
 workspace = true