diff --git a/CHANGELOG.md b/CHANGELOG.md index 11d24d09d..f53873a20 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,154 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Performance — FT.SEARCH intra-query worker pool + bounded bulk compaction + +- **`--ft-search-workers N` worker pool** (`src/vector/search_pool.rs`): the + per-segment HNSW searches of ONE KNN query fan out across a pool of searcher + threads while the shard task scans the mutable segment; replies are awaited + with `flume::recv_async`, so the shard event loop is never blocked. Default + **0 (off, opt-in)** — each segment pays the full resolved ef, so a pooled + N-segment query does ~N× the CPU work for its latency win: a 4.9× QPS win on + physical-core-rich boxes but a measured regression on SMT-constrained ones + (same posture as `--io-busy-poll-us`). Good opt-in size: physical cores − + shards, cap 8. Results are identical pooled or serial (enforced by + pooled-vs-serial identity + 8-thread concurrency stress tests); worker + panics are contained per-job (empty segment result + warn, never a hang). + Runtime-agnostic (monoio + tokio). +- **Bounded bulk compaction**: with a pool active and `COMPACT_THRESHOLD > 0`, + one compaction build freezes at most `max(threshold, len/8)` entries + (`MutableSegment::freeze_prefix`), so FT.COMPACT after a bulk load yields + several independently searchable segments instead of one giant graph — + bounded build memory and pool-parallel search. Pool-less deployments keep + single-segment builds (multi-segment serial search would be strictly + slower). Red/green: `test_force_compact_bulk_bounded_segments`, + `test_bg_compact_bulk_bounded_segments`. +- Measured (20k×384d clustered SQ8, single connection, post-FT.COMPACT, + R@10 = 1.0 in all rows): macOS 10-core p50 2.06 ms → **0.46 ms**, 473 → + **2,321 QPS (4.9×)**; GCE c3-standard-8 (4 physical cores) regresses at + default ef (1,732 → 1,165 QPS) — hence the opt-in default. The GCE probe + matrix also showed the per-index `EF_RUNTIME` knob alone closes most of the + vs-RediSearch clustered gap (ef 64: 4,493 QPS serial at R@10 = 1.0). Full + GCE vs-RediSearch tables in PR #214. + +### Fixed — update-churn no longer mass-deletes vectors at compact/merge install (soak-diagnostic find) + +- **Compact install: dead window entries no longer kill their own update.** + `snap_and_reconcile` treated ANY tombstoned entry in the frozen window as + "key deleted" and applied a key_hash-wide tombstone to the new immutable — + but since the VEC-1 update path tombstones the old copy in place, a key + updated before the freeze had a dead old copy AND a live new copy in the + same window, and the install deleted the new copy out of the segment. + Every key updated-then-compacted silently vanished from FT.SEARCH (32% of + live keys lost / live-set recall 0.985 → 0.685 in a 1-minute churn soak; + regression vs v0.5.1). A dead window entry now only proves deletion when + the key has no live window sibling. +- **Merge install: source tombstones are origin-gated.** The merge replay + applied each source segment's lifetime interior tombstone set key_hash-wide + to the merged output, so a key whose old copy was tombstoned-by-update in + one source killed its current copy merged in from a sibling segment. The + replay now only tombstones entries whose `global_id` originated in the + tombstone's own source; DEL/UNLINK tombstones land in every source's set + and still apply everywhere. +- **`VectorStore::insert_vector` allocates monotonic LSNs** (same allocator + as the wire path) instead of `mutable.len()+1`, which restarted after every + compaction and made merge dedup keep a stale copy over the current one. +- Red/green: `test_bg_compact_update_before_freeze_survives_install`, + `test_bg_merge_update_across_segments_survives`; end-to-end churn repro + (24k mixed ops): LOST 1054 → 0, live-set recall 0.685 → 0.98. + +### Fixed — DEL/UNLINK now unindexes vectors on every dispatch path (soak-diagnostic find) + +- **Deleted keys no longer resurface in FT.SEARCH** — the vector auto-delete + hook (`mark_deleted_for_key`) existed only on the cross-shard SPSC `Execute` + arm and the tokio sharded handler. The monoio conn-local path (the default + runtime's only path at `--shards 1`), `handler_single`, the MULTI/EXEC batch + paths, and the SPSC pipeline arms never tombstoned vectors on DEL/UNLINK, so + deleted keys kept matching KNN searches forever (20% of results after one + minute of mixed churn; live-set recall collapsed 0.985 → 0.735 in the + Bundle-5 soak diagnostic). All paths now share one `auto_delete_vectors` + parity helper; wire-level red/green coverage in `tests/vector_del_unindex.rs`. + +### Added — long-run vector reliability harness + +- `scripts/vector-validate.py` — on-target validation driver: recall/QPS + comparison between two moon binaries (SQ8 + TQ4, ground-truth brute force), + churn soak with live-set recall / RSS / resurrection sampling, and kill -9 + durability (settled-write survival + double-crash restart) under + `appendonly yes`. +- `scripts/gcloud-vector-soak.sh` — GCE orchestration (c4a ARM + c3 x86): + ships the working branch via `git bundle`, builds baseline + branch binaries, + runs the validation driver, fetches JSON results, tears down on exit. +- `scripts/bench-vector-vs-redisearch.py` — Moon vs RediSearch head-to-head + driver (same FT.* wire dialect for both engines): insert throughput, KNN-10 + QPS/p50/p99, and R@10 vs numpy ground truth on random-Gaussian and + clustered-mixture 384d datasets, with an `EF_RUNTIME` sweep so QPS is + compared at matched recall. + +### Fixed — vector search correctness bundle (deep-review VEC-1/XC-SHARD-1/XC-3/VEC-4/VEC-7) + +- **HSET update no longer duplicates a vector** — re-indexing an existing key + tombstones the old copy first (O(1) fast path in the mutable segment via the + key→global-id map, scan fallback across mutable + immutable segments), so KNN + totals and results no longer count both the stale and the fresh vector. +- **FT.INFO is now cluster-wide at `--shards N`** — previously answered from the + local shard only (~1/N of `num_docs`). Now scatter-gathers to every shard and + merges additively (top-level counters + per-field stats), on both the sharded + and monoio handlers. +- **Filtered FT.SEARCH on the cooperative-yield path honors `FilterStrategy`** — + the yielding search always graph-filtered; the post-filter strategy (selective + filters) now searches unfiltered with 3×k oversampling and bitmap post-filter, + matching the non-yielding path's recall behavior. +- **`FT.CREATE ... MERGE_MODE KEEP_RAW` is rejected fail-loud** — it silently + behaved as graph-union; now errors until the raw-vector sidecar is implemented + (the separate `KEEP_RAW ON` flag is unchanged). +- **MEMORY DOCTOR / Prometheus KV memory no longer report 0 under unlimited + `maxmemory`** — the per-shard KV memory publish on the 100ms eviction tick + had been gated on `maxmemory > 0` since the GAP-1 elastic-budget work, + permanently zeroing the `DashTable + entries` line for the default config. + The publish (an O(1) accumulator read per DB) now runs unconditionally; + elastic-budget recompute stays gated on a finite cap. + +### Added + +- **`FT.CONFIG SET/GET MERGE_RECALL_TOLERANCE <0.0..=1.0>`** — per-index + recall gate for unattended (background/vacuum) GraphUnion merges; default 0.70 + unchanged. + +### Added — exact rerank stage (deep-review HQ-1) + +- **FT.SEARCH distances on compacted segments are now (near-)exact.** Immutable + segments carry an f16 sidecar of the original vectors (built at compaction, + BFS-ordered); the top `4·k` beam candidates are re-scored with true metric + distances (L2: squared L2; Cosine/InnerProduct: normalized-pair squared L2) + before top-k truncation, replacing pure quantized ADC estimates — the + recall lever vs engines that keep full-precision vectors. SQ8, previously + ZERO-refinement, benefits most; TQ4 estimates improve from percent-level + error to f16 tolerance (~1e-3). +- The sidecar persists (`raw_f16.bin` per segment dir, missing file = no + sidecar, fully backward/forward compatible), survives GraphUnion merges + (all-or-nothing propagation), and is MEMORY DOCTOR-accounted. Memory cost: + +2·dim bytes per vector in both the mutable segment and compacted segments + (384d ≈ +768 B/vector); an opt-out knob is a follow-up. + +### Performance — SIMD SQ8 ADC kernels (deep-review HQ-2) + +- **SQ8 asymmetric distance is now SIMD-dispatched** (NEON / AVX2+FMA / + AVX-512F, scalar fallback) via an algebraic decomposition: per-query + constants `(Σq, Σq²)` computed once, per-candidate work reduced to three + fused widen-u8→f32 FMA sums `(Σq·c, Σc, Σc²)` combined in O(1). Wired into + HNSW beam search and both mutable-segment brute-force paths. Criterion + (aarch64 NEON): 2.7×/1.8×/1.6× faster per candidate at 128/384/768d. + Note: the pre-AVX2 x86 scalar fallback is ~1.65× slower than the old naive + loop (the decomposition only pays with SIMD); all supported targets + (aarch64 NEON, x86-64 AVX2+) are wins. + +### Performance — vector search hot-path quick wins + +- Copy-on-write `Arc` key map (no full `HashMap` clone per snapshot), hoisted + brute-force query prep, striped search metrics counters, stable aarch64 + `prfm` node prefetch, dead quantize removed from the insert path, SQ8 + `code_len` fix in compaction, and raw-buffer work skipped for SQ8 segments. ### Performance — coordinator local legs ride group commit under appendfsync=always (PR #TBD) - The cross-shard coordinator's LOCAL-leg persist (co-located MSET/MSETNX and diff --git a/CLAUDE.md b/CLAUDE.md index 4c0a626dd..389ae4499 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -199,10 +199,11 @@ orb run -m moon-dev bash -c 'sudo apt-get update -qq && sudo apt-get install -y Moon ships a native HNSW + TurboQuant vector engine exposed via a RediSearch-compatible subset (`FT.CREATE`, `FT.DROPINDEX`, `FT.INFO`, `FT.SEARCH`, `FT.COMPACT`). Source: `src/vector/` and `src/command/vector_search/`. -- **Per-index knobs:** `EF_RUNTIME` (recall/QPS trade-off), `COMPACT_THRESHOLD` (when to flush mutable → immutable segment). -- **Segment lifecycle:** auto-indexed on HSET → mutable segment (brute force) → compact → immutable segment (HNSW graph + TQ codes). Segments do not merge once immutable — lossy decode+re-encode accumulates quantization error (tested, recall collapsed 0.73 → 0.0005). -- **Key hash map:** `key_hash_to_key: HashMap` must be populated in `auto_index_hset` and propagated via `SearchResult.key_hash`; otherwise multi-segment search returns synthetic `vec:` instead of original keys. -- **FT.INFO `num_docs`** must sum across all segments (mutable + immutable), not just the mutable one. +- **Per-index knobs:** `EF_RUNTIME` (recall/QPS trade-off), `COMPACT_THRESHOLD` (when to flush mutable → immutable segment), `MERGE_RECALL_TOLERANCE` (FT.CONFIG; recall gate for unattended GraphUnion merges, default 0.70). +- **Segment lifecycle:** auto-indexed on HSET → mutable segment (brute force) → compact → immutable segment (HNSW graph + TQ codes). Immutable segments DO merge: `MERGE_MODE GRAPH_UNION` (the default) auto-merges at ≥16 segments or ≥20% dead, gated by a recall check (0.70 background / 0.90 manual FT.COMPACT). What is forbidden is *decode+re-encode* merging — re-quantizing accumulates error (tested, recall collapsed 0.73 → 0.0005); GraphUnion stitches graphs over the original codes instead. `MERGE_MODE KEEP_RAW` is rejected at FT.CREATE (unimplemented stub). +- **Key hash map:** `key_hash_to_key: Arc>` (copy-on-write via `Arc::make_mut`) must be populated in `auto_index_hset` and propagated via `SearchResult.key_hash`; otherwise multi-segment search returns synthetic `vec:` instead of original keys. +- **FT.INFO `num_docs`** must sum across all segments (mutable + immutable), not just the mutable one — and across all shards: FT.INFO scatter-gathers via `scatter_ft_info` + `merge_ft_info_responses` (additive counters, per-field merge); a local-only answer under-reports by ~1/N. +- **Exact rerank sidecar (HQ-1):** immutable segments keep an f16 copy of each original vector (`raw_f16.bin` on disk; built from the mutable segment's always-retained f16 buffer); the top 4·k beam candidates are re-scored with true metric distances before truncation. Segments reloaded from pre-sidecar dirs (or rebuilt without raw data) silently fall back to quantized ADC distances — recall degrades, doesn't break. Merge propagation is all-or-nothing. - **TQ4 at 384d loses recall** (concentration of distances + quantization noise). Use **SQ8** or FP32 HNSW for ≤384d workloads; TQ4 shines at 768d+. (There is no TQ8 — SQ8 is the real 8-bit option: per-vector affine scalar quant, normalizes for the unit-sphere metrics Cosine + InnerProduct, validated across the full lifecycle — search/merge/persistence — at ~0.90 R@10 on real MiniLM 384d. PR #166.) ## GPU / CUDA Acceleration diff --git a/Cargo.toml b/Cargo.toml index 458f562b4..957ea5f6f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -223,6 +223,10 @@ harness = false name = "distance_bench" harness = false +[[bench]] +name = "sq8_adc_bench" +harness = false + [[bench]] name = "hnsw_bench" harness = false diff --git a/benches/sq8_adc_bench.rs b/benches/sq8_adc_bench.rs new file mode 100644 index 000000000..6aa27c67a --- /dev/null +++ b/benches/sq8_adc_bench.rs @@ -0,0 +1,89 @@ +//! Criterion benchmarks for SQ8 asymmetric-distance-code (ADC) kernels. +//! +//! Validates finding HQ-2 (tmp/VECTOR-DEEP-REVIEW.md): SIMD-dispatched ADC +//! stats (NEON/AVX2/AVX-512 widen-u8-to-f32 + FMA) beat the naive per-element +//! scalar `sq8_l2_adc` at standard embedding dimensions (128/384/768). +//! +//! Three variants compared per dimension: +//! - `naive_scalar`: the original per-candidate scalar loop (baseline). +//! - `stats_scalar`: the algebraic stats decomposition, scalar stats kernel +//! (isolates the win from the ALGEBRA alone, no SIMD). +//! - `stats_dispatch`: the algebraic decomposition with the SIMD-dispatched +//! stats kernel (NEON on aarch64, AVX2/AVX-512 on x86_64) — this is what +//! the beam-search / brute-force hot paths now call. + +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use moon::vector::distance; +use moon::vector::turbo_quant::sq8::{ + sq8_candidate_stats_scalar, sq8_l2_adc, sq8_l2_from_stats, sq8_query_stats, +}; +use std::hint::black_box; + +fn make_f32_vector(dim: usize, seed: u64) -> Vec { + let mut s = seed as u32; + let mut v = Vec::with_capacity(dim); + for _ in 0..dim { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + v.push((s as f32) / (u32::MAX as f32) * 2.0 - 1.0); + } + v +} + +fn make_u8_codes(dim: usize, seed: u64) -> Vec { + let mut s = seed as u32; + let mut v = Vec::with_capacity(dim); + for _ in 0..dim { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + v.push((s >> 24) as u8); + } + v +} + +const DIMS: &[usize] = &[128, 384, 768]; + +fn bench_sq8_adc(c: &mut Criterion) { + distance::init(); + let mut group = c.benchmark_group("sq8_l2_adc"); + + for &dim in DIMS { + let query = make_f32_vector(dim, 42); + let codes = make_u8_codes(dim, 99); + let min = -0.5f32; + let scale = 0.004f32; + + group.bench_with_input(BenchmarkId::new("naive_scalar", dim), &dim, |bench, _| { + bench.iter(|| { + sq8_l2_adc( + black_box(&query), + black_box(&codes), + black_box(min), + black_box(scale), + ) + }); + }); + + group.bench_with_input(BenchmarkId::new("stats_scalar", dim), &dim, |bench, _| { + // Per-query stats hoisted outside the timed loop, matching how the + // real beam-search / brute-force call sites use it. + let (q_sum, q_sumsq) = sq8_query_stats(&query); + bench.iter(|| { + let (dot_qc, sum_c, sumsq_c) = + sq8_candidate_stats_scalar(black_box(&query), black_box(&codes)); + sq8_l2_from_stats(dim, min, scale, q_sum, q_sumsq, dot_qc, sum_c, sumsq_c) + }); + }); + + group.bench_with_input(BenchmarkId::new("stats_dispatch", dim), &dim, |bench, _| { + let stats_fn = distance::table().sq8_stats; + let (q_sum, q_sumsq) = sq8_query_stats(&query); + bench.iter(|| { + let (dot_qc, sum_c, sumsq_c) = stats_fn(black_box(&query), black_box(&codes)); + sq8_l2_from_stats(dim, min, scale, q_sum, q_sumsq, dot_qc, sum_c, sumsq_c) + }); + }); + } + group.finish(); +} + +criterion_group!(benches, bench_sq8_adc); +criterion_main!(benches); diff --git a/scripts/bench-vector-vs-redisearch.py b/scripts/bench-vector-vs-redisearch.py new file mode 100644 index 000000000..20d88ea0a --- /dev/null +++ b/scripts/bench-vector-vs-redisearch.py @@ -0,0 +1,360 @@ +#!/usr/bin/env python3 +"""bench-vector-vs-redisearch.py — Moon vs RediSearch (Redis Stack) vector head-to-head. + +Runs ON the target machine. Both engines speak the same FT.* dialect, so one +driver benches both: + - insert throughput (pipelined HSET, vectors auto-indexed) + - FT.SEARCH KNN-10 QPS + p50/p99 (single connection, 10s tight loop) + - R@10 vs numpy brute-force ground truth + +Datasets (384d unit vectors): + gaussian — iid random Gaussian, normalized (harsh: concentration of distances) + clustered — 200-center Gaussian mixture, sigma 0.25 (closer to real embeddings) + +Configs: + redisearch — FLOAT32 HNSW, EF_RUNTIME default(10) and 64 + moon SQ8 / TQ4 — Moon quantized HNSW (COMPACT_THRESHOLD 4096 + FT.COMPACT) + +Usage: + python3 bench-vector-vs-redisearch.py --moon-bin PATH --redis-bin PATH \ + [--n-vectors 20000] [--out results.json] + +`--redis-bin` is the redis-stack-server *redis-server* binary; the RediSearch +module is loaded via --loadmodule if REDISEARCH_SO is set, otherwise assumed +built in (redis-stack-server layout). +""" + +import argparse +import json +import os +import shutil +import signal +import socket +import subprocess +import tempfile +import time + +import numpy as np + +DIM = 384 +K = 10 + + +# ── RESP client (same shape as vector-validate.py) ── +class Resp: + def __init__(self, port, timeout=60): + self.sock = socket.create_connection(("127.0.0.1", port), timeout=timeout) + self.sock.settimeout(timeout) + self.buf = b"" + + def close(self): + try: + self.sock.close() + except Exception: + pass + + def _read_more(self): + chunk = self.sock.recv(1 << 16) + if not chunk: + raise ConnectionError("server closed connection") + self.buf += chunk + + def _read_line(self): + while b"\r\n" not in self.buf: + self._read_more() + line, self.buf = self.buf.split(b"\r\n", 1) + return line + + def _read_exact(self, n): + while len(self.buf) < n + 2: + self._read_more() + data, self.buf = self.buf[:n], self.buf[n + 2 :] + return data + + def _parse(self): + line = self._read_line() + t, rest = line[:1], line[1:] + if t == b"+": + return rest.decode() + if t == b"-": + return Exception(rest.decode()) + if t == b":": + return int(rest) + if t == b"$": + n = int(rest) + return None if n == -1 else self._read_exact(n) + if t == b"*": + n = int(rest) + return None if n == -1 else [self._parse() for _ in range(n)] + if t == b"%": + n = int(rest) + return [self._parse() for _ in range(2 * n)] + raise ValueError(f"unexpected RESP type {line!r}") + + @staticmethod + def encode(*args): + parts = [f"*{len(args)}\r\n".encode()] + for a in args: + b = a if isinstance(a, bytes) else str(a).encode() + parts += [f"${len(b)}\r\n".encode(), b, b"\r\n"] + return b"".join(parts) + + def cmd(self, *args): + self.sock.sendall(self.encode(*args)) + return self._parse() + + def pipeline(self, cmds): + self.sock.sendall(b"".join(self.encode(*c) for c in cmds)) + return [self._parse() for _ in cmds] + + +class Server: + """Spawn either moon or redis-server(+RediSearch). Kill -9 only.""" + + def __init__(self, argv, port, log_path): + self.port = port + self.log = open(log_path, "ab") + self.proc = subprocess.Popen(argv, stdout=self.log, stderr=self.log) + + def wait_ready(self, timeout=60): + deadline = time.time() + timeout + while time.time() < deadline: + try: + c = Resp(self.port, timeout=2) + if c.cmd("PING") == "PONG": + c.close() + return + c.close() + except Exception: + pass + if self.proc.poll() is not None: + raise RuntimeError(f"server exited rc={self.proc.returncode}") + time.sleep(0.2) + raise RuntimeError(f"no PING on {self.port}") + + def stop(self): + if self.proc and self.proc.poll() is None: + self.proc.send_signal(signal.SIGKILL) + self.proc.wait(10) + + +# ── Data ── +def gen_gaussian(n, seed): + rng = np.random.default_rng(seed) + v = rng.standard_normal((n, DIM), dtype=np.float32) + v /= np.linalg.norm(v, axis=1, keepdims=True) + return v + + +def gen_clustered(n, seed, centers=200, sigma=0.04, center_seed=1234): + # center_seed is FIXED so DB vectors and query vectors come from the SAME + # mixture (in-distribution queries — like real embedding workloads). + # sigma 0.04 at 384d puts members at ~0.8 cosine to their center (noise + # norm sigma*sqrt(d) ~= 0.78 vs unit center) — real-embedding-like + # neighborhoods. Larger sigmas drown the centers and degenerate to the + # random-Gaussian regime. + crng = np.random.default_rng(center_seed) + c = crng.standard_normal((centers, DIM), dtype=np.float32) + c /= np.linalg.norm(c, axis=1, keepdims=True) + rng = np.random.default_rng(seed) + which = rng.integers(0, centers, size=n) + v = c[which] + sigma * rng.standard_normal((n, DIM), dtype=np.float32) + v /= np.linalg.norm(v, axis=1, keepdims=True) + return v.astype(np.float32) + + +def cosine_gt(queries, db, k): + sims = queries @ db.T + return np.argsort(-sims, axis=1)[:, :k] + + +def knn_ids(c, idx, q, k=K, ef_runtime=None): + query = f"*=>[KNN {k} @vec $B" + (f" EF_RUNTIME {ef_runtime}" if ef_runtime else "") + "]" + r = c.cmd( + "FT.SEARCH", idx, query, "PARAMS", "2", "B", q.tobytes(), + "NOCONTENT", "DIALECT", "2", + ) + if isinstance(r, Exception): + # Engines differ on NOCONTENT / EF_RUNTIME support — retry plain. + r = c.cmd( + "FT.SEARCH", idx, f"*=>[KNN {k} @vec $B]", + "PARAMS", "2", "B", q.tobytes(), "DIALECT", "2", + ) + if isinstance(r, Exception): + raise RuntimeError(f"FT.SEARCH failed: {r}") + keys = [r[i] for i in range(1, len(r), 2)] + else: + # NOCONTENT reply: [total, key1, key2, ...]; with-content: keys at odd idx. + body = r[1:] + if body and isinstance(body[0], bytes) and len(body) >= 2 and isinstance(body[1], list): + keys = body[0::2] + else: + keys = body + out = [] + for kk in keys: + s = kk.decode() if isinstance(kk, bytes) else str(kk) + if ":" in s: + out.append(int(s.split(":", 1)[1])) + return out + + +def ft_num_docs(c, idx): + r = c.cmd("FT.INFO", idx) + if isinstance(r, Exception): + return -1 + for i in range(0, len(r) - 1, 2): + k = r[i].decode() if isinstance(r[i], bytes) else str(r[i]) + if k == "num_docs": + v = r[i + 1] + return int(v.decode() if isinstance(v, bytes) else v) + return -1 + + +def settle_index(c, idx, n, timeout=120): + """Wait until the index reports >= n docs (RediSearch backfill / Moon mutable).""" + deadline = time.time() + timeout + while time.time() < deadline: + if ft_num_docs(c, idx) >= n: + return + time.sleep(0.5) + + +def bench_engine(c, idx, vecs, queries, gt, ef_runtime=None, insert=True): + n = len(vecs) + res = {} + if insert: + t0 = time.time() + for s in range(0, n, 500): + cmds = [ + ("HSET", f"{idx}:{i}", "vec", vecs[i].tobytes()) + for i in range(s, min(s + 500, n)) + ] + for r in c.pipeline(cmds): + if isinstance(r, Exception): + raise RuntimeError(f"HSET failed: {r}") + res["insert_secs"] = round(time.time() - t0, 2) + res["insert_vps"] = round(n / max(1e-9, time.time() - t0)) + settle_index(c, idx, n) + + # warmup (drives Moon background-build installs; JITs RediSearch caches) + for qi in range(20): + knn_ids(c, idx, queries[qi % len(queries)], ef_runtime=ef_runtime) + + # recall + latency + lat = [] + hits = 0 + for qi in range(len(queries)): + t0 = time.time() + got = knn_ids(c, idx, queries[qi], ef_runtime=ef_runtime) + lat.append(time.time() - t0) + hits += len(set(got[:K]) & set(gt[qi].tolist())) + res["recall_at_10"] = round(hits / (len(queries) * K), 4) + lat.sort() + res["p50_ms"] = round(lat[len(lat) // 2] * 1000, 3) + res["p99_ms"] = round(lat[int(len(lat) * 0.99)] * 1000, 3) + + t_end = time.time() + 10 + nq = 0 + while time.time() < t_end: + knn_ids(c, idx, queries[nq % len(queries)], ef_runtime=ef_runtime) + nq += 1 + res["qps"] = round(nq / 10.0, 1) + return res + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--moon-bin") + ap.add_argument("--redis-bin") + ap.add_argument("--n-vectors", type=int, default=20000) + ap.add_argument("--port", type=int, default=6460) + ap.add_argument("--out", default="vector-vs-redisearch.json") + ap.add_argument("--only", choices=["moon", "redisearch"], + help="bench a single engine (smoke tests)") + ap.add_argument("--moon-args", default="", + help="extra CLI args appended to the moon server command " + "(e.g. '--ft-search-workers 0' for A/B runs)") + args = ap.parse_args() + if args.only != "redisearch" and not args.moon_bin: + ap.error("--moon-bin required") + if args.only != "moon" and not args.redis_bin: + ap.error("--redis-bin required") + + report = {"machine": os.uname().machine, "n_vectors": args.n_vectors} + queries_per_ds = 200 + + for ds_name, gen in (("gaussian", gen_gaussian), ("clustered", gen_clustered)): + vecs = gen(args.n_vectors, seed=42) + queries = gen(queries_per_ds, seed=7) + gt = cosine_gt(queries, vecs, K) + + # ── RediSearch (FLOAT32 HNSW) ── + if args.only != "moon": + d = tempfile.mkdtemp(prefix="redisearch-") + argv = [args.redis_bin, "--port", str(args.port), "--dir", d, + "--save", "", "--appendonly", "no"] + so = os.environ.get("REDISEARCH_SO") + if so: + argv += ["--loadmodule", so] + srv = Server(argv, args.port, os.path.join(d, "log")) + try: + srv.wait_ready() + c = Resp(args.port) + r = c.cmd( + "FT.CREATE", "r", "ON", "HASH", "PREFIX", "1", "r:", + "SCHEMA", "vec", "VECTOR", "HNSW", "6", + "TYPE", "FLOAT32", "DIM", DIM, "DISTANCE_METRIC", "COSINE", + ) + if isinstance(r, Exception): + raise RuntimeError(f"redisearch FT.CREATE: {r}") + base = bench_engine(c, "r", vecs, queries, gt) + report[f"{ds_name}/redisearch/ef-default"] = base + print(f"[{ds_name}] redisearch/ef-default: {base}", flush=True) + # Sweep EF_RUNTIME to find RediSearch's recall-matched operating + # point — QPS is only comparable at equal recall. + for ef in (64, 128, 256, 512): + hi = bench_engine(c, "r", vecs, queries, gt, ef_runtime=ef, insert=False) + report[f"{ds_name}/redisearch/ef-{ef}"] = hi + print(f"[{ds_name}] redisearch/ef-{ef}: {hi}", flush=True) + c.close() + finally: + srv.stop() + shutil.rmtree(d, ignore_errors=True) + + # ── Moon SQ8 / TQ4 ── + for quant in ("SQ8", "TQ4") if args.only != "redisearch" else (): + d = tempfile.mkdtemp(prefix=f"moon-{quant}-") + argv = [args.moon_bin, "--port", str(args.port), "--shards", "1", + "--admin-port", "0", "--appendonly", "no", "--dir", d] + argv += args.moon_args.split() + srv = Server(argv, args.port, os.path.join(d, "log")) + try: + srv.wait_ready() + c = Resp(args.port) + r = c.cmd( + "FT.CREATE", "m", "ON", "HASH", "PREFIX", "1", "m:", + "SCHEMA", "vec", "VECTOR", "HNSW", "10", + "TYPE", "FLOAT32", "DIM", DIM, "DISTANCE_METRIC", "COSINE", + "QUANTIZATION", quant, "COMPACT_THRESHOLD", "4096", + ) + if isinstance(r, Exception): + raise RuntimeError(f"moon FT.CREATE: {r}") + res = bench_engine(c, "m", vecs, queries, gt) + c.cmd("FT.COMPACT", "m") + time.sleep(2) + post = bench_engine(c, "m", vecs, queries, gt, insert=False) + res.update({f"compacted_{k}": v for k, v in post.items()}) + report[f"{ds_name}/moon/{quant}"] = res + print(f"[{ds_name}] moon/{quant}: {res}", flush=True) + c.close() + finally: + srv.stop() + shutil.rmtree(d, ignore_errors=True) + + with open(args.out, "w") as f: + json.dump(report, f, indent=2) + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/scripts/gcloud-vector-soak.sh b/scripts/gcloud-vector-soak.sh new file mode 100755 index 000000000..19273c69f --- /dev/null +++ b/scripts/gcloud-vector-soak.sh @@ -0,0 +1,200 @@ +#!/usr/bin/env bash +# gcloud-vector-soak.sh — vector-search reliability/stability/durability validation on GCloud, cross-arch. +# +# Runs scripts/vector-validate.py (recall/QPS vs baseline + long soak + kill-9 durability) on +# fresh GCE instances, comparing the CURRENT (possibly unpushed) branch against the v0.5.1 tag. +# The branch is shipped via `git bundle` so nothing needs to be pushed to GitHub first. +# +# Pattern follows gcloud-kv-scale-bench.sh (provision / SSH-wait / toolchain / teardown trap), +# plus the standing gotchas: ELF-magic assert (Mach-O trap), kill -9 only (SIGTERM+SO_REUSEPORT +# hang), distinct binary basenames (moon-base / moon-branch) so pkill backstops can't cross-kill. +# +# Subcommands: +# --self-test local gates: refs exist, bundle round-trips, driver compiles (NO GCloud, NO cost) +# --remote INNER; assumes it runs ON a provisioned Linux instance +# --one OUTER for a single $GCE_MACHINE +# --gcloud (default) OUTER sweep over $MACHINES sequentially +set -uo pipefail + +BRANCH="${BRANCH:-perf/vector-search-optimization}" +BASE_REF="${BASE_REF:-v0.5.1}" +SOAK_MINUTES="${SOAK_MINUTES:-20}" +N_VECTORS="${N_VECTORS:-20000}" +GCE_MACHINE="${GCE_MACHINE:-c4a-standard-8}" +GCE_NAME="${GCE_NAME:-moon-vec-soak}" +GCE_ZONE="${GCE_ZONE:-us-central1-a}" +MACHINES="${MACHINES:-c4a-standard-8 c3-standard-8}" + +log(){ printf '%s\n' "$*" >&2; } +die(){ log "FATAL: $*"; exit 1; } + +# ============================================================================= self-test (no GCloud) +self_test(){ + local fails=0 tmp + _ok(){ printf ' ok %s\n' "$1" >&2; } + _bad(){ printf ' FAIL %s\n' "$1" >&2; fails=$((fails+1)); } + log "=== gcloud-vector-soak --self-test (refs + bundle + driver; no GCloud) ===" + + /usr/bin/git rev-parse -q --verify "$BRANCH" >/dev/null && _ok "branch $BRANCH exists" || _bad "branch $BRANCH missing" + /usr/bin/git rev-parse -q --verify "$BASE_REF" >/dev/null && _ok "base $BASE_REF exists" || _bad "base $BASE_REF missing" + + python3 -m py_compile scripts/vector-validate.py 2>/dev/null \ + && _ok "vector-validate.py compiles" || _bad "vector-validate.py does not compile" + + # Bundle round-trip: both refs must be clonable from the bundle. + tmp=$(mktemp -d) + if /usr/bin/git bundle create "$tmp/m.bundle" "$BASE_REF" "$BRANCH" >/dev/null 2>&1 \ + && /usr/bin/git clone -q -b "${BRANCH##*/}" "$tmp/m.bundle" "$tmp/clone" 2>/dev/null \ + && /usr/bin/git -C "$tmp/clone" rev-parse -q --verify "$BASE_REF" >/dev/null; then + _ok "bundle carries $BASE_REF + $BRANCH" + else + # Bundles register branch refs verbatim; retry with the full ref name. + if /usr/bin/git clone -q -b "$BRANCH" "$tmp/m.bundle" "$tmp/clone2" 2>/dev/null \ + && /usr/bin/git -C "$tmp/clone2" rev-parse -q --verify "$BASE_REF" >/dev/null; then + _ok "bundle carries $BASE_REF + $BRANCH (full ref name)" + else + _bad "bundle round-trip failed" + fi + fi + rm -rf "$tmp" + + # Result-parse gate: the driver's pass/fail JSON must be machine-readable. + echo '{"pass": true, "failures": []}' | python3 -c 'import json,sys; d=json.load(sys.stdin); sys.exit(0 if d["pass"] else 1)' \ + && _ok "result JSON parse" || _bad "result JSON parse" + + [[ "$fails" -eq 0 ]] && { log "=== self-test PASS ==="; return 0; } + log "=== self-test FAIL ($fails) ==="; return 1 +} + +# ============================================================================= inner: on the instance +elf_assert(){ + local magic; magic=$(od -An -tx1 -N4 "$1" | tr -d ' ') + [[ "$magic" == "7f454c46" ]] || die "$1 not ELF (magic=$magic) — stale Mach-O?" +} + +remote_run(){ + source "$HOME/.cargo/env" 2>/dev/null || true + command -v cargo >/dev/null || die "cargo missing" + [[ -f "$HOME/moon.bundle" ]] || die "moon.bundle not shipped" + + # Stale-server guard (leaked busy-poller gotcha): nothing moon-like may be running. + pkill -9 -f 'moon-base|moon-branch' 2>/dev/null || true + + local src="$HOME/moon-src" + if [[ ! -d "$src/.git" ]]; then + git clone -q -b "$BRANCH" "$HOME/moon.bundle" "$src" 2>/dev/null \ + || git clone -q -b "${BRANCH##*/}" "$HOME/moon.bundle" "$src" \ + || die "clone from bundle failed" + fi + cd "$src" + log "=== source: branch $(git rev-parse --short HEAD) / base $(git rev-parse --short "$BASE_REF") ===" + + # Build BASELINE first, park the binary, then build the branch. + if [[ ! -x "$HOME/moon-base" ]]; then + git checkout -q "$BASE_REF" || die "checkout $BASE_REF" + log "=== building baseline ($BASE_REF) ===" + cargo build --release >/dev/null 2>&1 || die "baseline build failed" + cp target/release/moon "$HOME/moon-base" + fi + git checkout -q "$BRANCH" 2>/dev/null || git checkout -q "${BRANCH##*/}" || die "checkout $BRANCH" + if [[ ! -x "$HOME/moon-branch" ]]; then + log "=== building branch ($BRANCH) ===" + cargo build --release >/dev/null 2>&1 || die "branch build failed" + cp target/release/moon "$HOME/moon-branch" + fi + elf_assert "$HOME/moon-base" + elf_assert "$HOME/moon-branch" + + log "=== running vector-validate.py (recall + ${SOAK_MINUTES}m soak + durability) ===" + python3 "$HOME/vector-validate.py" \ + --moon-bin "$HOME/moon-branch" --baseline-bin "$HOME/moon-base" \ + --soak-minutes "$SOAK_MINUTES" --n-vectors "$N_VECTORS" \ + --out "$HOME/results.json" + local rc=$? + pkill -9 -f 'moon-base|moon-branch' 2>/dev/null || true + log "=== validate exited rc=$rc ===" + return "$rc" +} + +# ============================================================================= gcloud orchestration +arch_image_family(){ case "$1" in c4a-*|t2a-*|*arm*|*arm64*) echo ubuntu-2404-lts-arm64;; *) echo ubuntu-2404-lts-amd64;; esac; } +# c4a (Axion) rejects pd-ssd — hyperdisk-balanced is its only SSD-class boot disk. +boot_disk_type(){ case "$1" in c4a-*|c4-*) echo hyperdisk-balanced;; *) echo pd-ssd;; esac; } + +gcloud_run_one(){ + command -v gcloud >/dev/null || die "gcloud CLI not found" + local imgfam; imgfam=$(arch_image_family "$GCE_MACHINE") + local GSSH=(gcloud compute ssh "$GCE_NAME" --zone="$GCE_ZONE" --quiet + --ssh-flag=-oStrictHostKeyChecking=no --ssh-flag=-oConnectTimeout=15) + + mkdir -p tmp + log "=== bundling $BASE_REF + $BRANCH ===" + /usr/bin/git bundle create "tmp/moon-vec-$GCE_NAME.bundle" "$BASE_REF" "$BRANCH" || die "git bundle failed" + + log "=== provisioning $GCE_NAME ($GCE_MACHINE, $GCE_ZONE, $imgfam) ===" + gcloud compute instances create "$GCE_NAME" \ + --machine-type="$GCE_MACHINE" --zone="$GCE_ZONE" \ + --image-family="$imgfam" --image-project=ubuntu-os-cloud \ + --boot-disk-size=50GB --boot-disk-type="$(boot_disk_type "$GCE_MACHINE")" --quiet \ + || die "instance create failed" + + # Pin the actual name into the teardown trap NOW (env-prefix reverts GCE_NAME at fire time). + local _inst="$GCE_NAME" _zone="$GCE_ZONE" + trap "log '=== tearing down $_inst ==='; gcloud compute instances delete '$_inst' --zone='$_zone' -q 2>/dev/null || true" EXIT INT TERM + + log "=== waiting for SSH ===" + local tries=0 + until "${GSSH[@]}" --command='echo up' >/dev/null 2>&1; do + tries=$((tries+1)); [[ $tries -ge 40 ]] && die "SSH never came up"; sleep 5 + done + + log "=== provisioning toolchain (rust + numpy) ===" + "${GSSH[@]}" --command=' + set -e + sudo apt-get update -qq + sudo apt-get install -y -qq build-essential pkg-config libssl-dev git curl ca-certificates python3-numpy + command -v cargo >/dev/null || curl --proto "=https" --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain 1.94.1 + ' || die "instance provisioning failed" + + log "=== pushing bundle + harness ===" + gcloud compute scp "tmp/moon-vec-$GCE_NAME.bundle" scripts/vector-validate.py scripts/gcloud-vector-soak.sh \ + "$GCE_NAME":~/ --zone="$GCE_ZONE" --quiet --scp-flag=-oStrictHostKeyChecking=no || die "scp failed" + "${GSSH[@]}" --command="mv ~/moon-vec-$GCE_NAME.bundle ~/moon.bundle" || die "bundle rename failed" + + local rawfile="tmp/vec-soak-${GCE_MACHINE}.log" + log "=== running --remote on $GCE_MACHINE (build x2 + validate; the long part) -> $rawfile ===" + "${GSSH[@]}" --command=" + BRANCH='$BRANCH' BASE_REF='$BASE_REF' SOAK_MINUTES='$SOAK_MINUTES' N_VECTORS='$N_VECTORS' \ + bash ~/gcloud-vector-soak.sh --remote + " 2>&1 | tee "$rawfile" + local rc=${PIPESTATUS[0]} + + log "=== fetching results.json ===" + gcloud compute scp "$GCE_NAME":~/results.json "tmp/vec-soak-${GCE_MACHINE}.json" \ + --zone="$GCE_ZONE" --quiet --scp-flag=-oStrictHostKeyChecking=no \ + || log "WARN: results.json fetch failed (validate rc=$rc)" + log "=== $GCE_MACHINE done (rc=$rc); results tmp/vec-soak-${GCE_MACHINE}.json; teardown follows (trap) ===" + return "$rc" +} + +gcloud_sweep(){ + log "=== VECTOR-SOAK SWEEP: $MACHINES ===" + local m short overall=0 + for m in $MACHINES; do + short="${m%%-*}" + log ""; log "##################### machine: $m #####################" + ( GCE_MACHINE="$m"; GCE_NAME="moon-vec-soak-${short}"; gcloud_run_one ) \ + || { overall=1; log "WARN: machine $m run failed (see log); continuing"; } + done + log "=== sweep complete; per-machine results: tmp/vec-soak-.json ===" + return "$overall" +} + +# ============================================================================= dispatch +case "${1:---gcloud}" in + --self-test) self_test ;; + --remote) remote_run ;; + --gcloud) gcloud_sweep ;; + --one) gcloud_run_one ;; + *) die "unknown subcommand: $1 (use --self-test | --remote | --gcloud | --one)" ;; +esac diff --git a/scripts/vector-validate.py b/scripts/vector-validate.py new file mode 100644 index 000000000..19f3c844b --- /dev/null +++ b/scripts/vector-validate.py @@ -0,0 +1,538 @@ +#!/usr/bin/env python3 +""" +vector-validate.py — long-run reliability / stability / durability validation +for Moon's vector search (FT.*), plus recall/QPS comparison between two moon +binaries (baseline vs branch). + +Runs ON the target machine (GCloud instance or Linux VM). Manages moon server +processes itself (spawn / kill -9 / restart) so the durability phase needs no +outer orchestration. + +Phases: + recall — per binary: fresh server, MiniLM-like 384d unit vectors, + SQ8 + TQ4 indexes, compaction, R@10 vs numpy brute force, + query latency p50/p99 + QPS. + soak — branch binary, appendonly yes: update/insert/delete/search + churn for SOAK_MINUTES; samples RSS, MEMORY DOCTOR, FT.INFO, + live-set recall every SAMPLE_SECS; flags RSS runaway, recall + drift, deleted-key resurrection. + durability — kill -9 mid-churn, restart on the same dir, verify recovery: + keys settled >FSYNC_MARGIN before the kill must survive + (appendonly everysec), survivor recall must stay >= floor. + +Usage: + python3 vector-validate.py --moon-bin PATH [--baseline-bin PATH] + [--phases recall,soak,durability] [--soak-minutes 20] [--port 6470] + [--n-vectors 20000] [--out results.json] + +Exit code 0 = all hard assertions passed. JSON report to --out. +""" + +import argparse +import json +import math +import os +import shutil +import signal +import socket +import subprocess +import sys +import tempfile +import time + +import numpy as np + +DIM = 384 +K = 10 +FSYNC_MARGIN = 5.0 # seconds a write must predate kill -9 to be "must survive" + + +# ── RESP client (proper incremental parser — FT.SEARCH replies are nested) ── +class Resp: + def __init__(self, port, timeout=30): + self.sock = socket.create_connection(("127.0.0.1", port), timeout=timeout) + self.sock.settimeout(timeout) + self.buf = b"" + + def close(self): + try: + self.sock.close() + except Exception: + pass + + def _read_more(self): + chunk = self.sock.recv(1 << 16) + if not chunk: + raise ConnectionError("server closed connection") + self.buf += chunk + + def _read_line(self): + while b"\r\n" not in self.buf: + self._read_more() + line, self.buf = self.buf.split(b"\r\n", 1) + return line + + def _read_exact(self, n): + while len(self.buf) < n + 2: + self._read_more() + data, self.buf = self.buf[:n], self.buf[n + 2 :] + return data + + def _parse(self): + line = self._read_line() + t, rest = line[:1], line[1:] + if t == b"+": + return rest.decode() + if t == b"-": + return Exception(rest.decode()) + if t == b":": + return int(rest) + if t == b"$": + n = int(rest) + return None if n == -1 else self._read_exact(n) + if t == b"*": + n = int(rest) + return None if n == -1 else [self._parse() for _ in range(n)] + if t == b"%": # RESP3 map + n = int(rest) + return [self._parse() for _ in range(2 * n)] + raise ValueError(f"unexpected RESP type {line!r}") + + @staticmethod + def encode(*args): + parts = [f"*{len(args)}\r\n".encode()] + for a in args: + b = a if isinstance(a, bytes) else str(a).encode() + parts += [f"${len(b)}\r\n".encode(), b, b"\r\n"] + return b"".join(parts) + + def cmd(self, *args): + self.sock.sendall(self.encode(*args)) + return self._parse() + + def pipeline(self, cmds): + self.sock.sendall(b"".join(self.encode(*c) for c in cmds)) + return [self._parse() for _ in cmds] + + +# ── Server lifecycle ── +class Moon: + def __init__(self, binary, port, data_dir, appendonly="no", extra=()): + self.binary, self.port, self.data_dir = binary, port, data_dir + self.appendonly = appendonly + self.extra = list(extra) + self.proc = None + self.log = open(os.path.join(data_dir, "moon.log"), "ab") + + def start(self, wait=30): + args = [ + self.binary, + "--port", str(self.port), + "--shards", "1", + "--admin-port", "0", + "--appendonly", self.appendonly, + "--dir", self.data_dir, + ] + self.extra + self.proc = subprocess.Popen(args, stdout=self.log, stderr=self.log) + deadline = time.time() + wait + while time.time() < deadline: + try: + c = Resp(self.port, timeout=2) + if c.cmd("PING") == "PONG": + c.close() + return + c.close() + except Exception: + pass + if self.proc.poll() is not None: + raise RuntimeError(f"moon exited early rc={self.proc.returncode}") + time.sleep(0.2) + raise RuntimeError(f"moon did not answer PING on {self.port} within {wait}s") + + def kill9(self): + if self.proc and self.proc.poll() is None: + self.proc.send_signal(signal.SIGKILL) + self.proc.wait(10) + + def stop(self): + self.kill9() # tests always hard-kill (SIGTERM+SO_REUSEPORT hang gotcha) + + +# ── Data ── +def gen_unit(n, seed): + rng = np.random.default_rng(seed) + v = rng.standard_normal((n, DIM), dtype=np.float32) + v /= np.linalg.norm(v, axis=1, keepdims=True) + return v + + +def cosine_gt(queries, db_matrix, k): + """Ground-truth top-k ids by cosine distance (rows of db_matrix are unit).""" + sims = queries @ db_matrix.T + return np.argsort(-sims, axis=1)[:, :k] + + +def ft_create(c, idx, quant): + r = c.cmd( + "FT.CREATE", idx, "ON", "HASH", "PREFIX", "1", f"{idx}:", + "SCHEMA", "vec", "VECTOR", "HNSW", "10", + "TYPE", "FLOAT32", "DIM", DIM, "DISTANCE_METRIC", "COSINE", + "QUANTIZATION", quant, "COMPACT_THRESHOLD", "4096", + ) + if isinstance(r, Exception): + raise r + + +def insert_batch(c, idx, ids, vecs): + cmds = [ + ("HSET", f"{idx}:{i}", "vec", vecs[j].tobytes()) + for j, i in enumerate(ids) + ] + for r in c.pipeline(cmds): + if isinstance(r, Exception): + raise r + + +def knn(c, idx, q, k=K, timeout_note=""): + r = c.cmd( + "FT.SEARCH", idx, f"*=>[KNN {k} @vec $BLOB]", + "PARAMS", "2", "BLOB", q.tobytes(), "DIALECT", "2", + ) + if isinstance(r, Exception): + raise RuntimeError(f"FT.SEARCH failed{timeout_note}: {r}") + # reply: [total, key1, fields1, key2, fields2, ...] + keys = [r[i].decode() for i in range(1, len(r), 2)] + return [int(kk.split(":", 1)[1]) for kk in keys if ":" in kk] + + +def ft_info(c, idx): + r = c.cmd("FT.INFO", idx) + if isinstance(r, Exception): + return {} + out = {} + i = 0 + while i + 1 < len(r): + key = r[i].decode() if isinstance(r[i], bytes) else str(r[i]) + val = r[i + 1] + if isinstance(val, bytes): + val = val.decode() + if not isinstance(val, list): + out[key] = val + i += 2 + return out + + +def rss_mb(c): + r = c.cmd("MEMORY", "DOCTOR") + if isinstance(r, (bytes, str)): + text = r.decode() if isinstance(r, bytes) else r + for line in text.splitlines(): + if "RSS:" in line: + parts = line.split() + try: + val = float(parts[-2]) + unit = parts[-1] + return val * (1024 if unit == "GB" else 1) if unit in ("MB", "GB") else val / 1024 + except (ValueError, IndexError): + pass + return -1.0 + + +# ── Phase: recall/QPS ── +def phase_recall(binary, label, port, n_vectors, report): + db = gen_unit(n_vectors, seed=42) + queries = gen_unit(200, seed=7) + gt = cosine_gt(queries, db, K) + + for quant in ("SQ8", "TQ4"): + d = tempfile.mkdtemp(prefix=f"moon-recall-{label}-{quant}-") + srv = Moon(binary, port, d) + srv.start() + try: + c = Resp(port, timeout=60) + idx = f"r{quant.lower()}" + ft_create(c, idx, quant) + + t0 = time.time() + for s in range(0, n_vectors, 500): + ids = range(s, min(s + 500, n_vectors)) + insert_batch(c, idx, ids, db[s:]) + insert_secs = time.time() - t0 + c.cmd("FT.COMPACT", idx) + time.sleep(2) # let background build install + + # Recall + single-query latency + lat = [] + hits = 0 + for qi in range(len(queries)): + t0 = time.time() + got = knn(c, idx, queries[qi]) + lat.append(time.time() - t0) + hits += len(set(got[:K]) & set(gt[qi].tolist())) + recall = hits / (len(queries) * K) + + # QPS: 10s tight loop + t_end = time.time() + 10 + nq = 0 + while time.time() < t_end: + knn(c, idx, queries[nq % len(queries)]) + nq += 1 + qps = nq / 10.0 + + lat.sort() + info = ft_info(c, idx) + report[f"recall/{label}/{quant}"] = { + "recall_at_10": round(recall, 4), + "qps": round(qps, 1), + "p50_ms": round(lat[len(lat) // 2] * 1000, 3), + "p99_ms": round(lat[int(len(lat) * 0.99)] * 1000, 3), + "insert_secs": round(insert_secs, 1), + "num_docs": info.get("num_docs"), + } + print(f"[recall] {label}/{quant}: R@10={recall:.4f} qps={qps:.0f} " + f"p50={lat[len(lat)//2]*1000:.2f}ms", flush=True) + c.close() + finally: + srv.stop() + shutil.rmtree(d, ignore_errors=True) + + +# ── Phase: soak ── +def phase_soak(binary, port, minutes, n_vectors, report): + d = tempfile.mkdtemp(prefix="moon-soak-") + srv = Moon(binary, port, d, appendonly="yes") + srv.start() + failures = [] + warnings = [] + samples = [] + try: + c = Resp(port, timeout=60) + idx = "soak" + ft_create(c, idx, "SQ8") + + rng = np.random.default_rng(99) + live = {} # id -> vector row (numpy) + id_pool = [] # O(1) random sampling; kept in sync with `live` + deleted = set() + + # Seed data + seed_vecs = gen_unit(n_vectors, seed=11) + for s in range(0, n_vectors, 500): + ids = list(range(s, min(s + 500, n_vectors))) + insert_batch(c, idx, ids, seed_vecs[s:]) + for i in ids: + live[i] = seed_vecs[i] + id_pool.append(i) + next_id = n_vectors + probe_q = gen_unit(20, seed=3) + + t_end = time.time() + minutes * 60 + t_sample = 0.0 + ops = 0 + while time.time() < t_end: + r = rng.random() + if r < 0.60: # search + knn(c, idx, probe_q[ops % len(probe_q)]) + elif r < 0.85 and id_pool: # update existing (tombstone pressure) + i = id_pool[int(rng.integers(len(id_pool)))] + v = gen_unit(1, seed=int(rng.integers(1 << 30)))[0] + cr = c.cmd("HSET", f"{idx}:{i}", "vec", v.tobytes()) + if isinstance(cr, Exception): + failures.append(f"HSET update failed: {cr}") + else: + live[i] = v + elif r < 0.95: # insert new + v = gen_unit(1, seed=int(rng.integers(1 << 30)))[0] + cr = c.cmd("HSET", f"{idx}:{next_id}", "vec", v.tobytes()) + if not isinstance(cr, Exception): + live[next_id] = v + id_pool.append(next_id) + next_id += 1 + elif id_pool: # delete (swap-remove from the pool) + j = int(rng.integers(len(id_pool))) + i = id_pool[j] + id_pool[j] = id_pool[-1] + id_pool.pop() + c.cmd("DEL", f"{idx}:{i}") + live.pop(i, None) + deleted.add(i) + ops += 1 + + now = time.time() + if now - t_sample >= 60: + t_sample = now + mat = np.stack(list(live.values())) + id_list = np.array(list(live.keys())) + gt = cosine_gt(probe_q, mat, K) + hits, resurrections = 0, 0 + for qi in range(len(probe_q)): + got = knn(c, idx, probe_q[qi]) + truth = set(id_list[gt[qi]].tolist()) + hits += len(set(got[:K]) & truth) + resurrections += sum(1 for g in got if g in deleted) + recall = hits / (len(probe_q) * K) + info = ft_info(c, idx) + mem = rss_mb(c) + sample = { + "t": round(now - (t_end - minutes * 60), 1), + "ops": ops, + "live": len(live), + "recall": round(recall, 4), + "rss_mb": round(mem, 1), + "num_docs": info.get("num_docs"), + "resurrected": resurrections, + } + samples.append(sample) + print(f"[soak] {json.dumps(sample)}", flush=True) + if resurrections: + failures.append(f"deleted keys resurfaced in search: {sample}") + if recall < 0.60: # catastrophic collapse guard + failures.append(f"recall collapsed below 0.60: {sample}") + elif recall < 0.85: + # Warning only: on random-Gaussian 384d, live-set recall + # legitimately declines as N grows (concentration of + # distances). Data LOSS is judged by the self-recall probe. + warnings.append(f"recall below 0.85 (grew to {sample['live']} live): {sample}") + + # Self-recall LOST probe: every sampled live key, queried by its own + # CURRENT vector, must appear in its own top-10. A key that fails is + # GONE from the index (the direct mass-loss detector — recall drift + # can be dataset noise; lost keys cannot). + probe_ids = list(live.keys()) + if len(probe_ids) > 1000: + probe_ids = [probe_ids[int(i)] for i in + rng.choice(len(probe_ids), size=1000, replace=False)] + lost = sum(1 for i in probe_ids if i not in knn(c, idx, live[i])) + lost_frac = lost / max(1, len(probe_ids)) + report["soak_lost_probe"] = {"sampled": len(probe_ids), "lost": lost} + print(f"[soak] lost-probe: {lost}/{len(probe_ids)} sampled live keys missing", + flush=True) + if lost_frac > 0.005: + failures.append( + f"index lost {lost}/{len(probe_ids)} sampled live keys (>0.5%)") + + # RSS runaway check: last sample vs first, adjusted for growth in live set + if len(samples) >= 2 and samples[0]["rss_mb"] > 0: + growth = samples[-1]["rss_mb"] / samples[0]["rss_mb"] + live_growth = max(1.0, samples[-1]["live"] / max(1, samples[0]["live"])) + if growth > 3.0 * live_growth: + failures.append( + f"RSS runaway: {samples[0]['rss_mb']}MB -> {samples[-1]['rss_mb']}MB " + f"(live set only grew {live_growth:.2f}x)" + ) + c.close() + finally: + srv.stop() + shutil.rmtree(d, ignore_errors=True) + report["soak"] = {"samples": samples, "failures": failures, "warnings": warnings} + return failures + + +# ── Phase: durability ── +def phase_durability(binary, port, report): + d = tempfile.mkdtemp(prefix="moon-dur-") + failures = [] + srv = Moon(binary, port, d, appendonly="yes") + srv.start() + try: + c = Resp(port, timeout=60) + idx = "dur" + ft_create(c, idx, "SQ8") + vecs = gen_unit(5000, seed=21) + for s in range(0, 5000, 500): + insert_batch(c, idx, range(s, s + 500), vecs[s:]) + settled = 5000 # ids 0..4999 written now + time.sleep(FSYNC_MARGIN) # everysec fsync window + margin + + # churn a bit more (these may or may not survive), then SIGKILL + extra = gen_unit(500, seed=22) + insert_batch(c, idx, range(5000, 5500), extra) + c.close() + srv.kill9() + print("[durability] killed -9 mid-churn", flush=True) + + # restart on same dir + srv2 = Moon(binary, port, d, appendonly="yes") + srv2.start(wait=60) + c = Resp(port, timeout=60) + + # every settled key must be readable + missing = 0 + for s in range(0, settled, 500): + cmds = [("HEXISTS", f"{idx}:{i}", "vec") for i in range(s, s + 500)] + missing += sum(1 for r in c.pipeline(cmds) if r != 1) + if missing: + failures.append(f"{missing}/{settled} settled keys missing after kill -9 + restart") + + # recall of survivors must hold (validates index rebuild incl. rerank) + queries = gen_unit(50, seed=5) + gt = cosine_gt(queries, vecs[:settled], K) + hits = 0 + for qi in range(len(queries)): + got = knn(c, idx, queries[qi], timeout_note=" (post-recovery)") + hits += len(set(g for g in got[:K] if g < settled) & set(gt[qi].tolist())) + recall = hits / (len(queries) * K) + if recall < 0.85: + failures.append(f"post-recovery recall {recall:.4f} < 0.85") + print(f"[durability] post-recovery: missing={missing} recall={recall:.4f}", flush=True) + + # double-crash: kill again right after recovery, restart once more + c.close() + srv2.kill9() + srv3 = Moon(binary, port, d, appendonly="yes") + srv3.start(wait=60) + c = Resp(port, timeout=60) + if c.cmd("PING") != "PONG": + failures.append("server unresponsive after double crash-recovery") + c.close() + srv3.stop() + report["durability"] = { + "settled": settled, + "missing": missing, + "post_recovery_recall": round(recall, 4), + "failures": failures, + } + finally: + srv.kill9() + # -x (exact comm match), NOT -f: the driver's own cmdline contains the + # binary path via --moon-bin, so -f SIGKILLs the driver itself (rc=137). + subprocess.run(["pkill", "-9", "-x", os.path.basename(binary)], check=False) + shutil.rmtree(d, ignore_errors=True) + return failures + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--moon-bin", required=True) + ap.add_argument("--baseline-bin") + ap.add_argument("--phases", default="recall,soak,durability") + ap.add_argument("--soak-minutes", type=float, default=20) + ap.add_argument("--port", type=int, default=6470) + ap.add_argument("--n-vectors", type=int, default=20000) + ap.add_argument("--out", default="vector-validate-results.json") + args = ap.parse_args() + + phases = set(args.phases.split(",")) + report = {"host": os.uname().nodename, "machine": os.uname().machine} + all_failures = [] + + if "recall" in phases: + if args.baseline_bin: + phase_recall(args.baseline_bin, "baseline", args.port, args.n_vectors, report) + phase_recall(args.moon_bin, "branch", args.port, args.n_vectors, report) + if "soak" in phases: + all_failures += phase_soak(args.moon_bin, args.port, args.soak_minutes, + args.n_vectors, report) + if "durability" in phases: + all_failures += phase_durability(args.moon_bin, args.port, report) + + report["failures"] = all_failures + report["pass"] = not all_failures + with open(args.out, "w") as f: + json.dump(report, f, indent=2) + print(json.dumps({"pass": report["pass"], "failures": all_failures}, indent=2)) + sys.exit(0 if report["pass"] else 1) + + +if __name__ == "__main__": + main() diff --git a/src/command/connection.rs b/src/command/connection.rs index 794193531..694fcc7b8 100644 --- a/src/command/connection.rs +++ b/src/command/connection.rs @@ -240,9 +240,9 @@ pub fn info(db: &Database, _args: &[Frame]) -> Frame { vector_compaction_duration_ms:{}\r\n\ vector_mutable_segment_bytes:{}\r\n", crate::vector::metrics::VECTOR_INDEXES.load(std::sync::atomic::Ordering::Relaxed), - crate::vector::metrics::VECTOR_TOTAL_VECTORS.load(std::sync::atomic::Ordering::Relaxed), + crate::vector::metrics::total_vectors(), crate::vector::metrics::VECTOR_MEMORY_BYTES.load(std::sync::atomic::Ordering::Relaxed), - crate::vector::metrics::VECTOR_SEARCH_TOTAL.load(std::sync::atomic::Ordering::Relaxed), + crate::vector::metrics::search_total(), crate::vector::metrics::VECTOR_SEARCH_LATENCY_US.load(std::sync::atomic::Ordering::Relaxed), crate::vector::metrics::VECTOR_COMPACTION_COUNT.load(std::sync::atomic::Ordering::Relaxed), crate::vector::metrics::VECTOR_COMPACTION_DURATION_MS diff --git a/src/command/vector_search/ft_config.rs b/src/command/vector_search/ft_config.rs index 9816dd876..9cb5a875e 100644 --- a/src/command/vector_search/ft_config.rs +++ b/src/command/vector_search/ft_config.rs @@ -131,6 +131,31 @@ fn ft_config_set( } Err(e) => Frame::Error(Bytes::from(format!("ERR {e}").into_bytes())), } + } else if param.eq_ignore_ascii_case(b"MERGE_RECALL_TOLERANCE") { + // VEC-4: recall gate for UNATTENDED (background/vacuum) GraphUnion + // merges. Default 0.70 catches only catastrophic collapse; operators + // running recall-sensitive workloads can tighten toward the manual + // FT.COMPACT gate (0.90). Raising it can make auto-merge abort + // repeatedly on small indexes (segments stay > threshold) — informed + // trade-off, hence a knob rather than a new default. + let parsed: f32 = match std::str::from_utf8(value) + .ok() + .and_then(|s| s.parse::().ok()) + { + Some(v) => v, + None => { + return Frame::Error(Bytes::from_static( + b"ERR MERGE_RECALL_TOLERANCE must be a number", + )); + } + }; + if !(0.0..=1.0).contains(&parsed) { + return Frame::Error(Bytes::from_static( + b"ERR MERGE_RECALL_TOLERANCE must be between 0.0 and 1.0", + )); + } + idx.merge_recall_tolerance = parsed; + Frame::SimpleString(Bytes::from_static(b"OK")) } else { Frame::Error(Bytes::from_static(b"ERR unknown config parameter")) } @@ -171,6 +196,11 @@ fn ft_config_get( use std::fmt::Write as _; let _ = write!(buf, "{}", idx.compaction_weight()); Frame::BulkString(Bytes::from(buf)) + } else if param.eq_ignore_ascii_case(b"MERGE_RECALL_TOLERANCE") { + let mut buf = String::with_capacity(8); + use std::fmt::Write as _; + let _ = write!(buf, "{}", idx.merge_recall_tolerance); + Frame::BulkString(Bytes::from(buf)) } else { Frame::Error(Bytes::from_static(b"ERR unknown config parameter")) } diff --git a/src/command/vector_search/ft_create.rs b/src/command/vector_search/ft_create.rs index 049c9e76a..0b2f9ead6 100644 --- a/src/command/vector_search/ft_create.rs +++ b/src/command/vector_search/ft_create.rs @@ -692,6 +692,15 @@ fn parse_vector_field_params(args: &[Frame], pos: &mut usize) -> Result Frame { + const ADDITIVE_TOP: &[&[u8]] = &[ + b"num_docs", + b"num_terms", + b"total_inverted_index_size", + b"vector_version_token", + b"text_version_token", + ]; + const ADDITIVE_FIELD: &[&[u8]] = &[b"num_docs", b"mutable_vectors", b"immutable_segments"]; + + if matches!(local, Frame::Error(_)) { + return local; + } + if let Some(err) = remotes.iter().find(|r| matches!(r, Frame::Error(_))) { + return err.clone(); + } + let mut items: Vec = match &local { + Frame::Array(a) => a.to_vec(), + _ => return local, + }; + + // Value-index of `key` in an alternating key/value frame list. + fn value_idx(items: &[Frame], key: &[u8]) -> Option { + let mut i = 0; + while i + 1 < items.len() { + if let Frame::BulkString(k) = &items[i] { + if k.as_ref() == key { + return Some(i + 1); + } + } + i += 2; + } + None + } + + fn int_at(items: &[Frame], idx: usize) -> Option { + match items.get(idx) { + Some(Frame::Integer(n)) => Some(*n), + _ => None, + } + } + + // `field_name` value of one per-field entry array. + fn entry_field_name(entry: &Frame) -> Option { + if let Frame::Array(pairs) = entry { + if let Some(vi) = value_idx(pairs, b"field_name") { + if let Some(Frame::BulkString(name)) = pairs.get(vi) { + return Some(name.clone()); + } + } + } + None + } + + for remote in remotes { + let r_items: &[Frame] = match remote { + Frame::Array(a) => a, + _ => continue, + }; + for key in ADDITIVE_TOP { + if let (Some(li), Some(ri)) = (value_idx(&items, key), value_idx(r_items, key)) { + if let (Some(lv), Some(rv)) = (int_at(&items, li), int_at(r_items, ri)) { + items[li] = Frame::Integer(lv.saturating_add(rv)); + } + } + } + for list_key in [b"vector_fields".as_slice(), b"text_fields".as_slice()] { + let (Some(li), Some(ri)) = (value_idx(&items, list_key), value_idx(r_items, list_key)) + else { + continue; + }; + let remote_entries: Vec = match r_items.get(ri) { + Some(Frame::Array(a)) => a.to_vec(), + _ => continue, + }; + let Some(Frame::Array(local_entries)) = items.get(li) else { + continue; + }; + let mut local_entries: Vec = local_entries.to_vec(); + for le in local_entries.iter_mut() { + let Some(name) = entry_field_name(le) else { + continue; + }; + let Some(re) = remote_entries + .iter() + .find(|re| entry_field_name(re).as_deref() == Some(name.as_ref())) + else { + continue; + }; + let (Frame::Array(lp), Frame::Array(rp)) = (&*le, re) else { + continue; + }; + let mut pairs: Vec = lp.to_vec(); + for key in ADDITIVE_FIELD { + if let (Some(lpi), Some(rpi)) = (value_idx(&pairs, key), value_idx(rp, key)) { + if let (Some(lv), Some(rv)) = (int_at(&pairs, lpi), int_at(rp, rpi)) { + pairs[lpi] = Frame::Integer(lv.saturating_add(rv)); + } + } + } + *le = Frame::Array(pairs.into()); + } + items[li] = Frame::Array(local_entries.into()); + } + } + + Frame::Array(items.into()) +} + /// Full FT.INFO response for TEXT-only indexes. /// /// Returns index_name, num_docs, num_terms, per-field stats (num_docs, @@ -285,3 +410,105 @@ fn ft_info_text_only( Frame::Array(items.into()) } + +#[cfg(test)] +mod merge_tests { + use super::*; + + fn bs(s: &[u8]) -> Frame { + Frame::BulkString(Bytes::copy_from_slice(s)) + } + + fn info_frame(num_docs: i64, field_docs: i64, mutable: i64, imms: i64) -> Frame { + let field_entry = Frame::Array( + vec![ + bs(b"field_name"), + bs(b"vec"), + bs(b"dimension"), + Frame::Integer(8), + bs(b"num_docs"), + Frame::Integer(field_docs), + bs(b"mutable_vectors"), + Frame::Integer(mutable), + bs(b"immutable_segments"), + Frame::Integer(imms), + ] + .into(), + ); + Frame::Array( + vec![ + bs(b"index_name"), + bs(b"idx"), + bs(b"num_docs"), + Frame::Integer(num_docs), + bs(b"dimension"), + Frame::Integer(8), + bs(b"vector_fields"), + Frame::Array(vec![field_entry].into()), + bs(b"vector_version_token"), + Frame::Integer(7), + ] + .into(), + ) + } + + fn get_int(frame: &Frame, key: &[u8]) -> i64 { + let Frame::Array(items) = frame else { + panic!("not an array") + }; + let mut i = 0; + while i + 1 < items.len() { + if let Frame::BulkString(k) = &items[i] { + if k.as_ref() == key { + if let Frame::Integer(n) = &items[i + 1] { + return *n; + } + panic!("value for {key:?} not Integer"); + } + } + i += 2; + } + panic!("key {key:?} not found"); + } + + #[test] + fn sums_additive_fields_across_shards() { + let local = info_frame(10, 10, 4, 1); + let remotes = [info_frame(5, 5, 2, 1), info_frame(3, 3, 3, 0)]; + let merged = merge_ft_info_responses(local, &remotes); + assert_eq!(get_int(&merged, b"num_docs"), 18); + assert_eq!(get_int(&merged, b"vector_version_token"), 21); + // Config fields untouched. + assert_eq!(get_int(&merged, b"dimension"), 8); + // Nested per-field additivity. + let Frame::Array(items) = &merged else { + unreachable!() + }; + let vf_idx = items + .iter() + .position(|f| matches!(f, Frame::BulkString(b) if b.as_ref() == b"vector_fields")) + .unwrap(); + let Frame::Array(entries) = &items[vf_idx + 1] else { + panic!("vector_fields not array") + }; + assert_eq!(get_int(&entries[0], b"num_docs"), 18); + assert_eq!(get_int(&entries[0], b"mutable_vectors"), 9); + assert_eq!(get_int(&entries[0], b"immutable_segments"), 2); + } + + #[test] + fn propagates_remote_error() { + let local = info_frame(10, 10, 4, 1); + let remotes = [Frame::Error(Bytes::from_static(b"ERR boom"))]; + let merged = merge_ft_info_responses(local, &remotes); + assert!(matches!(merged, Frame::Error(_))); + } + + #[test] + fn no_remotes_is_identity() { + let local = info_frame(10, 10, 4, 1); + let merged = merge_ft_info_responses(local.clone(), &[]); + assert_eq!(get_int(&merged, b"num_docs"), 10); + assert_eq!(get_int(&merged, b"vector_version_token"), 7); + } +} diff --git a/src/command/vector_search/ft_search/dispatch.rs b/src/command/vector_search/ft_search/dispatch.rs index 961d6e385..5ecafd6fd 100644 --- a/src/command/vector_search/ft_search/dispatch.rs +++ b/src/command/vector_search/ft_search/dispatch.rs @@ -637,10 +637,15 @@ fn capture_dense_knn_snapshot( (base * dim_factor / 2).clamp(200, 1000) }; - let filter_bitmap = filter.map(|f| { - let total = idx.segments.total_vectors(); - idx.payload_index.evaluate_bitmap(f, total) - }); + let total_vectors = idx.segments.total_vectors(); + let filter_bitmap = filter.map(|f| idx.payload_index.evaluate_bitmap(f, total_vectors)); + // XC-3: resolve the selectivity-based strategy at capture, mirroring the + // sync path's `select_strategy` dispatch (holder.rs `search_filtered`). The + // yield refactor (PR #189) originally hardcoded ACORN-filtered search for + // every filtered query, losing the >80%-selectivity oversample+post-filter + // branch. + let filter_strategy = + crate::vector::filter::selectivity::select_strategy(filter_bitmap.as_ref(), total_vectors); let segments = idx.segments.load_full(); let mutable_len = segments.mutable.len(); @@ -651,6 +656,7 @@ fn capture_dense_knn_snapshot( k, ef_search, filter_bitmap, + filter_strategy, snapshot_lsn: as_of_lsn, my_txn_id: 0, committed, diff --git a/src/command/vector_search/ft_search/execute.rs b/src/command/vector_search/ft_search/execute.rs index 75d6b2911..cb8656f7e 100644 --- a/src/command/vector_search/ft_search/execute.rs +++ b/src/command/vector_search/ft_search/execute.rs @@ -25,7 +25,7 @@ use super::response::build_search_response; pub(super) enum SearchRawResult { Ok { results: SmallVec<[SearchResult; 32]>, - key_hash_to_key: std::collections::HashMap, + key_hash_to_key: std::sync::Arc>, }, Error(Frame), } diff --git a/src/command/vector_search/hybrid.rs b/src/command/vector_search/hybrid.rs index 97aad41fd..9401dd466 100644 --- a/src/command/vector_search/hybrid.rs +++ b/src/command/vector_search/hybrid.rs @@ -513,7 +513,13 @@ pub(super) fn run_dense_knn( k: usize, as_of_lsn: u64, committed: &roaring::RoaringTreemap, -) -> Result<(Vec, std::collections::HashMap), Frame> { +) -> Result< + ( + Vec, + std::sync::Arc>, + ), + Frame, +> { let field_opt = if field_name.is_empty() { None } else { diff --git a/src/command/vector_search/mod.rs b/src/command/vector_search/mod.rs index 8e3f8cc5f..1139efb6a 100644 --- a/src/command/vector_search/mod.rs +++ b/src/command/vector_search/mod.rs @@ -35,7 +35,7 @@ pub use ft_admin::{ft_compact, ft_dropindex, ft_list}; pub use ft_aggregate::ft_aggregate; pub use ft_config::ft_config; pub use ft_create::ft_create; -pub use ft_info::ft_info; +pub use ft_info::{ft_info, merge_ft_info_responses}; #[cfg(feature = "text-index")] pub use ft_invalidate_range::ft_invalidate_range; #[cfg(feature = "graph")] diff --git a/src/command/vector_search/tests.rs b/src/command/vector_search/tests.rs index 3c33aefe3..5aaff66ed 100644 --- a/src/command/vector_search/tests.rs +++ b/src/command/vector_search/tests.rs @@ -683,8 +683,7 @@ fn test_end_to_end_create_insert_search() { for (i, v) in vectors.iter().enumerate() { let mut sq = vec![0i8; dim]; quantize_f32_to_sq(v, &mut sq); - let norm = v.iter().map(|x| x * x).sum::().sqrt(); - snap.mutable.append(i as u64, v, &sq, norm, i as u64); + snap.mutable.append(i as u64, v, i as u64); } drop(snap); @@ -957,7 +956,7 @@ fn test_vector_metrics_increment_decrement() { // FT.SEARCH should increment VECTOR_SEARCH_TOTAL crate::vector::distance::init(); - let before_search = crate::vector::metrics::VECTOR_SEARCH_TOTAL.load(Ordering::Relaxed); + let before_search = crate::vector::metrics::search_total(); let query_vec: Vec = vec![0u8; 128 * 4]; let search_args = vec![ bulk(b"myidx"), @@ -968,7 +967,7 @@ fn test_vector_metrics_increment_decrement() { Frame::BulkString(Bytes::from(query_vec)), ]; ft_search(&mut store, &search_args, None, None, 0); - let after_search = crate::vector::metrics::VECTOR_SEARCH_TOTAL.load(Ordering::Relaxed); + let after_search = crate::vector::metrics::search_total(); assert_eq!( after_search, before_search + 1, @@ -2420,8 +2419,7 @@ fn test_ft_search_field_targeting() { for (i, v) in title_vecs.iter().enumerate() { let mut sq = vec![0i8; 4]; quantize_f32_to_sq(v, &mut sq); - let norm = v.iter().map(|x| x * x).sum::().sqrt(); - snap.mutable.append(i as u64, v, &sq, norm, i as u64); + snap.mutable.append(i as u64, v, i as u64); } drop(snap); @@ -2435,8 +2433,7 @@ fn test_ft_search_field_targeting() { for (i, v) in body_vecs.iter().enumerate() { let mut sq = vec![0i8; 8]; quantize_f32_to_sq(v, &mut sq); - let norm = v.iter().map(|x| x * x).sum::().sqrt(); - snap.mutable.append(i as u64, v, &sq, norm, i as u64); + snap.mutable.append(i as u64, v, i as u64); } } @@ -2526,8 +2523,7 @@ fn test_ft_search_default_field_compat() { for (i, v) in vectors.iter().enumerate() { let mut sq = vec![0i8; 4]; quantize_f32_to_sq(v, &mut sq); - let norm = v.iter().map(|x| x * x).sum::().sqrt(); - snap.mutable.append(i as u64, v, &sq, norm, i as u64); + snap.mutable.append(i as u64, v, i as u64); } drop(snap); @@ -2813,13 +2809,11 @@ fn insert_hybrid_doc( let snap = idx.segments.load(); let mut sq = vec![0i8; dim]; quantize_f32_to_sq(dense_vec, &mut sq); - let norm = dense_vec.iter().map(|x| x * x).sum::().sqrt(); - snap.mutable.append(key_hash, dense_vec, &sq, norm, 0); + snap.mutable.append(key_hash, dense_vec, 0); drop(snap); // Record key mapping - idx.key_hash_to_key - .insert(key_hash, Bytes::from(key.to_vec())); + std::sync::Arc::make_mut(&mut idx.key_hash_to_key).insert(key_hash, Bytes::from(key.to_vec())); // Insert sparse vector if let Some(ss) = idx.sparse_stores.get_mut(b"sparse_vec".as_ref()) { @@ -3157,8 +3151,7 @@ fn test_range_filter_l2_search() { for (i, v) in vectors.iter().enumerate() { let mut sq = vec![0i8; dim]; quantize_f32_to_sq(v, &mut sq); - let norm = v.iter().map(|x| x * x).sum::().sqrt(); - snap.mutable.append(i as u64, v, &sq, norm, i as u64); + snap.mutable.append(i as u64, v, i as u64); } } @@ -3356,10 +3349,9 @@ fn test_recommend_basic_with_vectors() { let snap = idx.segments.load(); let mut sq = vec![0i8; dim]; quantize_f32_to_sq(v, &mut sq); - let norm = v.iter().map(|x| x * x).sum::().sqrt(); - snap.mutable.append(key_hash, v, &sq, norm, i as u64); + snap.mutable.append(key_hash, v, i as u64); drop(snap); - idx.key_hash_to_key + std::sync::Arc::make_mut(&mut idx.key_hash_to_key) .insert(key_hash, Bytes::from(key.to_vec())); } } @@ -3529,8 +3521,8 @@ fn test_ft_dropindex_dd_deletes_docs() { if let Some(idx) = store.get_index_mut(b"ddtest") { let h1 = xxhash_rust::xxh64::xxh64(&key1, 0); let h2 = xxhash_rust::xxh64::xxh64(&key2, 0); - idx.key_hash_to_key.insert(h1, key1.clone()); - idx.key_hash_to_key.insert(h2, key2.clone()); + std::sync::Arc::make_mut(&mut idx.key_hash_to_key).insert(h1, key1.clone()); + std::sync::Arc::make_mut(&mut idx.key_hash_to_key).insert(h2, key2.clone()); } // Verify keys exist in database @@ -3598,7 +3590,7 @@ fn test_ft_dropindex_preserves_docs() { // Register key in vector index if let Some(idx) = store.get_index_mut(b"preservetest") { let h1 = xxhash_rust::xxh64::xxh64(&key1, 0); - idx.key_hash_to_key.insert(h1, key1.clone()); + std::sync::Arc::make_mut(&mut idx.key_hash_to_key).insert(h1, key1.clone()); } // Drop index WITHOUT DD flag (using None for db since we don't need it) @@ -3650,7 +3642,7 @@ fn test_ft_dropindex_dd_case_insensitive() { let key = Bytes::from_static(b"c1:doc"); db.set(key.clone(), crate::storage::entry::Entry::new_hash()); if let Some(idx) = store.get_index_mut(b"casetest1") { - idx.key_hash_to_key + std::sync::Arc::make_mut(&mut idx.key_hash_to_key) .insert(xxhash_rust::xxh64::xxh64(&key, 0), key.clone()); } @@ -3700,7 +3692,7 @@ fn test_ft_dropindex_dd_case_insensitive() { let key = Bytes::from_static(b"c2:doc"); db.set(key.clone(), crate::storage::entry::Entry::new_hash()); if let Some(idx) = store.get_index_mut(b"casetest2") { - idx.key_hash_to_key + std::sync::Arc::make_mut(&mut idx.key_hash_to_key) .insert(xxhash_rust::xxh64::xxh64(&key, 0), key.clone()); } diff --git a/src/config.rs b/src/config.rs index 7c7fea87d..1e3b12e64 100644 --- a/src/config.rs +++ b/src/config.rs @@ -279,6 +279,19 @@ pub struct ServerConfig { #[arg(long = "io-busy-poll-us", default_value_t = 0)] pub io_busy_poll_us: u64, + /// FT.SEARCH intra-query worker threads: per-segment HNSW searches of one + /// KNN query fan out across this pool, cutting single-query latency on + /// multi-segment indexes (the pool also serves concurrent queries). + /// Default 0 = disabled (serial per-segment loop; results are identical + /// either way). Opt in on boxes with spare PHYSICAL cores — a good size is + /// cores minus shards, capped at 8 (see search_pool::auto_workers). + /// Measured (20k×384d clustered SQ8, single conn, R@10=1.0): macOS + /// 10-core 473→2,321 QPS (4.9×); GCE c3-standard-8 (4 physical cores) + /// REGRESSES at default ef — each segment pays the full resolved ef, so a + /// pooled N-segment query does ~N× the CPU work for its latency win. + #[arg(long = "ft-search-workers")] + pub ft_search_workers: Option, + // ── MoonStore v2: Disk Offload ────────────────────────────────── /// Enable disk offload (tiered storage: RAM -> mmap -> NVMe) #[arg(long = "disk-offload", default_value = "enable")] @@ -1238,6 +1251,24 @@ mod tests { assert!(ServerConfig::try_parse_from(["moon", "--io-driver", "iouring"]).is_err()); } + #[test] + fn test_ft_search_workers_flag() { + let config = ServerConfig::parse_from::<[&str; 0], &str>([]); + assert_eq!( + config.ft_search_workers, None, + "default must be auto (None)" + ); + let config = ServerConfig::parse_from(["moon", "--ft-search-workers", "4"]); + assert_eq!(config.ft_search_workers, Some(4)); + let config = ServerConfig::parse_from(["moon", "--ft-search-workers", "0"]); + assert_eq!( + config.ft_search_workers, + Some(0), + "0 must parse (explicit off)" + ); + assert!(ServerConfig::try_parse_from(["moon", "--ft-search-workers", "x"]).is_err()); + } + #[test] fn test_io_busy_poll_flag_parses_with_zero_default() { let config = ServerConfig::parse_from::<[&str; 0], &str>([]); diff --git a/src/main.rs b/src/main.rs index 922b6e7a6..2675e876e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -392,6 +392,20 @@ fn main() -> anyhow::Result<()> { info!("Starting with {} shards", num_shards); + // FT.SEARCH intra-query worker pool: fan per-segment HNSW searches of one + // query across workers (threads spawn eagerly here — the pool is tiny and + // parks on its channel when vector search is unused). Default OFF: each + // segment is searched at the full resolved ef, so the pool trades ~Nseg× + // CPU work for the latency win — a regression on SMT-constrained boxes + // (GCE c3-standard-8 clustered SQ8: 1,732 → 1,165 QPS) while a real win + // on physical-core-rich ones (macOS 10-core: 473 → 2,321 QPS, 4.9×). + // Same posture as --io-busy-poll-us: measured opt-in, never a silent tax. + let ft_workers = config.ft_search_workers.unwrap_or(0); + moon::vector::search_pool::init_global(ft_workers); + if ft_workers > 0 { + info!("FT.SEARCH worker pool: {ft_workers} threads"); + } + // Checked cast: --shards is bounded by clap's value_parser, but `as u16` // would silently wrap for values > 65535. Fail loudly instead. // ALLOW: panic is appropriate here — this is `main`, not library code. diff --git a/src/server/conn/handler_monoio/ft.rs b/src/server/conn/handler_monoio/ft.rs index 62614db3c..3adba024c 100644 --- a/src/server/conn/handler_monoio/ft.rs +++ b/src/server/conn/handler_monoio/ft.rs @@ -271,9 +271,17 @@ pub(super) async fn try_handle_ft_command( return true; } if cmd.eq_ignore_ascii_case(b"FT.INFO") { - let response = crate::shard::slice::with_shard(|s| { - crate::command::vector_search::ft_info(&s.vector_store, &s.text_store, cmd_args) - }); + // Scatter + sum additive stats across shards (XC-SHARD-1): the local + // shard's index holds only its key-hash partition of the documents. + let response = crate::shard::coordinator::scatter_ft_info( + std::sync::Arc::new(frame.clone()), + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await; responses.push(response); return true; } diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index 4e7330369..eb1b013ce 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -1322,6 +1322,20 @@ pub(crate) async fn handle_connection_sharded_monoio< smallvec::SmallVec::new() }; + // Auto-delete vectors on DEL/UNLINK (conn-local write path). + // Parity with the SPSC Execute arm and the tokio sharded + // handler — without this, deleted keys keep matching + // FT.SEARCH at shards=1 (soak-diagnostic resurrection bug). + if !is_error + && (cmd.eq_ignore_ascii_case(b"DEL") + || cmd.eq_ignore_ascii_case(b"UNLINK")) + { + crate::shard::spsc_handler::auto_delete_vectors( + &mut s.vector_store, + cmd_args, + ); + } + // Blocking wakeup: re-borrow db by index (NLL) if !is_error { let needs_wake = cmd.eq_ignore_ascii_case(b"LPUSH") diff --git a/src/server/conn/handler_sharded/ft.rs b/src/server/conn/handler_sharded/ft.rs index 29f8a2a4a..d08d41f6a 100644 --- a/src/server/conn/handler_sharded/ft.rs +++ b/src/server/conn/handler_sharded/ft.rs @@ -278,6 +278,22 @@ pub(super) async fn try_handle_ft_command( ))); return true; } + // FT.INFO: stats are per-shard partitions — scatter to every shard and + // sum the additive fields (XC-SHARD-1); broadcast_vector_command would + // return only the local shard's counts (~1/N of the truth). + if cmd.eq_ignore_ascii_case(b"FT.INFO") { + let response = crate::shard::coordinator::scatter_ft_info( + std::sync::Arc::new(frame.clone()), + ctx.shard_id, + ctx.num_shards, + &ctx.shard_databases, + &ctx.dispatch_tx, + &ctx.spsc_notifiers, + ) + .await; + responses.push(response); + return true; + } let response = crate::shard::coordinator::broadcast_vector_command( std::sync::Arc::new(frame.clone()), ctx.shard_id, diff --git a/src/server/conn/handler_single.rs b/src/server/conn/handler_single.rs index ea06591fe..1d4d8c8d2 100644 --- a/src/server/conn/handler_single.rs +++ b/src/server/conn/handler_single.rs @@ -1138,6 +1138,17 @@ pub async fn handle_connection( ); } } + } else if (c.eq_ignore_ascii_case(b"DEL") + || c.eq_ignore_ascii_case(b"UNLINK")) + && i < txn_results.len() + && !matches!(txn_results[i], Frame::Error(_)) + { + // Auto-delete vectors (parity with + // the HSET auto-index arm above). + crate::shard::spsc_handler::auto_delete_vectors( + &mut vs.lock(), + a, + ); } } } @@ -2222,6 +2233,20 @@ pub async fn handle_connection( } } + // Auto-delete vectors on DEL/UNLINK (parity with + // the HSET auto-index hook above). + if !matches!(&response, Frame::Error(_)) + && (d_cmd.eq_ignore_ascii_case(b"DEL") + || d_cmd.eq_ignore_ascii_case(b"UNLINK")) + { + if let Some(ref vs) = vector_store { + crate::shard::spsc_handler::auto_delete_vectors( + &mut vs.lock(), + d_args, + ); + } + } + // Invalidate tracked key on successful write if !matches!(&response, Frame::Error(_)) { if let Some(key) = d_args.first().and_then(|f| extract_bytes(f)) { diff --git a/src/server/conn/shared.rs b/src/server/conn/shared.rs index 82685e431..2ef3b93fe 100644 --- a/src/server/conn/shared.rs +++ b/src/server/conn/shared.rs @@ -237,6 +237,15 @@ pub(crate) fn execute_transaction_sharded( } } + // Auto-delete vectors on DEL/UNLINK (parity with the HSET hook above). + if !matches!(response, Frame::Error(_)) + && (cmd.eq_ignore_ascii_case(b"DEL") || cmd.eq_ignore_ascii_case(b"UNLINK")) + { + crate::shard::slice::with_shard(|s| { + crate::shard::spsc_handler::auto_delete_vectors(&mut s.vector_store, cmd_args); + }); + } + results.push(response); } diff --git a/src/shard/coordinator.rs b/src/shard/coordinator.rs index 6ef9fb6bd..96dc79efd 100644 --- a/src/shard/coordinator.rs +++ b/src/shard/coordinator.rs @@ -1925,6 +1925,67 @@ pub async fn scatter_invalidate_range( } } +/// Scatter `FT.INFO` to all shards and merge the per-shard stats (XC-SHARD-1). +/// +/// Vector data is key-hash partitioned: each shard's index holds only the +/// vectors whose keys route there, so a single shard's FT.INFO reports ~1/N of +/// the true document count. This helper collects every shard's response and +/// sums the additive fields via +/// [`crate::command::vector_search::merge_ft_info_responses`]; config fields +/// come from the local response (identical everywhere by FT.CREATE broadcast). +/// +/// # Lock safety +/// Local execution is synchronous inside `with_shard` (no `.await` while the +/// shard slice is borrowed). +pub async fn scatter_ft_info( + command: std::sync::Arc, + my_shard: usize, + num_shards: usize, + shard_databases: &Arc, + dispatch_tx: &Rc>>>, + spsc_notifiers: &[Arc], +) -> Frame { + let _ = shard_databases; + let mut receivers = Vec::with_capacity(num_shards.saturating_sub(1)); + for target in 0..num_shards { + if target == my_shard { + continue; + } + let (reply_tx, reply_rx) = channel::oneshot(); + let msg = ShardMessage::VectorCommand { + command: command.clone(), + reply_tx, + }; + spsc_send(dispatch_tx, my_shard, target, msg, spsc_notifiers).await; + receivers.push(reply_rx); + } + + let mut remote_responses: Vec = Vec::with_capacity(receivers.len()); + for rx in receivers { + match rx.recv().await { + Ok(frame) => remote_responses.push(frame), + Err(_) => { + return Frame::Error(Bytes::from_static( + b"ERR FT.INFO: cross-shard reply channel closed", + )); + } + } + } + + let local = crate::shard::slice::with_shard(|s| { + crate::shard::spsc_handler::dispatch_vector_command( + &mut s.vector_store, + &mut s.text_store, + #[cfg(feature = "graph")] + Some(&s.graph_store), + &command, + None, + ) + }); + + crate::command::vector_search::merge_ft_info_responses(local, &remote_responses) +} + /// Two-phase DFS scatter-gather for globally accurate BM25 text search (per D-04). /// /// **Phase 1** — DocFreq scatter: collect (term, df) + total N from every shard, diff --git a/src/shard/persistence_tick.rs b/src/shard/persistence_tick.rs index a435c3794..ebf1403a3 100644 --- a/src/shard/persistence_tick.rs +++ b/src/shard/persistence_tick.rs @@ -351,16 +351,20 @@ pub(crate) fn run_eviction_tick( // static scheme already has between eviction passes. { let rt = runtime_config.read(); + // C5 / Phase 3: compute per-shard KV memory via ShardSlice without + // lock acquisitions (avoids per-DB read locks; estimated_memory() is + // an O(1) accumulator read). Published unconditionally: MEMORY DOCTOR + // and the Prometheus KV gauge read this atomic even when maxmemory is + // unlimited — gating it on maxmemory > 0 left them at a permanent 0. + let used = crate::shard::slice::with_shard(|s| { + s.databases + .iter() + .map(|db| db.estimated_memory()) + .sum::() + }); + shard_databases.publish_memory(shard_id, used); + // Elastic budgets only exist under a finite maxmemory cap. if rt.maxmemory > 0 { - // C5 / Phase 3: compute per-shard KV memory via ShardSlice without - // lock acquisitions (avoids per-DB read locks). - let used = crate::shard::slice::with_shard(|s| { - s.databases - .iter() - .map(|db| db.estimated_memory()) - .sum::() - }); - shard_databases.publish_memory(shard_id, used); shard_databases.recompute_elastic_budget(shard_id, &rt); } } diff --git a/src/shard/spsc_handler.rs b/src/shard/spsc_handler.rs index b63c635b1..75e1c4a1a 100644 --- a/src/shard/spsc_handler.rs +++ b/src/shard/spsc_handler.rs @@ -918,6 +918,14 @@ pub(crate) fn handle_shard_message_shared( } } + // Auto-delete vectors on DEL/UNLINK (parity with the HSET + // hook above and the Execute arm's auto-delete). + if !matches!(frame, crate::protocol::Frame::Error(_)) + && (cmd.eq_ignore_ascii_case(b"DEL") || cmd.eq_ignore_ascii_case(b"UNLINK")) + { + auto_delete_vectors(&mut s.vector_store, args); + } + // Post-dispatch wakeup hooks for producer commands (cross-shard blocking) if !matches!(frame, crate::protocol::Frame::Error(_)) { let needs_wake = cmd.eq_ignore_ascii_case(b"LPUSH") @@ -1281,6 +1289,14 @@ pub(crate) fn handle_shard_message_shared( } } + // Auto-delete vectors on DEL/UNLINK (parity with the HSET + // hook above and the Execute arm's auto-delete). + if !matches!(frame, crate::protocol::Frame::Error(_)) + && (cmd.eq_ignore_ascii_case(b"DEL") || cmd.eq_ignore_ascii_case(b"UNLINK")) + { + auto_delete_vectors(&mut s.vector_store, args); + } + if !matches!(frame, crate::protocol::Frame::Error(_)) { let needs_wake = cmd.eq_ignore_ascii_case(b"LPUSH") || cmd.eq_ignore_ascii_case(b"RPUSH") @@ -2131,6 +2147,19 @@ pub fn auto_index_hset_public( auto_index_hset(vector_store, text_store, key, args, 0) } +/// Tombstone auto-indexed vectors for every key argument of a successful +/// DEL/UNLINK. Wire-parity requirement: every dispatch path that runs +/// `auto_index_hset*` on HSET must run this on DEL/UNLINK, or deleted keys +/// keep matching FT.SEARCH forever (resurrection + live-set recall collapse; +/// found by the Bundle-5 soak diagnostic at shards=1). +pub fn auto_delete_vectors(vector_store: &mut VectorStore, args: &[crate::protocol::Frame]) { + for arg in args { + if let Some(key) = crate::server::connection::extract_bytes(arg) { + vector_store.mark_deleted_for_key(key.as_ref()); + } + } +} + /// TXN-aware variant: tags each inserted vector entry with `txn_id` so /// non-transactional readers (snapshot_lsn == 0) see it as uncommitted and /// exclude it until TXN.COMMIT calls `txn_manager.commit(txn_id)`. @@ -2329,13 +2358,9 @@ fn handle_vector_insert( for chunk in blob.chunks_exact(4) { f32_vec.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])); } - // SQ quantize - let mut sq_vec = vec![0i8; dim]; - vector_search::quantize_f32_to_sq(&f32_vec, &mut sq_vec); - // Compute norm - let norm: f32 = f32_vec.iter().map(|x| x * x).sum::().sqrt(); - // Record original Redis key for FT.SEARCH response. - idx.key_hash_to_key + // Record original Redis key for FT.SEARCH response. COW via make_mut: + // clones only if a search snapshot holds the map concurrently (QP-1). + std::sync::Arc::make_mut(&mut idx.key_hash_to_key) .entry(key_hash) .or_insert_with(|| bytes::Bytes::copy_from_slice(key)); // Append to mutable segment. `insert_lsn` is the monotonic LSN allocated @@ -2344,12 +2369,39 @@ fn handle_vector_insert( // TXN snapshot isolation. When inside a TXN (txn_id != 0), use the // transactional variant so non-TXN readers see the entry as uncommitted. let snap = idx.segments.load(); + // VEC-1: an HSET on an already-indexed key is an UPDATE — tombstone the + // prior version BEFORE appending, or the index accumulates stale + // duplicates (doc returned twice, num_docs inflating under churn). + // Non-txn path only: a txn's tombstone must not leak to other readers + // before commit (txn vector updates keep prior append-only behavior). + if txn_id == 0 { + if let Some(&old_gid) = idx.key_hash_to_global_id.get(&key_hash) { + let base = snap.mutable.global_id_base(); + if old_gid >= base + && snap + .mutable + .mark_deleted_if_key(old_gid - base, key_hash, insert_lsn) + { + // O(1) fast path: old version still in the mutable segment, + // MVCC-tombstoned at the new version's LSN (older snapshots + // keep seeing the old vector; new snapshots see only the new). + } else { + // Old version was compacted (or the gid mapping was stale): + // steady-state interior tombstone across immutable segments — + // the same path DEL/UNLINK takes via `mark_deleted_for_key` — + // plus a defensive mutable scan for the stale-mapping case. + snap.mutable.mark_deleted_by_key_hash(key_hash, insert_lsn); + for imm in snap.immutable.iter() { + imm.mark_deleted_by_key_hash(key_hash); + } + } + } + } let internal_id = if txn_id != 0 { snap.mutable - .append_transactional(key_hash, &f32_vec, &sq_vec, norm, insert_lsn, txn_id) + .append_transactional(key_hash, &f32_vec, insert_lsn, txn_id) } else { - snap.mutable - .append(key_hash, &f32_vec, &sq_vec, norm, insert_lsn) + snap.mutable.append(key_hash, &f32_vec, insert_lsn) }; // Use global_id for payload index so filter bitmaps match // search results after compaction advances global_id_base. @@ -2399,14 +2451,9 @@ fn handle_vector_insert_field( for chunk in blob.chunks_exact(4) { f32_vec.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])); } - // SQ quantize - let mut sq_vec = vec![0i8; dim]; - vector_search::quantize_f32_to_sq(&f32_vec, &mut sq_vec); - // Compute norm - let norm: f32 = f32_vec.iter().map(|x| x * x).sum::().sqrt(); - - // Record original Redis key (shared across all fields) - idx.key_hash_to_key + // Record original Redis key (shared across all fields). COW via make_mut: + // clones only if a search snapshot holds the map concurrently (QP-1). + std::sync::Arc::make_mut(&mut idx.key_hash_to_key) .entry(key_hash) .or_insert_with(|| bytes::Bytes::copy_from_slice(key)); @@ -2419,12 +2466,20 @@ fn handle_vector_insert_field( // so both fields share one logical write event (Phase 165 MVCC contract). // When inside a TXN (txn_id != 0), tag with txn_id for uncommitted visibility. let snap = fs.segments.load(); + // VEC-1 (additional fields): tombstone the prior version on update. Field + // segments have no `key_hash → global_id` map, so this is the scan path + // (mutable is bounded by compact_threshold; immutables are set lookups). + if txn_id == 0 { + snap.mutable.mark_deleted_by_key_hash(key_hash, insert_lsn); + for imm in snap.immutable.iter() { + imm.mark_deleted_by_key_hash(key_hash); + } + } let _internal_id = if txn_id != 0 { snap.mutable - .append_transactional(key_hash, &f32_vec, &sq_vec, norm, insert_lsn, txn_id) + .append_transactional(key_hash, &f32_vec, insert_lsn, txn_id) } else { - snap.mutable - .append(key_hash, &f32_vec, &sq_vec, norm, insert_lsn) + snap.mutable.append(key_hash, &f32_vec, insert_lsn) }; crate::vector::metrics::add_vectors(1); // Note: global_id and payload_index are NOT updated here. diff --git a/src/vector/distance/avx2.rs b/src/vector/distance/avx2.rs index a92420177..5b83cc708 100644 --- a/src/vector/distance/avx2.rs +++ b/src/vector/distance/avx2.rs @@ -301,6 +301,81 @@ pub unsafe fn cosine_f32(a: &[f32], b: &[f32]) -> f32 { 1.0 - dot_sum / (norm_a * norm_b) } +/// SQ8 ADC (HQ-2) per-candidate statistics (AVX2+FMA): `(Σ(q_i·c_i), Σc_i, Σc_i²)`. +/// +/// Widens 16 u8 codes per iteration to f32 (`VPMOVZXBD` via +/// `_mm256_cvtepu8_epi32`, 8 lanes/call, 2x unrolled, then `_mm256_cvtepi32_ps`) +/// and FMAs against the f32 query. See `turbo_quant::sq8` module docs for how +/// these three running sums combine (O(1), architecture-independent) into +/// the final asymmetric L2/inner-product ADC distance — that combine step +/// deliberately has no SIMD variant, only this stats pass does. +/// +/// # Safety +/// Caller must ensure AVX2 and FMA CPU features are available. +#[cfg(target_arch = "x86_64")] +#[inline] +#[target_feature(enable = "avx2,fma")] +pub unsafe fn sq8_stats(query: &[f32], codes: &[u8]) -> (f32, f32, f32) { + debug_assert_eq!(query.len(), codes.len(), "sq8_stats: dimension mismatch"); + + let n = query.len(); + let mut dot0 = _mm256_setzero_ps(); + let mut dot1 = _mm256_setzero_ps(); + let mut sumc0 = _mm256_setzero_ps(); + let mut sumc1 = _mm256_setzero_ps(); + let mut sumsqc0 = _mm256_setzero_ps(); + let mut sumsqc1 = _mm256_setzero_ps(); + + let pq = query.as_ptr(); + let pc = codes.as_ptr(); + + let chunks = n / 16; + let mut i = 0usize; + + for _ in 0..chunks { + // SAFETY: i + 16 <= n guaranteed by chunks = n / 16. `_mm_loadl_epi64` + // reads 8 bytes (unaligned); `_mm256_cvtepu8_epi32` zero-extends the + // low 8 lanes to i32, which `_mm256_cvtepi32_ps` converts exactly + // (u8 values 0..=255 have exact f32 representations). + let c0_u8 = _mm_loadl_epi64(pc.add(i) as *const __m128i); + let c0_f32 = _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(c0_u8)); + let q0 = _mm256_loadu_ps(pq.add(i)); + dot0 = _mm256_fmadd_ps(q0, c0_f32, dot0); + sumc0 = _mm256_add_ps(sumc0, c0_f32); + sumsqc0 = _mm256_fmadd_ps(c0_f32, c0_f32, sumsqc0); + + let c1_u8 = _mm_loadl_epi64(pc.add(i + 8) as *const __m128i); + let c1_f32 = _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(c1_u8)); + let q1 = _mm256_loadu_ps(pq.add(i + 8)); + dot1 = _mm256_fmadd_ps(q1, c1_f32, dot1); + sumc1 = _mm256_add_ps(sumc1, c1_f32); + sumsqc1 = _mm256_fmadd_ps(c1_f32, c1_f32, sumsqc1); + + i += 16; + } + + dot0 = _mm256_add_ps(dot0, dot1); + sumc0 = _mm256_add_ps(sumc0, sumc1); + sumsqc0 = _mm256_add_ps(sumsqc0, sumsqc1); + + // SAFETY: hsum_f32_avx2 requires AVX2, which we have via target_feature. + let mut dot_sum = hsum_f32_avx2(dot0); + let mut sum_c_sum = hsum_f32_avx2(sumc0); + let mut sumsq_c_sum = hsum_f32_avx2(sumsqc0); + + // Scalar tail — safe indexing; bounds-checked slices cost nothing here + // (at most one sub-vector-width pass) and keep the unsafe surface to the + // intrinsics above (UNSAFE_POLICY). + for (&c, &qv) in codes[i..n].iter().zip(query[i..n].iter()) { + let cf = c as f32; + dot_sum += qv * cf; + sum_c_sum += cf; + sumsq_c_sum += cf * cf; + } + + (dot_sum, sum_c_sum, sumsq_c_sum) +} + #[cfg(test)] #[cfg(target_arch = "x86_64")] mod tests { @@ -458,4 +533,75 @@ mod tests { assert_eq!(l2_i8(ai, bi), 0); } } + + // ── HQ-2: SQ8 ADC stats (AVX2 vs scalar reference) ────────────────── + // + // Gated on `has_avx2_fma()` like every other test in this file: on this + // aarch64 dev host these compile (proving the cfg discipline holds) but + // no-op at runtime. They exercise for real on any x86_64 host/CI runner + // with AVX2+FMA. + + fn gen_u8(len: usize, seed: u32) -> Vec { + let mut v = Vec::with_capacity(len); + let mut s = seed; + for _ in 0..len { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + v.push((s >> 24) as u8); + } + v + } + + #[test] + fn test_sq8_stats_matches_scalar() { + if !has_avx2_fma() { + return; + } + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_scalar; + let q = gen_f32(768, 42); + let c = gen_u8(768, 99); + let expected = sq8_candidate_stats_scalar(&q, &c); + // SAFETY: AVX2+FMA verified above. + let got = unsafe { sq8_stats(&q, &c) }; + let rel_dot = (got.0 - expected.0).abs() / expected.0.abs().max(1.0); + let rel_sum = (got.1 - expected.1).abs() / expected.1.abs().max(1.0); + let rel_sq = (got.2 - expected.2).abs() / expected.2.abs().max(1.0); + assert!( + rel_dot < 1e-3 && rel_sum < 1e-3 && rel_sq < 1e-3, + "sq8_stats mismatch: scalar={expected:?}, avx2={got:?}" + ); + } + + #[test] + fn test_sq8_stats_tail_handling() { + if !has_avx2_fma() { + return; + } + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_scalar; + for len in [0, 1, 3, 7, 13, 15, 16, 17, 31, 33, 100] { + let q = gen_f32(len, 42); + let c = gen_u8(len, 99); + let expected = sq8_candidate_stats_scalar(&q, &c); + // SAFETY: AVX2+FMA verified above. + let got = unsafe { sq8_stats(&q, &c) }; + let rel_dot = (got.0 - expected.0).abs() / expected.0.abs().max(1.0); + let rel_sum = (got.1 - expected.1).abs() / expected.1.abs().max(1.0); + let rel_sq = (got.2 - expected.2).abs() / expected.2.abs().max(1.0); + assert!( + rel_dot < 1e-3 && rel_sum < 1e-3 && rel_sq < 1e-3, + "sq8_stats tail len={len}: scalar={expected:?}, avx2={got:?}" + ); + } + } + + #[test] + fn test_sq8_stats_empty() { + if !has_avx2_fma() { + return; + } + let q: &[f32] = &[]; + let c: &[u8] = &[]; + // SAFETY: AVX2+FMA verified above. + let got = unsafe { sq8_stats(q, c) }; + assert_eq!(got, (0.0, 0.0, 0.0)); + } } diff --git a/src/vector/distance/avx512.rs b/src/vector/distance/avx512.rs index 328ceb914..626cd9f36 100644 --- a/src/vector/distance/avx512.rs +++ b/src/vector/distance/avx512.rs @@ -244,6 +244,80 @@ pub unsafe fn cosine_f32(a: &[f32], b: &[f32]) -> f32 { 1.0 - dot_sum / (norm_a * norm_b) } +/// SQ8 ADC (HQ-2) per-candidate statistics (AVX-512F): `(Σ(q_i·c_i), Σc_i, Σc_i²)`. +/// +/// Widens 32 u8 codes per iteration to f32 (`VPMOVZXBD` via +/// `_mm512_cvtepu8_epi32`, 16 lanes/call, 2x unrolled, then +/// `_mm512_cvtepi32_ps`) and FMAs against the f32 query. See +/// `turbo_quant::sq8` module docs for how these three running sums combine +/// (O(1), architecture-independent) into the final ADC distance. +/// +/// # Safety +/// Caller must ensure AVX-512F CPU feature is available. +#[cfg(target_arch = "x86_64")] +#[inline] +#[target_feature(enable = "avx512f")] +pub unsafe fn sq8_stats(query: &[f32], codes: &[u8]) -> (f32, f32, f32) { + debug_assert_eq!(query.len(), codes.len(), "sq8_stats: dimension mismatch"); + + let n = query.len(); + let mut dot0 = _mm512_setzero_ps(); + let mut dot1 = _mm512_setzero_ps(); + let mut sumc0 = _mm512_setzero_ps(); + let mut sumc1 = _mm512_setzero_ps(); + let mut sumsqc0 = _mm512_setzero_ps(); + let mut sumsqc1 = _mm512_setzero_ps(); + + let pq = query.as_ptr(); + let pc = codes.as_ptr(); + + let chunks = n / 32; + let mut i = 0usize; + + for _ in 0..chunks { + // SAFETY: i + 32 <= n guaranteed by chunks = n / 32. `_mm_loadu_si128` + // reads 16 bytes (unaligned); `_mm512_cvtepu8_epi32` zero-extends all + // 16 lanes to i32, which `_mm512_cvtepi32_ps` converts exactly (u8 + // values 0..=255 have exact f32 representations). + let c0_u8 = _mm_loadu_si128(pc.add(i) as *const __m128i); + let c0_f32 = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(c0_u8)); + let q0 = _mm512_loadu_ps(pq.add(i)); + dot0 = _mm512_fmadd_ps(q0, c0_f32, dot0); + sumc0 = _mm512_add_ps(sumc0, c0_f32); + sumsqc0 = _mm512_fmadd_ps(c0_f32, c0_f32, sumsqc0); + + let c1_u8 = _mm_loadu_si128(pc.add(i + 16) as *const __m128i); + let c1_f32 = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(c1_u8)); + let q1 = _mm512_loadu_ps(pq.add(i + 16)); + dot1 = _mm512_fmadd_ps(q1, c1_f32, dot1); + sumc1 = _mm512_add_ps(sumc1, c1_f32); + sumsqc1 = _mm512_fmadd_ps(c1_f32, c1_f32, sumsqc1); + + i += 32; + } + + dot0 = _mm512_add_ps(dot0, dot1); + sumc0 = _mm512_add_ps(sumc0, sumc1); + sumsqc0 = _mm512_add_ps(sumsqc0, sumsqc1); + + // SAFETY: _mm512_reduce_add_ps requires AVX-512F, verified via target_feature. + let mut dot_sum = _mm512_reduce_add_ps(dot0); + let mut sum_c_sum = _mm512_reduce_add_ps(sumc0); + let mut sumsq_c_sum = _mm512_reduce_add_ps(sumsqc0); + + // Scalar tail — safe indexing; bounds-checked slices cost nothing here + // (at most one sub-vector-width pass) and keep the unsafe surface to the + // intrinsics above (UNSAFE_POLICY). + for (&c, &qv) in codes[i..n].iter().zip(query[i..n].iter()) { + let cf = c as f32; + dot_sum += qv * cf; + sum_c_sum += cf; + sumsq_c_sum += cf * cf; + } + + (dot_sum, sum_c_sum, sumsq_c_sum) +} + #[cfg(test)] #[cfg(target_arch = "x86_64")] mod tests { @@ -374,4 +448,58 @@ mod tests { ); } } + + // ── HQ-2: SQ8 ADC stats (AVX-512 vs scalar reference) ─────────────── + + fn gen_u8(len: usize, seed: u32) -> Vec { + let mut v = Vec::with_capacity(len); + let mut s = seed; + for _ in 0..len { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + v.push((s >> 24) as u8); + } + v + } + + #[test] + fn test_sq8_stats_matches_scalar() { + if !has_avx512f() { + return; + } + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_scalar; + let q = gen_f32(768, 42); + let c = gen_u8(768, 99); + let expected = sq8_candidate_stats_scalar(&q, &c); + // SAFETY: AVX-512F verified above. + let got = unsafe { sq8_stats(&q, &c) }; + let rel_dot = (got.0 - expected.0).abs() / expected.0.abs().max(1.0); + let rel_sum = (got.1 - expected.1).abs() / expected.1.abs().max(1.0); + let rel_sq = (got.2 - expected.2).abs() / expected.2.abs().max(1.0); + assert!( + rel_dot < 1e-3 && rel_sum < 1e-3 && rel_sq < 1e-3, + "sq8_stats mismatch: scalar={expected:?}, avx512={got:?}" + ); + } + + #[test] + fn test_sq8_stats_tail_handling() { + if !has_avx512f() { + return; + } + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_scalar; + for len in [0, 1, 3, 7, 13, 15, 31, 32, 33, 63, 100] { + let q = gen_f32(len, 42); + let c = gen_u8(len, 99); + let expected = sq8_candidate_stats_scalar(&q, &c); + // SAFETY: AVX-512F verified above. + let got = unsafe { sq8_stats(&q, &c) }; + let rel_dot = (got.0 - expected.0).abs() / expected.0.abs().max(1.0); + let rel_sum = (got.1 - expected.1).abs() / expected.1.abs().max(1.0); + let rel_sq = (got.2 - expected.2).abs() / expected.2.abs().max(1.0); + assert!( + rel_dot < 1e-3 && rel_sum < 1e-3 && rel_sq < 1e-3, + "sq8_stats tail len={len}: scalar={expected:?}, avx512={got:?}" + ); + } + } } diff --git a/src/vector/distance/mod.rs b/src/vector/distance/mod.rs index 6a5fe569f..3d23ed028 100644 --- a/src/vector/distance/mod.rs +++ b/src/vector/distance/mod.rs @@ -33,6 +33,17 @@ pub struct DistanceTable { /// Centroids must be dimension-scaled (from CollectionMetadata.codebook_16()). /// All tiers use scalar ADC for now; AVX2/AVX-512 VPERMPS ADC is Phase 61+ work. pub tq_l2: fn(&[f32], &[u8], f32, &[f32; 16]) -> f32, + /// SQ8 asymmetric-distance-code (ADC) per-candidate statistics: + /// `(query, codes) -> (Σ(q_i·c_i), Σc_i, Σc_i²)`. + /// + /// This is the SIMD-accelerated inner loop of the HQ-2 fix (see + /// `turbo_quant::sq8` module docs for the algebraic decomposition): u8 + /// codes are widened to f32 and FMA'd against the query in one pass. + /// Combine with `turbo_quant::sq8::sq8_l2_from_stats` / + /// `sq8_ip_from_stats` (using per-query `turbo_quant::sq8::sq8_query_stats`) + /// to get the final ADC distance — that combine step is O(1) arithmetic, + /// architecture-independent, and deliberately NOT part of this table. + pub sq8_stats: fn(&[f32], &[u8]) -> (f32, f32, f32), } static DISTANCE_TABLE: OnceLock = OnceLock::new(); @@ -74,6 +85,10 @@ pub fn init() { unsafe { avx512::cosine_f32(a, b) } }, tq_l2: crate::vector::turbo_quant::tq_adc::tq_l2_adc_scaled, + sq8_stats: |q, c| { + // SAFETY: AVX-512F verified by is_x86_feature_detected! above. + unsafe { avx512::sq8_stats(q, c) } + }, }; } if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") { @@ -95,6 +110,10 @@ pub fn init() { unsafe { avx2::cosine_f32(a, b) } }, tq_l2: crate::vector::turbo_quant::tq_adc::tq_l2_adc_scaled, + sq8_stats: |q, c| { + // SAFETY: AVX2+FMA verified by is_x86_feature_detected! above. + unsafe { avx2::sq8_stats(q, c) } + }, }; } } @@ -121,6 +140,10 @@ pub fn init() { unsafe { neon::cosine_f32(a, b) } }, tq_l2: crate::vector::turbo_quant::tq_adc::tq_l2_adc_scaled, + sq8_stats: |q, c| { + // SAFETY: NEON is guaranteed on AArch64. + unsafe { neon::sq8_stats(q, c) } + }, }; } @@ -132,6 +155,7 @@ pub fn init() { dot_f32: scalar::dot_f32, cosine_f32: scalar::cosine_f32, tq_l2: crate::vector::turbo_quant::tq_adc::tq_l2_adc_scaled, + sq8_stats: crate::vector::turbo_quant::sq8::sq8_candidate_stats_scalar, } }); } @@ -184,6 +208,16 @@ mod tests { let centroids = crate::vector::turbo_quant::codebook::scaled_centroids(8); let dist = (t.tq_l2)(&q, &code, 1.0, ¢roids); assert!(dist >= 0.0, "tq_l2 should be non-negative, got {dist}"); + + // SQ8 ADC stats smoke test (HQ-2): dot=1*10+2*20+3*30=140, sum_c=60, sumsq_c=1400. + let codes: [u8; 3] = [10, 20, 30]; + let (dot, sum_c, sumsq_c) = (t.sq8_stats)(&a[..3], &codes); + assert!((dot - 140.0).abs() < 1e-4, "sq8_stats dot={dot}"); + assert!((sum_c - 60.0).abs() < 1e-4, "sq8_stats sum_c={sum_c}"); + assert!( + (sumsq_c - 1400.0).abs() < 1e-4, + "sq8_stats sumsq_c={sumsq_c}" + ); } #[test] @@ -391,4 +425,90 @@ mod integration_tests { let bi = [-10i8]; assert_eq!(scalar::l2_i8(&ai, &bi), (t.l2_i8)(&ai, &bi)); } + + /// Deterministic u8 codes via LCG PRNG, full 0..=255 range. + fn deterministic_u8(dim: usize, seed: u64) -> Vec { + let mut v = Vec::with_capacity(dim); + let mut s = seed as u32; + for _ in 0..dim { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + v.push((s >> 24) as u8); + } + v + } + + /// HQ-2: the SIMD-dispatched `sq8_stats` kernel must match the scalar + /// reference (`turbo_quant::sq8::sq8_candidate_stats_scalar`) at every + /// dimension, including SIMD-width tail remainders. On x86_64 this + /// exercises whichever tier `is_x86_feature_detected!` selected at + /// `init()` time (AVX-512 > AVX2+FMA); on aarch64, NEON (always + /// available); elsewhere, the scalar fallback trivially matches itself. + #[test] + fn test_simd_matches_scalar_sq8_stats() { + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_scalar; + + init(); + let t = table(); + for &dim in TEST_DIMS { + let q = deterministic_f32(dim, 42); + let codes = deterministic_u8(dim, 99); + let expected = sq8_candidate_stats_scalar(&q, &codes); + let got = (t.sq8_stats)(&q, &codes); + assert!( + approx_eq_f32(expected.0, got.0, 1e-3), + "sq8_stats dot mismatch at dim={dim}: scalar={:?} dispatch={:?}", + expected, + got + ); + assert!( + approx_eq_f32(expected.1, got.1, 1e-3), + "sq8_stats sum_c mismatch at dim={dim}: scalar={:?} dispatch={:?}", + expected, + got + ); + assert!( + approx_eq_f32(expected.2, got.2, 1e-3), + "sq8_stats sumsq_c mismatch at dim={dim}: scalar={:?} dispatch={:?}", + expected, + got + ); + } + } + + /// End-to-end: dispatched `sq8_stats` combined via + /// `sq8_l2_from_stats`/`sq8_ip_from_stats` must reproduce the original + /// naive per-element `sq8_l2_adc`/`sq8_ip_adc` ADC (same math, + /// reassociated for vectorization). + #[test] + fn test_sq8_dispatched_adc_matches_naive() { + use crate::vector::turbo_quant::sq8::{ + sq8_ip_adc, sq8_ip_from_stats, sq8_l2_adc, sq8_l2_from_stats, sq8_query_stats, + }; + + init(); + let t = table(); + for &dim in TEST_DIMS { + let q = deterministic_f32(dim, 7); + let codes = deterministic_u8(dim, 13); + let min = -0.37f32; + let scale = 0.0123f32; + + let naive_l2 = sq8_l2_adc(&q, &codes, min, scale); + let (q_sum, q_sumsq) = sq8_query_stats(&q); + let (dot_qc, sum_c, sumsq_c) = (t.sq8_stats)(&q, &codes); + let fast_l2 = + sq8_l2_from_stats(dim, min, scale, q_sum, q_sumsq, dot_qc, sum_c, sumsq_c); + assert!( + approx_eq_f32(naive_l2, fast_l2, 1e-3), + "l2 mismatch at dim={dim}: naive={naive_l2} dispatched={fast_l2}" + ); + + let naive_ip = sq8_ip_adc(&q, &codes, min, scale); + let fast_ip = sq8_ip_from_stats(min, scale, q_sum, dot_qc); + assert!( + approx_eq_f32(naive_ip, fast_ip, 1e-3), + "ip mismatch at dim={dim}: naive={naive_ip} dispatched={fast_ip}" + ); + } + } } diff --git a/src/vector/distance/neon.rs b/src/vector/distance/neon.rs index f1ba9b2f5..290d3fa42 100644 --- a/src/vector/distance/neon.rs +++ b/src/vector/distance/neon.rs @@ -275,6 +275,88 @@ pub unsafe fn cosine_f32(a: &[f32], b: &[f32]) -> f32 { 1.0 - dot_sum / (norm_a * norm_b) } +/// SQ8 ADC (HQ-2) per-candidate statistics (NEON): `(Σ(q_i·c_i), Σc_i, Σc_i²)`. +/// +/// Widens 16 u8 codes per iteration to f32 (`u8 -> u16 -> u32 -> f32` via +/// `vmovl`/`vcvtq_f32_u32`) and FMAs against the f32 query with three +/// independent accumulators. See `turbo_quant::sq8` module docs for how +/// these three running sums combine (O(1), architecture-independent) into +/// the final asymmetric L2/inner-product ADC distance — that combine step +/// deliberately has no SIMD variant, only this stats pass does. +/// +/// # Safety +/// Caller must ensure the CPU supports NEON (baseline on all AArch64). +#[cfg(target_arch = "aarch64")] +#[inline] +#[target_feature(enable = "neon")] +pub unsafe fn sq8_stats(query: &[f32], codes: &[u8]) -> (f32, f32, f32) { + debug_assert_eq!(query.len(), codes.len(), "sq8_stats: dimension mismatch"); + + let n = query.len(); + let mut dot = vdupq_n_f32(0.0); + let mut sum_c = vdupq_n_f32(0.0); + let mut sumsq_c = vdupq_n_f32(0.0); + + let pq = query.as_ptr(); + let pc = codes.as_ptr(); + + let chunks = n / 16; + let mut i = 0usize; + + for _ in 0..chunks { + // SAFETY: i + 16 <= n guaranteed by chunks = n / 16. Pointers are + // valid for 16 elements (f32 query / u8 codes) at this offset. + let c_u8 = vld1q_u8(pc.add(i)); + let c_u16_lo = vmovl_u8(vget_low_u8(c_u8)); + let c_u16_hi = vmovl_u8(vget_high_u8(c_u8)); + + // Widen each group of 4 lanes u16 -> u32 -> f32. + let c0 = vcvtq_f32_u32(vmovl_u16(vget_low_u16(c_u16_lo))); + let c1 = vcvtq_f32_u32(vmovl_u16(vget_high_u16(c_u16_lo))); + let c2 = vcvtq_f32_u32(vmovl_u16(vget_low_u16(c_u16_hi))); + let c3 = vcvtq_f32_u32(vmovl_u16(vget_high_u16(c_u16_hi))); + + let q0 = vld1q_f32(pq.add(i)); + let q1 = vld1q_f32(pq.add(i + 4)); + let q2 = vld1q_f32(pq.add(i + 8)); + let q3 = vld1q_f32(pq.add(i + 12)); + + dot = vfmaq_f32(dot, q0, c0); + dot = vfmaq_f32(dot, q1, c1); + dot = vfmaq_f32(dot, q2, c2); + dot = vfmaq_f32(dot, q3, c3); + + sum_c = vaddq_f32(sum_c, c0); + sum_c = vaddq_f32(sum_c, c1); + sum_c = vaddq_f32(sum_c, c2); + sum_c = vaddq_f32(sum_c, c3); + + sumsq_c = vfmaq_f32(sumsq_c, c0, c0); + sumsq_c = vfmaq_f32(sumsq_c, c1, c1); + sumsq_c = vfmaq_f32(sumsq_c, c2, c2); + sumsq_c = vfmaq_f32(sumsq_c, c3, c3); + + i += 16; + } + + // SAFETY: vaddvq_f32 requires NEON, which we have via target_feature. + let mut dot_sum = vaddvq_f32(dot); + let mut sum_c_sum = vaddvq_f32(sum_c); + let mut sumsq_c_sum = vaddvq_f32(sumsq_c); + + // Scalar tail — safe indexing; bounds-checked slices cost nothing here + // (at most one sub-vector-width pass) and keep the unsafe surface to the + // intrinsics above (UNSAFE_POLICY). + for (&c, &qv) in codes[i..n].iter().zip(query[i..n].iter()) { + let cf = c as f32; + dot_sum += qv * cf; + sum_c_sum += cf; + sumsq_c_sum += cf * cf; + } + + (dot_sum, sum_c_sum, sumsq_c_sum) +} + #[cfg(test)] #[cfg(target_arch = "aarch64")] mod tests { @@ -407,4 +489,61 @@ mod tests { assert_eq!(l2_i8(ai, bi), 0); } } + + // ── HQ-2: SQ8 ADC stats (NEON vs scalar reference) ────────────────── + + fn gen_u8(len: usize, seed: u32) -> Vec { + let mut v = Vec::with_capacity(len); + let mut s = seed; + for _ in 0..len { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + v.push((s >> 24) as u8); + } + v + } + + #[test] + fn test_sq8_stats_matches_scalar() { + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_scalar; + let q = gen_f32(768, 42); + let c = gen_u8(768, 99); + let expected = sq8_candidate_stats_scalar(&q, &c); + // SAFETY: NEON is baseline on AArch64. + let got = unsafe { sq8_stats(&q, &c) }; + let rel_dot = (got.0 - expected.0).abs() / expected.0.abs().max(1.0); + let rel_sum = (got.1 - expected.1).abs() / expected.1.abs().max(1.0); + let rel_sq = (got.2 - expected.2).abs() / expected.2.abs().max(1.0); + assert!( + rel_dot < 1e-3 && rel_sum < 1e-3 && rel_sq < 1e-3, + "sq8_stats mismatch: scalar={expected:?}, neon={got:?}" + ); + } + + #[test] + fn test_sq8_stats_tail_handling() { + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_scalar; + for len in [0, 1, 3, 7, 13, 15, 16, 17, 31, 33, 100] { + let q = gen_f32(len, 42); + let c = gen_u8(len, 99); + let expected = sq8_candidate_stats_scalar(&q, &c); + // SAFETY: NEON is baseline on AArch64. + let got = unsafe { sq8_stats(&q, &c) }; + let rel_dot = (got.0 - expected.0).abs() / expected.0.abs().max(1.0); + let rel_sum = (got.1 - expected.1).abs() / expected.1.abs().max(1.0); + let rel_sq = (got.2 - expected.2).abs() / expected.2.abs().max(1.0); + assert!( + rel_dot < 1e-3 && rel_sum < 1e-3 && rel_sq < 1e-3, + "sq8_stats tail len={len}: scalar={expected:?}, neon={got:?}" + ); + } + } + + #[test] + fn test_sq8_stats_empty() { + let q: &[f32] = &[]; + let c: &[u8] = &[]; + // SAFETY: NEON is baseline on AArch64. + let got = unsafe { sq8_stats(q, c) }; + assert_eq!(got, (0.0, 0.0, 0.0)); + } } diff --git a/src/vector/f16.rs b/src/vector/f16.rs new file mode 100644 index 000000000..ac3f6236b --- /dev/null +++ b/src/vector/f16.rs @@ -0,0 +1,224 @@ +//! Minimal IEEE 754 binary16 (half-precision) conversion. +//! +//! Used by the exact-rerank sidecar (deep-review HQ-1): immutable segments +//! keep an optional f16 copy of each original vector so the top-of-beam +//! candidates can be re-scored with (near-)exact distances instead of pure +//! quantized ADC estimates. f16 halves the sidecar footprint vs f32 while its +//! ~1e-3 relative error is far below SQ8/TQ4 quantization error. +//! +//! Hand-rolled instead of the `half` crate: two total functions, no new +//! dependency on the hot path. Conversion follows IEEE 754-2019 +//! round-to-nearest-even, with subnormal, overflow→infinity, and NaN +//! handling — each pinned by a unit test below. + +/// Convert an f32 to IEEE 754 binary16 bits (round-to-nearest-even). +#[inline] +pub fn f32_to_f16(value: f32) -> u16 { + let bits = value.to_bits(); + let sign = ((bits >> 16) & 0x8000) as u16; + let exp = ((bits >> 23) & 0xFF) as i32; + let mant = bits & 0x007F_FFFF; + + if exp == 0xFF { + // Inf / NaN. Preserve NaN-ness with a quiet-NaN mantissa bit. + return if mant == 0 { + sign | 0x7C00 + } else { + sign | 0x7E00 + }; + } + + // Re-bias: f32 bias 127 -> f16 bias 15. + let unbiased = exp - 127; + if unbiased > 15 { + // Overflows f16 range -> infinity. + return sign | 0x7C00; + } + if unbiased >= -14 { + // Normal f16. Keep top 10 mantissa bits, RNE on the dropped 13. + let exp16 = (unbiased + 15) as u32; + let mant16 = mant >> 13; + let rest = mant & 0x1FFF; + let mut out = (exp16 << 10) | mant16; + // Round up on >half, or exactly half with odd LSB (ties-to-even). + if rest > 0x1000 || (rest == 0x1000 && (mant16 & 1) == 1) { + out += 1; // Mantissa overflow correctly carries into the exponent. + } + return sign | (out as u16); + } + if unbiased >= -25 { + // Subnormal f16: implicit leading 1 becomes explicit, shifted right. + let full = mant | 0x0080_0000; + let shift = (-14 - unbiased) as u32 + 13; + let mant16 = full >> shift; + let rest = full & ((1u32 << shift) - 1); + let half = 1u32 << (shift - 1); + let mut out = mant16; + if rest > half || (rest == half && (mant16 & 1) == 1) { + out += 1; + } + return sign | (out as u16); + } + // Underflows to signed zero. + sign +} + +/// Convert IEEE 754 binary16 bits to f32 (exact — every f16 is representable). +#[inline] +pub fn f16_to_f32(bits: u16) -> f32 { + let sign = ((bits & 0x8000) as u32) << 16; + let exp = ((bits >> 10) & 0x1F) as u32; + let mant = (bits & 0x03FF) as u32; + + let out = if exp == 0 { + if mant == 0 { + sign // Signed zero. + } else { + // Subnormal: value is mant * 2^-24. Normalize: with `lead` zero + // bits above the 10-bit field, the leading 1 sits at bit + // p = 10 - lead, so mant = 1.frac * 2^p and the value is + // 1.frac * 2^(p - 24). Biased f32 exponent: 127 + p - 24 + // = 113 - lead; shifting by `lead` moves the leading 1 to + // bit 10, where the mask drops it (it becomes implicit). + let lead = mant.leading_zeros() - 21; + let exp32 = 113 - lead; + let mant32 = (mant << lead) & 0x03FF; + sign | (exp32 << 23) | (mant32 << 13) + } + } else if exp == 0x1F { + // Inf / NaN. + sign | 0x7F80_0000 | (mant << 13) + } else { + sign | ((exp + 127 - 15) << 23) | (mant << 13) + }; + f32::from_bits(out) +} + +/// Encode a full f32 slice into f16 bits, appending to `out`. +#[inline] +pub fn encode_f16_slice(src: &[f32], out: &mut Vec) { + out.reserve(src.len()); + for &v in src { + out.push(f32_to_f16(v)); + } +} + +/// Squared L2 between an f32 query and an f16-encoded vector. +#[inline] +pub fn l2_sq_f16(query: &[f32], vec_f16: &[u16]) -> f32 { + debug_assert_eq!(query.len(), vec_f16.len()); + let mut sum = 0.0f32; + for (q, &h) in query.iter().zip(vec_f16.iter()) { + let d = q - f16_to_f32(h); + sum += d * d; + } + sum +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn known_values_roundtrip_exact() { + // Values exactly representable in f16 must roundtrip bit-perfectly. + for &v in &[ + 0.0f32, -0.0, 1.0, -1.0, 0.5, 2.0, 65504.0, -65504.0, 0.25, 1024.0, + ] { + let back = f16_to_f32(f32_to_f16(v)); + assert_eq!(v.to_bits(), back.to_bits(), "value {v}"); + } + } + + #[test] + fn known_bit_patterns() { + assert_eq!(f32_to_f16(1.0), 0x3C00); + assert_eq!(f32_to_f16(-2.0), 0xC000); + assert_eq!(f32_to_f16(65504.0), 0x7BFF); // f16::MAX + assert_eq!(f32_to_f16(f32::INFINITY), 0x7C00); + assert_eq!(f32_to_f16(f32::NEG_INFINITY), 0xFC00); + assert_eq!(f32_to_f16(65520.0), 0x7C00); // Overflow -> inf. + assert_eq!(f32_to_f16(6.10352e-5), 0x0400); // Smallest normal. + assert_eq!(f32_to_f16(5.96046e-8), 0x0001); // Smallest subnormal. + assert_eq!(f32_to_f16(1e-9), 0x0000); // Underflow -> zero. + } + + #[test] + fn nan_preserved() { + assert!(f16_to_f32(f32_to_f16(f32::NAN)).is_nan()); + } + + #[test] + fn all_subnormals_decode_exact() { + // Every f16 subnormal (exp=0, mant 1..=1023) is exactly mant * 2^-24. + // Regression: an exponent off-by-one decoded every subnormal to HALF + // its value (2^-25 scale), breaking f16_to_f32's "exact" contract. + for mant in 1u16..=0x03FF { + let expected = mant as f32 * (-24f32).exp2(); + let got = f16_to_f32(mant); + assert_eq!( + got.to_bits(), + expected.to_bits(), + "subnormal mant={mant}: got {got:e}, expected {expected:e}" + ); + // Negative counterpart. + let got_neg = f16_to_f32(0x8000 | mant); + assert_eq!(got_neg.to_bits(), (-expected).to_bits()); + } + } + + #[test] + fn all_finite_patterns_roundtrip() { + // f16 -> f32 is exact, so re-encoding must reproduce the identical + // bit pattern for every finite f16 (an incorrectly-scaled decode + // cannot survive this). + for bits in 0u16..=0xFFFF { + let exp = (bits >> 10) & 0x1F; + if exp == 0x1F { + continue; // inf/NaN handled elsewhere + } + let back = f32_to_f16(f16_to_f32(bits)); + assert_eq!(back, bits, "bits={bits:#06x}"); + } + } + + #[test] + fn round_to_nearest_even() { + // 1.0 + 2^-11 is exactly halfway between f16(1.0) and the next f16 up; + // ties-to-even keeps the even mantissa (1.0). + let halfway = f32::from_bits(0x3F80_1000); + assert_eq!(f32_to_f16(halfway), 0x3C00); + // Just above halfway rounds up. + let above = f32::from_bits(0x3F80_1001); + assert_eq!(f32_to_f16(above), 0x3C01); + } + + #[test] + fn relative_error_bounded() { + // RNE guarantees relative error <= 2^-11 for normal-range values. + let mut s = 0x2545F491u32; + for _ in 0..10_000 { + s ^= s << 13; + s ^= s >> 17; + s ^= s << 5; + let v = (s as f32 / u32::MAX as f32) * 200.0 - 100.0; + let back = f16_to_f32(f32_to_f16(v)); + if v.abs() > 6.2e-5 { + let rel = ((v - back) / v).abs(); + assert!(rel <= 4.9e-4, "v={v} back={back} rel={rel}"); + } + } + } + + #[test] + fn l2_sq_matches_f32_within_tolerance() { + let q: Vec = (0..64).map(|i| (i as f32) * 0.031 - 1.0).collect(); + let x: Vec = (0..64).map(|i| (i as f32) * -0.017 + 0.5).collect(); + let mut enc = Vec::new(); + encode_f16_slice(&x, &mut enc); + let exact: f32 = q.iter().zip(x.iter()).map(|(a, b)| (a - b) * (a - b)).sum(); + let approx = l2_sq_f16(&q, &enc); + let rel = ((exact - approx) / exact).abs(); + assert!(rel < 1e-3, "exact={exact} approx={approx} rel={rel}"); + } +} diff --git a/src/vector/hnsw/graph.rs b/src/vector/hnsw/graph.rs index b7be59102..63d8d893e 100644 --- a/src/vector/hnsw/graph.rs +++ b/src/vector/hnsw/graph.rs @@ -274,8 +274,33 @@ impl HnswGraph { #[cfg(target_arch = "aarch64")] { - // No-op on AArch64 for now (PRFM requires nightly intrinsics). - let _ = (neighbor_offset, vector_offset); + // Stable-Rust PRFM via inline asm (the stdarch `_prefetch` intrinsic + // is nightly-only, but `asm!` is stable). Mirrors the x86_64 hint + // pattern above: 2 neighbor cache lines + 3 TQ-code cache lines. + let nptr = self.layer0_neighbors.as_ptr(); + let vptr = _vectors_tq.as_ptr(); + // Addresses are formed with `wrapping_add` so `pointer::add`'s + // in-bounds contract is never invoked; the asm only materializes + // each address in a register. + // SAFETY: PRFM PLDL1KEEP is an architectural hint — it never + // faults, never architecturally reads or writes memory, and + // silently ignores invalid/out-of-bounds addresses. + unsafe { + use core::arch::asm; + asm!( + "prfm pldl1keep, [{n0}]", + "prfm pldl1keep, [{n1}]", + "prfm pldl1keep, [{v0}]", + "prfm pldl1keep, [{v1}]", + "prfm pldl1keep, [{v2}]", + n0 = in(reg) nptr.wrapping_add(neighbor_offset), + n1 = in(reg) nptr.wrapping_add(neighbor_offset + 16), + v0 = in(reg) vptr.wrapping_add(vector_offset), + v1 = in(reg) vptr.wrapping_add(vector_offset + 64), + v2 = in(reg) vptr.wrapping_add(vector_offset + 128), + options(nostack, preserves_flags, readonly) + ); + } } #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))] diff --git a/src/vector/hnsw/search.rs b/src/vector/hnsw/search.rs index 603c04f35..cef11b8ac 100644 --- a/src/vector/hnsw/search.rs +++ b/src/vector/hnsw/search.rs @@ -9,9 +9,10 @@ use smallvec::SmallVec; use super::graph::{HnswGraph, SENTINEL}; use crate::vector::aligned_buffer::AlignedBuffer; +use crate::vector::distance; use crate::vector::turbo_quant::collection::{CollectionMetadata, QuantizationConfig}; use crate::vector::turbo_quant::fwht; -use crate::vector::turbo_quant::sq8::{sq8_l2_adc, sq8_params}; +use crate::vector::turbo_quant::sq8::{sq8_l2_from_stats, sq8_params, sq8_query_stats}; use crate::vector::types::{DistanceMetric, SearchResult, VectorId}; /// Bit vector for O(1) visited tracking. 64x more cache-efficient than HashSet @@ -274,6 +275,16 @@ pub fn hnsw_search_filtered( } else { SmallVec::new() }; + // HQ-2: resolve the SIMD-dispatched ADC stats kernel ONCE per query (not + // per beam candidate) and precompute the query-only constants (Σq_i, + // Σq_i²) it feeds `sq8_l2_from_stats` with — see `turbo_quant::sq8` + // module docs for the algebraic decomposition. + let sq8_stats_fn = distance::table().sq8_stats; + let (sq8_q_sum, sq8_q_sumsq) = if is_sq8 { + sq8_query_stats(&sq8_query) + } else { + (0.0, 0.0) + }; let q_rot = scratch.query_rotated.as_mut_slice(); // Copy query and zero-pad @@ -390,10 +401,22 @@ pub fn hnsw_search_filtered( let dist_bfs = |bfs_pos: u32| -> f32 { let offset = bfs_pos as usize * bytes_per_code; if is_sq8 { - // SQ8: decode (min, scale) from the slot trailer, ADC squared-L2. + // SQ8: decode (min, scale) from the slot trailer, SIMD-dispatched + // ADC squared-L2 (HQ-2). `sq8_stats_fn` was resolved once above, + // outside this per-candidate closure. let slot = &vectors_tq[offset..offset + bytes_per_code]; let (min, scale) = sq8_params(slot, dim); - return sq8_l2_adc(&sq8_query, &slot[..dim], min, scale); + let (dot_qc, sum_c, sumsq_c) = sq8_stats_fn(&sq8_query, &slot[..dim]); + return sq8_l2_from_stats( + dim, + min, + scale, + sq8_q_sum, + sq8_q_sumsq, + dot_qc, + sum_c, + sumsq_c, + ); } let code_only = &vectors_tq[offset..offset + code_len]; let norm_bytes = &vectors_tq[offset + code_len..offset + bytes_per_code]; @@ -529,10 +552,21 @@ pub fn hnsw_search_filtered( let offset = bfs_pos as usize * bytes_per_code; if is_sq8 { // SQ8 ADC is cheap and exact; ignore the budget early-exit. + // SIMD-dispatched (HQ-2), same as `dist_bfs` above. let _ = budget; let slot = &vectors_tq[offset..offset + bytes_per_code]; let (min, scale) = sq8_params(slot, dim); - return sq8_l2_adc(&sq8_query, &slot[..dim], min, scale); + let (dot_qc, sum_c, sumsq_c) = sq8_stats_fn(&sq8_query, &slot[..dim]); + return sq8_l2_from_stats( + dim, + min, + scale, + sq8_q_sum, + sq8_q_sumsq, + dot_qc, + sum_c, + sumsq_c, + ); } let code_only = &vectors_tq[offset..offset + code_len]; let norm_bytes = &vectors_tq[offset + code_len..offset + bytes_per_code]; diff --git a/src/vector/metrics.rs b/src/vector/metrics.rs index a252733cc..52e4704f9 100644 --- a/src/vector/metrics.rs +++ b/src/vector/metrics.rs @@ -18,18 +18,66 @@ pub static MOONSTORE_DISK_OFFLOAD_ENABLED: AtomicBool = AtomicBool::new(false); /// Number of active vector indexes (incremented on FT.CREATE, decremented on FT.DROPINDEX). pub static VECTOR_INDEXES: AtomicU64 = AtomicU64::new(0); -/// Total vectors inserted across all indexes. -pub static VECTOR_TOTAL_VECTORS: AtomicU64 = AtomicU64::new(0); - /// Approximate total memory usage of vector data in bytes. pub static VECTOR_MEMORY_BYTES: AtomicU64 = AtomicU64::new(0); -/// Total number of FT.SEARCH operations executed. -pub static VECTOR_SEARCH_TOTAL: AtomicU64 = AtomicU64::new(0); - /// Rolling last-search latency in microseconds (last-writer-wins). +/// Plain relaxed `store` (no RMW) — cheap enough to stay unstriped. pub static VECTOR_SEARCH_LATENCY_US: AtomicU64 = AtomicU64::new(0); +// -- Striped hot-path counters (XC-5) -- +// +// `increment_search`/`add_vectors` fire on every FT.SEARCH / vector insert from +// every shard thread. A single global `fetch_add` cache line ping-pongs across +// cores under thread-per-core load, taxing exactly the throughput the counter +// observes. Each stripe is cache-line padded; a thread picks a stripe once +// (round-robin at first use) and always RMWs its own line. Reads (INFO — cold +// path) sum all stripes. + +const METRIC_STRIPES: usize = 16; + +#[repr(align(64))] +struct PaddedAtomicU64(AtomicU64); + +#[allow(clippy::declare_interior_mutable_const)] +const PADDED_ZERO: PaddedAtomicU64 = PaddedAtomicU64(AtomicU64::new(0)); + +/// Total FT.SEARCH operations, striped. Read via [`search_total`]. +static VECTOR_SEARCH_TOTAL_STRIPES: [PaddedAtomicU64; METRIC_STRIPES] = + [PADDED_ZERO; METRIC_STRIPES]; + +/// Total vectors inserted, striped. Read via [`total_vectors`]. +static VECTOR_TOTAL_VECTORS_STRIPES: [PaddedAtomicU64; METRIC_STRIPES] = + [PADDED_ZERO; METRIC_STRIPES]; + +/// Stripe slot for the calling thread: assigned round-robin on first use, +/// cached in a thread-local thereafter (one TLS read per metric bump). +#[inline] +fn stripe_index() -> usize { + use std::sync::atomic::AtomicUsize; + static NEXT: AtomicUsize = AtomicUsize::new(0); + thread_local! { + static SLOT: usize = NEXT.fetch_add(1, Ordering::Relaxed) % METRIC_STRIPES; + } + SLOT.with(|s| *s) +} + +/// Sum of all search-counter stripes (cold path: INFO/stats only). +pub fn search_total() -> u64 { + VECTOR_SEARCH_TOTAL_STRIPES + .iter() + .map(|s| s.0.load(Ordering::Relaxed)) + .sum() +} + +/// Sum of all vector-count stripes (cold path: INFO/stats only). +pub fn total_vectors() -> u64 { + VECTOR_TOTAL_VECTORS_STRIPES + .iter() + .map(|s| s.0.load(Ordering::Relaxed)) + .sum() +} + /// Total number of compaction operations completed. pub static VECTOR_COMPACTION_COUNT: AtomicU64 = AtomicU64::new(0); @@ -41,10 +89,12 @@ pub static VECTOR_MUTABLE_SEGMENT_BYTES: AtomicU64 = AtomicU64::new(0); // -- Helper functions (zero-allocation, pure atomics) -- -/// Increment the search counter by 1. +/// Increment the search counter by 1 (striped: RMW on this thread's own line). #[inline] pub fn increment_search() { - VECTOR_SEARCH_TOTAL.fetch_add(1, Ordering::Relaxed); + VECTOR_SEARCH_TOTAL_STRIPES[stripe_index()] + .0 + .fetch_add(1, Ordering::Relaxed); } /// Store the latest search latency in microseconds (last-writer-wins). @@ -70,10 +120,12 @@ pub fn decrement_indexes() { .ok(); } -/// Add to total vector count (called on vector insertion). +/// Add to total vector count (striped: RMW on this thread's own line). #[inline] pub fn add_vectors(count: u64) { - VECTOR_TOTAL_VECTORS.fetch_add(count, Ordering::Relaxed); + VECTOR_TOTAL_VECTORS_STRIPES[stripe_index()] + .0 + .fetch_add(count, Ordering::Relaxed); } /// Update the memory usage gauge (relaxed store). diff --git a/src/vector/mod.rs b/src/vector/mod.rs index 67050b309..a9cee4335 100644 --- a/src/vector/mod.rs +++ b/src/vector/mod.rs @@ -4,6 +4,7 @@ pub mod aligned_buffer; pub mod background_compact; pub mod diskann; pub mod distance; +pub mod f16; pub mod filter; pub mod fusion; pub mod hnsw; @@ -11,6 +12,7 @@ pub mod index_persist; pub mod metrics; pub mod mvcc; pub mod persistence; +pub mod search_pool; pub mod segment; pub mod sparse; pub mod store; diff --git a/src/vector/persistence/recovery.rs b/src/vector/persistence/recovery.rs index a24ad4289..aed9fd691 100644 --- a/src/vector/persistence/recovery.rs +++ b/src/vector/persistence/recovery.rs @@ -150,9 +150,9 @@ fn replay_vector_wal(records: &[VectorWalRecord]) -> (HashMap { let dim = f32_vector.len() as u32; @@ -173,13 +173,11 @@ fn replay_vector_wal(records: &[VectorWalRecord]) -> (HashMap> = match fs::read(seg_dir.join("raw_f16.bin")) { + Ok(bytes) if bytes.len() == mvcc.len() * dim * 2 => Some( + bytes + .chunks_exact(2) + .map(|c| u16::from_le_bytes([c[0], c[1]])) + .collect(), + ), + Ok(bytes) => { + tracing::warn!( + "segment-{segment_id}: raw_f16.bin has {} bytes, expected {} — \ + ignoring sidecar (search degrades to quantized distances)", + bytes.len(), + mvcc.len() * dim * 2 + ); + None + } + Err(_) => None, + }; + let segment = ImmutableSegment::new( graph, vectors_tq, @@ -451,7 +488,8 @@ pub fn read_immutable_segment( collection.clone(), meta.live_count, meta.total_count, - ); + ) + .with_raw_f16(raw_f16); Ok((segment, collection)) } @@ -719,7 +757,7 @@ mod tests { for i in 0..n { let mut v = lcg_f32(dim, (i * 7 + 13) as u32); normalize(&mut v); - seg.append(i as u64, &v, &[], 1.0, i as u64 + 1); + seg.append(i as u64, &v, i as u64 + 1); db.push(v); } let frozen = seg.freeze(); diff --git a/src/vector/search_pool.rs b/src/vector/search_pool.rs new file mode 100644 index 000000000..7abf9fee1 --- /dev/null +++ b/src/vector/search_pool.rs @@ -0,0 +1,402 @@ +//! Intra-query FT.SEARCH worker pool — parallel fan-out of per-segment graph +//! searches within ONE query. +//! +//! ## Why intra-query (not per-query) parallelism +//! +//! Moon's post-compact index is a LIST of immutable segments; a KNN query +//! searches every segment and merges. The serial per-segment loop is the +//! dominant single-query latency term once ≥2 segments exist (each segment is +//! an independent HNSW beam search at the full resolved `ef_search`). Fanning +//! the per-segment searches across worker threads cuts p50 directly — which is +//! what closes a single-connection QPS gap (QPS = 1/latency there). A +//! RediSearch-style pool that only runs *different* queries concurrently +//! cannot move that number. +//! +//! ## Design +//! +//! - Plain `std::thread` workers + `flume` channels (runtime-agnostic: the +//! shard side awaits replies with `recv_async`, which works under monoio and +//! tokio alike and never blocks the event loop). Mirrors +//! [`crate::vector::background_compact::BackgroundCompactor`]. +//! - Jobs carry `Arc` segment references (`SegmentList` holds +//! `Vec>`) — O(1) refcount bumps, no data copies. The +//! segment types are already proven `Send + Sync` by the compaction pool, +//! which ships them to its own threads. +//! - Each worker owns a cached [`SearchScratch`] (grown when a larger +//! `padded_dim` arrives) — zero steady-state allocation per job. +//! - Mutable-segment scanning, cold (DiskANN) and IVF tiers stay on the +//! calling task: mutable needs MVCC context, cold/IVF are rare tiers. The +//! caller overlaps its own mutable scan with the workers. +//! - **Failure containment**: a panicking segment search is caught per-job; +//! the worker survives and the caller receives an empty result for that +//! segment (fail-degraded, never hung — the reply channel is always +//! answered or dropped, and a dropped reply surfaces as `RecvError`). +//! +//! ## Identity guarantee +//! +//! Pooled and serial searches produce identical results: per-segment searches +//! are independent reads of committed-by-definition segments, and the caller +//! merges with `sort_unstable` under `SearchResult`'s total order (distance, +//! then id), so accumulation order cannot affect the final top-k. + +use std::sync::Arc; +use std::sync::OnceLock; + +use roaring::RoaringBitmap; +use smallvec::SmallVec; + +use crate::vector::hnsw::search::SearchScratch; +use crate::vector::persistence::warm_search::WarmSearchSegment; +use crate::vector::segment::immutable::ImmutableSegment; +use crate::vector::turbo_quant::encoder::padded_dimension; +use crate::vector::types::SearchResult; + +/// Hard cap on auto-sized worker count — beyond ~8 workers the per-query +/// segment count (merge auto-triggers at 16 segments) no longer fills the +/// pool, and extra idle threads only cost memory. +const MAX_AUTO_WORKERS: usize = 8; + +/// A graph-tier segment reference a worker can search independently. +/// Immutable and warm segments are committed-by-definition — no MVCC context +/// crosses the thread boundary. +pub enum GraphSegmentRef { + Immutable(Arc), + Warm(Arc), +} + +/// One per-segment search unit. `reply` is answered exactly once (empty on +/// caught panic); if the worker dies anyway, the dropped sender surfaces as +/// `RecvError` on the caller side — never a hang. +pub struct SegmentSearchJob { + pub segment: GraphSegmentRef, + /// Owned query copy shared across this query's jobs (one alloc per query). + pub query: Arc<[f32]>, + pub fetch_k: usize, + pub ef_search: usize, + /// `Some` = ACORN-filtered traversal (allow-list). Post-filter oversampling + /// mode ships `None` here; the caller applies the bitmap to the results. + pub filter: Option>, + pub reply: flume::Sender>, +} + +/// Worker pool for intra-query segment fan-out. Create once per process +/// (see [`init_global`]); drop disconnects the channel and workers exit. +pub struct SearchWorkerPool { + job_tx: flume::Sender, + num_workers: usize, + _workers: Vec>, +} + +impl SearchWorkerPool { + /// Spawn `num_workers` (≥1) searcher threads. + pub fn new(num_workers: usize) -> Self { + assert!(num_workers >= 1, "at least one worker required"); + // Unbounded: job volume is (graph segments per in-flight query) — small + // and naturally bounded by connection count; the shard thread must + // NEVER block on submit. + let (job_tx, job_rx) = flume::unbounded::(); + let workers: Vec<_> = (0..num_workers) + .map(|i| { + let rx = job_rx.clone(); + std::thread::Builder::new() + .name(format!("moon-vec-search-{i}")) + .spawn(move || worker_loop(&rx)) + .expect("failed to spawn vector search worker thread") + }) + .collect(); + Self { + job_tx, + num_workers, + _workers: workers, + } + } + + /// Submit a job. Returns `false` if the pool is shut down (caller must + /// then run the segment search inline — design-for-failure fallback). + #[must_use] + pub fn submit(&self, job: SegmentSearchJob) -> bool { + self.job_tx.send(job).is_ok() + } + + pub fn num_workers(&self) -> usize { + self.num_workers + } +} + +fn worker_loop(rx: &flume::Receiver) { + // Cached scratch, rebuilt when the job's padded_dim differs — amortized + // zero allocation per job for a homogeneous index workload. EXACT match is + // required (not ≥): `query_rotated.len()` acts as the search's padded_dim + // (LUT sizing + code_len invariants in hnsw/search.rs key off it), so an + // oversized buffer is as wrong as an undersized one. + let mut scratch = SearchScratch::new(0, 0); + let mut scratch_pdim: u32 = 0; + while let Ok(job) = rx.recv() { + let pdim = padded_dimension(job.query.len() as u32); + if pdim != scratch_pdim { + scratch = SearchScratch::new(0, pdim); + scratch_pdim = pdim; + } + // Contain a panicking segment search to this job: reply empty and keep + // the worker alive. Scratch is rebuilt afterwards — a panic can leave + // it in an arbitrary (but memory-safe) state. + let outcome = + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| run_job(&job, &mut scratch))); + let results = match outcome { + Ok(r) => r, + Err(_) => { + tracing::warn!( + "vector search worker caught a per-segment search panic; \ + returning empty results for this segment" + ); + scratch = SearchScratch::new(0, scratch_pdim); + SmallVec::new() + } + }; + // Receiver may already be gone (client disconnected mid-search). + let _ = job.reply.send(results); + } + // job_tx dropped → channel disconnected → exit cleanly. +} + +fn run_job(job: &SegmentSearchJob, scratch: &mut SearchScratch) -> SmallVec<[SearchResult; 32]> { + let filter = job.filter.as_deref(); + match &job.segment { + GraphSegmentRef::Immutable(seg) => match filter { + Some(bm) => { + seg.search_filtered(&job.query, job.fetch_k, job.ef_search, scratch, Some(bm)) + } + None => seg.search(&job.query, job.fetch_k, job.ef_search, scratch), + }, + GraphSegmentRef::Warm(seg) => match filter { + Some(bm) => { + seg.search_filtered(&job.query, job.fetch_k, job.ef_search, scratch, Some(bm)) + } + None => seg.search(&job.query, job.fetch_k, job.ef_search, scratch), + }, + } +} + +// ── Process-global pool ───────────────────────────────────────────────────── + +/// `None` inside = explicitly disabled (0 workers). Uninitialized = disabled: +/// unit tests and embedded uses see the exact serial path unless they opt in. +static GLOBAL_POOL: OnceLock> = OnceLock::new(); + +/// Initialize the process-global pool. `num_workers == 0` disables pooling. +/// First call wins (idempotent); called once from server startup. +pub fn init_global(num_workers: usize) { + let _ = GLOBAL_POOL.set(if num_workers == 0 { + None + } else { + Some(SearchWorkerPool::new(num_workers)) + }); +} + +/// The process-global pool, or `None` when disabled/uninitialized. +pub fn global() -> Option<&'static SearchWorkerPool> { + GLOBAL_POOL.get().and_then(|p| p.as_ref()) +} + +/// Auto-sizing: leave the shard threads their cores, cap at +/// [`MAX_AUTO_WORKERS`]. `cores <= shards` (the KV-tuned deployment shape) +/// yields 0 — pooling disabled, zero interference with KV latency. +pub fn auto_workers(cores: usize, shards: usize) -> usize { + cores.saturating_sub(shards).min(MAX_AUTO_WORKERS) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::vector::distance; + use crate::vector::segment::holder::{SearchSnapshot, SegmentHolder, YieldBudget}; + + /// Yield-free budget: unit tests drive the async search with + /// `futures::executor::block_on`, where monoio's `cooperative_yield` + /// (scoped-TLS) would panic. With MAX caps no yield point is ever taken; + /// the pooled path's `recv_async` is executor-agnostic. + fn no_yield_budget() -> YieldBudget { + YieldBudget { + max_segments_per_chunk: usize::MAX, + max_graph_nodes_per_chunk: usize::MAX, + max_brute_force_vecs_per_chunk: usize::MAX, + } + } + use crate::vector::store::VectorStore; + use bytes::Bytes; + + fn random_vec(dim: usize, seed: u64) -> Vec { + let mut state = seed.wrapping_add(1); + (0..dim) + .map(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((state >> 33) as f32) / (u32::MAX as f32) * 2.0 - 1.0 + }) + .collect() + } + + /// Store with one index, `segs` immutable segments of `per_seg` vectors + /// each, plus `tail` vectors left in the mutable segment. + fn build_store(dim: u32, segs: usize, per_seg: usize, tail: usize) -> VectorStore { + distance::init(); + let mut store = VectorStore::new(); + store + .create_index(crate::vector::store::test_index_meta(dim)) + .unwrap(); + let mut n = 0u64; + for _ in 0..segs { + for _ in 0..per_seg { + let key = format!("doc:{n}"); + let hash = xxhash_rust::xxh64::xxh64(key.as_bytes(), 0); + store + .insert_vector(b"idx", &random_vec(dim as usize, n), hash, Bytes::from(key)) + .unwrap(); + n += 1; + } + store.force_compact_index(b"idx").unwrap(); + } + for _ in 0..tail { + let key = format!("doc:{n}"); + let hash = xxhash_rust::xxh64::xxh64(key.as_bytes(), 0); + store + .insert_vector(b"idx", &random_vec(dim as usize, n), hash, Bytes::from(key)) + .unwrap(); + n += 1; + } + store + } + + fn snapshot_for(store: &mut VectorStore, query: Vec, k: usize) -> SearchSnapshot { + let committed = store.txn_manager().committed_treemap().clone(); + let snapshot_lsn = store.txn_manager().current_lsn(); + let idx = store.get_index_mut(b"idx").unwrap(); + let segments = idx.segments.load_full(); + let mutable_len = segments.mutable.len(); + let dim = query.len() as u32; + SearchSnapshot { + segments, + query_f32: query, + k, + ef_search: 100, + filter_bitmap: None, + filter_strategy: crate::vector::filter::selectivity::FilterStrategy::Unfiltered, + snapshot_lsn, + my_txn_id: 0, + committed, + dimension: dim, + mutable_len, + scratch: SearchScratch::new(0, padded_dimension(dim)), + key_hash_to_key: idx.key_hash_to_key.clone(), + } + } + + fn run_search( + store: &mut VectorStore, + query: &[f32], + k: usize, + pool: Option<&SearchWorkerPool>, + ) -> Vec<(u32, u64)> { + let mut snap = snapshot_for(store, query.to_vec(), k); + let results = futures::executor::block_on(SegmentHolder::search_mvcc_yielding_with_pool( + &mut snap, + no_yield_budget(), + pool, + )); + results.iter().map(|r| (r.id.0, r.key_hash)).collect() + } + + /// G-IDENTITY for the pooled path: multi-segment + mutable-tail index must + /// return exactly the serial path's results. + #[test] + fn test_pooled_search_identical_to_serial() { + let dim = 32u32; + let mut store = build_store(dim, 4, 60, 25); + let pool = SearchWorkerPool::new(3); + for q in 0..20u64 { + let query = random_vec(dim as usize, 10_000 + q); + let serial = run_search(&mut store, &query, 10, None); + let pooled = run_search(&mut store, &query, 10, Some(&pool)); + assert_eq!(serial, pooled, "query {q}: pooled != serial"); + assert!(!serial.is_empty(), "query {q}: no results"); + } + } + + /// Thread-safety under concurrency: many OS threads driving pooled + /// searches against the same segments must each match serial. + #[test] + fn test_pooled_search_concurrent_stress() { + let dim = 32u32; + let mut store = build_store(dim, 5, 40, 0); + let pool = Arc::new(SearchWorkerPool::new(4)); + + // Serial ground truth per query. + let queries: Vec> = (0..32u64) + .map(|q| random_vec(dim as usize, 77_000 + q)) + .collect(); + let expected: Vec<_> = queries + .iter() + .map(|q| run_search(&mut store, q, 10, None)) + .collect(); + + // Owned snapshot ingredients (SegmentList Arc is Send+Sync). + let idx = store.get_index_mut(b"idx").unwrap(); + let segments = idx.segments.load_full(); + let key_map = idx.key_hash_to_key.clone(); + let committed = store.txn_manager().committed_treemap().clone(); + let snapshot_lsn = store.txn_manager().current_lsn(); + + let handles: Vec<_> = (0..8) + .map(|t| { + let pool = Arc::clone(&pool); + let segments = Arc::clone(&segments); + let key_map = key_map.clone(); + let committed = committed.clone(); + let queries = queries.clone(); + let expected = expected.clone(); + std::thread::spawn(move || { + for (qi, query) in queries.iter().enumerate() { + let mut snap = SearchSnapshot { + segments: Arc::clone(&segments), + query_f32: query.clone(), + k: 10, + ef_search: 100, + filter_bitmap: None, + filter_strategy: + crate::vector::filter::selectivity::FilterStrategy::Unfiltered, + snapshot_lsn, + my_txn_id: 0, + committed: committed.clone(), + dimension: dim, + mutable_len: segments.mutable.len(), + scratch: SearchScratch::new(0, padded_dimension(dim)), + key_hash_to_key: key_map.clone(), + }; + let results = futures::executor::block_on( + SegmentHolder::search_mvcc_yielding_with_pool( + &mut snap, + no_yield_budget(), + Some(&pool), + ), + ); + let got: Vec<(u32, u64)> = + results.iter().map(|r| (r.id.0, r.key_hash)).collect(); + assert_eq!(got, expected[qi], "thread {t} query {qi} diverged"); + } + }) + }) + .collect(); + for h in handles { + h.join().unwrap(); + } + } + + #[test] + fn test_auto_workers_sizing() { + assert_eq!(auto_workers(8, 1), 7); + assert_eq!(auto_workers(8, 8), 0, "cores==shards must disable pooling"); + assert_eq!(auto_workers(4, 12), 0); + assert_eq!(auto_workers(32, 4), MAX_AUTO_WORKERS, "cap at 8"); + } +} diff --git a/src/vector/segment/compaction.rs b/src/vector/segment/compaction.rs index b1961b3a1..234233919 100644 --- a/src/vector/segment/compaction.rs +++ b/src/vector/segment/compaction.rs @@ -551,7 +551,11 @@ pub fn compact( } else { &[0.0; 16] }; - let code_len = bytes_per_code - 4; + // SQ8 slots are `dim` u8 codes + 8 params bytes (not the TQ `padded/2` nibble + // layout + 4-byte norm) — same ternary as the read sites at ~:1079/:1450. + // Without the branch this is `dim + 4` for SQ8: latent today (every SQ8 read + // is guarded by `codebook_opt.is_none() → continue`) but a slice-OOB landmine. + let code_len = if is_sq8 { dim } else { bytes_per_code - 4 }; let has_raw = !frozen.raw_f32.is_empty(); let dim = frozen.dimension as usize; @@ -721,7 +725,11 @@ pub fn compact( // Sign bit = 1 if original >= centroid (upper sub-bin), 0 if below. let sub_bpv = (padded + 7) / 8; let mut sub_signs_bfs = vec![0u8; n * sub_bpv]; - if has_raw { + // SQ8 has no sub-centroid refinement (no codebook): both inner branches + // below `continue` unconditionally, so without this gate the loop spends + // O(n · padded log padded) on normalize+FWHT whose results are discarded. + // The zero-filled buffer is exactly the SQ8 contract. + if has_raw && !is_sq8 { // Use raw f32 → FWHT rotate → compare against centroid per TQ index let mut work = vec![0.0f32; padded]; for bfs_pos in 0..n { @@ -796,7 +804,7 @@ pub fn compact( } } } - } else if need_cpu_build && !frozen.sub_centroid_signs.is_empty() { + } else if need_cpu_build && !is_sq8 && !frozen.sub_centroid_signs.is_empty() { // Light mode with insert-time sub-centroid signs: remap to BFS order. // graph.to_original(bfs_pos) returns the builder's sequential ID (0..n-1), // which is the index into live_entries. Use it directly, not as internal_id. @@ -834,6 +842,24 @@ pub fn compact( let total_count = frozen.entries.len() as u32; let live_count = n as u32; + // HQ-1: exact-rerank sidecar — f16 copies of the original vectors in BFS + // order (same permutation as tq_bfs), lifted verbatim from the mutable + // segment's f16 buffer (kept in BOTH build modes; unlike raw_f32, which is + // Exact-only). Disk-reloaded rebuilds have no f16 buffer and fall back to + // quantized ADC distances. + let expected_f16 = frozen.entries.len() * dim; + let raw_f16_bfs: Option> = if frozen.raw_f16.len() == expected_f16 && n > 0 { + let mut buf: Vec = Vec::with_capacity(n * dim); + for bfs_pos in 0..n { + let orig_id = graph.to_original(bfs_pos as u32) as usize; + let src = live_entries[orig_id].internal_id as usize * dim; + buf.extend_from_slice(&frozen.raw_f16[src..src + dim]); + } + Some(buf) + } else { + None + }; + let segment = ImmutableSegment::new( graph, AlignedBuffer::from_vec(tq_bfs), @@ -846,7 +872,8 @@ pub fn compact( collection.clone(), live_count, total_count, - ); + ) + .with_raw_f16(raw_f16_bfs); // Step 7 (continued): persist to disk if requested if let Some((dir, segment_id)) = persist { @@ -1080,10 +1107,11 @@ fn merge_graph_union( // ── Step 1: Collect live entries, deduplicate by key_hash ──────────────── // Map key_hash → (insert_lsn, global_id, tq_code_bytes, qjl_bytes, residual_norm, - // sub_centroid_bytes) + // sub_centroid_bytes, raw_f16_bytes) + #[allow(clippy::type_complexity)] let mut by_key_hash: std::collections::HashMap< u64, - (u64, u32, Vec, Vec, f32, Vec), + (u64, u32, Vec, Vec, f32, Vec, Vec), > = std::collections::HashMap::new(); let qjl_bpv = { @@ -1104,9 +1132,15 @@ fn merge_graph_union( }; let sub_bpv = (padded + 7) / 8; + // HQ-1: the merged segment keeps the exact-rerank sidecar only when every + // surviving entry can supply its f16 vector (all-or-nothing — a partial + // sidecar would silently mix exact and ADC distances within one segment). + let mut all_have_raw = true; + for seg in segments { let tq_buf = seg.vectors_tq().as_slice(); let headers = seg.mvcc_headers(); + let seg_raw = seg.raw_f16(); for hdr in headers { // Skip tombstoned entries. @@ -1141,6 +1175,16 @@ fn merge_graph_union( // Sub-centroid sign bytes. let sub_bytes = seg.sub_centroid_bytes_for(bfs_pos, sub_bpv); + // Exact-rerank sidecar slice for this entry (HQ-1). + let raw_bytes: Vec = + match seg_raw.and_then(|r| r.get(bfs_pos * dim..(bfs_pos + 1) * dim)) { + Some(slice) => slice.to_vec(), + None => { + all_have_raw = false; + Vec::new() + } + }; + // Deduplicate: keep highest insert_lsn. let entry = by_key_hash.entry(hdr.key_hash).or_insert(( 0, @@ -1149,6 +1193,7 @@ fn merge_graph_union( Vec::new(), 0.0, Vec::new(), + Vec::new(), )); if hdr.insert_lsn >= entry.0 { *entry = ( @@ -1158,6 +1203,7 @@ fn merge_graph_union( qjl_bytes, norm, sub_bytes, + raw_bytes, ); } } @@ -1180,9 +1226,12 @@ fn merge_graph_union( // ── Step 2: Lay out entries in deterministic order ─────────────────────── // Sort by (insert_lsn asc, key_hash asc) for determinism. - let mut entries: Vec<(u64, u32, Vec, Vec, f32, Vec, u64)> = by_key_hash + #[allow(clippy::type_complexity)] + let mut entries: Vec<(u64, u32, Vec, Vec, f32, Vec, u64, Vec)> = by_key_hash .into_iter() - .map(|(kh, (lsn, gid, code, qjl, norm, sub))| (lsn, gid, code, qjl, norm, sub, kh)) + .map(|(kh, (lsn, gid, code, qjl, norm, sub, raw))| { + (lsn, gid, code, qjl, norm, sub, kh, raw) + }) .collect(); entries.sort_by(|a, b| a.0.cmp(&b.0).then(a.6.cmp(&b.6))); @@ -1192,9 +1241,17 @@ fn merge_graph_union( let mut residual_norms: Vec = Vec::with_capacity(n); let mut sub_orig: Vec = Vec::with_capacity(n * sub_bpv); let mut mvcc_orig: Vec = Vec::with_capacity(n); + let mut raw_orig: Vec = if all_have_raw { + Vec::with_capacity(n * dim) + } else { + Vec::new() + }; - for (i, (lsn, gid, code, qjl, _norm, sub, kh)) in entries.iter().enumerate() { + for (i, (lsn, gid, code, qjl, _norm, sub, kh, raw)) in entries.iter().enumerate() { tq_buffer_orig.extend_from_slice(code); + if all_have_raw { + raw_orig.extend_from_slice(raw); + } if qjl_bpv > 0 { if qjl.len() == qjl_bpv { qjl_orig.extend_from_slice(qjl); @@ -1365,6 +1422,19 @@ fn merge_graph_union( mvcc_bfs.push(hdr); } + // BFS-reorder the exact-rerank sidecar (HQ-1), same permutation as tq_bfs. + let raw_f16_bfs: Option> = if all_have_raw && raw_orig.len() == n * dim { + let mut buf = vec![0u16; n * dim]; + for bfs_pos in 0..n { + let orig_id = graph.to_original(bfs_pos as u32) as usize; + buf[bfs_pos * dim..(bfs_pos + 1) * dim] + .copy_from_slice(&raw_orig[orig_id * dim..(orig_id + 1) * dim]); + } + Some(buf) + } else { + None + }; + // ── Step 7: Recall verification ────────────────────────────────────────── // Sample queries from the merged TQ codes and compare against fan-out // search across the original segments. @@ -1393,7 +1463,8 @@ fn merge_graph_union( collection.clone(), n as u32, n as u32, - ); + ) + .with_raw_f16(raw_f16_bfs); Ok(merged) } @@ -1612,11 +1683,7 @@ mod tests { for i in 0..n { let mut f32_v = lcg_f32(dim, (i * 7 + 13) as u32); normalize(&mut f32_v); - let sq_v: Vec = f32_v - .iter() - .map(|&x| (x * 127.0).clamp(-128.0, 127.0) as i8) - .collect(); - seg.append(i as u64, &f32_v, &sq_v, 1.0, i as u64 + 1); + seg.append(i as u64, &f32_v, i as u64 + 1); } // Mark some as deleted @@ -1665,7 +1732,7 @@ mod tests { for i in 0..n { let mut v = lcg_f32(dim, (i * 7 + 13) as u32); normalize(&mut v); - seg.append(i as u64, &v, &[], 1.0, i as u64 + 1); + seg.append(i as u64, &v, i as u64 + 1); db.push(v); } let frozen = seg.freeze(); @@ -1734,7 +1801,7 @@ mod tests { let gid = id_base + i; let mut v = lcg_f32(dim, (gid * 7 + 13) as u32); normalize(&mut v); - seg.append(gid as u64, &v, &[], 1.0, gid as u64 + 1); + seg.append(gid as u64, &v, gid as u64 + 1); db.push(v); } let frozen = seg.freeze(); diff --git a/src/vector/segment/holder.rs b/src/vector/segment/holder.rs index d3c65aed1..1429fd6af 100644 --- a/src/vector/segment/holder.rs +++ b/src/vector/segment/holder.rs @@ -122,6 +122,10 @@ pub struct SearchSnapshot { pub ef_search: usize, /// Pre-evaluated payload/numeric filter bitmap (owned), or None. pub filter_bitmap: Option, + /// Selectivity-based filter strategy resolved at capture (XC-3), matching + /// the sync path's `select_strategy` dispatch. `Unfiltered` when + /// `filter_bitmap` is None. + pub filter_strategy: FilterStrategy, /// MVCC snapshot LSN captured at entry — governs visibility across yields. pub snapshot_lsn: u64, /// Active txn id (0 for non-transactional reads). @@ -138,8 +142,9 @@ pub struct SearchSnapshot { pub scratch: SearchScratch, /// Key-hash → key map captured at START (§3 C1) so a mid-search delete cannot /// drop an entry this search still needs to resolve. Used by the response - /// builder, not the segment scan. - pub key_hash_to_key: std::collections::HashMap, + /// builder, not the segment scan. `Arc` snapshot: capture is O(1); writers + /// copy-on-write via `Arc::make_mut` (QP-1). + pub key_hash_to_key: std::sync::Arc>, } /// Lock-free segment holder. Searches load() once at query start and hold @@ -596,6 +601,24 @@ impl SegmentHolder { pub async fn search_mvcc_yielding( snap: &mut SearchSnapshot, budget: YieldBudget, + ) -> SmallVec<[SearchResult; 32]> { + Self::search_mvcc_yielding_with_pool(snap, budget, crate::vector::search_pool::global()) + .await + } + + /// [`Self::search_mvcc_yielding`] with an explicit worker pool: when + /// `pool` is `Some` and the index holds ≥2 graph-tier (immutable/warm) + /// segments, the per-segment HNSW searches fan out to the pool while THIS + /// task runs the mutable MVCC scan — replies are awaited via + /// `recv_async`, which parks the task, never the shard event loop. With + /// `pool = None` (or <2 graph segments) the body is the exact serial path. + /// Results are identical either way: segment searches are independent + /// reads and the final `sort_unstable` under `SearchResult`'s total order + /// (distance, then id) makes accumulation order immaterial. + pub async fn search_mvcc_yielding_with_pool( + snap: &mut SearchSnapshot, + budget: YieldBudget, + pool: Option<&crate::vector::search_pool::SearchWorkerPool>, ) -> SmallVec<[SearchResult; 32]> { // Capture-before-yield: move all read-only inputs into owned locals so // the per-chunk loops touch only `snap.scratch` (mutably) — no aliasing @@ -609,6 +632,21 @@ impl SegmentHolder { let filter_ref = filter_bitmap.as_ref(); let k = snap.k; let ef_search = snap.ef_search; + // XC-3: high-selectivity filters (>80% of vectors pass) run the graph + // UNFILTERED with 3×k oversampling and post-filter the results — + // mirrors the sync path's `FilterStrategy::HnswPostFilter` branch. + let post_filter = + filter_ref.is_some() && matches!(snap.filter_strategy, FilterStrategy::HnswPostFilter); + let (fetch_k, graph_filter) = if post_filter { + (k * 3, None) + } else { + (k, filter_ref) + }; + let graph_ef = if post_filter { + ef_search.max(k * 3) + } else { + ef_search + }; let snapshot_lsn = snap.snapshot_lsn; let my_txn_id = snap.my_txn_id; let mutable_len = snap.mutable_len; @@ -630,74 +668,223 @@ impl SegmentHolder { let mut all: SmallVec<[SearchResult; 32]> = SmallVec::new(); + // 0. Intra-query fan-out (search_pool): submit every graph-tier + // (immutable/warm) segment search to the worker pool BEFORE the + // mutable scan so workers overlap with it. Filter semantics mirror + // the serial loops below: ACORN mode ships the allow-list bitmap to + // the worker; HnswPostFilter mode searches unfiltered at fetch_k and + // the bitmap is applied to the collected results (step 2). + let graph_jobs = segments.immutable.len() + segments.warm.len(); + let pooled = pool.filter(|_| graph_jobs >= 2); + let mut pending_replies = 0usize; + let mut reply_rx = None; + if let Some(pool) = pooled { + let (tx, rx) = flume::bounded::>(graph_jobs); + // One owned query copy + optional bitmap clone per query — shared + // across this query's jobs via Arc (not per-segment copies). + let query_arc: std::sync::Arc<[f32]> = std::sync::Arc::from(query_f32); + let filter_arc = graph_filter.map(|bm| std::sync::Arc::new(bm.clone())); + // On submit failure (pool shut down at process teardown) the + // segment is searched inline — the query still answers correctly. + for seg in &segments.immutable { + let job = crate::vector::search_pool::SegmentSearchJob { + segment: crate::vector::search_pool::GraphSegmentRef::Immutable( + std::sync::Arc::clone(seg), + ), + query: std::sync::Arc::clone(&query_arc), + fetch_k, + ef_search: graph_ef, + filter: filter_arc.clone(), + reply: tx.clone(), + }; + if pool.submit(job) { + pending_replies += 1; + } else if graph_filter.is_some() { + all.extend(seg.search_filtered( + query_f32, + fetch_k, + graph_ef, + &mut snap.scratch, + graph_filter, + )); + } else { + let results = seg.search(query_f32, fetch_k, graph_ef, &mut snap.scratch); + if post_filter { + if let Some(bm) = filter_ref { + all.extend(results.into_iter().filter(|r| bm.contains(r.id.0))); + } + } else { + all.extend(results); + } + } + } + for seg in &segments.warm { + let job = crate::vector::search_pool::SegmentSearchJob { + segment: crate::vector::search_pool::GraphSegmentRef::Warm( + std::sync::Arc::clone(seg), + ), + query: std::sync::Arc::clone(&query_arc), + fetch_k, + ef_search: graph_ef, + filter: filter_arc.clone(), + reply: tx.clone(), + }; + if pool.submit(job) { + pending_replies += 1; + } else if graph_filter.is_some() { + all.extend(seg.search_filtered( + query_f32, + fetch_k, + graph_ef, + &mut snap.scratch, + graph_filter, + )); + } else { + let results = seg.search(query_f32, fetch_k, graph_ef, &mut snap.scratch); + if post_filter { + if let Some(bm) = filter_ref { + all.extend(results.into_iter().filter(|r| bm.contains(r.id.0))); + } + } else { + all.extend(results); + } + } + } + reply_rx = Some(rx); + } + // 1. MVCC brute-force over the captured append-only range [0, mutable_len), - // chunked + cooperatively yielded between chunks. Each chunk's top-k - // merges into the same global top-k a single full scan produces. + // chunked + cooperatively yielded between chunks. Query prep (FWHT + // rotation / SQ8 normalize) and the top-k heap are hoisted out of the + // chunk loop (QP-4): one prepare per query, one shared heap — the + // accumulated result is exactly the global top-k a single full scan + // produces. let chunk = budget.max_brute_force_vecs_per_chunk.max(1); - let mut start = 0usize; - while start < mutable_len { - let end = (start + chunk).min(mutable_len); - let part = segments.mutable.brute_force_search_mvcc( + if mutable_len > 0 { + // fetch_k oversamples under HnswPostFilter (filter still applied — + // the mutable scan is linear, filtering there is free). + let mut bf_query = segments.mutable.prepare_brute_force_query( query_f32, - query_state.as_ref(), - k, - filter_ref, - snapshot_lsn, - my_txn_id, - &committed, - start, - end, + query_state.is_some(), + fetch_k, ); - all.extend(part); - start = end; - if start < mutable_len { - crate::admin::metrics_setup::bump_ft_search_cooperative_yield(); - crate::runtime::cooperative_yield().await; + let mut start = 0usize; + while start < mutable_len { + let end = (start + chunk).min(mutable_len); + segments.mutable.brute_force_scan_mvcc_chunk( + &mut bf_query, + query_state.as_ref(), + fetch_k, + filter_ref, + snapshot_lsn, + my_txn_id, + &committed, + start, + end, + ); + start = end; + if start < mutable_len { + crate::admin::metrics_setup::bump_ft_search_cooperative_yield(); + crate::runtime::cooperative_yield().await; + } } + all.extend(bf_query.into_results()); } let seg_cap = budget.max_segments_per_chunk.max(1); let mut since_yield = 0usize; - // 2. HNSW search on immutable segments (committed by definition). - for imm in &segments.immutable { - if filter_ref.is_some() { - all.extend(imm.search_filtered( - query_f32, - k, - ef_search, - &mut snap.scratch, - filter_ref, - )); - } else { - all.extend(imm.search(query_f32, k, ef_search, &mut snap.scratch)); + // 2. Graph-tier (immutable/warm) results. + // + // Pooled: collect the worker replies submitted in step 0 — one + // `recv_async` await per job parks this task (never the event loop) + // until that segment's results land. A dropped reply (worker death — + // catch_unwind already contains per-job panics) degrades to missing + // segment results with a warning, never a hang. HnswPostFilter mode + // applies the bitmap here, mirroring the serial branch below. + let pooled_graph = reply_rx.is_some(); + if let Some(rx) = reply_rx { + for _ in 0..pending_replies { + match rx.recv_async().await { + Ok(results) => { + if post_filter { + if let Some(bm) = filter_ref { + all.extend(results.into_iter().filter(|r| bm.contains(r.id.0))); + } + } else { + all.extend(results); + } + } + Err(_) => { + tracing::warn!( + "vector search worker reply dropped; a segment's results \ + are missing from this query" + ); + } + } } - since_yield += 1; - if since_yield >= seg_cap { - since_yield = 0; - crate::admin::metrics_setup::bump_ft_search_cooperative_yield(); - crate::runtime::cooperative_yield().await; + } + + // Serial path (no pool / <2 graph segments): identical to search_mvcc. + // Strategy dispatch (XC-3): `graph_filter` is None under HnswPostFilter + // (unfiltered traversal at fetch_k = 3×k, bitmap applied to results) — + // otherwise the ACORN-filtered traversal, same as the sync path. + if !pooled_graph { + for imm in &segments.immutable { + if graph_filter.is_some() { + all.extend(imm.search_filtered( + query_f32, + fetch_k, + graph_ef, + &mut snap.scratch, + graph_filter, + )); + } else { + let results = imm.search(query_f32, fetch_k, graph_ef, &mut snap.scratch); + if post_filter { + if let Some(bm) = filter_ref { + all.extend(results.into_iter().filter(|r| bm.contains(r.id.0))); + } + } else { + all.extend(results); + } + } + since_yield += 1; + if since_yield >= seg_cap { + since_yield = 0; + crate::admin::metrics_setup::bump_ft_search_cooperative_yield(); + crate::runtime::cooperative_yield().await; + } } } // 2a. Warm segment search (committed by definition, same as immutable). - for warm_seg in &segments.warm { - if filter_ref.is_some() { - all.extend(warm_seg.search_filtered( - query_f32, - k, - ef_search, - &mut snap.scratch, - filter_ref, - )); - } else { - all.extend(warm_seg.search(query_f32, k, ef_search, &mut snap.scratch)); - } - since_yield += 1; - if since_yield >= seg_cap { - since_yield = 0; - crate::admin::metrics_setup::bump_ft_search_cooperative_yield(); - crate::runtime::cooperative_yield().await; + if !pooled_graph { + for warm_seg in &segments.warm { + if graph_filter.is_some() { + all.extend(warm_seg.search_filtered( + query_f32, + fetch_k, + graph_ef, + &mut snap.scratch, + graph_filter, + )); + } else { + let results = warm_seg.search(query_f32, fetch_k, graph_ef, &mut snap.scratch); + if post_filter { + if let Some(bm) = filter_ref { + all.extend(results.into_iter().filter(|r| bm.contains(r.id.0))); + } + } else { + all.extend(results); + } + } + since_yield += 1; + if since_yield >= seg_cap { + since_yield = 0; + crate::admin::metrics_setup::bump_ft_search_cooperative_yield(); + crate::runtime::cooperative_yield().await; + } } } @@ -813,13 +1000,13 @@ mod tests { // Insert into original mutable { let snap = holder.load(); - snap.mutable.append(1, &[0.0f32; 128], &[0i8; 128], 1.0, 1); + snap.mutable.append(1, &[0.0f32; 128], 1); } // Swap with a new list let new_mutable = Arc::new(MutableSegment::new(128, collection)); - new_mutable.append(2, &[1.0f32; 128], &[1i8; 128], 1.0, 2); - new_mutable.append(3, &[2.0f32; 128], &[2i8; 128], 1.0, 3); + new_mutable.append(2, &[1.0f32; 128], 2); + new_mutable.append(3, &[2.0f32; 128], 3); holder.swap(SegmentList { mutable: new_mutable, @@ -844,9 +1031,8 @@ mod tests { { let snap = holder.load(); for i in 0..5u32 { - let sq = make_sq_vector(dim, i * 13 + 1); let f32_v = vec![0.0f32; dim]; - snap.mutable.append(i as u64, &f32_v, &sq, 1.0, i as u64); + snap.mutable.append(i as u64, &f32_v, i as u64); } } @@ -870,9 +1056,8 @@ mod tests { { let snap = holder.load(); for i in 0..5u32 { - let sq = make_sq_vector(dim, i * 13 + 1); let f32_v = vec![0.0f32; dim]; - snap.mutable.append(i as u64, &f32_v, &sq, 1.0, i as u64); + snap.mutable.append(i as u64, &f32_v, i as u64); } } let _query_sq = make_sq_vector(dim, 1); @@ -896,9 +1081,8 @@ mod tests { { let snap = holder.load(); for i in 0..5u32 { - let sq = make_sq_vector(dim, i * 13 + 1); let f32_v = vec![0.0f32; dim]; - snap.mutable.append(i as u64, &f32_v, &sq, 1.0, i as u64); + snap.mutable.append(i as u64, &f32_v, i as u64); } } let _query_sq = make_sq_vector(dim, 1); @@ -932,9 +1116,8 @@ mod tests { { let snap = holder.load(); for i in 0..5u32 { - let sq = make_sq_vector(dim as usize, i * 13 + 1); let f32_v = vec![0.0f32; dim as usize]; - snap.mutable.append(i as u64, &f32_v, &sq, 1.0, i as u64); + snap.mutable.append(i as u64, &f32_v, i as u64); } } let _query_sq = make_sq_vector(dim as usize, 1); @@ -968,9 +1151,9 @@ mod tests { { let snap = holder.load(); // insert_lsn=1, visible to snapshot=5 - snap.mutable.append(0, &[0.0f32; 4], &[0i8; 4], 1.0, 1); + snap.mutable.append(0, &[0.0f32; 4], 1); // insert_lsn=10, NOT visible to snapshot=5 - snap.mutable.append(1, &[0.0f32; 4], &[1i8; 4], 1.0, 10); + snap.mutable.append(1, &[0.0f32; 4], 10); } let _query_sq = vec![0i8; dim as usize]; let query_f32 = vec![0.0f32; dim as usize]; @@ -1000,8 +1183,7 @@ mod tests { { let snap = holder.load(); // One existing entry far from query (f32 L2 distance) - snap.mutable - .append(0, &[100.0f32; 4], &[100i8, 100, 100, 100], 1.0, 1); + snap.mutable.append(0, &[100.0f32; 4], 1); } let _query_sq = vec![0i8; dim]; let query_f32 = vec![0.0f32; dim]; @@ -1062,9 +1244,8 @@ mod tests { { let snap = holder.load(); for i in 0..5u32 { - let sq = make_sq_vector(dim as usize, i * 13 + 1); let f32_v = vec![0.0f32; dim as usize]; - snap.mutable.append(i as u64, &f32_v, &sq, 1.0, i as u64); + snap.mutable.append(i as u64, &f32_v, i as u64); } } let _query_sq = make_sq_vector(dim as usize, 1); @@ -1107,14 +1288,12 @@ mod tests { assert_eq!(snap_before.mutable.len(), 0); // Insert into mutable (through original snapshot's Arc) - snap_before - .mutable - .append(1, &[0.0f32; 128], &[0i8; 128], 1.0, 1); + snap_before.mutable.append(1, &[0.0f32; 128], 1); // Swap with completely new list let new_mutable = Arc::new(MutableSegment::new(128, collection)); - new_mutable.append(2, &[1.0f32; 128], &[1i8; 128], 1.0, 2); - new_mutable.append(3, &[2.0f32; 128], &[2i8; 128], 1.0, 3); + new_mutable.append(2, &[1.0f32; 128], 2); + new_mutable.append(3, &[2.0f32; 128], 3); holder.swap(SegmentList { mutable: new_mutable, immutable: Vec::new(), @@ -1189,9 +1368,8 @@ mod tests { { let snap = holder.load(); for i in 0..5u32 { - let sq = make_sq_vector(dim, i * 13 + 1); let f32_v = vec![0.0f32; dim]; - snap.mutable.append(i as u64, &f32_v, &sq, 1.0, i as u64); + snap.mutable.append(i as u64, &f32_v, i as u64); } } diff --git a/src/vector/segment/immutable.rs b/src/vector/segment/immutable.rs index 868fb8ccf..4a28fcec4 100644 --- a/src/vector/segment/immutable.rs +++ b/src/vector/segment/immutable.rs @@ -85,6 +85,14 @@ pub struct ImmutableSegment { // limitation — count becomes a lower bound, not exact). has_tombstones: AtomicBool, tombstoned_keys: parking_lot::RwLock>, + + /// Exact-rerank sidecar (HQ-1): f16 copy of each ORIGINAL vector, + /// BFS-ordered like `vectors_tq`, `dimension` halves per entry. When + /// present, beam candidates are re-scored with (near-)exact distances + /// before top-k truncation — the returned distances are then true metric + /// values to f16 tolerance instead of quantized ADC estimates. `None` for + /// segments built without raw vectors (pre-sidecar disk segments). + raw_f16: Option>, } impl ImmutableSegment { @@ -117,7 +125,104 @@ impl ImmutableSegment { created_at: Instant::now(), has_tombstones: AtomicBool::new(false), tombstoned_keys: parking_lot::RwLock::new(HashSet::new()), + raw_f16: None, + } + } + + /// Attach the exact-rerank sidecar (HQ-1): BFS-ordered f16 copies of the + /// original vectors, `dimension` halves per entry. Builder-style so the + /// many `new()` call sites without raw vectors stay untouched. + #[must_use] + pub fn with_raw_f16(mut self, raw_f16: Option>) -> Self { + if let Some(ref buf) = raw_f16 { + debug_assert_eq!( + buf.len(), + self.mvcc.len() * self.collection_meta.dimension as usize, + "raw_f16 sidecar must hold dimension halves per BFS entry" + ); + } + self.raw_f16 = raw_f16; + self + } + + /// The exact-rerank sidecar, if this segment carries one (BFS-ordered, + /// `dimension` u16 halves per entry). Used by segment persistence and + /// GraphUnion merge to propagate the sidecar. + pub fn raw_f16(&self) -> Option<&[u16]> { + self.raw_f16.as_deref() + } + + /// Exact rerank (HQ-1): re-score `candidates` with (near-)exact distances + /// decoded from the f16 sidecar, replacing their quantized ADC estimates, + /// then re-sort ascending. No-op when the segment has no sidecar. + /// + /// Distance conventions match the quantized paths so cross-segment merge + /// stays consistent: + /// - L2: true squared L2 on the original vectors. + /// - Cosine / InnerProduct (unit-sphere metrics, mirroring SQ8's encode + /// normalization): squared L2 between the normalized pair, + /// `2 − 2·⟨q̂,x⟩/‖x‖`. + /// + /// Candidates still carry per-segment internal ids (pre + /// `remap_to_global_ids`); `graph.to_bfs` maps them into the sidecar. + /// + /// Cost control: only the top `4·k` ADC-ranked candidates are re-scored + /// (candidates arrive nearest-first from the beam). The true top-k landing + /// outside a 4× ADC oversample is rare; re-scoring the full ef-wide beam + /// costs ~ef·dim f16 decodes per segment for negligible recall beyond that. + fn rerank_exact(&self, candidates: &mut SmallVec<[SearchResult; 32]>, query: &[f32], k: usize) { + let Some(raw) = self.raw_f16.as_deref() else { + return; + }; + if candidates.is_empty() { + return; + } + let rerank_n = (4 * k.max(1)).min(candidates.len()); + let dim = self.collection_meta.dimension as usize; + let is_l2 = self.collection_meta.metric == crate::vector::types::DistanceMetric::L2; + + // Unit-sphere metrics: normalize the query once per call. + let mut q_unit: Vec = Vec::new(); + let q_ref: &[f32] = if is_l2 { + query + } else { + let norm: f32 = query.iter().map(|x| x * x).sum::().sqrt(); + if norm > 0.0 { + let inv = 1.0 / norm; + q_unit.extend(query.iter().map(|x| x * inv)); + } else { + q_unit.extend_from_slice(query); + } + &q_unit + }; + + for result in candidates[..rerank_n].iter_mut() { + let bfs_pos = self.graph.to_bfs(result.id.0) as usize; + let start = bfs_pos * dim; + let Some(vec_f16) = raw.get(start..start + dim) else { + continue; // Out-of-range id: keep the ADC estimate. + }; + if is_l2 { + result.distance = crate::vector::f16::l2_sq_f16(q_ref, vec_f16); + } else { + // One pass: ⟨q̂,x⟩ and ‖x‖² from the f16-decoded vector. + let mut dot = 0.0f32; + let mut xsq = 0.0f32; + for (q, &h) in q_ref.iter().zip(vec_f16.iter()) { + let x = crate::vector::f16::f16_to_f32(h); + dot += q * x; + xsq += x * x; + } + if xsq > 0.0 { + // f16 rounding can push cos slightly outside [-1, 1]; + // clamp so distances stay in the metric's [0, 4] range. + let cos = (dot / xsq.sqrt()).clamp(-1.0, 1.0); + result.distance = 2.0 - 2.0 * cos; + } + // Zero vector: normalized form undefined — keep ADC estimate. + } } + candidates.sort_unstable(); } /// Mark a key as deleted via interior mutability. @@ -208,16 +313,24 @@ impl ImmutableSegment { ef_search, scratch, ); - // Fallback: rerank with TQ_prod when no sub-centroid data - self.rerank_with_prod(&mut cands, query); + // Fallback: rerank with TQ_prod when no sub-centroid data AND no + // exact sidecar (rerank_exact below supersedes the estimator). + if self.raw_f16.is_none() { + self.rerank_with_prod(&mut cands, query); + } cands }; - // Filter deleted entries before truncating so that k live results are - // returned even when some candidates are tombstoned. + // Filter deleted entries first so tombstones neither consume the + // exact-rerank 4·k budget nor leave stale ADC scores mixed into the + // post-rerank ordering. candidates.retain(|c| { let bfs = self.graph.to_bfs(c.id.0); self.is_live_bfs(bfs) }); + // HQ-1: exact rerank of the live beam (ef candidates) from the f16 + // sidecar — replaces quantized estimates with true metric distances + // before top-k truncation. No-op without a sidecar. + self.rerank_exact(&mut candidates, query, k); candidates.truncate(k); self.remap_to_global_ids(&mut candidates); candidates @@ -248,15 +361,18 @@ impl ImmutableSegment { ); // When sub-centroid signs are used in beam, no rerank needed. - // Only rerank if beam used standard 16-level scoring. - if self.sub_centroid_signs.is_empty() { + // Only rerank if beam used standard 16-level scoring (and no exact + // sidecar — rerank_exact below supersedes the estimator). + if self.sub_centroid_signs.is_empty() && self.raw_f16.is_none() { self.rerank_with_prod(&mut candidates, query); } - // Filter deleted entries before truncating. + // Filter deleted entries first (see comment in search()). candidates.retain(|c| { let bfs = self.graph.to_bfs(c.id.0); self.is_live_bfs(bfs) }); + // HQ-1: exact rerank of the live beam from the f16 sidecar. + self.rerank_exact(&mut candidates, query, k); candidates.truncate(k); self.remap_to_global_ids(&mut candidates); candidates @@ -500,7 +616,11 @@ impl ImmutableSegment { let norms = self.residual_norms.len() * std::mem::size_of::(); let sub = self.sub_centroid_signs.len(); let mvcc = self.mvcc.len() * std::mem::size_of::(); - graph + tq + qjl + norms + sub + mvcc + let sidecar = self + .raw_f16 + .as_ref() + .map_or(0, |v| v.len() * std::mem::size_of::()); + graph + tq + qjl + norms + sub + mvcc + sidecar } /// Fraction of dead entries: (total - live) / total. @@ -677,6 +797,31 @@ impl ImmutableSegment { count } + /// Origin-filtered variant of [`mark_deleted_by_key_hash_install`]: only + /// tombstones entries whose `global_id` is in `source_gids`. + /// + /// Used by the merge-install tombstone replay. A source segment's interior + /// tombstone was recorded against the copies THAT source held; replaying it + /// key_hash-wide onto the merged output would also kill a NEWER same-key + /// copy merged in from a sibling segment (the update-then-compact case — + /// mass index loss under churn). Real DEL/UNLINK tombstones land in every + /// source's interior set, so gating by origin still kills them everywhere. + pub fn mark_deleted_by_key_hash_install_from( + &mut self, + key_hash: u64, + source_gids: &std::collections::HashSet, + ) -> u32 { + let mut count = 0u32; + for h in self.mvcc.iter_mut() { + if h.key_hash == key_hash && h.delete_lsn == 0 && source_gids.contains(&h.global_id) { + h.delete_lsn = 1; // sentinel: deleted during reconciliation + self.live_count = self.live_count.saturating_sub(1); + count += 1; + } + } + count + } + // ── Merge-support accessors (P2) ───────────────────────────────────────── /// Clone the set of key_hashes tombstoned via steady-state interior deletion. diff --git a/src/vector/segment/mutable.rs b/src/vector/segment/mutable.rs index c3514760e..d38342c44 100644 --- a/src/vector/segment/mutable.rs +++ b/src/vector/segment/mutable.rs @@ -11,13 +11,16 @@ use parking_lot::RwLock; use roaring::RoaringBitmap; use smallvec::SmallVec; +use crate::vector::distance; use crate::vector::mvcc::visibility::is_visible; use crate::vector::turbo_quant::collection::{CollectionMetadata, QuantizationConfig}; use crate::vector::turbo_quant::encoder::{ encode_tq_mse_a2, encode_tq_mse_scaled, encode_tq_mse_scaled_with_signs, padded_dimension, }; use crate::vector::turbo_quant::fwht; -use crate::vector::turbo_quant::sq8::{SQ8_PARAMS_BYTES, encode_sq8_into, sq8_l2_adc, sq8_params}; +use crate::vector::turbo_quant::sq8::{ + SQ8_PARAMS_BYTES, encode_sq8_into, sq8_l2_from_stats, sq8_params, sq8_query_stats, +}; use crate::vector::turbo_quant::tq_adc::tq_l2_adc_scaled; use crate::vector::types::{DistanceMetric, SearchResult, VectorId}; @@ -65,6 +68,35 @@ impl<'a> Sq8Query<'a> { } } +/// Per-query state for the chunked MVCC brute-force scan (QP-4): the prepared +/// query buffer and the top-k heap survive across yield chunks so neither is +/// redone per chunk. Build via `MutableSegment::prepare_brute_force_query`. +pub struct BruteForceQuery { + /// SQ8: (possibly normalized) owned query copy; TQ-ADC: FWHT-rotated padded + /// query; TQ-prod: empty (the caller's `TqProdQueryState` carries the prep). + prepared: Vec, + /// Whether the TQ-ADC distance path applies (resolved once at prepare). + use_tq_adc: bool, + /// SQ8 (HQ-2): per-query ADC constants `(Σq_i, Σq_i²)` over the PREPARED + /// (possibly normalized) query — computed once here, combined per + /// candidate via `sq8_l2_from_stats`. Zero for non-SQ8 collections. + sq8_q_sum: f32, + sq8_q_sumsq: f32, + /// Shared top-k accumulator across all chunks. + heap: BinaryHeap, +} + +impl BruteForceQuery { + /// Drain the accumulated global top-k, ascending by distance. + pub fn into_results(self) -> SmallVec<[SearchResult; 32]> { + self.heap + .into_sorted_vec() + .into_iter() + .map(|DistF32(d, id, kh)| SearchResult::with_key_hash(d, VectorId(id), kh)) + .collect() + } +} + /// 48 bytes. MVCC fields prepared for Phase 65. #[repr(C)] pub struct MutableEntry { @@ -89,6 +121,9 @@ pub struct FrozenSegment { /// Raw f32 vectors for exact pairwise distance during HNSW build. /// Layout: dim floats per vector, contiguous. Dropped after compaction. pub raw_f32: Vec, + /// f16 originals for the exact-rerank sidecar (HQ-1). Present in BOTH + /// build modes; dim halves per vector, mutable-internal-id order. + pub raw_f16: Vec, /// Sub-centroid sign bits per vector (ceil(padded_dim/8) bytes each). /// Computed at insert time from pre-quantization FWHT values. pub sub_centroid_signs: Vec, @@ -115,6 +150,11 @@ struct MutableSegmentInner { /// Raw f32 vectors retained for deferred QJL encoding at freeze time. /// Layout: dim floats per vector, contiguous. raw_f32: Vec, + /// f16 copies of the original vectors (HQ-1 exact-rerank sidecar source). + /// Layout: dim halves per vector, contiguous — retained in BOTH build + /// modes (unlike raw_f32, Exact-only) at 2·dim B/vector so compaction can + /// hand the immutable segment its rerank sidecar by permutation alone. + raw_f16: Vec, /// Sub-centroid sign bits computed at insert time. sub_centroid_signs: Vec, sub_sign_bytes_per_vec: usize, @@ -209,6 +249,7 @@ impl MutableSegment { qjl_signs: Vec::new(), residual_norms: Vec::new(), raw_f32: Vec::new(), + raw_f16: Vec::new(), sub_centroid_signs: Vec::new(), sub_sign_bytes_per_vec, entries: Vec::new(), @@ -228,14 +269,7 @@ impl MutableSegment { /// Fast path: only FWHT + quantize + nibble pack (O(d log d)). /// QJL encoding (O(M×d²)) is deferred to freeze() when the segment compacts. /// Mutable brute-force search uses TQ-MSE-only distance (no QJL correction). - pub fn append( - &self, - key_hash: u64, - vector_f32: &[f32], - _vector_sq: &[i8], - _norm: f32, - insert_lsn: u64, - ) -> u32 { + pub fn append(&self, key_hash: u64, vector_f32: &[f32], insert_lsn: u64) -> u32 { let mut inner = self.inner.write(); let internal_id = inner.entries.len() as u32; let dim = inner.dimension as usize; @@ -258,14 +292,15 @@ impl MutableSegment { let is_exact = self.collection.build_mode == crate::vector::turbo_quant::collection::BuildMode::Exact; - let mut extra_bytes = 0usize; + crate::vector::f16::encode_f16_slice(vector_f32, &mut inner.raw_f16); + let mut extra_bytes = dim * 2; // f16 sidecar source if is_exact { let qjl_bpv = inner.qjl_bytes_per_vec; let new_qjl_len = inner.qjl_signs.len() + qjl_bpv; inner.qjl_signs.resize(new_qjl_len, 0u8); inner.residual_norms.push(0.0); inner.raw_f32.extend_from_slice(vector_f32); - extra_bytes = qjl_bpv + 4 + dim * 4; + extra_bytes += qjl_bpv + 4 + dim * 4; } inner.entries.push(MutableEntry { @@ -336,14 +371,15 @@ impl MutableSegment { // Light mode: skip both — saves 1,536 B/vec + avoids O(M×d²) at freeze. let is_exact = self.collection.build_mode == crate::vector::turbo_quant::collection::BuildMode::Exact; - let mut extra_bytes = 0usize; + crate::vector::f16::encode_f16_slice(vector_f32, &mut inner.raw_f16); + let mut extra_bytes = dim * 2; // f16 sidecar source if is_exact { let qjl_bpv = inner.qjl_bytes_per_vec; let new_qjl_len = inner.qjl_signs.len() + qjl_bpv; inner.qjl_signs.resize(new_qjl_len, 0u8); inner.residual_norms.push(0.0); inner.raw_f32.extend_from_slice(vector_f32); - extra_bytes = qjl_bpv + 4 + dim * 4; + extra_bytes += qjl_bpv + 4 + dim * 4; } inner.entries.push(MutableEntry { @@ -393,6 +429,11 @@ impl MutableSegment { // Heap-free query prep: borrow for L2, inline-normalize otherwise. let q = Sq8Query::prepare(query_f32, self.collection.metric); let q = q.as_slice(); + // HQ-2: resolve the SIMD-dispatched ADC stats kernel and the + // per-query constants (Σq_i, Σq_i²) ONCE, before the candidate + // loop below — never per candidate. + let sq8_stats_fn = distance::table().sq8_stats; + let (q_sum, q_sumsq) = sq8_query_stats(q); let mut heap: BinaryHeap = BinaryHeap::with_capacity(k + 1); for entry in &inner.entries { if entry.delete_lsn != 0 { @@ -408,7 +449,9 @@ impl MutableSegment { let off = id * bytes_per_code; let slot = &inner.tq_codes[off..off + bytes_per_code]; let (min, scale) = sq8_params(slot, dim); - let dist = sq8_l2_adc(q, &slot[..dim], min, scale); + let (dot_qc, sum_c, sumsq_c) = sq8_stats_fn(q, &slot[..dim]); + let dist = + sq8_l2_from_stats(dim, min, scale, q_sum, q_sumsq, dot_qc, sum_c, sumsq_c); let global_id = inner.global_id_base + entry.internal_id; if heap.len() < k { heap.push(DistF32(dist, global_id, entry.key_hash)); @@ -564,6 +607,12 @@ impl MutableSegment { /// a yield land at indices ≥ `end` (invisible to this scan), and deletes set /// `delete_lsn > snapshot_lsn` (still visible to this snapshot via /// `is_visible`). Full-scan callers pass `0..len`. (ft-search-off-eventloop) + /// + /// Chunked callers (the yielding path) should instead call + /// [`Self::prepare_brute_force_query`] once and + /// [`Self::brute_force_scan_mvcc_chunk`] per chunk, so the query + /// rotation/normalization and the top-k heap are NOT redone per chunk (QP-4). + #[allow(clippy::too_many_arguments)] pub fn brute_force_search_mvcc( &self, query_f32: &[f32], @@ -576,19 +625,118 @@ impl MutableSegment { start: usize, end: usize, ) -> SmallVec<[SearchResult; 32]> { + let mut q = self.prepare_brute_force_query(query_f32, query_state.is_some(), k); + self.brute_force_scan_mvcc_chunk( + &mut q, + query_state, + k, + allow_bitmap, + snapshot_lsn, + my_txn_id, + committed, + start, + end, + ); + q.into_results() + } + + /// One-time per-query setup for the chunked MVCC brute-force scan (QP-4): + /// prepares the (possibly normalized/FWHT-rotated) query buffer and the + /// shared top-k heap that persist across yield chunks. A single per-query + /// allocation at capture — never per-chunk (G-HOTPATH SAFETY-NET clause). + pub fn prepare_brute_force_query( + &self, + query_f32: &[f32], + have_query_state: bool, + k: usize, + ) -> BruteForceQuery { + let dim = query_f32.len(); + let padded = self.collection.padded_dimension as usize; + let prepared: Vec; + let use_tq_adc: bool; + if self.collection.quantization == QuantizationConfig::Sq8 { + // Mirrors `Sq8Query::prepare`, owned: copy for L2, normalize otherwise. + use_tq_adc = false; + let mut q = query_f32.to_vec(); + if self.collection.metric != DistanceMetric::L2 { + let n: f32 = q.iter().map(|x| x * x).sum::().sqrt(); + if n > 0.0 { + let inv = 1.0 / n; + for v in q.iter_mut() { + *v *= inv; + } + } + } + prepared = q; + } else { + let is_a2 = self.collection.quantization == QuantizationConfig::TurboQuant4A2; + use_tq_adc = !is_a2 + && (!have_query_state + || self.collection.build_mode + == crate::vector::turbo_quant::collection::BuildMode::Light); + if use_tq_adc { + let mut buf = vec![0.0f32; padded]; + buf[..dim].copy_from_slice(query_f32); + let norm: f32 = query_f32.iter().map(|x| x * x).sum::().sqrt(); + if norm > 0.0 { + let inv = 1.0 / norm; + for v in buf[..dim].iter_mut() { + *v *= inv; + } + } + fwht::fwht(&mut buf, self.collection.fwht_sign_flips.as_slice()); + prepared = buf; + } else { + prepared = Vec::new(); + } + } + // HQ-2: per-query ADC constants over the prepared (normalized) query — + // the same buffer the chunk scan feeds the stats kernel. + let (sq8_q_sum, sq8_q_sumsq) = if self.collection.quantization == QuantizationConfig::Sq8 { + sq8_query_stats(&prepared) + } else { + (0.0, 0.0) + }; + BruteForceQuery { + prepared, + use_tq_adc, + sq8_q_sum, + sq8_q_sumsq, + heap: BinaryHeap::with_capacity(k + 1), + } + } + + /// Scan one `[start, end)` chunk, accumulating into the query's shared + /// top-k heap. Isolation contract identical to + /// [`Self::brute_force_search_mvcc`]; the inner read lock is taken per + /// chunk so cooperative yields between chunks never hold it. + #[allow(clippy::too_many_arguments)] + pub fn brute_force_scan_mvcc_chunk( + &self, + q: &mut BruteForceQuery, + query_state: Option<&crate::vector::turbo_quant::inner_product::TqProdQueryState>, + k: usize, + allow_bitmap: Option<&RoaringBitmap>, + snapshot_lsn: u64, + my_txn_id: u64, + committed: &roaring::RoaringTreemap, + start: usize, + end: usize, + ) { let inner = self.inner.read(); let hi = end.min(inner.entries.len()); let lo = start.min(hi); let dim = inner.dimension as usize; - let padded = inner.padded_dimension as usize; let bytes_per_code = inner.bytes_per_code; // SQ8: per-vector affine decode + true squared-L2 ADC (MVCC-visible scan). if self.collection.quantization == QuantizationConfig::Sq8 { - // Heap-free query prep: borrow for L2, inline-normalize otherwise. - let q = Sq8Query::prepare(query_f32, self.collection.metric); - let q = q.as_slice(); - let mut heap: BinaryHeap = BinaryHeap::with_capacity(k + 1); + let q_slice = q.prepared.as_slice(); + // HQ-2: SIMD-dispatched stats kernel resolved once per chunk (a fn + // pointer read), per-query constants carried in `BruteForceQuery`. + let sq8_stats_fn = distance::table().sq8_stats; + let (q_sum, q_sumsq) = (q.sq8_q_sum, q.sq8_q_sumsq); + let heap = &mut q.heap; for entry in &inner.entries[lo..hi] { if !is_visible( entry.insert_lsn, @@ -610,7 +758,9 @@ impl MutableSegment { let off = id * bytes_per_code; let slot = &inner.tq_codes[off..off + bytes_per_code]; let (min, scale) = sq8_params(slot, dim); - let dist = sq8_l2_adc(q, &slot[..dim], min, scale); + let (dot_qc, sum_c, sumsq_c) = sq8_stats_fn(q_slice, &slot[..dim]); + let dist = + sq8_l2_from_stats(dim, min, scale, q_sum, q_sumsq, dot_qc, sum_c, sumsq_c); let global_id = inner.global_id_base + entry.internal_id; if heap.len() < k { heap.push(DistF32(dist, global_id, entry.key_hash)); @@ -621,11 +771,7 @@ impl MutableSegment { } } } - return heap - .into_sorted_vec() - .into_iter() - .map(|DistF32(d, id, kh)| SearchResult::with_key_hash(d, VectorId(id), kh)) - .collect(); + return; } let code_len = bytes_per_code - 4; @@ -637,27 +783,9 @@ impl MutableSegment { self.collection.codebook_16() }; - let use_tq_adc = !is_a2 - && (query_state.is_none() - || self.collection.build_mode - == crate::vector::turbo_quant::collection::BuildMode::Light); - let q_rotated: Vec = if use_tq_adc { - let mut buf = vec![0.0f32; padded]; - buf[..dim].copy_from_slice(query_f32); - let norm: f32 = query_f32.iter().map(|x| x * x).sum::().sqrt(); - if norm > 0.0 { - let inv = 1.0 / norm; - for v in buf[..dim].iter_mut() { - *v *= inv; - } - } - fwht::fwht(&mut buf, self.collection.fwht_sign_flips.as_slice()); - buf - } else { - Vec::new() - }; - - let mut heap: BinaryHeap = BinaryHeap::with_capacity(k + 1); + let use_tq_adc = q.use_tq_adc; + let q_rotated = q.prepared.as_slice(); + let heap = &mut q.heap; for entry in &inner.entries[lo..hi] { if !is_visible( @@ -681,7 +809,7 @@ impl MutableSegment { let tq_code = &inner.tq_codes[tq_offset..tq_offset + code_len]; let dist = if use_tq_adc { - tq_l2_adc_scaled(&q_rotated, tq_code, entry.norm, centroids) + tq_l2_adc_scaled(q_rotated, tq_code, entry.norm, centroids) } else { let qs = query_state.unwrap(); let qjl_bpv = inner.qjl_bytes_per_vec; @@ -712,11 +840,6 @@ impl MutableSegment { } } } - - heap.into_sorted_vec() - .into_iter() - .map(|DistF32(d, id, kh)| SearchResult::with_key_hash(d, VectorId(id), kh)) - .collect() } /// Append within a transaction context. @@ -724,8 +847,6 @@ impl MutableSegment { &self, key_hash: u64, vector_f32: &[f32], - _vector_sq: &[i8], - _norm: f32, insert_lsn: u64, txn_id: u64, ) -> u32 { @@ -753,14 +874,15 @@ impl MutableSegment { let is_exact = self.collection.build_mode == crate::vector::turbo_quant::collection::BuildMode::Exact; - let mut extra_bytes = 0usize; + crate::vector::f16::encode_f16_slice(vector_f32, &mut inner.raw_f16); + let mut extra_bytes = dim * 2; // f16 sidecar source if is_exact { let qjl_bpv = inner.qjl_bytes_per_vec; let new_qjl_len = inner.qjl_signs.len() + qjl_bpv; inner.qjl_signs.resize(new_qjl_len, 0u8); inner.residual_norms.push(0.0); inner.raw_f32.extend_from_slice(vector_f32); - extra_bytes = qjl_bpv + 4 + dim * 4; + extra_bytes += qjl_bpv + 4 + dim * 4; } inner.entries.push(MutableEntry { @@ -793,14 +915,15 @@ impl MutableSegment { let is_exact = self.collection.build_mode == crate::vector::turbo_quant::collection::BuildMode::Exact; - let mut extra_bytes = 0usize; + crate::vector::f16::encode_f16_slice(vector_f32, &mut inner.raw_f16); + let mut extra_bytes = dim * 2; // f16 sidecar source if is_exact { let qjl_bpv = inner.qjl_bytes_per_vec; let new_qjl_len = inner.qjl_signs.len() + qjl_bpv; inner.qjl_signs.resize(new_qjl_len, 0u8); inner.residual_norms.push(0.0); inner.raw_f32.extend_from_slice(vector_f32); - extra_bytes = qjl_bpv + 4 + dim * 4; + extra_bytes += qjl_bpv + 4 + dim * 4; } inner.entries.push(MutableEntry { @@ -872,6 +995,23 @@ impl MutableSegment { } } + /// Tombstone the entry at `internal_id` iff its key_hash matches (VEC-1 + /// HSET-update fast path). The key check makes the O(1) index lookup safe + /// even if the caller's `key_hash → global_id` mapping is stale (e.g. + /// remapped by a concurrent compaction install) — on mismatch the caller + /// falls back to the O(n) `mark_deleted_by_key_hash` scan. + /// Returns `true` if an entry was tombstoned. + pub fn mark_deleted_if_key(&self, internal_id: u32, key_hash: u64, delete_lsn: u64) -> bool { + let mut inner = self.inner.write(); + if let Some(entry) = inner.entries.get_mut(internal_id as usize) { + if entry.key_hash == key_hash && entry.delete_lsn == 0 { + entry.delete_lsn = delete_lsn; + return true; + } + } + false + } + /// Mark all entries matching a key_hash as deleted. pub fn mark_deleted_by_key_hash(&self, key_hash: u64, delete_lsn: u64) -> u32 { let mut inner = self.inner.write(); @@ -970,10 +1110,40 @@ impl MutableSegment { /// Freeze: snapshot TQ codes and entries for compaction. pub fn freeze(&self) -> FrozenSegment { + self.freeze_prefix(usize::MAX) + } + + /// Freeze only the first `n` entries (clamped to len) for a **bounded** + /// compaction build. Bulk loads compact into several threshold-sized + /// segments instead of one giant graph — bounding build memory/latency and + /// giving the intra-query search pool independent segments to fan out + /// over. The frozen window is the prefix `[0, n)`, so entry + /// `vector_offset`s (absolute from 0) stay valid; the tail survives via + /// `clone_suffix(n)` at install, exactly like a mid-build append. + pub fn freeze_prefix(&self, n: usize) -> FrozenSegment { let inner = self.inner.read(); + let n = n.min(inner.entries.len()); + let dim = inner.dimension as usize; + // SQ8 has no QJL/residual side data; its codes are not TQ-decodable, so + // the recompute paths (which assume TQ layout) must be skipped entirely. + let exact_tq = self.collection.build_mode + == crate::vector::turbo_quant::collection::BuildMode::Exact + && self.collection.quantization != QuantizationConfig::Sq8; + // Recompute is whole-buffer; truncate to the frozen window afterwards. + let mut qjl_signs = if exact_tq { + self.recompute_qjl_signs(&inner) + } else { + Vec::new() + }; + qjl_signs.truncate(n * inner.qjl_bytes_per_vec); + let mut residual_norms = if exact_tq { + self.recompute_residual_norms(&inner) + } else { + Vec::new() + }; + residual_norms.truncate(n); FrozenSegment { - entries: inner - .entries + entries: inner.entries[..n] .iter() .map(|e| MutableEntry { internal_id: e.internal_id, @@ -985,27 +1155,25 @@ impl MutableSegment { txn_id: e.txn_id, }) .collect(), - tq_codes: inner.tq_codes.clone(), - // SQ8 has no QJL/residual side data; its codes are not TQ-decodable, so - // the recompute paths (which assume TQ layout) must be skipped entirely. - qjl_signs: if self.collection.build_mode - == crate::vector::turbo_quant::collection::BuildMode::Exact - && self.collection.quantization != QuantizationConfig::Sq8 - { - self.recompute_qjl_signs(&inner) - } else { + tq_codes: inner.tq_codes[..n * inner.bytes_per_code].to_vec(), + qjl_signs, + residual_norms, + // empty in Light mode (nothing was appended) + raw_f32: if inner.raw_f32.is_empty() { Vec::new() + } else { + inner.raw_f32[..n * dim].to_vec() }, - residual_norms: if self.collection.build_mode - == crate::vector::turbo_quant::collection::BuildMode::Exact - && self.collection.quantization != QuantizationConfig::Sq8 - { - self.recompute_residual_norms(&inner) + raw_f16: if inner.raw_f16.is_empty() { + Vec::new() } else { + inner.raw_f16[..n * dim].to_vec() + }, + sub_centroid_signs: if inner.sub_centroid_signs.is_empty() { Vec::new() + } else { + inner.sub_centroid_signs[..n * inner.sub_sign_bytes_per_vec].to_vec() }, - raw_f32: inner.raw_f32.clone(), // empty in Light mode (nothing was appended) - sub_centroid_signs: inner.sub_centroid_signs.clone(), sub_sign_bytes_per_vec: inner.sub_sign_bytes_per_vec, bytes_per_code: inner.bytes_per_code, qjl_bytes_per_vec: inner.qjl_bytes_per_vec, @@ -1071,6 +1239,12 @@ impl MutableSegment { let rs = start * dim; inner.raw_f32[rs..].to_vec() }; + let raw_f16 = if inner.raw_f16.is_empty() { + Vec::new() + } else { + let rs = start * dim; + inner.raw_f16[rs..].to_vec() + }; // ── Entries: rebase internal_id and vector_offset to 0-based ───────── let entries: Vec = inner.entries[start..] @@ -1103,6 +1277,11 @@ impl MutableSegment { count * dim * 4 } else { 0 + }) + + (if !raw_f16.is_empty() { + count * dim * 2 + } else { + 0 }); let new_inner = MutableSegmentInner { @@ -1110,6 +1289,7 @@ impl MutableSegment { qjl_signs, residual_norms, raw_f32, + raw_f16, sub_centroid_signs, sub_sign_bytes_per_vec: sub_bpv, entries, @@ -1347,8 +1527,8 @@ mod tests { let seg = MutableSegment::new(128, col); let v1 = make_f32_vector(128, 1); let v2 = make_f32_vector(128, 2); - assert_eq!(seg.append(100, &v1, &[], 1.0, 1), 0); - assert_eq!(seg.append(200, &v2, &[], 1.0, 2), 1); + assert_eq!(seg.append(100, &v1, 1), 0); + assert_eq!(seg.append(200, &v2, 2), 1); assert_eq!(seg.len(), 2); } @@ -1363,7 +1543,7 @@ mod tests { .map(|i| make_f32_vector(dim, i * 7 + 1)) .collect(); for (i, v) in vectors.iter().enumerate() { - seg.append(i as u64, v, &[], 1.0, i as u64); + seg.append(i as u64, v, i as u64); } let _q_rot = rotate_query(&vectors[0], &col); @@ -1392,7 +1572,7 @@ mod tests { let seg = MutableSegment::new(dim as u32, collection); let db: Vec> = (0..50u32).map(|i| make_f32_vector(dim, 100 + i)).collect(); for (i, v) in db.iter().enumerate() { - seg.append(i as u64, v, &[], 0.0, 1); + seg.append(i as u64, v, 1); } // Query == db[7]; SQ8 must rank that vector first (exact-match invariant @@ -1454,7 +1634,7 @@ mod tests { for x in v.iter_mut() { *x *= scale; } - seg.append(i as u64, &v, &[], 0.0, 1); + seg.append(i as u64, &v, 1); db.push(v); } // Query == db[12] scaled by a different factor: identical direction (cos = 1), @@ -1492,7 +1672,7 @@ mod tests { let seg = MutableSegment::new(dim as u32, collection); let db: Vec> = (0..50u32).map(|i| make_f32_vector(dim, 100 + i)).collect(); for (i, v) in db.iter().enumerate() { - seg.append_transactional(i as u64, v, &[], 0.0, 1, 7); + seg.append_transactional(i as u64, v, 1, 7); } // Direct stride check: n slots of exactly dim + SQ8_PARAMS_BYTES bytes. @@ -1547,7 +1727,7 @@ mod tests { for x in v.iter_mut() { *x *= scale; } - seg.append(i as u64, &v, &[], 0.0, 1); + seg.append(i as u64, &v, 1); db.push(v); } // Query == db[12] at a very different magnitude (same direction). With the @@ -1575,9 +1755,9 @@ mod tests { let v0 = make_f32_vector(dim, 1); let v1 = make_f32_vector(dim, 2); let v2 = make_f32_vector(dim, 3); - seg.append(0, &v0, &[], 1.0, 1); - seg.append(1, &v1, &[], 1.0, 2); - seg.append(2, &v2, &[], 1.0, 3); + seg.append(0, &v0, 1); + seg.append(1, &v1, 2); + seg.append(2, &v2, 3); seg.mark_deleted(0, 10); @@ -1594,8 +1774,8 @@ mod tests { let seg = MutableSegment::new(128, col); let v1 = make_f32_vector(128, 1); let v2 = make_f32_vector(128, 2); - seg.append(100, &v1, &[], 1.5, 1); - seg.append(200, &v2, &[], 2.5, 2); + seg.append(100, &v1, 1); + seg.append(200, &v2, 2); let frozen = seg.freeze(); assert_eq!(frozen.entries.len(), 2); @@ -1613,7 +1793,7 @@ mod tests { distance::init(); let col = make_collection(128); let seg = MutableSegment::new(128, col); - seg.append(1, &make_f32_vector(128, 1), &[], 1.0, 1); + seg.append(1, &make_f32_vector(128, 1), 1); seg.mark_deleted(0, 42); let frozen = seg.freeze(); assert_eq!(frozen.entries[0].delete_lsn, 42); @@ -1630,7 +1810,7 @@ mod tests { .map(|i| make_f32_vector(dim, i * 7 + 1)) .collect(); for (i, v) in vectors.iter().enumerate() { - seg.append(i as u64, v, &[], 1.0, i as u64); + seg.append(i as u64, v, i as u64); } let _q_rot = rotate_query(&vectors[0], &col); diff --git a/src/vector/store.rs b/src/vector/store.rs index c0ab802a8..cb3d5dc61 100644 --- a/src/vector/store.rs +++ b/src/vector/store.rs @@ -189,7 +189,12 @@ pub struct VectorIndex { /// return the original Redis key (e.g., `doc:1755`) instead of the internal /// `vec:` form. Survives compaction and segment merging because /// it's keyed by the stable `key_hash`, not the volatile internal ID. - pub key_hash_to_key: std::collections::HashMap, + /// + /// `Arc`-wrapped so search snapshots capture it in O(1) (QP-1: the previous + /// owned `HashMap` was deep-cloned per query — one entry per indexed vector). + /// Writers mutate via [`Arc::make_mut`]: copy-on-write triggers only when a + /// snapshot is concurrently alive, at most once per snapshot lifetime. + pub key_hash_to_key: Arc>, /// Maps `key_hash` → `global_id` for metadata-only updates. /// /// When `HSET doc:1 category "science"` is called without a vector blob, @@ -200,6 +205,11 @@ pub struct VectorIndex { /// Set to false via FT.CONFIG SET idx AUTOCOMPACT OFF for bulk ingestion. /// Manual FT.COMPACT always works regardless of this flag. pub autocompact_enabled: bool, + /// Recall-gate tolerance for UNATTENDED (background/vacuum) GraphUnion + /// merges (VEC-4). Default 0.70 (catastrophic-collapse guard only); the + /// manual FT.COMPACT merge path uses 0.90. Tunable per index via + /// `FT.CONFIG SET MERGE_RECALL_TOLERANCE <0.0..=1.0>`. + pub merge_recall_tolerance: f32, /// Per-index compaction priority weight for the autovacuum scheduler (W3-deep). /// /// Multiplies the raw `dead_bytes_rate` before comparison in `CompactionScheduler`. @@ -347,7 +357,7 @@ impl VectorIndex { let fs_len = fs.segments.load().mutable.len(); if fs_len >= threshold { let dim = fs.collection.dimension; - Self::compact_segments(&mut fs.segments, &mut fs.scratch, &fs.collection, dim); + Self::compact_segments(&mut fs.segments, &mut fs.scratch, &fs.collection, dim, 0); } } } @@ -399,7 +409,13 @@ impl VectorIndex { // Compact additional fields inline (no in-flight for those). for (_, fs) in &mut self.field_segments { let dim = fs.collection.dimension; - Self::compact_segments(&mut fs.segments, &mut fs.scratch, &fs.collection, dim); + Self::compact_segments( + &mut fs.segments, + &mut fs.scratch, + &fs.collection, + dim, + 0, + ); } return; } @@ -412,60 +428,97 @@ impl VectorIndex { &mut self.scratch, &self.collection, self.meta.dimension, + self.meta.compact_threshold as usize, ); - // Compact additional fields + // Compact additional fields (legacy unbounded semantics: threshold 0). for (_, fs) in &mut self.field_segments { let dim = fs.collection.dimension; - Self::compact_segments(&mut fs.segments, &mut fs.scratch, &fs.collection, dim); + Self::compact_segments(&mut fs.segments, &mut fs.scratch, &fs.collection, dim, 0); } } - /// Compact a single field's mutable segment into an immutable HNSW segment. + /// Compact a field's mutable segment into immutable HNSW segment(s). + /// + /// With a non-zero `compact_threshold` the mutable is drained in + /// `bulk_freeze_cap`-bounded prefix builds — a bulk load yields several + /// independently searchable segments (see `bulk_freeze_cap`). The tail + /// survives each install via `clone_suffix`, preserving the global ID + /// space, exactly like the background install path. fn compact_segments( segments: &mut SegmentHolder, scratch: &mut SearchScratch, collection: &Arc, dimension: u32, + compact_threshold: usize, ) { - let mutable_len = segments.load().mutable.len(); - if mutable_len == 0 { - return; - } - - let frozen = segments.load().mutable.freeze(); + let _ = dimension; let seed = collection.collection_id.wrapping_mul(6364136223846793005); - - match compaction::compact(&frozen, collection, seed, None) { - Ok(immutable) => { - let num_nodes = immutable.graph().num_nodes(); - let padded = collection.padded_dimension; - *scratch = SearchScratch::new(num_nodes, padded); - - let old = segments.load(); - let next_global = old.mutable.next_global_id(); - let mut imm_list = old.immutable.clone(); - imm_list.push(Arc::new(immutable)); - let new_mutable = Arc::new(crate::vector::segment::mutable::MutableSegment::new( - dimension, - collection.clone(), - )); - new_mutable.set_global_id_base(next_global); - let new_list = SegmentList { - mutable: new_mutable, - immutable: imm_list, - ivf: old.ivf.clone(), - warm: old.warm.clone(), - cold: old.cold.clone(), - }; - segments.swap(new_list); + loop { + let snap = segments.load(); + let mutable_len = snap.mutable.len(); + if mutable_len == 0 { + return; } - Err(_e) => { - // Compaction failed (recall too low, etc.) — fall back to brute force + let frozen_len = mutable_len.min(bulk_freeze_cap(mutable_len, compact_threshold)); + let frozen = snap.mutable.freeze_prefix(frozen_len); + drop(snap); + + match compaction::compact(&frozen, collection, seed, None) { + Ok(immutable) => { + let num_nodes = immutable.graph().num_nodes(); + let padded = collection.padded_dimension; + *scratch = SearchScratch::new(num_nodes, padded); + + let old = segments.load(); + let tail_mutable = old.mutable.clone_suffix(frozen_len); + let mut imm_list = old.immutable.clone(); + imm_list.push(Arc::new(immutable)); + let new_list = SegmentList { + mutable: tail_mutable, + immutable: imm_list, + ivf: old.ivf.clone(), + warm: old.warm.clone(), + cold: old.cold.clone(), + }; + segments.swap(new_list); + if frozen_len == mutable_len { + return; // drained + } + } + Err(_e) => { + // Compaction failed (recall too low, etc.) — leave the + // rest in brute-force mutable. + return; + } } } } } +/// Max segments one bulk-loaded mutable is split into when its compact +/// threshold can't bound the build count sensibly (huge loads). Matches the +/// search pool's worker cap — more segments than workers adds merge overhead +/// without more parallelism. +const MAX_BULK_SEGMENTS: usize = 8; + +/// Bounded-freeze cap for one compaction build. `compact_threshold == 0` +/// (auto-compact disabled — legacy/test indexes) keeps the historical +/// whole-mutable single-segment semantics; otherwise a bulk-loaded mutable is +/// compacted in `max(threshold, len/MAX_BULK_SEGMENTS)`-sized builds, so +/// FT.COMPACT after a bulk load yields several independently searchable +/// segments (bounded build memory; intra-query pool fan-out) instead of one +/// giant graph. +fn bulk_freeze_cap(mutable_len: usize, compact_threshold: usize) -> usize { + // Only split when an intra-query pool exists: multiple segments searched + // SERIALLY are strictly slower than one graph (each segment pays the full + // resolved ef beam), so pool-less deployments keep single-segment builds. + if compact_threshold == 0 || crate::vector::search_pool::global().is_none() { + mutable_len + } else { + compact_threshold.max(mutable_len.div_ceil(MAX_BULK_SEGMENTS)) + } +} + /// Walk the window `[0..frozen_len)` of `segments.mutable` and apply /// post-freeze tombstones to `immutable` before it is wrapped in `Arc`. /// @@ -491,9 +544,29 @@ fn snap_and_reconcile( tail_keys.insert(key_hash); }); + // Key_hashes that still have a LIVE copy inside the window. A dead window + // entry whose key also has a live window sibling is an UPDATE leftover + // (VEC-1 tombstones the old copy in place before appending the new one) — + // compact() already filtered the dead copy, and the live sibling IS the + // current version inside `immutable`. Key_hash-wide tombstoning on that + // evidence would delete the current version: every key updated before the + // freeze vanished from search (32% of live keys in the churn soak). + // Only a dead entry with NO live window sibling proves the key is gone + // (DEL/UNLINK marks ALL copies dead; a post-freeze update lands its new + // copy in the tail, which the `tail_keys` arm handles). + let mut live_window_keys: std::collections::HashSet = std::collections::HashSet::new(); snap.mutable .for_each_window_entry(frozen_len, |key_hash, delete_lsn| { - if delete_lsn != 0 || tail_keys.contains(&key_hash) { + if delete_lsn == 0 { + live_window_keys.insert(key_hash); + } + }); + + snap.mutable + .for_each_window_entry(frozen_len, |key_hash, delete_lsn| { + if (delete_lsn != 0 && !live_window_keys.contains(&key_hash)) + || tail_keys.contains(&key_hash) + { immutable.mark_deleted_by_key_hash_install(key_hash); } }); @@ -520,12 +593,19 @@ impl VectorIndex { return false; } let snap = self.segments.load(); - let frozen_len = snap.mutable.len(); - if frozen_len == 0 { + let mutable_len = snap.mutable.len(); + if mutable_len == 0 { return false; } + // Bounded build: a bulk-loaded mutable compacts in threshold-sized + // chunks (the due-gate re-fires while len >= threshold, so the tail + // drains across successive begin/install cycles). + let frozen_len = mutable_len.min(bulk_freeze_cap( + mutable_len, + self.meta.compact_threshold as usize, + )); let frozen_global_base = snap.mutable.global_id_base(); - let frozen = snap.mutable.freeze(); + let frozen = snap.mutable.freeze_prefix(frozen_len); drop(snap); let seed = self @@ -717,9 +797,10 @@ impl VectorIndex { .collection_id .wrapping_mul(6364136223846793005); let mode = self.meta.merge_mode; - // Use 0.70 tolerance (same as vacuum_pass): catch catastrophic recall - // collapse without false-positives on small/medium indexes. - let tolerance = 0.70; + // Default 0.70 (same as vacuum_pass): catch catastrophic recall + // collapse without false-positives on small/medium indexes. Per-index + // override: FT.CONFIG SET MERGE_RECALL_TOLERANCE (VEC-4). + let tolerance = self.merge_recall_tolerance; match compactor.submit_merge(segs.clone(), self.collection.clone(), seed, mode, tolerance) { Ok(reply_rx) => { @@ -804,10 +885,22 @@ impl VectorIndex { // merge_immutable already dropped entries with mvcc.delete_lsn != 0 // at snapshot time. Any `mark_deleted_by_key_hash` call that landed // AFTER the worker snapshot only wrote to the source Arc's interior - // `tombstoned_keys` set. Apply those to the merged output. + // `tombstoned_keys` set. Apply those to the merged output — but gated + // by ORIGIN: a source's tombstone may only kill merged entries whose + // global_id came from that source. An HSET update interior-tombstones + // the OLD copy's home segment while the NEW copy lives on (mutable or + // a sibling segment); a hash-wide replay would kill the new copy too + // (mass loss under update churn). Real DEL/UNLINK tombstones are + // recorded in EVERY segment's interior set, so they still apply. for src in &inflight.merged_sources { - for kh in src.tombstoned_key_hashes() { - merged.mark_deleted_by_key_hash_install(kh); + let tombs = src.tombstoned_key_hashes(); + if tombs.is_empty() { + continue; + } + let src_gids: std::collections::HashSet = + src.mvcc_headers().iter().map(|h| h.global_id).collect(); + for kh in tombs { + merged.mark_deleted_by_key_hash_install_from(kh, &src_gids); } } @@ -1283,9 +1376,10 @@ impl VectorStore { scratch, collection, payload_index: PayloadIndex::new(), - key_hash_to_key: std::collections::HashMap::new(), + key_hash_to_key: Arc::new(std::collections::HashMap::new()), key_hash_to_global_id: std::collections::HashMap::new(), autocompact_enabled: true, + merge_recall_tolerance: 0.70, compaction_weight: COMPACTION_WEIGHT_DEFAULT, field_segments: extra_fields, sparse_stores: HashMap::new(), @@ -1412,7 +1506,7 @@ impl VectorStore { // they track LIVE keys, not historical inserts — without this // they grow monotonically under key churn (~1GB / 24M deletes). // A re-insert of the same key repopulates both maps. - idx.key_hash_to_key.remove(&key_hash); + Arc::make_mut(&mut idx.key_hash_to_key).remove(&key_hash); idx.key_hash_to_global_id.remove(&key_hash); any_deleted = true; } @@ -1773,20 +1867,24 @@ impl VectorStore { key_hash: u64, key: bytes::Bytes, ) -> Result<(), &'static str> { - let idx = self.indexes.get_mut(index_name).ok_or("index not found")?; - let snap = idx.segments.load(); - let insert_lsn = snap.mutable.len() as u64 + 1; - drop(snap); - let sq_vec: Vec = vector - .iter() - .map(|&x| (x * 127.0).clamp(-128.0, 127.0) as i8) - .collect(); - let norm: f32 = vector.iter().map(|x| x * x).sum::().sqrt(); + if !self.indexes.contains_key(index_name) { + return Err("index not found"); + } + // Monotonic store-wide LSN, same allocator as the wire path + // (auto_index_hset). The previous `mutable.len() + 1` restarted after + // every compaction, so a RE-inserted key could carry a LOWER lsn than + // its compacted predecessor — merge dedup (keep highest insert_lsn) + // then kept the stale copy and dropped the current one. + let insert_lsn = self.txn_manager_mut().allocate_lsn(); + let idx = match self.indexes.get_mut(index_name) { + Some(idx) => idx, + None => return Err("index not found"), + }; idx.segments .load() .mutable - .append(key_hash, vector, &sq_vec, norm, insert_lsn); - idx.key_hash_to_key.insert(key_hash, key); + .append(key_hash, vector, insert_lsn); + Arc::make_mut(&mut idx.key_hash_to_key).insert(key_hash, key); Ok(()) } @@ -1969,9 +2067,15 @@ impl VectorStore { let mut stats = VacuumPassStats::default(); for name in names { if self.needs_merge(&name) == Some(true) { - // Use 0.70 tolerance for vacuum: catch catastrophic recall collapse - // without false-positives on small/medium indexes. - match self.force_merge_index_with_tolerance(&name, 0.70) { + // Default 0.70: catch catastrophic recall collapse without + // false-positives on small/medium indexes. Per-index override: + // FT.CONFIG SET MERGE_RECALL_TOLERANCE (VEC-4). + let tolerance = self + .indexes + .get(&name) + .map(|i| i.merge_recall_tolerance) + .unwrap_or(0.70); + match self.force_merge_index_with_tolerance(&name, tolerance) { Ok(ms) => { stats.indexes_merged += 1; stats.total_merged += ms.segments_merged; @@ -2074,6 +2178,30 @@ fn enforce_segment_holder_budget( stats.segments_evicted } +/// Minimal single-field index meta for cross-module unit tests (search_pool +/// identity tests build multi-segment stores through the public store API). +#[cfg(test)] +pub(crate) fn test_index_meta(dim: u32) -> IndexMeta { + IndexMeta { + name: Bytes::from_static(b"idx"), + dimension: dim, + padded_dimension: padded_dimension(dim), + metric: DistanceMetric::L2, + hnsw_m: 8, + hnsw_ef_construction: 50, + hnsw_ef_runtime: 0, + compact_threshold: 0, + source_field: Bytes::from_static(b"vec"), + key_prefixes: vec![Bytes::from_static(b"doc:")], + quantization: QuantizationConfig::TurboQuant4, + build_mode: BuildMode::Light, + vector_fields: Vec::new(), + schema_fields: Vec::new(), + merge_mode: MergeMode::GraphUnion, + keep_raw: false, + } +} + #[cfg(test)] mod tests { use super::*; @@ -2529,6 +2657,95 @@ mod bg_compact_tests { results.iter().map(|r| r.key_hash).collect() } + // ── Bounded bulk compaction (search_pool enabler) ──────────────────────── + + /// A bulk load compacted via FT.COMPACT (force path) must produce + /// ceil(n/threshold) threshold-sized immutable segments, not one giant + /// graph — multiple segments are what the intra-query worker pool fans + /// out over. All keys must remain findable (self-recall probe). + #[test] + fn test_force_compact_bulk_bounded_segments() { + distance::init(); + // Bounded bulk builds are gated on an active intra-query search pool. + crate::vector::search_pool::init_global(1); + let dim = 16u32; + let mut store = VectorStore::new(); + let mut meta = make_idx(dim); + meta.compact_threshold = 100; + store.create_index(meta).unwrap(); + let n = 500u64; + for i in 0..n { + insert( + &mut store, + format!("doc:{i}").as_bytes(), + random_vec(dim as usize, i), + ); + } + store.force_compact_index(b"idx").unwrap(); + + let idx = store.indexes.get_mut(b"idx".as_ref()).unwrap(); + let snap = idx.segments.load_full(); + assert_eq!( + snap.mutable.len(), + 0, + "force compact must drain the mutable" + ); + assert_eq!( + snap.immutable.len(), + 5, + "500 vectors at threshold 100 must yield 5 bounded segments" + ); + + // Self-recall: every key still findable by its own vector. + for i in (0..n).step_by(7) { + let hash = xxhash_rust::xxh64::xxh64(format!("doc:{i}").as_bytes(), 0); + let got = search_key_hashes(&mut store, &random_vec(dim as usize, i), 3); + assert!( + got.contains(&hash), + "doc:{i} lost after bounded bulk compact" + ); + } + } + + /// Background path: successive begin/poll cycles over a bulk-loaded + /// mutable must also chip away in threshold-bounded builds. + #[test] + fn test_bg_compact_bulk_bounded_segments() { + distance::init(); + // Bounded bulk builds are gated on an active intra-query search pool. + crate::vector::search_pool::init_global(1); + let dim = 16u32; + let compactor = BackgroundCompactor::new(1); + let mut store = VectorStore::new(); + let mut meta = make_idx(dim); + meta.compact_threshold = 100; + store.create_index(meta).unwrap(); + for i in 0..500u64 { + insert( + &mut store, + format!("doc:{i}").as_bytes(), + random_vec(dim as usize, i), + ); + } + // Drive begin+install until the mutable drains (bounded per build). + for _ in 0..64 { + store.begin_background_compactions(&compactor); + poll_until_installed(&mut store, 400); + let idx = store.indexes.get_mut(b"idx".as_ref()).unwrap(); + if idx.segments.load().mutable.len() == 0 { + break; + } + } + let idx = store.indexes.get_mut(b"idx".as_ref()).unwrap(); + let snap = idx.segments.load_full(); + assert_eq!(snap.mutable.len(), 0, "bg compaction must eventually drain"); + assert_eq!( + snap.immutable.len(), + 5, + "500 vectors at threshold 100 must yield 5 bounded segments (bg path)" + ); + } + /// Like [`make_idx`] but with a caller-chosen index name (and a matching /// key prefix), so a test can create several independent indexes. fn make_idx_named(name: Bytes, dim: u32) -> IndexMeta { @@ -2824,6 +3041,49 @@ mod bg_compact_tests { ); } + /// A key UPDATED (tombstone old + append new, VEC-1 semantics) BEFORE + /// begin_background_compact() must survive the install. Both the dead old + /// copy and the live new copy sit inside the frozen window; compact() + /// already filters the dead copy, so the install reconcile must NOT + /// key_hash-wide-tombstone the new copy out of the immutable. + /// + /// RED before the fix: snap_and_reconcile treated ANY dead window entry as + /// "key deleted" and killed the key's live compacted copy — every key + /// updated-then-compacted vanished from FT.SEARCH (32% of live keys lost + /// in the Bundle-5 churn soak; regression introduced with VEC-1). + #[test] + fn test_bg_compact_update_before_freeze_survives_install() { + distance::init(); + let compactor = BackgroundCompactor::new(1); + let mut store = VectorStore::new(); + store.create_index(make_idx(64)).unwrap(); + + const T: usize = 20; + for i in 0..T { + let key = format!("doc:{i}"); + insert(&mut store, key.as_bytes(), random_vec(64, i as u64)); + } + + // UPDATE doc:5 BEFORE dispatch: dead old + live new, both in-window. + store.mark_deleted_for_key(b"doc:5"); + let updated_hash = xxhash_rust::xxh64::xxh64(b"doc:5", 0); + let new_vec = random_vec(64, 555); + store + .insert_vector(b"idx", &new_vec, updated_hash, Bytes::from_static(b"doc:5")) + .unwrap(); + + assert_eq!(store.begin_background_compactions(&compactor), 1); + assert!(poll_until_installed(&mut store, 200), "must install"); + + // The updated key must still be findable by its NEW vector, exactly once. + let results = search_key_hashes(&mut store, &new_vec, T + 5); + let count = results.iter().filter(|&&h| h == updated_hash).count(); + assert_eq!( + count, 1, + "updated-then-compacted key must survive install (0=lost, 2=duplicate), got {count}" + ); + } + // ── Test 5: steady-state HDEL tombstones installed immutable ───────────── /// mark_deleted_for_key on an already-installed immutable segment must @@ -3041,6 +3301,60 @@ mod bg_compact_tests { ); } + /// A key UPDATED across segments must survive a merge: old copy in seg1 + /// (interior-tombstoned by the update), new copy compacted into seg2, then + /// seg1+seg2 merged. + /// + /// RED before the fix: `poll_install_merge` replayed seg1's interior + /// tombstone set key_hash-WIDE onto the merged output, killing the NEW + /// copy that came from seg2 — the merge-side twin of the + /// `snap_and_reconcile` update bug (657 keys lost in the churn soak with + /// merges enabled even after the compact-install fix). + #[test] + fn test_bg_merge_update_across_segments_survives() { + distance::init(); + let compactor = BackgroundCompactor::new(1); + let mut store = VectorStore::new(); + store.create_index(make_idx(64)).unwrap(); + + const T: usize = 15; + + // seg1: doc:0..T, including the soon-to-be-updated doc:5. + for i in 0..T { + let key = format!("doc:{i}"); + insert(&mut store, key.as_bytes(), random_vec(64, i as u64)); + } + store.force_compact_index(b"idx").unwrap(); + + // UPDATE doc:5 (VEC-1 semantics): interior-tombstone the old copy in + // the Arc'd seg1, append the new vector to the mutable segment. + let updated_hash = xxhash_rust::xxh64::xxh64(b"doc:5", 0); + store.mark_deleted_for_key(b"doc:5"); + let new_vec = random_vec(64, 555); + store + .insert_vector(b"idx", &new_vec, updated_hash, Bytes::from_static(b"doc:5")) + .unwrap(); + + // seg2: padding + the new doc:5 copy, sealed. + for i in T..2 * T { + let key = format!("doc:{i}"); + insert(&mut store, key.as_bytes(), random_vec(64, i as u64)); + } + store.force_compact_index(b"idx").unwrap(); + + // Merge seg1+seg2 — seg1's tombstone must NOT kill seg2's new copy. + let idx = store.get_index_mut(b"idx").unwrap(); + assert!(idx.begin_background_merge(&compactor), "merge dispatched"); + assert!(poll_until_merged(&mut store, 500), "merge installed"); + + let results = search_key_hashes(&mut store, &new_vec, 2 * T + 5); + let count = results.iter().filter(|&&h| h == updated_hash).count(); + assert_eq!( + count, 1, + "updated-then-merged key must survive install (0=lost, 2=duplicate), got {count}" + ); + } + // ── Merge test 3 ───────────────────────────────────────────────────────── /// A key deleted via steady-state interior tombstone on an Arc'd immutable diff --git a/src/vector/turbo_quant/sq8.rs b/src/vector/turbo_quant/sq8.rs index 6e7211dae..87ca9fe35 100644 --- a/src/vector/turbo_quant/sq8.rs +++ b/src/vector/turbo_quant/sq8.rs @@ -154,6 +154,129 @@ pub fn sq8_ip_adc(query: &[f32], codes: &[u8], min: f32, scale: f32) -> f32 { min * q_sum + scale * dot_code } +// ── SIMD-ready ADC decomposition (HQ-2) ───────────────────────────────── +// +// `sq8_l2_adc`/`sq8_ip_adc` above touch every dimension with a query-and-min +// dependent expression (`query[i] - (min + code[i]*scale)`), which vectorizes +// fine on its own but doesn't factor apart the parts that are pure +// per-QUERY constants — those get recomputed for every beam-search candidate. +// +// Algebraic expansion (finding HQ-2, tmp/VECTOR-DEEP-REVIEW.md): let +// `a_i = q_i - min`, `s = scale`. +// +// d = Σ(q_i - min - s·c_i)² +// = Σa_i² - 2s·Σ(a_i·c_i) + s²·Σc_i² +// = Σa_i² - 2s·(Σq_i·c_i - min·Σc_i) + s²·Σc_i² +// +// Σa_i² = Σ(q_i - min)² = Σq_i² - 2·min·Σq_i + n·min² +// +// `Σq_i` and `Σq_i²` depend ONLY on the query — compute once per query via +// [`sq8_query_stats`]. The remaining per-candidate work is exactly THREE +// running sums over `dim` elements — `Σ(q_i·c_i)`, `Σc_i`, `Σc_i²` — computed +// together in [`sq8_candidate_stats_scalar`] (SIMD-dispatched via +// `distance::table().sq8_stats`). [`sq8_l2_from_stats`] does the final O(1) +// combine. Same asymptotic work as the naive loop (one pass over `dim` per +// candidate), but every summand is now a clean widen+FMA reduction — the +// shape NEON/AVX2/LLVM vectorize well — instead of an interleaved +// subtract-then-square chain that re-touches `min` and `scale` per element. + +/// Per-query precomputation for the fast ADC decomposition: `(Σq_i, Σq_i²)`. +/// +/// Call **once per query** (not per candidate) and thread the result through +/// [`sq8_l2_from_stats`] / [`sq8_ip_from_stats`] for every candidate in the +/// beam. `O(dim)`, but paid once instead of once per beam candidate. +#[inline] +pub fn sq8_query_stats(query: &[f32]) -> (f32, f32) { + let mut sum = 0.0f32; + let mut sumsq = 0.0f32; + for &q in query { + sum += q; + sumsq += q * q; + } + (sum, sumsq) +} + +/// Per-candidate ADC statistics: `(Σ(q_i·c_i), Σc_i, Σc_i²)`. +/// +/// This is the **only** per-candidate pass over `dim` elements in the fast +/// ADC path — cleanly vectorizable (u8 codes widened to f32, multiplied +/// against the f32 query with FMA, three independent running sums). This +/// scalar version is both the portable fallback (installed into +/// `DistanceTable::sq8_stats` when no SIMD tier is available) and the +/// correctness oracle the NEON/AVX2 kernels are checked against. +/// +/// # Panics (debug only) +/// `debug_assert_eq!(query.len(), codes.len())`. +#[inline] +pub fn sq8_candidate_stats_scalar(query: &[f32], codes: &[u8]) -> (f32, f32, f32) { + debug_assert_eq!(query.len(), codes.len()); + let (mut dot0, mut dot1, mut dot2, mut dot3) = (0.0f32, 0.0f32, 0.0f32, 0.0f32); + let (mut sc0, mut sc1, mut sc2, mut sc3) = (0.0f32, 0.0f32, 0.0f32, 0.0f32); + let (mut sq0, mut sq1, mut sq2, mut sq3) = (0.0f32, 0.0f32, 0.0f32, 0.0f32); + let n = codes.len(); + let chunks = n / 4; + for c in 0..chunks { + let i = c * 4; + let c0 = codes[i] as f32; + let c1 = codes[i + 1] as f32; + let c2 = codes[i + 2] as f32; + let c3 = codes[i + 3] as f32; + dot0 += query[i] * c0; + dot1 += query[i + 1] * c1; + dot2 += query[i + 2] * c2; + dot3 += query[i + 3] * c3; + sc0 += c0; + sc1 += c1; + sc2 += c2; + sc3 += c3; + sq0 += c0 * c0; + sq1 += c1 * c1; + sq2 += c2 * c2; + sq3 += c3 * c3; + } + for i in (chunks * 4)..n { + let c0 = codes[i] as f32; + dot0 += query[i] * c0; + sc0 += c0; + sq0 += c0 * c0; + } + ( + (dot0 + dot1) + (dot2 + dot3), + (sc0 + sc1) + (sc2 + sc3), + (sq0 + sq1) + (sq2 + sq3), + ) +} + +/// Combine per-query stats ([`sq8_query_stats`]) with per-candidate stats +/// ([`sq8_candidate_stats_scalar`] or a SIMD equivalent) into the asymmetric +/// squared-L2 ADC distance. `O(1)` — see the module-level algebra derivation +/// above. Architecture-independent: no SIMD needed here, only in the stats +/// pass. +#[inline] +#[allow(clippy::too_many_arguments)] +pub fn sq8_l2_from_stats( + dim: usize, + min: f32, + scale: f32, + q_sum: f32, + q_sumsq: f32, + dot_qc: f32, + sum_c: f32, + sumsq_c: f32, +) -> f32 { + let n = dim as f32; + let a_sumsq = q_sumsq - 2.0 * min * q_sum + n * min * min; + a_sumsq - 2.0 * scale * (dot_qc - min * sum_c) + scale * scale * sumsq_c +} + +/// Combine per-query `Σq_i` with per-candidate `Σ(q_i·c_i)` into the +/// asymmetric inner product ADC: ` = min·Σq_i + scale·Σ(q_i·c_i)`. +/// `O(1)`. +#[inline] +pub fn sq8_ip_from_stats(min: f32, scale: f32, q_sum: f32, dot_qc: f32) -> f32 { + min * q_sum + scale * dot_qc +} + #[cfg(test)] mod tests { use super::*; @@ -299,4 +422,185 @@ mod tests { assert_eq!(&slot[..v.len()], &codes[..]); assert_eq!(sq8_params(&slot, v.len()), (min, scale)); } + + // ── HQ-2: stats-decomposition parity (scalar reference) ───────────── + // + // These prove the ALGEBRA is correct in isolation from any SIMD + // concerns: `sq8_l2_from_stats(sq8_query_stats(q), sq8_candidate_stats_scalar(q, c))` + // must reproduce the original naive `sq8_l2_adc` (same math, reassociated + // for vectorization) to tight float tolerance. NEON/AVX2 kernels are + // checked against `sq8_candidate_stats_scalar` directly in + // `distance::neon` / `distance::avx2`, so passing here is a precondition + // for those to mean anything. + + fn l2_via_stats(query: &[f32], codes: &[u8], min: f32, scale: f32) -> f32 { + let (q_sum, q_sumsq) = sq8_query_stats(query); + let (dot_qc, sum_c, sumsq_c) = sq8_candidate_stats_scalar(query, codes); + sq8_l2_from_stats( + query.len(), + min, + scale, + q_sum, + q_sumsq, + dot_qc, + sum_c, + sumsq_c, + ) + } + + fn ip_via_stats(query: &[f32], codes: &[u8], min: f32, scale: f32) -> f32 { + let (q_sum, _) = sq8_query_stats(query); + let (dot_qc, _, _) = sq8_candidate_stats_scalar(query, codes); + sq8_ip_from_stats(min, scale, q_sum, dot_qc) + } + + #[test] + fn test_stats_l2_matches_naive_adc() { + for &dim in &[1usize, 2, 3, 7, 8, 15, 16, 31, 32, 63, 100, 128, 384, 768] { + let q = pseudo_vec(1000 + dim as u64, dim); + let v = pseudo_vec(2000 + dim as u64, dim); + let slot = encode_sq8(&v); + let (min, scale) = sq8_params(&slot, dim); + let naive = sq8_l2_adc(&q, &slot[..dim], min, scale); + let fast = l2_via_stats(&q, &slot[..dim], min, scale); + let rel = (naive - fast).abs() / naive.max(1e-6); + assert!( + rel < 1e-3, + "dim={dim}: naive={naive} fast(stats)={fast} rel={rel}" + ); + } + } + + #[test] + fn test_stats_negative_min_and_extreme_scale_matches_naive() { + // sq8_params stores the vector's literal min/max — for zero-centered + // embeddings (MiniLM-style, the real workload per CLAUDE.md) `min` is + // routinely negative. Also stress the `scale` extremes: near-zero + // (near-constant vector -> tiny max-min spread) and large (values + // scaled up 1000x -> big spread), since `scale` is squared in the + // `s²·Σc²` term of the expansion and errors there compound fastest. + let dim = 384usize; + let base = pseudo_vec(5000, dim); // negative-min by construction ([-1,1) range) + let q = pseudo_vec(6000, dim); + + let cases: [(&str, Vec); 3] = [ + ("normal_negative_min", base.clone()), + ("large_scale", base.iter().map(|&x| x * 1000.0).collect()), + ( + "tiny_scale_near_constant", + base.iter().map(|&x| 5.0 + x * 1e-4).collect(), + ), + ]; + + for (name, v) in &cases { + let slot = encode_sq8(v); + let (min, scale) = sq8_params(&slot, dim); + assert!( + min < 0.0 || *name != "normal_negative_min", + "case {name}: expected negative min, got {min}" + ); + let naive = sq8_l2_adc(&q, &slot[..dim], min, scale); + let fast = l2_via_stats(&q, &slot[..dim], min, scale); + let rel = (naive - fast).abs() / naive.max(1e-6); + assert!( + rel < 1e-3, + "case {name}: naive={naive} fast(stats)={fast} rel={rel} min={min} scale={scale}" + ); + } + } + + #[test] + fn test_stats_ip_matches_naive_adc() { + for &dim in &[1usize, 7, 16, 100, 256, 384] { + let q = pseudo_vec(3000 + dim as u64, dim); + let v = pseudo_vec(4000 + dim as u64, dim); + let slot = encode_sq8(&v); + let (min, scale) = sq8_params(&slot, dim); + let naive = sq8_ip_adc(&q, &slot[..dim], min, scale); + let fast = ip_via_stats(&q, &slot[..dim], min, scale); + let rel = (naive - fast).abs() / naive.abs().max(1.0); + assert!( + rel < 1e-3, + "dim={dim}: naive={naive} fast(stats)={fast} rel={rel}" + ); + } + } + + #[test] + fn test_stats_self_distance_near_zero_matches_naive_absolute() { + // Advisor must-fix: the algebraic expansion `d = Σq² - 2·min·Σq + + // n·min² - 2s(Σqc - min·Σc) + s²·Σc²` is a difference of large + // near-equal terms — catastrophic cancellation is worst exactly + // where the true distance is near zero, i.e. for the near-neighbor + // candidates whose ranking determines recall. The naive + // `sq8_l2_adc` (direct sum-of-squares) is unconditionally + // well-conditioned there. All other stats-path parity tests above + // use far-apart vectors + *relative* tolerance, which structurally + // hides this failure mode. Assert *absolute* agreement in the + // exact-match regime instead, mirroring + // `test_exact_match_distance_near_zero`'s bound. + for &dim in &[16usize, 64, 128, 256, 384, 768] { + let v = pseudo_vec(9000 + dim as u64, dim); + let slot = encode_sq8(&v); + let (min, scale) = sq8_params(&slot, dim); + let naive = sq8_l2_adc(&v, &slot[..dim], min, scale); + let fast = l2_via_stats(&v, &slot[..dim], min, scale); + // Same quantization bound as test_exact_match_distance_near_zero, + // plus slack for the expansion's extra cancellation error. + let bound = dim as f32 * (scale * 0.5) * (scale * 0.5) + 1e-3; + assert!( + fast <= bound, + "dim={dim}: stats self-distance {fast} exceeds bound {bound} (naive={naive})" + ); + assert!( + fast >= -1e-3, + "dim={dim}: stats self-distance {fast} is meaningfully negative (naive={naive})" + ); + let abs_err = (naive - fast).abs(); + assert!( + abs_err < 1e-2, + "dim={dim}: naive={naive} fast(stats)={fast} abs_err={abs_err} exceeds absolute bound" + ); + } + } + + #[test] + fn test_stats_zero_dim_is_zero() { + let (q_sum, q_sumsq) = sq8_query_stats(&[]); + assert_eq!((q_sum, q_sumsq), (0.0, 0.0)); + let (dot, sc, sq) = sq8_candidate_stats_scalar(&[], &[]); + assert_eq!((dot, sc, sq), (0.0, 0.0, 0.0)); + assert_eq!(sq8_l2_from_stats(0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), 0.0); + } + + #[test] + fn test_stats_nearest_neighbor_ranking_preserved() { + // Same property as `test_nearest_neighbor_ranking_preserved` but for + // the stats-decomposed distance — this is what the beam search will + // actually rank candidates with once dispatched. + let dim = 128; + let q = pseudo_vec(100, dim); + let db: Vec> = (0..50).map(|i| pseudo_vec(200 + i, dim)).collect(); + let exact_best = (0..db.len()) + .min_by(|&a, &b| l2_sq(&q, &db[a]).total_cmp(&l2_sq(&q, &db[b]))) + .unwrap(); + let (q_sum, q_sumsq) = sq8_query_stats(&q); + let sq8_best = (0..db.len()) + .min_by(|&a, &b| { + let sa = encode_sq8(&db[a]); + let sb = encode_sq8(&db[b]); + let (mna, sca) = sq8_params(&sa, dim); + let (mnb, scb) = sq8_params(&sb, dim); + let (dot_a, sc_a, sq_a) = sq8_candidate_stats_scalar(&q, &sa[..dim]); + let (dot_b, sc_b, sq_b) = sq8_candidate_stats_scalar(&q, &sb[..dim]); + sq8_l2_from_stats(dim, mna, sca, q_sum, q_sumsq, dot_a, sc_a, sq_a).total_cmp( + &sq8_l2_from_stats(dim, mnb, scb, q_sum, q_sumsq, dot_b, sc_b, sq_b), + ) + }) + .unwrap(); + assert_eq!( + exact_best, sq8_best, + "stats-decomposed SQ8 nearest neighbor diverged from exact" + ); + } } diff --git a/tests/moonstore_warm_e2e.rs b/tests/moonstore_warm_e2e.rs index 18a94c030..61d0fc26b 100644 --- a/tests/moonstore_warm_e2e.rs +++ b/tests/moonstore_warm_e2e.rs @@ -63,9 +63,7 @@ fn test_warm_transition_end_to_end() { let snap = idx.segments.load(); for i in 0..150u32 { let f32_vec: Vec = (0..128).map(|d| (i * 128 + d) as f32 * 0.001).collect(); - let sq_vec: Vec = f32_vec.iter().map(|v| (v * 100.0) as i8).collect(); - snap.mutable - .append(i as u64, &f32_vec, &sq_vec, 1.0, i as u64); + snap.mutable.append(i as u64, &f32_vec, i as u64); } } @@ -199,9 +197,7 @@ fn test_warm_transition_respects_age_threshold() { let snap = idx.segments.load(); for i in 0..150u32 { let f32_vec: Vec = (0..128).map(|d| (i * 128 + d) as f32 * 0.001).collect(); - let sq_vec: Vec = f32_vec.iter().map(|v| (v * 100.0) as i8).collect(); - snap.mutable - .append(i as u64, &f32_vec, &sq_vec, 1.0, i as u64); + snap.mutable.append(i as u64, &f32_vec, i as u64); } } { @@ -265,9 +261,7 @@ fn test_warm_transition_search_still_works_on_mutable() { let snap = idx.segments.load(); for i in 0..150u32 { let f32_vec: Vec = (0..128).map(|d| (i * 128 + d) as f32 * 0.001).collect(); - let sq_vec: Vec = f32_vec.iter().map(|v| (v * 100.0) as i8).collect(); - snap.mutable - .append(i as u64, &f32_vec, &sq_vec, 1.0, i as u64); + snap.mutable.append(i as u64, &f32_vec, i as u64); } } { @@ -291,9 +285,7 @@ fn test_warm_transition_search_still_works_on_mutable() { let snap = idx.segments.load(); for i in 200..210u32 { let f32_vec: Vec = (0..128).map(|d| (i * 128 + d) as f32 * 0.001).collect(); - let sq_vec: Vec = f32_vec.iter().map(|v| (v * 100.0) as i8).collect(); - snap.mutable - .append(i as u64, &f32_vec, &sq_vec, 1.0, i as u64); + snap.mutable.append(i as u64, &f32_vec, i as u64); } // Mutable segment should have the new vectors assert!( diff --git a/tests/mq_integration.rs b/tests/mq_integration.rs index f289f70c4..1f4fa9d38 100644 --- a/tests/mq_integration.rs +++ b/tests/mq_integration.rs @@ -85,6 +85,7 @@ async fn start_mq_server(num_shards: usize) -> (u16, CancellationToken) { uring_sqpoll_ms: None, io_driver: "auto".to_string(), io_busy_poll_us: 0, + ft_search_workers: None, admin_port: 0, slowlog_log_slower_than: 10000, slowlog_max_len: 128, diff --git a/tests/txn_kv_wiring.rs b/tests/txn_kv_wiring.rs index 64ebb943e..35215dd88 100644 --- a/tests/txn_kv_wiring.rs +++ b/tests/txn_kv_wiring.rs @@ -90,6 +90,7 @@ async fn start_txn_server(num_shards: usize, persistence_dir: &str) -> (u16, Can uring_sqpoll_ms: None, io_driver: "auto".to_string(), io_busy_poll_us: 0, + ft_search_workers: None, admin_port: 0, slowlog_log_slower_than: 10000, slowlog_max_len: 128, diff --git a/tests/vector_del_unindex.rs b/tests/vector_del_unindex.rs new file mode 100644 index 000000000..4335ac56c --- /dev/null +++ b/tests/vector_del_unindex.rs @@ -0,0 +1,381 @@ +//! DEL/UNLINK must remove auto-indexed vectors from FT.* indexes on EVERY +//! dispatch path — not just the cross-shard SPSC `Execute` arm. +//! +//! Found by the Bundle-5 soak diagnostic (scripts/vector-validate.py): after +//! one minute of mixed churn at --shards 1, 20% of FT.SEARCH results were +//! DELETED keys (resurrection) and live-set recall collapsed to 0.735, +//! because the conn-local write path never called +//! `VectorStore::mark_deleted_for_key`. Classic three-dispatch-paths gap: +//! the hook existed on the SPSC `Execute` arm and the tokio sharded handler, +//! but not on the monoio conn-local path (the default runtime's ONLY path at +//! shards=1), handler_single, the MULTI batch path, or the SPSC pipeline arms. +//! +//! Wire-level on purpose: store-level tests cannot catch dispatch wiring. +//! +//! Run alone with: +//! MOON_BIN=$PWD/target/release/moon cargo test --test vector_del_unindex + +#![allow(clippy::unwrap_used)] + +use std::io::{BufReader, Read, Write}; +use std::net::{TcpStream, ToSocketAddrs}; +use std::process::{Child, Command}; +use std::time::{Duration, Instant}; + +const DIM: usize = 8; + +// --------------------------------------------------------------------------- +// Binary resolution + server spawn (pattern: tests/shardslice_live.rs) +// --------------------------------------------------------------------------- + +fn find_moon_binary() -> std::path::PathBuf { + if let Ok(bin) = std::env::var("MOON_BIN") { + let p = std::path::PathBuf::from(bin); + if p.exists() { + return p; + } + } + let manifest = env!("CARGO_MANIFEST_DIR"); + let release = std::path::PathBuf::from(format!("{manifest}/target/release/moon")); + if release.exists() { + return release; + } + let debug = std::path::PathBuf::from(format!("{manifest}/target/debug/moon")); + if debug.exists() { + return debug; + } + panic!("No moon binary found. Build first or set MOON_BIN=/path/to/moon."); +} + +fn free_port() -> u16 { + let l = std::net::TcpListener::bind("127.0.0.1:0").expect("bind :0"); + let p = l.local_addr().expect("local_addr").port(); + drop(l); + p +} + +struct ServerGuard(Child); + +impl Drop for ServerGuard { + fn drop(&mut self) { + let _ = self.0.kill(); + let _ = self.0.wait(); + } +} + +fn spawn_moon(port: u16, dir: &std::path::Path, shards: u32) -> ServerGuard { + let child = Command::new(find_moon_binary()) + .args([ + "--port", + &port.to_string(), + "--dir", + &dir.to_string_lossy(), + "--shards", + &shards.to_string(), + "--appendonly", + "no", + ]) + .stdout(std::fs::File::create(dir.join("moon.stdout.log")).expect("stdout log")) + .stderr(std::fs::File::create(dir.join("moon.stderr.log")).expect("stderr log")) + .spawn() + .expect("spawn moon"); + ServerGuard(child) +} + +// --------------------------------------------------------------------------- +// Minimal RESP client (binary-safe args, full-frame parser) +// --------------------------------------------------------------------------- + +#[derive(Debug, PartialEq)] +enum V { + Simple(String), + Err(String), + Int(i64), + Bulk(Vec), + Arr(Vec), + Null, +} + +struct Client { + reader: BufReader, + writer: TcpStream, +} + +impl Client { + fn connect(port: u16) -> Self { + let addr = format!("127.0.0.1:{port}") + .to_socket_addrs() + .unwrap() + .next() + .unwrap(); + let start = Instant::now(); + let stream = loop { + match TcpStream::connect_timeout(&addr, Duration::from_millis(200)) { + Ok(s) => break s, + Err(_) if start.elapsed() < Duration::from_secs(30) => { + std::thread::sleep(Duration::from_millis(50)); + } + Err(e) => panic!("server never accepted on port {port}: {e}"), + } + }; + stream + .set_read_timeout(Some(Duration::from_secs(15))) + .unwrap(); + let writer = stream.try_clone().unwrap(); + Client { + reader: BufReader::new(stream), + writer, + } + } + + fn encode(args: &[&[u8]]) -> Vec { + let mut out = format!("*{}\r\n", args.len()).into_bytes(); + for a in args { + out.extend_from_slice(format!("${}\r\n", a.len()).as_bytes()); + out.extend_from_slice(a); + out.extend_from_slice(b"\r\n"); + } + out + } + + fn read_line(&mut self) -> String { + let mut line = Vec::new(); + let mut b = [0u8; 1]; + loop { + self.reader.read_exact(&mut b).expect("read byte"); + if b[0] == b'\n' { + break; + } + if b[0] != b'\r' { + line.push(b[0]); + } + } + String::from_utf8_lossy(&line).into_owned() + } + + fn parse(&mut self) -> V { + let line = self.read_line(); + let (t, rest) = line.split_at(1); + match t { + "+" => V::Simple(rest.to_string()), + "-" => V::Err(rest.to_string()), + ":" => V::Int(rest.parse().expect("int")), + "$" => { + let n: i64 = rest.parse().expect("bulk len"); + if n < 0 { + return V::Null; + } + let mut buf = vec![0u8; n as usize + 2]; + self.reader.read_exact(&mut buf).expect("bulk body"); + buf.truncate(n as usize); + V::Bulk(buf) + } + "*" => { + let n: i64 = rest.parse().expect("arr len"); + if n < 0 { + return V::Null; + } + V::Arr((0..n).map(|_| self.parse()).collect()) + } + other => panic!("unexpected RESP type {other:?} (line {line:?})"), + } + } + + fn cmd(&mut self, args: &[&[u8]]) -> V { + self.writer.write_all(&Self::encode(args)).expect("send"); + self.parse() + } + + /// Send all commands in ONE write (a wire pipeline), then read all replies. + fn pipeline(&mut self, cmds: &[Vec>]) -> Vec { + let mut buf = Vec::new(); + for c in cmds { + let refs: Vec<&[u8]> = c.iter().map(|a| a.as_slice()).collect(); + buf.extend_from_slice(&Self::encode(&refs)); + } + self.writer.write_all(&buf).expect("send pipeline"); + cmds.iter().map(|_| self.parse()).collect() + } +} + +fn wait_ready(port: u16) -> Client { + let mut c = Client::connect(port); + let start = Instant::now(); + loop { + match c.cmd(&[b"PING"]) { + V::Simple(s) if s == "PONG" => return c, + _ if start.elapsed() < Duration::from_secs(30) => { + std::thread::sleep(Duration::from_millis(100)); + } + other => panic!("server never answered PING: {other:?}"), + } + } +} + +// --------------------------------------------------------------------------- +// Test fixtures +// --------------------------------------------------------------------------- + +fn vec_blob(seed: u32) -> Vec { + // Distinct, deterministic unit-ish vectors; exact values don't matter. + let mut out = Vec::with_capacity(DIM * 4); + for i in 0..DIM { + let v = ((seed * 31 + i as u32 * 7) % 97) as f32 / 97.0 + 0.01; + out.extend_from_slice(&v.to_le_bytes()); + } + out +} + +fn ft_create(c: &mut Client) { + let r = c.cmd(&[ + b"FT.CREATE", + b"idx", + b"ON", + b"HASH", + b"PREFIX", + b"1", + b"d:", + b"SCHEMA", + b"vec", + b"VECTOR", + b"HNSW", + b"6", + b"TYPE", + b"FLOAT32", + b"DIM", + b"8", + b"DISTANCE_METRIC", + b"L2", + ]); + assert_eq!(r, V::Simple("OK".into()), "FT.CREATE failed"); +} + +fn hset_vectors(c: &mut Client, ids: std::ops::Range) { + for i in ids { + let key = format!("d:{i}"); + let blob = vec_blob(i); + let r = c.cmd(&[b"HSET", key.as_bytes(), b"vec", &blob]); + assert!(matches!(r, V::Int(_)), "HSET d:{i} failed: {r:?}"); + } +} + +/// Returns the set of keys FT.SEARCH finds for a KNN-k probe. +fn search_keys(c: &mut Client, k: u32, probe_seed: u32) -> Vec { + let query = format!("*=>[KNN {k} @vec $B]"); + let blob = vec_blob(probe_seed); + let r = c.cmd(&[ + b"FT.SEARCH", + b"idx", + query.as_bytes(), + b"PARAMS", + b"2", + b"B", + &blob, + b"DIALECT", + b"2", + ]); + let V::Arr(items) = r else { + panic!("FT.SEARCH reply not array: {r:?}"); + }; + // Reply shape: [total, key1, fields1, key2, fields2, ...] + items[1..] + .iter() + .step_by(2) + .filter_map(|v| match v { + V::Bulk(b) => Some(String::from_utf8_lossy(b).into_owned()), + _ => None, + }) + .collect() +} + +fn assert_absent(keys: &[String], dead: &str, ctx: &str) { + assert!( + !keys.iter().any(|k| k == dead), + "{ctx}: deleted key {dead} resurfaced in FT.SEARCH results {keys:?}" + ); +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[test] +fn test_del_unindexes_vector_conn_local() { + let dir = tempfile::tempdir().expect("tempdir"); + let port = free_port(); + let _guard = spawn_moon(port, dir.path(), 1); + let mut c = wait_ready(port); + + ft_create(&mut c); + hset_vectors(&mut c, 0..6); + let before = search_keys(&mut c, 6, 0); + assert!( + before.iter().any(|k| k == "d:1"), + "d:1 must be indexed before DEL (got {before:?})" + ); + + assert_eq!(c.cmd(&[b"DEL", b"d:1"]), V::Int(1), "DEL d:1"); + let after = search_keys(&mut c, 6, 0); + assert_absent(&after, "d:1", "conn-local DEL (shards=1)"); +} + +#[test] +fn test_unlink_unindexes_vector_conn_local() { + let dir = tempfile::tempdir().expect("tempdir"); + let port = free_port(); + let _guard = spawn_moon(port, dir.path(), 1); + let mut c = wait_ready(port); + + ft_create(&mut c); + hset_vectors(&mut c, 0..6); + assert_eq!(c.cmd(&[b"UNLINK", b"d:2"]), V::Int(1), "UNLINK d:2"); + let after = search_keys(&mut c, 6, 0); + assert_absent(&after, "d:2", "conn-local UNLINK (shards=1)"); +} + +#[test] +fn test_multi_exec_del_unindexes_vector() { + let dir = tempfile::tempdir().expect("tempdir"); + let port = free_port(); + let _guard = spawn_moon(port, dir.path(), 1); + let mut c = wait_ready(port); + + ft_create(&mut c); + hset_vectors(&mut c, 0..6); + + assert_eq!(c.cmd(&[b"MULTI"]), V::Simple("OK".into())); + assert_eq!(c.cmd(&[b"DEL", b"d:3"]), V::Simple("QUEUED".into())); + let exec = c.cmd(&[b"EXEC"]); + assert!( + matches!(&exec, V::Arr(rs) if rs.first() == Some(&V::Int(1))), + "EXEC should report DEL=1: {exec:?}" + ); + + let after = search_keys(&mut c, 6, 0); + assert_absent(&after, "d:3", "MULTI/EXEC DEL (shards=1)"); +} + +#[test] +fn test_pipelined_del_unindexes_vector_multishard() { + let dir = tempfile::tempdir().expect("tempdir"); + let port = free_port(); + let _guard = spawn_moon(port, dir.path(), 4); + let mut c = wait_ready(port); + + ft_create(&mut c); + hset_vectors(&mut c, 0..8); + + // One wire write carrying several DELs: exercises the batched/pipelined + // dispatch arms at shards=4 (whichever arm handles it, the vector must go). + let dels: Vec>> = (4..7) + .map(|i| vec![b"DEL".to_vec(), format!("d:{i}").into_bytes()]) + .collect(); + for (i, r) in c.pipeline(&dels).iter().enumerate() { + assert_eq!(*r, V::Int(1), "pipelined DEL #{i}"); + } + + let after = search_keys(&mut c, 8, 0); + for dead in ["d:4", "d:5", "d:6"] { + assert_absent(&after, dead, "pipelined DEL (shards=4)"); + } +} diff --git a/tests/vector_edge_cases.rs b/tests/vector_edge_cases.rs index 56bb2208b..a6c2c82a9 100644 --- a/tests/vector_edge_cases.rs +++ b/tests/vector_edge_cases.rs @@ -104,9 +104,8 @@ fn test_zero_vector_insert_and_search() { let collection = make_test_collection(dim as u32); let seg = MutableSegment::new(dim as u32, collection); let zeros_f32 = vec![0.0f32; dim]; - let zeros_sq = vec![0i8; dim]; - seg.append(1, &zeros_f32, &zeros_sq, 0.0, 1); + seg.append(1, &zeros_f32, 1); let results = seg.brute_force_search(&zeros_f32, None, 1); assert_eq!(results.len(), 1, "should find the zero vector"); @@ -131,8 +130,7 @@ fn test_max_dimension_3072() { sq_vec.push((val.clamp(-1.0, 1.0) * 127.0) as i8); } - let norm = f32_vec.iter().map(|x| x * x).sum::().sqrt(); - seg.append(1, &f32_vec, &sq_vec, norm, 1); + seg.append(1, &f32_vec, 1); assert_eq!(seg.len(), 1); let results = seg.brute_force_search(&f32_vec, None, 1); @@ -164,8 +162,7 @@ fn test_search_k_zero() { let collection = make_test_collection(dim as u32); let seg = MutableSegment::new(dim as u32, collection); let f32_v = vec![1.0f32; dim]; - let sq_v = vec![1i8; dim]; - seg.append(1, &f32_v, &sq_v, 1.0, 1); + seg.append(1, &f32_v, 1); let results = seg.brute_force_search(&f32_v, None, 0); assert!(results.is_empty(), "k=0 should return empty results"); @@ -182,8 +179,7 @@ fn test_search_k_larger_than_index() { let f32_v: Vec = (0..dim) .map(|d| (i * 10 + d as u32) as f32 / 100.0) .collect(); - let sq_v = make_sq_vec(&f32_v); - seg.append(i as u64, &f32_v, &sq_v, 1.0, i as u64); + seg.append(i as u64, &f32_v, i as u64); } let query = vec![0.0f32; dim]; diff --git a/tests/vector_exact_rerank.rs b/tests/vector_exact_rerank.rs new file mode 100644 index 000000000..c4151f344 --- /dev/null +++ b/tests/vector_exact_rerank.rs @@ -0,0 +1,448 @@ +//! HQ-1 (vector deep review 2026-07-05): exact rerank stage. +//! +//! Compacted (immutable) segments rank candidates purely by quantized ADC +//! distance — SQ8/TQ4 estimates with per-coordinate error around scale/2. +//! The exact-rerank sidecar keeps an f16 copy of each original vector and +//! re-scores the beam with (near-)exact distances before top-k truncation: +//! returned distances must match the true metric to f16 tolerance (~1e-3 +//! relative), and recall@k against exact ground truth must be ~1.0. +//! +//! RED before the sidecar lands: returned distances are ADC estimates whose +//! relative error is dominated by quantization (≫ 1e-3), and recall@10 on the +//! seeded workload is measurably below 1.0. + +use bytes::Bytes; + +use moon::vector::distance; +use moon::vector::segment::compaction::MergeMode; +use moon::vector::store::{IndexMeta, VectorStore}; +use moon::vector::turbo_quant::collection::{BuildMode, QuantizationConfig}; +use moon::vector::turbo_quant::encoder::padded_dimension; +use moon::vector::types::{DistanceMetric, SearchResult}; + +const DIM: usize = 32; + +/// xorshift64* PRNG — deterministic, no rng dependency. +struct Rng(u64); +impl Rng { + fn new(seed: u64) -> Self { + Rng(seed.max(1)) + } + fn next_f32(&mut self) -> f32 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + ((self.0 >> 40) as f32 / (1u64 << 24) as f32) * 2.0 - 1.0 + } + fn vec(&mut self, dim: usize) -> Vec { + (0..dim).map(|_| self.next_f32()).collect() + } +} + +fn make_meta(name: &str, quant: QuantizationConfig, metric: DistanceMetric) -> IndexMeta { + IndexMeta { + name: Bytes::from(name.to_owned()), + dimension: DIM as u32, + padded_dimension: padded_dimension(DIM as u32), + metric, + hnsw_m: 16, + hnsw_ef_construction: 200, + hnsw_ef_runtime: 0, + compact_threshold: 0, + source_field: Bytes::from_static(b"vec"), + key_prefixes: vec![Bytes::from_static(b"doc:")], + quantization: quant, + build_mode: BuildMode::Light, + vector_fields: Vec::new(), + schema_fields: Vec::new(), + merge_mode: MergeMode::GraphUnion, + keep_raw: false, + } +} + +/// Build a store with `n` seeded vectors in one compacted immutable segment. +fn build_compacted( + name: &str, + quant: QuantizationConfig, + metric: DistanceMetric, + n: usize, + seed: u64, +) -> (VectorStore, Vec>) { + distance::init(); + let mut store = VectorStore::new(); + store + .create_index(make_meta(name, quant, metric)) + .expect("create_index"); + let mut rng = Rng::new(seed); + let mut vecs = Vec::with_capacity(n); + for i in 0..n { + let v = rng.vec(DIM); + let key = Bytes::from(format!("doc:{i}")); + let key_hash = xxhash_rust::xxh64::xxh64(&key, 0); + store + .insert_vector(name.as_bytes(), &v, key_hash, key) + .expect("insert"); + vecs.push(v); + } + store + .force_compact_index(name.as_bytes()) + .expect("force_compact"); + assert_eq!( + store.immutable_segment_count(name.as_bytes()), + Some(1), + "expected exactly one immutable segment" + ); + (store, vecs) +} + +fn l2_sq(a: &[f32], b: &[f32]) -> f32 { + a.iter().zip(b.iter()).map(|(x, y)| (x - y) * (x - y)).sum() +} + +fn normalize(v: &[f32]) -> Vec { + let n: f32 = v.iter().map(|x| x * x).sum::().sqrt(); + if n > 0.0 { + v.iter().map(|x| x / n).collect() + } else { + v.to_vec() + } +} + +fn segment_search(store: &mut VectorStore, name: &str, q: &[f32], k: usize) -> Vec { + let idx = store.get_index_mut(name.as_bytes()).expect("index"); + idx.segments + .search(q, k, 128, &mut idx.scratch) + .into_iter() + .collect() +} + +// ── 1. L2: returned distances are exact (f16 tolerance) ───────────────────── + +#[test] +fn sq8_compacted_l2_distances_are_exact() { + let (mut store, vecs) = build_compacted( + "xr_l2", + QuantizationConfig::Sq8, + DistanceMetric::L2, + 64, + 0xA11CE, + ); + let mut rng = Rng::new(0xBEEF); + for _ in 0..5 { + let q = rng.vec(DIM); + let results = segment_search(&mut store, "xr_l2", &q, 10); + assert!(!results.is_empty()); + for r in &results { + let truth = l2_sq(&q, &vecs[r.id.0 as usize]); + let rel = (r.distance - truth).abs() / truth.max(1e-6); + assert!( + rel < 1.5e-3, + "id={} returned={} exact={} rel={rel} — distance is a quantized \ + estimate, exact rerank missing", + r.id.0, + r.distance, + truth + ); + } + } +} + +// ── 2. Cosine: distances are exact normalized-pair L2² ────────────────────── + +#[test] +fn sq8_compacted_cosine_distances_are_exact() { + let (mut store, vecs) = build_compacted( + "xr_cos", + QuantizationConfig::Sq8, + DistanceMetric::Cosine, + 64, + 0xC051, + ); + let mut rng = Rng::new(0xF00D); + for _ in 0..5 { + let q = rng.vec(DIM); + let qn = normalize(&q); + let results = segment_search(&mut store, "xr_cos", &q, 10); + assert!(!results.is_empty()); + for r in &results { + let vn = normalize(&vecs[r.id.0 as usize]); + let truth = l2_sq(&qn, &vn); + let rel = (r.distance - truth).abs() / truth.max(1e-6); + assert!( + rel < 1.5e-3, + "id={} returned={} exact={} rel={rel} (cosine convention: \ + normalized-pair squared L2)", + r.id.0, + r.distance, + truth + ); + } + } +} + +// ── 2b. TQ4: the coarser 4-bit codes make the red dramatic (~percent-level) ── + +#[test] +fn tq4_compacted_l2_distances_are_exact() { + let (mut store, vecs) = build_compacted( + "xr_tq4", + QuantizationConfig::TurboQuant4, + DistanceMetric::L2, + 64, + 0x71B4, + ); + let mut rng = Rng::new(0xCAFE); + for _ in 0..5 { + let q = rng.vec(DIM); + let results = segment_search(&mut store, "xr_tq4", &q, 10); + assert!(!results.is_empty()); + for r in &results { + let truth = l2_sq(&q, &vecs[r.id.0 as usize]); + let rel = (r.distance - truth).abs() / truth.max(1e-6); + assert!( + rel < 5e-3, + "TQ4 id={} returned={} exact={} rel={rel} — distance is a \ + quantized estimate, exact rerank missing", + r.id.0, + r.distance, + truth + ); + } + } +} + +// ── 3. Recall@10 vs exact ground truth ─────────────────────────────────────── + +#[test] +fn sq8_rerank_recall_at_10_near_perfect() { + // 200 vectors / 20 queries, all seeded. Pre-rerank SQ8 ADC misranks some + // near-ties (quantization error scale/2 per coord); with exact rerank the + // only residual error is f16 rounding (~1e-3 relative), far below typical + // neighbor gaps on this workload. + let (mut store, vecs) = build_compacted( + "xr_rec", + QuantizationConfig::Sq8, + DistanceMetric::L2, + 200, + 0x5EED, + ); + let mut rng = Rng::new(0xDEAD); + let mut hit = 0usize; + let mut total = 0usize; + for _ in 0..20 { + let q = rng.vec(DIM); + let mut truth: Vec<(f32, u32)> = vecs + .iter() + .enumerate() + .map(|(i, v)| (l2_sq(&q, v), i as u32)) + .collect(); + truth.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + let truth_ids: std::collections::HashSet = + truth[..10].iter().map(|(_, i)| *i).collect(); + let results = segment_search(&mut store, "xr_rec", &q, 10); + total += 10; + hit += results + .iter() + .filter(|r| truth_ids.contains(&r.id.0)) + .count(); + } + let recall = hit as f32 / total as f32; + assert!( + recall >= 0.99, + "recall@10 = {recall} (pre-rerank SQ8 ADC baseline measured ~0.9x on \ + this seed; exact rerank must reach ≥0.99)" + ); +} + +// ── 4. Persistence: sidecar survives segment write/read roundtrip ─────────── + +#[test] +fn rerank_sidecar_survives_persistence_roundtrip() { + use moon::vector::persistence::segment_io::{read_immutable_segment, write_immutable_segment}; + + let (mut store, vecs) = build_compacted( + "xr_per", + QuantizationConfig::Sq8, + DistanceMetric::L2, + 64, + 0x9E12, + ); + + let dir = std::env::temp_dir().join(format!("moon-xr-per-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + + // Write the compacted segment, then read it back fresh. + { + let idx = store.get_index_mut(b"xr_per").expect("index"); + let snap = idx.segments.load(); + write_immutable_segment(&dir, 7, &snap.immutable[0], &idx.collection) + .expect("write segment"); + } + let (seg, _meta) = read_immutable_segment(&dir, 7).expect("read segment"); + + let mut rng = Rng::new(0x0DDB); + let q = rng.vec(DIM); + let mut scratch = moon::vector::hnsw::search::SearchScratch::new(seg.total_count(), 64); + let results = seg.search(&q, 10, 128, &mut scratch); + assert!(!results.is_empty()); + for r in &results { + let truth = l2_sq(&q, &vecs[r.id.0 as usize]); + let rel = (r.distance - truth).abs() / truth.max(1e-6); + assert!( + rel < 1.5e-3, + "after disk roundtrip: id={} returned={} exact={} rel={rel} — \ + sidecar not persisted", + r.id.0, + r.distance, + truth + ); + } + + let _ = std::fs::remove_dir_all(&dir); +} + +// ── 5. Legacy segment dirs (no sidecar file) still load and search ────────── + +#[test] +fn legacy_segment_without_sidecar_still_searches() { + use moon::vector::persistence::segment_io::{read_immutable_segment, write_immutable_segment}; + + let (mut store, _vecs) = build_compacted( + "xr_leg", + QuantizationConfig::Sq8, + DistanceMetric::L2, + 32, + 0x1E64, + ); + + let dir = std::env::temp_dir().join(format!("moon-xr-leg-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + { + let idx = store.get_index_mut(b"xr_leg").expect("index"); + let snap = idx.segments.load(); + write_immutable_segment(&dir, 3, &snap.immutable[0], &idx.collection).expect("write"); + } + // Simulate a pre-sidecar segment directory. + let _ = std::fs::remove_file(dir.join("segment-3").join("raw_f16.bin")); + + let (seg, _meta) = read_immutable_segment(&dir, 3).expect("read legacy segment"); + let mut scratch = moon::vector::hnsw::search::SearchScratch::new(seg.total_count(), 64); + let mut rng = Rng::new(0x7E57); + let q = rng.vec(DIM); + let results = seg.search(&q, 5, 64, &mut scratch); + assert!( + !results.is_empty(), + "legacy segment without sidecar must still search (ADC distances)" + ); + + let _ = std::fs::remove_dir_all(&dir); +} + +// ── 6. GraphUnion merge preserves the sidecar ──────────────────────────────── + +#[test] +fn graph_union_merge_preserves_exact_distances() { + distance::init(); + let mut store = VectorStore::new(); + store + .create_index(make_meta( + "xr_mrg", + QuantizationConfig::Sq8, + DistanceMetric::L2, + )) + .expect("create_index"); + + let mut rng = Rng::new(0x4E46); + let mut vecs = Vec::new(); + for seg in 0..3 { + for i in 0..40 { + let v = rng.vec(DIM); + let key = Bytes::from(format!("doc:{}", seg * 40 + i)); + let key_hash = xxhash_rust::xxh64::xxh64(&key, 0); + store + .insert_vector(b"xr_mrg", &v, key_hash, key) + .expect("insert"); + vecs.push(v); + } + store.force_compact_index(b"xr_mrg").expect("compact"); + } + assert_eq!(store.immutable_segment_count(b"xr_mrg"), Some(3)); + + store + .force_merge_index_with_tolerance(b"xr_mrg", 0.5) + .expect("merge"); + assert_eq!(store.immutable_segment_count(b"xr_mrg"), Some(1)); + + // Post-merge ids are remapped; match results by key_hash instead. + let key_hash_to_pos: std::collections::HashMap = (0..vecs.len()) + .map(|i| { + let key = format!("doc:{i}"); + (xxhash_rust::xxh64::xxh64(key.as_bytes(), 0), i) + }) + .collect(); + + let mut qrng = Rng::new(0xAB1E); + let q = qrng.vec(DIM); + let results = segment_search(&mut store, "xr_mrg", &q, 10); + assert!(!results.is_empty()); + for r in &results { + let pos = key_hash_to_pos[&r.key_hash]; + let truth = l2_sq(&q, &vecs[pos]); + let rel = (r.distance - truth).abs() / truth.max(1e-6); + assert!( + rel < 1.5e-3, + "after merge: key pos={pos} returned={} exact={} rel={rel} — merged \ + segment lost the rerank sidecar", + r.distance, + truth + ); + } +} + +// ── 8. Tombstones must not consume the exact-rerank budget ────────────────── + +#[test] +fn tombstoned_candidates_do_not_eat_rerank_budget() { + // Regression: rerank_exact used to run BEFORE the liveness filter, so the + // top-4·k ADC candidates it re-scored could all be tombstones — the live + // results that survived the filter kept quantized ADC estimates (and the + // post-rerank sort mixed exact and ADC scores). Deleting the 30 nearest + // vectors with k=5 (budget 20) makes that failure deterministic. + let (mut store, vecs) = build_compacted( + "xr_tomb", + QuantizationConfig::Sq8, + DistanceMetric::L2, + 64, + 0x70B5, + ); + let mut rng = Rng::new(0xD00D); + let q = rng.vec(DIM); + + let mut order: Vec = (0..vecs.len()).collect(); + order.sort_by(|&a, &b| l2_sq(&q, &vecs[a]).total_cmp(&l2_sq(&q, &vecs[b]))); + let deleted: std::collections::HashSet = order[..30].iter().copied().collect(); + for &i in &deleted { + store.mark_deleted_for_key(format!("doc:{i}").as_bytes()); + } + + let results = segment_search(&mut store, "xr_tomb", &q, 5); + assert_eq!(results.len(), 5, "expected k live results"); + for r in &results { + let id = r.id.0 as usize; + assert!(!deleted.contains(&id), "tombstoned id {id} returned"); + let truth = l2_sq(&q, &vecs[id]); + let rel = (r.distance - truth).abs() / truth.max(1e-6); + assert!( + rel < 1.5e-3, + "id={id} returned={} exact={truth} rel={rel} — live candidate kept \ + its ADC estimate: tombstones consumed the rerank budget", + r.distance + ); + } + // The 5 results must be the true 5 nearest live vectors, in order. + let expect: Vec = order[30..35].to_vec(); + let got: Vec = results.iter().map(|r| r.id.0 as usize).collect(); + assert_eq!(got, expect, "live top-k mismatch after tombstone filtering"); +} diff --git a/tests/vector_insert_bench.rs b/tests/vector_insert_bench.rs index da488c33c..f5d5cf690 100644 --- a/tests/vector_insert_bench.rs +++ b/tests/vector_insert_bench.rs @@ -54,8 +54,7 @@ fn bench_raw_append_128d() { let start = Instant::now(); for i in 0..n { - let norm: f32 = vectors[i].iter().map(|x| x * x).sum::().sqrt(); - seg.append(i as u64, &vectors[i], &sq_vecs[i], norm, 0); + seg.append(i as u64, &vectors[i], 0); } let elapsed = start.elapsed(); @@ -109,8 +108,7 @@ fn bench_raw_append_768d() { let start = Instant::now(); for i in 0..n { - let norm: f32 = vectors[i].iter().map(|x| x * x).sum::().sqrt(); - seg.append(i as u64, &vectors[i], &sq_vecs[i], norm, 0); + seg.append(i as u64, &vectors[i], 0); } let elapsed = start.elapsed(); @@ -183,7 +181,6 @@ fn bench_full_insert_pipeline_128d() { let mut sq_vec = vec![0i8; dim as usize]; vector_search::quantize_f32_to_sq(&f32_vec, &mut sq_vec); // Norm - let norm: f32 = f32_vec.iter().map(|x| x * x).sum::().sqrt(); // Key hash let key = format!("doc:{i}"); let key_hash = xxhash_rust::xxh64::xxh64(key.as_bytes(), 0); @@ -192,7 +189,7 @@ fn bench_full_insert_pipeline_128d() { .get_index_mut(&bytes::Bytes::from_static(b"idx")) .unwrap(); let snap = idx.segments.load(); - snap.mutable.append(key_hash, &f32_vec, &sq_vec, norm, 0); + snap.mutable.append(key_hash, &f32_vec, 0); } let elapsed = start.elapsed(); @@ -258,14 +255,13 @@ fn bench_full_insert_pipeline_768d() { } let mut sq_vec = vec![0i8; dim as usize]; vector_search::quantize_f32_to_sq(&f32_vec, &mut sq_vec); - let norm: f32 = f32_vec.iter().map(|x| x * x).sum::().sqrt(); let key = format!("doc:{i}"); let key_hash = xxhash_rust::xxh64::xxh64(key.as_bytes(), 0); let idx = store .get_index_mut(&bytes::Bytes::from_static(b"idx")) .unwrap(); let snap = idx.segments.load(); - snap.mutable.append(key_hash, &f32_vec, &sq_vec, norm, 0); + snap.mutable.append(key_hash, &f32_vec, 0); } let elapsed = start.elapsed(); diff --git a/tests/vector_memory_audit.rs b/tests/vector_memory_audit.rs index f288e3849..d657dfdf6 100644 --- a/tests/vector_memory_audit.rs +++ b/tests/vector_memory_audit.rs @@ -173,7 +173,7 @@ fn test_per_vector_overhead_breakdown() { f32_v.push((s as f32) / (u32::MAX as f32) * 2.0 - 1.0); sq_v.push((s >> 24) as i8); } - seg.append(i as u64, &f32_v, &sq_v, 1.0, i as u64); + seg.append(i as u64, &f32_v, i as u64); } assert_eq!(seg.len(), n); diff --git a/tests/vector_segment_merge.rs b/tests/vector_segment_merge.rs index 3dcb9ccca..a38091fbc 100644 --- a/tests/vector_segment_merge.rs +++ b/tests/vector_segment_merge.rs @@ -555,15 +555,10 @@ fn test_merge_overlapping_ids_highest_lsn_wins() { let mut rng = Rng::new(42); let v = random_unit_vec(&mut rng, DIM as usize); let key_hash: u64 = 0xDEAD; - let sq_vec: Vec = v - .iter() - .map(|&x| (x * 127.0).clamp(-128.0, 127.0) as i8) - .collect(); - let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); // Insert into both with different insert_lsn (simulated via key_hash collision). - seg_a.append(key_hash, &v, &sq_vec, norm, 1); - seg_b.append(key_hash, &v, &sq_vec, norm, 2); // higher LSN in seg_b + seg_a.append(key_hash, &v, 1); + seg_b.append(key_hash, &v, 2); // higher LSN in seg_b let frozen_a = seg_a.freeze(); let frozen_b = seg_b.freeze(); @@ -800,3 +795,50 @@ fn test_ft_create_keep_raw_default_is_false() { "Default KEEP_RAW must be false when not specified" ); } + +// ── VEC-7: MERGE_MODE KEEP_RAW is an unimplemented stub — FT.CREATE must +// reject it fail-loud instead of silently falling back to graph-union. ── +#[test] +fn test_ft_create_merge_mode_keep_raw_rejected() { + use moon::command::vector_search::ft_create; + use moon::protocol::Frame; + + fn bulk(b: &[u8]) -> Frame { + Frame::BulkString(Bytes::copy_from_slice(b)) + } + + let mut store = VectorStore::new(); + let mut text_store = moon::text::store::TextStore::new(); + + let args: Vec = vec![ + bulk(b"idx_mm_kr"), + bulk(b"ON"), + bulk(b"HASH"), + bulk(b"PREFIX"), + bulk(b"1"), + bulk(b"doc:"), + bulk(b"SCHEMA"), + bulk(b"vec"), + bulk(b"VECTOR"), + bulk(b"HNSW"), + bulk(b"8"), + bulk(b"TYPE"), + bulk(b"FLOAT32"), + bulk(b"DIM"), + bulk(b"64"), + bulk(b"DISTANCE_METRIC"), + bulk(b"L2"), + bulk(b"MERGE_MODE"), + bulk(b"KEEP_RAW"), + ]; + + let result = ft_create(&mut store, &mut text_store, &args); + assert!( + matches!(result, Frame::Error(_)), + "MERGE_MODE KEEP_RAW must be rejected until implemented: {result:?}" + ); + assert!( + store.get_index(b"idx_mm_kr").is_none(), + "rejected FT.CREATE must not leave a partial index behind" + ); +} diff --git a/tests/vector_stress.rs b/tests/vector_stress.rs index 7677cdca1..9d90c3b5a 100644 --- a/tests/vector_stress.rs +++ b/tests/vector_stress.rs @@ -112,8 +112,7 @@ fn test_stress_10k_interleaved_operations() { if op < 40 { // INSERT (40%) fill_vectors(&mut rng, &mut f32_buf, &mut sq_buf, DIM); - let norm = f32_buf.iter().map(|x| x * x).sum::().sqrt(); - let id = mutable.append(i as u64, &f32_buf, &sq_buf, norm, i as u64); + let id = mutable.append(i as u64, &f32_buf, i as u64); inserted_ids.push(id); } else if op < 70 { // SEARCH (30%) @@ -198,8 +197,7 @@ fn test_stress_interleaved_search_during_compaction() { let insert_count = 5000; for i in 0..insert_count { fill_vectors(&mut rng, &mut f32_buf, &mut sq_buf, dim); - let norm = f32_buf.iter().map(|x| x * x).sum::().sqrt(); - seg.append(i as u64, &f32_buf, &sq_buf, norm, i as u64); + seg.append(i as u64, &f32_buf, i as u64); } assert_eq!(seg.len(), insert_count); diff --git a/tests/vector_update_tombstone.rs b/tests/vector_update_tombstone.rs new file mode 100644 index 000000000..76983b878 --- /dev/null +++ b/tests/vector_update_tombstone.rs @@ -0,0 +1,130 @@ +//! VEC-1 (vector deep review 2026-07-05): updating an existing key via HSET +//! must tombstone the previously indexed vector — both while it still lives in +//! the mutable segment AND after it has been compacted into an immutable +//! segment. Without the tombstone the index accumulates stale duplicates: +//! FT.SEARCH returns the same doc twice (once with the old embedding) and +//! num_docs inflates monotonically under update churn. + +use bytes::Bytes; + +use moon::command::vector_search::{ft_create, ft_search}; +use moon::protocol::Frame; +use moon::shard::spsc_handler::auto_index_hset_public; +use moon::text::store::TextStore; +use moon::vector::distance; +use moon::vector::store::VectorStore; + +fn bulk(s: &[u8]) -> Frame { + Frame::BulkString(Bytes::from(s.to_vec())) +} + +fn f32_blob(v: &[f32]) -> Frame { + let mut b = Vec::with_capacity(v.len() * 4); + for x in v { + b.extend_from_slice(&x.to_le_bytes()); + } + Frame::BulkString(Bytes::from(b)) +} + +fn ft_create_args(name: &str, dim: u32) -> Vec { + vec![ + bulk(name.as_bytes()), + bulk(b"ON"), + bulk(b"HASH"), + bulk(b"PREFIX"), + bulk(b"1"), + bulk(b"doc:"), + bulk(b"SCHEMA"), + bulk(b"vec"), + bulk(b"VECTOR"), + bulk(b"HNSW"), + bulk(b"6"), + bulk(b"TYPE"), + bulk(b"FLOAT32"), + bulk(b"DIM"), + bulk(dim.to_string().as_bytes()), + bulk(b"DISTANCE_METRIC"), + bulk(b"L2"), + ] +} + +fn hset(vs: &mut VectorStore, ts: &mut TextStore, key: &[u8], vec: &[f32]) { + let args = vec![bulk(key), bulk(b"vec"), f32_blob(vec)]; + auto_index_hset_public(vs, ts, key, &args); +} + +/// KNN k=10 search; returns the total-result count (first array element). +fn search_total(vs: &mut VectorStore, index: &str, query: &[f32]) -> i64 { + let mut qb = Vec::with_capacity(query.len() * 4); + for x in query { + qb.extend_from_slice(&x.to_le_bytes()); + } + let args = vec![ + bulk(index.as_bytes()), + bulk(b"*=>[KNN 10 @vec $query]"), + bulk(b"PARAMS"), + bulk(b"2"), + bulk(b"query"), + Frame::BulkString(Bytes::from(qb)), + ]; + match ft_search(vs, &args, None, None, 0) { + Frame::Array(items) => match items.first() { + Some(Frame::Integer(n)) => *n, + other => panic!("expected Integer total, got {other:?}"), + }, + other => panic!("expected Array response, got {other:?}"), + } +} + +const DIM: usize = 8; + +#[test] +fn hset_update_in_mutable_segment_does_not_duplicate() { + distance::init(); + let mut vs = VectorStore::new(); + let mut ts = TextStore::new(); + let out = ft_create(&mut vs, &mut ts, &ft_create_args("upd_mut", DIM as u32)); + assert!(!matches!(out, Frame::Error(_)), "ft_create failed: {out:?}"); + + let v1: Vec = (0..DIM).map(|i| if i == 0 { 1.0 } else { 0.0 }).collect(); + let v2: Vec = (0..DIM).map(|i| if i == 1 { 1.0 } else { 0.0 }).collect(); + + hset(&mut vs, &mut ts, b"doc:1", &v1); + hset(&mut vs, &mut ts, b"doc:1", &v2); // update, still in mutable + + let total = search_total(&mut vs, "upd_mut", &v2); + assert_eq!( + total, 1, + "updated key must appear exactly once (stale mutable duplicate returned)" + ); +} + +#[test] +fn hset_update_after_compaction_tombstones_immutable_copy() { + distance::init(); + let mut vs = VectorStore::new(); + let mut ts = TextStore::new(); + let out = ft_create(&mut vs, &mut ts, &ft_create_args("upd_imm", DIM as u32)); + assert!(!matches!(out, Frame::Error(_)), "ft_create failed: {out:?}"); + + let v1: Vec = (0..DIM).map(|i| if i == 0 { 1.0 } else { 0.0 }).collect(); + let v2: Vec = (0..DIM).map(|i| if i == 1 { 1.0 } else { 0.0 }).collect(); + + hset(&mut vs, &mut ts, b"doc:1", &v1); + // Push a few fillers so the compacted segment is non-trivial. + for i in 2..6 { + let vf: Vec = (0..DIM).map(|j| (i * j) as f32 * 0.1).collect(); + hset(&mut vs, &mut ts, format!("doc:{i}").as_bytes(), &vf); + } + vs.force_compact_index(b"upd_imm") + .expect("force_compact_index"); + + // Update doc:1 AFTER its old copy was compacted into the immutable segment. + hset(&mut vs, &mut ts, b"doc:1", &v2); + + let total = search_total(&mut vs, "upd_imm", &v2); + assert_eq!( + total, 5, + "5 live docs expected; a stale immutable copy of doc:1 was resurrected" + ); +} diff --git a/tests/workspace_integration.rs b/tests/workspace_integration.rs index 39c4e38bc..adabe2ec9 100644 --- a/tests/workspace_integration.rs +++ b/tests/workspace_integration.rs @@ -78,6 +78,7 @@ async fn start_workspace_server(num_shards: usize) -> (u16, CancellationToken) { uring_sqpoll_ms: None, io_driver: "auto".to_string(), io_busy_poll_us: 0, + ft_search_workers: None, admin_port: 0, slowlog_log_slower_than: 10000, slowlog_max_len: 128, @@ -306,6 +307,7 @@ async fn start_workspace_server_with_auth( uring_sqpoll_ms: None, io_driver: "auto".to_string(), io_busy_poll_us: 0, + ft_search_workers: None, admin_port: 0, slowlog_log_slower_than: 10000, slowlog_max_len: 128,