Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
183 changes: 107 additions & 76 deletions laurus/src/vector/index/flat/reader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,9 @@ use crate::storage::Storage;
use crate::vector::core::distance::DistanceMetric;
use crate::vector::core::quantization::QuantizedVectorMeta;
use crate::vector::core::vector::Vector;
use crate::vector::index::format::{QuantHeader, VectorSegmentHeader, record_prefix_size};
use crate::vector::index::format::{
FieldInterner, QuantHeader, VectorSegmentHeader, record_prefix_size,
};
use crate::vector::index::quantized_io::quantized_record_payload_size;
use crate::vector::index::quantized_storage::QuantizedVectorPool;
use crate::vector::reader::{ValidationReport, VectorIndexMetadata, VectorStats};
Expand All @@ -22,7 +24,12 @@ use crate::vector::index::storage::VectorStorage;
#[derive(Debug)]
pub struct FlatVectorIndexReader {
vectors: VectorStorage,
vector_ids: Vec<(u64, String)>,
/// `(doc_id, field_id)` per record; ids index [`Self::field_dict`]
/// (Issue #633 PR-B — interned, no per-record heap `String`).
vector_ids: Vec<(u64, u16)>,
/// Per-segment field-name dictionary (synthesized at load for
/// v1/v2 segments, taken from the header for v3).
field_dict: Arc<[Arc<str>]>,
dimension: usize,
distance_metric: DistanceMetric,
deletion_bitmap: Option<Arc<DeletionBitmap>>,
Expand All @@ -34,19 +41,17 @@ pub struct FlatVectorIndexReader {

/// Group `vector_ids` by field name into refcount-shared slices.
fn build_vector_ids_by_field(
vector_ids: &[(u64, String)],
vector_ids: &[(u64, u16)],
field_dict: &[Arc<str>],
) -> std::collections::HashMap<String, Arc<[u64]>> {
let mut by_field: std::collections::HashMap<String, Vec<u64>> =
std::collections::HashMap::new();
for (doc_id, field_name) in vector_ids {
by_field
.entry(field_name.clone())
.or_default()
.push(*doc_id);
let mut by_field: Vec<Vec<u64>> = vec![Vec::new(); field_dict.len()];
for &(doc_id, fid) in vector_ids {
by_field[fid as usize].push(doc_id);
}
by_field
.into_iter()
.map(|(field, ids)| (field, Arc::<[u64]>::from(ids)))
field_dict
.iter()
.zip(by_field)
.map(|(field, ids)| (field.to_string(), Arc::<[u64]>::from(ids)))
.collect()
}

Expand Down Expand Up @@ -134,7 +139,11 @@ impl FlatVectorIndexReader {
let record_stride =
record_prefix_size(header.version) + quantized_record_payload_size(dimension) as u64;

let (vectors, vector_ids) = match storage.loading_mode() {
// Interned field ids (Issue #633 PR-B): one shared dictionary per
// segment instead of one heap `String` per record.
let mut interner = FieldInterner::from_header(&header);

let (vectors, vector_ids, field_dict) = match storage.loading_mode() {
crate::storage::LoadingMode::Eager => {
// Step 7 of #481 Stage 1: load vectors as int8 + meta
// directly into a QuantizedVectorPool.
Expand All @@ -153,8 +162,10 @@ impl FlatVectorIndexReader {
input.read_exact(&mut doc_id_buf)?;
let doc_id = u64::from_le_bytes(doc_id_buf);

// Field reference: dictionary id (v3+) or inline name.
let field_name = header.read_record_field(
// Field reference: interned id (v3 = dictionary id,
// v1/v2 = inline name interned on first appearance).
let fid = interner.read_record_field_id(
&header,
&mut input,
records_remaining,
"flat field_name_len",
Expand All @@ -172,11 +183,18 @@ impl FlatVectorIndexReader {
norm_q: f32::from_le_bytes(norm_q_buf),
};

vector_ids.push((doc_id, field_name.clone()));
records.push((doc_id, field_name, int8, meta));
vector_ids.push((doc_id, fid));
// The pool's build input keeps the String shape; this
// clone is transient (the pool retains only per-field
// grouping keys), so nothing per-record is kept.
records.push((doc_id, interner.name(fid).to_string(), int8, meta));
}
let pool = QuantizedVectorPool::build(params, dimension, records);
(VectorStorage::OwnedQuantized(Arc::new(pool)), vector_ids)
(
VectorStorage::OwnedQuantized(Arc::new(pool)),
vector_ids,
interner.into_dict(),
)
}
crate::storage::LoadingMode::Lazy => {
checked_capacity(
Expand Down Expand Up @@ -204,7 +222,8 @@ impl FlatVectorIndexReader {
input.read_exact(&mut doc_id_buf)?;
let doc_id = u64::from_le_bytes(doc_id_buf);

let field_name = header.read_record_field(
let fid = interner.read_record_field_id(
&header,
&mut input,
records_remaining,
"flat field_name_len",
Expand All @@ -214,32 +233,36 @@ impl FlatVectorIndexReader {
// record prefix), so `VectorStorage::get` seeks straight
// to the int8 data without re-parsing the prefix.
let payload_offset = input.stream_position().map_err(LaurusError::Io)?;
offsets.insert((doc_id, field_name.clone()), payload_offset);
vector_ids.push((doc_id, field_name));
offsets.insert((doc_id, fid), payload_offset);
vector_ids.push((doc_id, fid));

// Skip int8 payload + per-vector meta.
input
.seek(std::io::SeekFrom::Current(quant_payload_size))
.map_err(LaurusError::Io)?;
}

let field_dict = interner.into_dict();
(
VectorStorage::OnDemand {
storage: storage.clone(),
file_name: file_name.clone(),
offsets: Arc::new(offsets),
field_dict: field_dict.clone(),
quant_params: Some(params),
cached_input: Arc::new(std::sync::RwLock::new(None)),
},
vector_ids,
field_dict,
)
}
};

let vector_ids_by_field = build_vector_ids_by_field(&vector_ids);
let vector_ids_by_field = build_vector_ids_by_field(&vector_ids, &field_dict);
Ok(Self {
vectors,
vector_ids,
field_dict,
dimension,
distance_metric,
deletion_bitmap: None,
Expand Down Expand Up @@ -276,18 +299,18 @@ impl VectorIndexReader for FlatVectorIndexReader {
if self.is_deleted(doc_id) {
return Ok(None);
}
self.vectors
.get(&(doc_id, field_name.to_string()), self.dimension)
self.vectors.get(doc_id, field_name, self.dimension)
}

fn get_vectors_for_doc(&self, doc_id: u64) -> Result<Vec<(String, Vector)>> {
let mut result = Vec::new();
for (id, field) in &self.vector_ids {
if *id == doc_id
&& !self.is_deleted(*id)
&& let Some(vec) = self.vectors.get(&(*id, field.clone()), self.dimension)?
for &(id, fid) in &self.vector_ids {
let field = &self.field_dict[fid as usize];
if id == doc_id
&& !self.is_deleted(id)
&& let Some(vec) = self.vectors.get(id, field, self.dimension)?
{
result.push((field.clone(), vec));
result.push((field.to_string(), vec));
}
}
Ok(result)
Expand All @@ -299,14 +322,20 @@ impl VectorIndexReader for FlatVectorIndexReader {
if self.is_deleted(*id) {
result.push(None);
} else {
result.push(self.vectors.get(&(*id, field.clone()), self.dimension)?);
result.push(self.vectors.get(*id, field, self.dimension)?);
}
}
Ok(result)
}

fn vector_ids(&self) -> Result<Vec<(u64, String)>> {
Ok(self.vector_ids.clone())
// Rehydrated at the trait boundary (Issue #633 PR-B): internal
// state is `(u64, u16)` + the shared dictionary.
Ok(self
.vector_ids
.iter()
.map(|&(id, fid)| (id, self.field_dict[fid as usize].to_string()))
.collect())
}

fn doc_ids_for_field(&self, field_name: &str) -> Arc<[u64]> {
Expand Down Expand Up @@ -354,20 +383,7 @@ impl VectorIndexReader for FlatVectorIndexReader {
}

fn contains_vector(&self, doc_id: u64, field_name: &str) -> bool {
match &self.vectors {
VectorStorage::Owned(vectors) => {
vectors.contains_key(&(doc_id, field_name.to_string()))
}
VectorStorage::OwnedQuantized(pool) => pool.contains(doc_id, field_name),
VectorStorage::OwnedPq(pool) => pool.contains(doc_id, field_name),
#[cfg(feature = "pq-fastscan")]
VectorStorage::OwnedPqFastScan(_) => {
unreachable!("Flat reader rejects PQ FastScan at the segment header (HNSW-only)")
}
VectorStorage::OnDemand { offsets, .. } => {
offsets.contains_key(&(doc_id, field_name.to_string()))
}
}
self.vectors.contains(doc_id, field_name)
}

fn get_vector_range(
Expand All @@ -376,41 +392,47 @@ impl VectorIndexReader for FlatVectorIndexReader {
end_doc_id: u64,
) -> Result<Vec<(u64, String, Vector)>> {
let mut result = Vec::new();
for (id, field) in &self.vector_ids {
if *id >= start_doc_id
&& *id < end_doc_id
&& !self.is_deleted(*id)
&& let Some(vec) = self.vectors.get(&(*id, field.clone()), self.dimension)?
for &(id, fid) in &self.vector_ids {
let field = &self.field_dict[fid as usize];
if id >= start_doc_id
&& id < end_doc_id
&& !self.is_deleted(id)
&& let Some(vec) = self.vectors.get(id, field, self.dimension)?
{
result.push((*id, field.clone(), vec));
result.push((id, field.to_string(), vec));
}
}
Ok(result)
}

fn get_vectors_by_field(&self, field_name: &str) -> Result<Vec<(u64, Vector)>> {
// One dictionary resolve, then integer compares per record.
let Some(target) =
crate::vector::index::format::resolve_field_id(&self.field_dict, field_name)
else {
return Ok(Vec::new());
};
let mut result = Vec::new();
for (id, field) in &self.vector_ids {
if field == field_name
&& !self.is_deleted(*id)
&& let Some(vec) = self.vectors.get(&(*id, field.clone()), self.dimension)?
for &(id, fid) in &self.vector_ids {
if fid == target
&& !self.is_deleted(id)
&& let Some(vec) = self.vectors.get(id, field_name, self.dimension)?
{
result.push((*id, vec));
result.push((id, vec));
}
}
Ok(result)
}

fn field_names(&self) -> Result<Vec<String>> {
use std::collections::HashSet;
let fields: HashSet<String> = self.vector_ids.iter().map(|val| val.1.clone()).collect();
Ok(fields.into_iter().collect())
Ok(self.field_dict.iter().map(|f| f.to_string()).collect())
}

fn vector_iterator(&self) -> Result<Box<dyn VectorIterator>> {
Ok(Box::new(FlatVectorIterator {
storage: self.vectors.clone(),
keys: self.vector_ids.clone(),
field_dict: self.field_dict.clone(),
current: 0,
dimension: self.dimension,
deletion_bitmap: self.deletion_bitmap.clone(),
Expand Down Expand Up @@ -442,7 +464,9 @@ impl VectorIndexReader for FlatVectorIndexReader {

match &self.vectors {
VectorStorage::Owned(map) => {
for (id, field) in &self.vector_ids {
for &(id, fid) in &self.vector_ids {
let field = &self.field_dict[fid as usize];
let (id, field) = (&id, &field.to_string());
if let Some(vector) = map.get(&(*id, field.clone())) {
if vector.dimension() != self.dimension {
errors.push(format!(
Expand All @@ -468,7 +492,9 @@ impl VectorIndexReader for FlatVectorIndexReader {
}
}
VectorStorage::OwnedQuantized(pool) => {
for (id, field) in &self.vector_ids {
for &(id, fid) in &self.vector_ids {
let field = &self.field_dict[fid as usize];
let id = &id;
if !pool.contains(*id, field) {
errors.push(format!(
"Vector {}:{} found in keys but missing in quantized pool",
Expand All @@ -483,7 +509,9 @@ impl VectorIndexReader for FlatVectorIndexReader {
);
}
VectorStorage::OwnedPq(pool) => {
for (id, field) in &self.vector_ids {
for &(id, fid) in &self.vector_ids {
let field = &self.field_dict[fid as usize];
let id = &id;
if !pool.contains(*id, field) {
errors.push(format!(
"Vector {}:{} found in keys but missing in PQ pool",
Expand All @@ -502,8 +530,10 @@ impl VectorIndexReader for FlatVectorIndexReader {
unreachable!("Flat reader rejects PQ FastScan at the segment header (HNSW-only)")
}
VectorStorage::OnDemand { offsets, .. } => {
for (id, field) in &self.vector_ids {
if !offsets.contains_key(&(*id, field.clone())) {
for &(id, fid) in &self.vector_ids {
let field = &self.field_dict[fid as usize];
let id = &id;
if !offsets.contains_key(&(*id, fid)) {
errors.push(format!(
"Vector {}:{} in ids but missing in storage",
id, field
Expand All @@ -526,7 +556,8 @@ impl VectorIndexReader for FlatVectorIndexReader {
/// Iterator for flat vector index.
struct FlatVectorIterator {
storage: VectorStorage,
keys: Vec<(u64, String)>,
keys: Vec<(u64, u16)>,
field_dict: Arc<[Arc<str>]>,
current: usize,
dimension: usize,
deletion_bitmap: Option<Arc<DeletionBitmap>>,
Expand All @@ -537,22 +568,20 @@ impl VectorIterator for FlatVectorIterator {
// Use a loop instead of recursion to avoid stack overflow when
// many consecutive entries are deleted.
while self.current < self.keys.len() {
let (doc_id, field) = &self.keys[self.current];
let (doc_id, fid) = self.keys[self.current];
let field = &self.field_dict[fid as usize];

// Skip deleted entries
if let Some(bitmap) = &self.deletion_bitmap
&& bitmap.is_deleted(*doc_id)
&& bitmap.is_deleted(doc_id)
{
self.current += 1;
continue;
}

if let Some(vec) = self
.storage
.get(&(*doc_id, field.clone()), self.dimension)?
{
if let Some(vec) = self.storage.get(doc_id, field, self.dimension)? {
self.current += 1;
return Ok(Some((*doc_id, field.clone(), vec)));
return Ok(Some((doc_id, field.to_string(), vec)));
} else {
return Err(LaurusError::internal(format!(
"Vector {}:{} found in keys but missing in storage",
Expand All @@ -566,8 +595,9 @@ impl VectorIterator for FlatVectorIterator {

fn skip_to(&mut self, doc_id: u64, field_name: &str) -> Result<bool> {
while self.current < self.keys.len() {
let (id, field) = &self.keys[self.current];
if *id > doc_id || (*id == doc_id && field.as_str() >= field_name) {
let (id, fid) = self.keys[self.current];
let field = &self.field_dict[fid as usize];
if id > doc_id || (id == doc_id && field.as_ref() as &str >= field_name) {
return Ok(true);
}
self.current += 1;
Expand All @@ -577,7 +607,8 @@ impl VectorIterator for FlatVectorIterator {

fn position(&self) -> (u64, String) {
if self.current < self.keys.len() {
self.keys[self.current].clone()
let (id, fid) = self.keys[self.current];
(id, self.field_dict[fid as usize].to_string())
} else {
(u64::MAX, String::new())
}
Expand Down
Loading