feat: SPANN vector search impl

Author: rebasedming

Cargo.lock

@@ -19,6 +19,7 @@ oneshot = "0.1.13"
 base64 = "0.22.0"
 byteorder = "1.4.3"
 crc32fast = "1.3.2"
+fastrand = "2"
 once_cell = "1.10.0"
 regex = { version = "1.5.5", default-features = false, features = [
   "std",
@@ -221,3 +222,11 @@ harness = false
 [[bench]]
 name = "term_set_queries"
 harness = false
+
+[[bench]]
+name = "distance"
+harness = false
+
+[[bench]]
+name = "graph"
+harness = false

Cargo.toml

@@ -0,0 +1,129 @@
+// Benchmarks the SIMD distance kernels in src/vector/distance.rs.
+//
+// What's measured:
+// - Pairwise f32:    l2_squared / dot / cosine on two f32 slices
+// - Pairwise bytes:  query is f32, doc side is LE bytes (segment shape)
+//
+// Dimension sweep covers common embedding widths:
+//   128 (word2vec) / 384 (MiniLM) / 768 (BERT-base) / 1024 (e5-large)
+//   1536 (OpenAI ada-002) / 3072 (text-embedding-3-large)
+//
+// Throughput is reported in bytes/sec so criterion prints MB/s, which
+// is directly comparable against:
+//   - L1 bandwidth   (~1 TB/s on modern cores)
+//   - FMA peak       (AVX-512: 64 FLOPS/cycle/core; NEON Apple M-series: 16-32 FLOPS/cycle/core
+//     depending on generation)
+//
+// Run with:  cargo bench --bench distance
+
+use std::hint::black_box;
+
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
+use tantivy::vector::{cosine, cosine_bytes, dot, dot_bytes, l2_squared, l2_squared_bytes};
+
+const DIMS: &[usize] = &[1024, 2048];
+
+/// Guaranteed-scalar dot product. Serves as the "no SIMD at all"
+/// baseline so we can see how much the autovectorizer is buying us.
+///
+/// `black_box` on the accumulator each iteration is the trick: it
+/// makes LLVM treat `acc` as opaque between iterations, which prevents
+/// the loop vectorizer from recognizing the reduction pattern and
+/// emitting SIMD.
+#[inline(never)]
+fn dot_scalar(a: &[f32], b: &[f32]) -> f32 {
+    debug_assert_eq!(a.len(), b.len());
+    let mut acc = 0.0f32;
+    for i in 0..a.len() {
+        acc = black_box(acc + a[i] * b[i]);
+    }
+    acc
+}
+
+/// Deterministic pseudo-random f32 values in [-1, 1). A bare LCG is
+/// enough — we just need non-zero, non-equal data the optimizer can't
+/// fold to a constant.
+fn make_f32(dim: usize, seed: u64) -> Vec<f32> {
+    let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15);
+    (0..dim)
+        .map(|_| {
+            state = state
+                .wrapping_mul(6_364_136_223_846_793_005)
+                .wrapping_add(1_442_695_040_888_963_407);
+            let bits = (state >> 33) as u32;
+            (bits as f32 / u32::MAX as f32) * 2.0 - 1.0
+        })
+        .collect()
+}
+
+fn f32_to_le_bytes(v: &[f32]) -> Vec<u8> {
+    let mut out = Vec::with_capacity(v.len() * 4);
+    for x in v {
+        out.extend_from_slice(&x.to_le_bytes());
+    }
+    out
+}
+
+fn bench_pairwise_f32(c: &mut Criterion) {
+    let mut group = c.benchmark_group("pairwise_f32");
+    for &dim in DIMS {
+        let a = make_f32(dim, 1);
+        let b = make_f32(dim, 2);
+
+        // Sanity check: the scalar baseline must agree with the
+        // autovectorized kernel within fp tolerance. Different reduction
+        // orders mean we can't expect bit-exact equality.
+        {
+            let autovec = dot(&a, &b);
+            let scalar = dot_scalar(&a, &b);
+            let tol = 1e-3 * autovec.abs().max(1.0);
+            assert!(
+                (autovec - scalar).abs() < tol,
+                "dot_scalar mismatch at dim={dim}: autovec={autovec} scalar={scalar}"
+            );
+        }
+
+        // 2 vectors * 4 B/element.
+        group.throughput(Throughput::Bytes((dim * 8) as u64));
+        group.bench_with_input(BenchmarkId::new("dot_scalar", dim), &dim, |bn, _| {
+            bn.iter(|| dot_scalar(black_box(&a), black_box(&b)))
+        });
+        group.bench_with_input(BenchmarkId::new("dot", dim), &dim, |bn, _| {
+            bn.iter(|| dot(black_box(&a), black_box(&b)))
+        });
+        group.bench_with_input(BenchmarkId::new("l2_squared", dim), &dim, |bn, _| {
+            bn.iter(|| l2_squared(black_box(&a), black_box(&b)))
+        });
+        group.bench_with_input(BenchmarkId::new("cosine", dim), &dim, |bn, _| {
+            bn.iter(|| cosine(black_box(&a), black_box(&b)))
+        });
+    }
+    group.finish();
+}
+
+fn bench_pairwise_bytes(c: &mut Criterion) {
+    let mut group = c.benchmark_group("pairwise_bytes");
+    for &dim in DIMS {
+        let q = make_f32(dim, 1);
+        let d = make_f32(dim, 2);
+        let d_bytes = f32_to_le_bytes(&d);
+        group.throughput(Throughput::Bytes((dim * 8) as u64));
+        group.bench_with_input(BenchmarkId::new("l2_squared_bytes", dim), &dim, |bn, _| {
+            bn.iter(|| l2_squared_bytes::<f32>(black_box(&q), black_box(&d_bytes)))
+        });
+        group.bench_with_input(BenchmarkId::new("dot_bytes", dim), &dim, |bn, _| {
+            bn.iter(|| dot_bytes::<f32>(black_box(&q), black_box(&d_bytes)))
+        });
+        group.bench_with_input(BenchmarkId::new("cosine_bytes", dim), &dim, |bn, _| {
+            bn.iter(|| cosine_bytes::<f32>(black_box(&q), black_box(&d_bytes)))
+        });
+    }
+    group.finish();
+}
+
+criterion_group! {
+    name = benches;
+    config = Criterion::default().sample_size(100);
+    targets = bench_pairwise_f32, bench_pairwise_bytes
+}
+criterion_main!(benches);

benches/distance.rs

@@ -0,0 +1,91 @@
+// Benchmarks the search workspace of src/vector/ivf/graph.rs.
+//
+// workspace/visited_check — the `Workspace.visited` bitset, exercised
+// with the same contains-then-insert pattern `search` inlines. The
+// bitset was chosen over a Vec<u32> epoch-stamp array, which measured
+// ~8-10% slower at n=1M and is 32x larger. To evaluate another
+// representation, mirror the swap here and in `Workspace` and rerun —
+// the bench names are stable, so `-- --save-baseline <name>` before and
+// `-- --baseline <name>` after prints the per-case delta.
+//
+// Two axes:
+// - n: node count (the centroid graph size the set is sized to)
+// - q: queries served by one allocation
+//
+// One iteration = allocate the set, then run q queries, each of which
+// resets the set and inserts n uniform-random node ids (~63% land on
+// unvisited slots, so both branch outcomes are exercised). Every
+// measured operation is one check-and-insert, so with throughput set to
+// n * q elements, elem/s reads directly as checks/sec; q=1 additionally
+// carries the allocation cost (the fresh-workspace-per-query regime)
+// while larger q amortizes it and isolates the per-query reset.
+//
+// Run with:  cargo bench --bench graph
+
+use std::hint::black_box;
+use std::time::Duration;
+
+use common::BitSet;
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
+
+/// Node counts to sweep (the centroid graph size).
+const NODE_COUNTS: &[usize] = &[100_000, 300_000, 1_000_000];
+
+/// Queries served by a single allocation. Each query costs n inserts,
+/// so amortization of the allocation saturates quickly; larger values
+/// only inflate iteration time (n * q inserts each).
+const QUERIES_PER_ALLOC: &[usize] = &[1, 10, 100];
+
+/// Same LCG as benches/distance.rs; deterministic and cheap enough to be
+/// noise next to the visited access it feeds.
+#[inline]
+fn lcg(state: u64) -> u64 {
+    state
+        .wrapping_mul(6_364_136_223_846_793_005)
+        .wrapping_add(1_442_695_040_888_963_407)
+}
+
+/// One iteration: q queries over one allocation, n random inserts each.
+#[inline]
+fn run_batch(n: usize, queries: usize) -> usize {
+    let mut visited = BitSet::with_max_value(n as u32);
+    let mut state = 0x9E37_79B9_7F4A_7C15u64;
+    let mut first_visits = 0;
+    for _ in 0..queries {
+        visited.clear();
+        for _ in 0..n {
+            state = lcg(state);
+            let node = ((state >> 33) as usize % n) as u32;
+            if !visited.contains(node) {
+                visited.insert(node);
+                first_visits += 1;
+            }
+        }
+    }
+    first_visits
+}
+
+fn bench_visited(c: &mut Criterion) {
+    let mut group = c.benchmark_group("workspace/visited_check");
+    for &n in NODE_COUNTS {
+        for &q in QUERIES_PER_ALLOC {
+            group.throughput(Throughput::Elements((n * q) as u64));
+            group.bench_with_input(
+                BenchmarkId::new(format!("n={n}"), q),
+                &(n, q),
+                |bn, &(n, q)| bn.iter(|| run_batch(black_box(n), black_box(q))),
+            );
+        }
+    }
+    group.finish();
+}
+
+criterion_group! {
+    name = benches;
+    config = Criterion::default()
+        .sample_size(10)
+        .warm_up_time(Duration::from_secs(1))
+        .measurement_time(Duration::from_secs(3));
+    targets = bench_visited
+}
+criterion_main!(benches);

benches/graph.rs

@@ -12,7 +12,7 @@ use std::ops::Range;
 
 pub use merge::merge_column_index;
 pub(crate) use multivalued_index::SerializableMultivalueIndex;
-pub use optional_index::{OptionalIndex, Set};
+pub use optional_index::{OptionalIndex, Set, open_optional_index, serialize_optional_index};
 pub use serialize::{
     SerializableColumnIndex, SerializableOptionalIndex, open_column_index, serialize_column_index,
 };

columnar/src/column_index/mod.rs

@@ -13,6 +13,8 @@ use crate::collector::sort_key_top_collector::TopBySortKeyCollector;
 use crate::collector::top_collector::ComparableDoc;
 use crate::collector::{SegmentSortKeyComputer, SortKeyComputer};
 use crate::fastfield::FastValue;
+use crate::schema::Field;
+use crate::vector::{TopDocsByVectorSimilarity, VectorElement};
 use crate::{DocAddress, DocId, Order, Score, SegmentOrdinal, SegmentReader};
 
 /// The `TopDocs` collector keeps track of the top `K` documents
@@ -329,6 +331,39 @@ impl TopDocs {
         TopBySortKeyCollector::new(sort_key_computer, self.doc_range())
     }
 
+    /// Ranks documents by similarity to a query vector against a
+    /// [`FieldType::Vector`](crate::schema::FieldType::Vector) field.
+    ///
+    /// The metric is read from the field's
+    /// [`VectorOptions`](crate::schema::VectorOptions) at query time —
+    /// callers don't pick a metric here. All metrics are presented in
+    /// "higher is better" form (L2 is internally negated), so this is
+    /// always a descending sort: the closest docs come first.
+    ///
+    /// Only documents that have a vector for `field` are returned —
+    /// filter-matching docs without a vector are dropped. This differs
+    /// from `order_by_fast_field`'s `None`-trailing convention; the
+    /// constraint comes from IVF, which can't see vectorless docs.
+    ///
+    /// Generic over `T: VectorElement` — `T` is typically `f32` and is
+    /// inferred from the `query` argument; it must match the schema's
+    /// declared dtype.
+    ///
+    /// The driving query (passed to `searcher.search`) decides which docs
+    /// are candidates. For brute-force scan over every doc, pair this with
+    /// [`AllQuery`](crate::query::AllQuery); pair with a filter query to
+    /// restrict candidates.
+    pub fn order_by_similarity<T>(
+        self,
+        field: Field,
+        query: Vec<T>,
+    ) -> TopDocsByVectorSimilarity<T>
+    where
+        T: VectorElement,
+    {
+        TopDocsByVectorSimilarity::new(field, query, self.limit).and_offset(self.offset)
+    }
+
     /// Helper function to tweak the similarity score of documents using a function.
     /// (usually a closure).
     ///
@@ -671,6 +706,26 @@ where
         self.append_doc(doc, sort_key);
     }
 
+    /// Like [`push`](Self::push), but doesn't require pushes to arrive
+    /// in ascending `DocId`/`DocAddress` order.
+    ///
+    /// The default `push` short-circuits below-threshold candidates with
+    /// a strict `> threshold` test. That's only safe under ascending-D
+    /// pushes — equal-sort-key future pushes are guaranteed to have a
+    /// larger D and thus lose the tie-break, so they can be skipped.
+    /// Out-of-order callers (eg a vector IVF format which visits cluster
+    /// by cluster) can't make that guarantee, so this path relaxes the
+    /// skip to strictly less than the threshold.
+    #[inline]
+    pub fn push_unordered(&mut self, sort_key: TSortKey, doc: D) {
+        if let Some((last_median, _thresh_ord)) = &self.threshold {
+            if self.comparator.compare(&sort_key, last_median) == std::cmp::Ordering::Less {
+                return;
+            }
+        }
+        self.append_doc(doc, sort_key);
+    }
+
     // Append a document to the top n.
     //
     // At this point, we need to have established that the doc is above the threshold.
@@ -870,6 +925,62 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_topn_computer_push_unordered() {
+        // Same inputs as test_topn_computer, but pushed in a non-ascending
+        // DocId order — push_unordered should still produce the right top-K
+        // and (where ties exist) the smallest-DocId tie-break.
+        let mut computer: TopNComputer<u32, u32, NaturalComparator> =
+            TopNComputer::new_with_comparator(2, NaturalComparator);
+
+        computer.push_unordered(3u32, 3u32);
+        computer.push_unordered(1u32, 5u32);
+        computer.push_unordered(2u32, 2u32);
+        computer.push_unordered(2u32, 4u32);
+        computer.push_unordered(1u32, 1u32);
+        assert_eq!(
+            computer.into_sorted_vec(),
+            &[
+                ComparableDoc {
+                    sort_key: 3u32,
+                    doc: 3u32,
+                },
+                ComparableDoc {
+                    sort_key: 2u32,
+                    doc: 2u32,
+                }
+            ]
+        );
+    }
+
+    #[test]
+    fn test_topn_computer_push_unordered_tiebreak_on_equal_sort_key() {
+        // Tied sort keys must still resolve to the smallest DocId, even when
+        // doc ids arrive out of order (and even when the smaller-DocId
+        // candidate arrives *after* the larger-DocId one).
+        let mut computer: TopNComputer<u32, u32, NaturalComparator> =
+            TopNComputer::new_with_comparator(2, NaturalComparator);
+
+        computer.push_unordered(1u32, 100u32);
+        computer.push_unordered(1u32, 50u32);
+        computer.push_unordered(1u32, 200u32);
+        computer.push_unordered(1u32, 25u32);
+        computer.push_unordered(1u32, 175u32);
+        assert_eq!(
+            computer.into_sorted_vec(),
+            &[
+                ComparableDoc {
+                    sort_key: 1u32,
+                    doc: 25u32,
+                },
+                ComparableDoc {
+                    sort_key: 1u32,
+                    doc: 50u32,
+                }
+            ]
+        );
+    }
+
     #[test]
     fn test_topn_computer_no_panic() {
         for top_n in 0..10 {