use crate::dsl::Field;
use crate::segment::SegmentReader;
use crate::{DocId, Score, TERMINATED};
use super::combiner::MultiValueCombiner;
use crate::query::ScoredPosition;
use crate::query::traits::{CountFuture, MatchedPositions, Query, Scorer, ScorerFuture};
#[derive(Debug, Clone)]
pub struct SparseVectorQuery {
pub field: Field,
pub vector: Vec<(u32, f32)>,
pub combiner: MultiValueCombiner,
pub heap_factor: f32,
pub weight_threshold: f32,
pub max_query_dims: Option<usize>,
pub pruning: Option<f32>,
pub min_query_dims: usize,
pub over_fetch_factor: f32,
pub max_superblocks: usize,
pruned: Option<Vec<(u32, f32)>>,
}
impl std::fmt::Display for SparseVectorQuery {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let dims = self.pruned_dims();
write!(f, "Sparse({}, dims={}", self.field.0, dims.len())?;
if self.heap_factor < 1.0 {
write!(f, ", heap={}", self.heap_factor)?;
}
if self.vector.len() != dims.len() {
write!(f, ", orig={}", self.vector.len())?;
}
write!(f, ")")
}
}
impl SparseVectorQuery {
pub fn new(field: Field, vector: Vec<(u32, f32)>) -> Self {
let mut q = Self {
field,
vector,
combiner: MultiValueCombiner::LogSumExp { temperature: 0.7 },
heap_factor: 1.0,
weight_threshold: 0.0,
max_query_dims: Some(crate::query::MAX_QUERY_TERMS),
pruning: None,
min_query_dims: 4,
over_fetch_factor: 2.0,
max_superblocks: 0,
pruned: None,
};
q.pruned = Some(q.compute_pruned_vector());
q
}
pub(crate) fn pruned_dims(&self) -> &[(u32, f32)] {
self.pruned.as_deref().unwrap_or(&self.vector)
}
fn validate(&self, reader: &SegmentReader) -> crate::Result<()> {
let entry = reader
.schema()
.get_field_entry(self.field)
.ok_or_else(|| crate::Error::FieldNotFound(self.field.0.to_string()))?;
if entry.field_type != crate::dsl::FieldType::SparseVector {
return Err(crate::Error::InvalidFieldType {
expected: "sparse_vector".to_string(),
got: format!("{:?}", entry.field_type),
});
}
if self.vector.iter().any(|(_, weight)| !weight.is_finite()) {
return Err(crate::Error::Query(
"sparse query contains a non-finite weight".to_string(),
));
}
if self.pruned_dims().len() > crate::query::MAX_QUERY_TERMS {
return Err(crate::Error::Query(format!(
"sparse query contains more than {} effective dimensions",
crate::query::MAX_QUERY_TERMS
)));
}
if !self.heap_factor.is_finite() || !(0.0..=1.0).contains(&self.heap_factor) {
return Err(crate::Error::Query(format!(
"sparse heap_factor must be finite and in [0, 1], got {}",
self.heap_factor
)));
}
if !self.over_fetch_factor.is_finite() || self.over_fetch_factor < 1.0 {
return Err(crate::Error::Query(format!(
"sparse over_fetch_factor must be finite and at least 1, got {}",
self.over_fetch_factor
)));
}
self.combiner.validate().map_err(crate::Error::Query)
}
pub fn with_combiner(mut self, combiner: MultiValueCombiner) -> Self {
self.combiner = combiner;
self
}
pub fn with_over_fetch_factor(mut self, factor: f32) -> Self {
self.over_fetch_factor = factor.max(1.0);
self
}
pub fn with_heap_factor(mut self, heap_factor: f32) -> Self {
self.heap_factor = heap_factor.clamp(0.0, 1.0);
self
}
pub fn with_weight_threshold(mut self, threshold: f32) -> Self {
self.weight_threshold = threshold;
self.pruned = Some(self.compute_pruned_vector());
self
}
pub fn with_max_query_dims(mut self, max_dims: usize) -> Self {
self.max_query_dims = Some(max_dims.min(crate::query::MAX_QUERY_TERMS));
self.pruned = Some(self.compute_pruned_vector());
self
}
pub fn with_pruning(mut self, fraction: f32) -> Self {
self.pruning = Some(fraction.clamp(0.0, 1.0));
self.pruned = Some(self.compute_pruned_vector());
self
}
pub fn with_min_query_dims(mut self, min_dims: usize) -> Self {
self.min_query_dims = min_dims;
self.pruned = Some(self.compute_pruned_vector());
self
}
fn compute_pruned_vector(&self) -> Vec<(u32, f32)> {
let original_len = self.vector.len();
let mut v: Vec<(u32, f32)> =
if self.weight_threshold > 0.0 && self.vector.len() > self.min_query_dims {
self.vector
.iter()
.copied()
.filter(|(_, w)| w.abs() >= self.weight_threshold)
.collect()
} else {
self.vector.clone()
};
let after_threshold = v.len();
let mut sorted_by_weight = false;
if let Some(fraction) = self.pruning
&& fraction < 1.0
&& v.len() > self.min_query_dims
{
v.sort_unstable_by(|a, b| {
b.1.abs()
.partial_cmp(&a.1.abs())
.unwrap_or(std::cmp::Ordering::Equal)
});
sorted_by_weight = true;
let keep = ((v.len() as f64 * fraction as f64).ceil() as usize).max(1);
v.truncate(keep);
}
let after_pruning = v.len();
let max_dims = self
.max_query_dims
.unwrap_or(crate::query::MAX_QUERY_TERMS)
.min(crate::query::MAX_QUERY_TERMS);
if v.len() > max_dims {
if !sorted_by_weight {
v.sort_unstable_by(|a, b| {
b.1.abs()
.partial_cmp(&a.1.abs())
.unwrap_or(std::cmp::Ordering::Equal)
});
}
v.truncate(max_dims);
}
if v.len() < original_len && log::log_enabled!(log::Level::Debug) {
let src: Vec<_> = self
.vector
.iter()
.map(|(d, w)| format!("({},{:.4})", d, w))
.collect();
let pruned_fmt: Vec<_> = v.iter().map(|(d, w)| format!("({},{:.4})", d, w)).collect();
log::debug!(
"[sparse query] field={}: pruned {}->{} dims \
(threshold: {}->{}, pruning: {}->{}, max_dims: {}->{}), \
source=[{}], pruned=[{}]",
self.field.0,
original_len,
v.len(),
original_len,
after_threshold,
after_threshold,
after_pruning,
after_pruning,
v.len(),
src.join(", "),
pruned_fmt.join(", "),
);
}
v
}
pub fn from_indices_weights(field: Field, indices: Vec<u32>, weights: Vec<f32>) -> Self {
let vector: Vec<(u32, f32)> = indices.into_iter().zip(weights).collect();
Self::new(field, vector)
}
#[cfg(feature = "native")]
pub fn from_text(
field: Field,
text: &str,
tokenizer_name: &str,
weighting: crate::structures::QueryWeighting,
sparse_index: Option<&crate::segment::SparseIndex>,
) -> crate::Result<Self> {
use crate::structures::QueryWeighting;
use crate::tokenizer::tokenizer_cache;
let tokenizer = tokenizer_cache().get_or_load(tokenizer_name)?;
let token_ids = tokenizer.tokenize_unique(text)?;
let weights: Vec<f32> = match weighting {
QueryWeighting::One => vec![1.0f32; token_ids.len()],
QueryWeighting::Idf => {
if let Some(index) = sparse_index {
index.idf_weights(&token_ids)
} else {
vec![1.0f32; token_ids.len()]
}
}
QueryWeighting::IdfFile => {
use crate::tokenizer::idf_weights_cache;
if let Some(idf) = idf_weights_cache().get_or_load(tokenizer_name, None) {
token_ids.iter().map(|&id| idf.get(id)).collect()
} else {
vec![1.0f32; token_ids.len()]
}
}
};
let vector: Vec<(u32, f32)> = token_ids.into_iter().zip(weights).collect();
Ok(Self::new(field, vector))
}
#[cfg(feature = "native")]
pub fn from_text_with_stats(
field: Field,
text: &str,
tokenizer: &crate::tokenizer::HfTokenizer,
weighting: crate::structures::QueryWeighting,
global_stats: Option<&crate::query::GlobalStats>,
) -> crate::Result<Self> {
use crate::structures::QueryWeighting;
let token_ids = tokenizer.tokenize_unique(text)?;
let weights: Vec<f32> = match weighting {
QueryWeighting::One => vec![1.0f32; token_ids.len()],
QueryWeighting::Idf => {
if let Some(stats) = global_stats {
stats
.sparse_idf_weights(field, &token_ids)
.into_iter()
.map(|w| w.max(0.0))
.collect()
} else {
vec![1.0f32; token_ids.len()]
}
}
QueryWeighting::IdfFile => {
vec![1.0f32; token_ids.len()]
}
};
let vector: Vec<(u32, f32)> = token_ids.into_iter().zip(weights).collect();
Ok(Self::new(field, vector))
}
#[cfg(feature = "native")]
pub fn from_text_with_tokenizer_bytes(
field: Field,
text: &str,
tokenizer_bytes: &[u8],
weighting: crate::structures::QueryWeighting,
global_stats: Option<&crate::query::GlobalStats>,
) -> crate::Result<Self> {
use crate::structures::QueryWeighting;
use crate::tokenizer::HfTokenizer;
let tokenizer = HfTokenizer::from_bytes(tokenizer_bytes)?;
let token_ids = tokenizer.tokenize_unique(text)?;
let weights: Vec<f32> = match weighting {
QueryWeighting::One => vec![1.0f32; token_ids.len()],
QueryWeighting::Idf => {
if let Some(stats) = global_stats {
stats
.sparse_idf_weights(field, &token_ids)
.into_iter()
.map(|w| w.max(0.0))
.collect()
} else {
vec![1.0f32; token_ids.len()]
}
}
QueryWeighting::IdfFile => {
vec![1.0f32; token_ids.len()]
}
};
let vector: Vec<(u32, f32)> = token_ids.into_iter().zip(weights).collect();
Ok(Self::new(field, vector))
}
}
impl SparseVectorQuery {
fn sparse_infos(&self) -> Vec<crate::query::SparseTermQueryInfo> {
self.pruned_dims()
.iter()
.map(|&(dim_id, weight)| crate::query::SparseTermQueryInfo {
field: self.field,
dim_id,
weight,
heap_factor: self.heap_factor,
combiner: self.combiner,
over_fetch_factor: self.over_fetch_factor,
max_superblocks: self.max_superblocks,
})
.collect()
}
}
impl Query for SparseVectorQuery {
fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
let validation = self.validate(reader);
let infos = self.sparse_infos();
Box::pin(async move {
validation?;
if infos.is_empty() {
return Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer>);
}
if let Some((raw, info)) =
crate::query::planner::build_sparse_bmp_results(&infos, reader, limit)
{
return Ok(crate::query::planner::combine_sparse_results(
raw,
info.combiner,
info.field,
limit,
));
}
if let Some((executor, info)) =
crate::query::planner::build_sparse_maxscore_executor(&infos, reader, limit, None)
{
let raw = executor.execute().await?;
return Ok(crate::query::planner::combine_sparse_results(
raw,
info.combiner,
info.field,
limit,
));
}
Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer>)
})
}
#[cfg(feature = "sync")]
fn scorer_sync<'a>(
&self,
reader: &'a SegmentReader,
limit: usize,
) -> crate::Result<Box<dyn Scorer + 'a>> {
self.validate(reader)?;
let infos = self.sparse_infos();
if infos.is_empty() {
return Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer + 'a>);
}
if let Some((raw, info)) =
crate::query::planner::build_sparse_bmp_results(&infos, reader, limit)
{
return Ok(crate::query::planner::combine_sparse_results(
raw,
info.combiner,
info.field,
limit,
));
}
if let Some((executor, info)) =
crate::query::planner::build_sparse_maxscore_executor(&infos, reader, limit, None)
{
let raw = executor.execute_sync()?;
return Ok(crate::query::planner::combine_sparse_results(
raw,
info.combiner,
info.field,
limit,
));
}
Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer + 'a>)
}
fn count_estimate<'a>(&self, _reader: &'a SegmentReader) -> CountFuture<'a> {
Box::pin(async move { Ok(u32::MAX) })
}
fn decompose(&self) -> crate::query::QueryDecomposition {
let infos = self.sparse_infos();
if infos.is_empty() {
crate::query::QueryDecomposition::Opaque
} else {
crate::query::QueryDecomposition::SparseTerms(infos)
}
}
}
#[derive(Debug, Clone)]
pub struct SparseTermQuery {
pub field: Field,
pub dim_id: u32,
pub weight: f32,
pub heap_factor: f32,
pub combiner: MultiValueCombiner,
pub over_fetch_factor: f32,
}
impl std::fmt::Display for SparseTermQuery {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"SparseTerm({}, dim={}, w={:.3})",
self.field.0, self.dim_id, self.weight
)
}
}
impl SparseTermQuery {
pub fn new(field: Field, dim_id: u32, weight: f32) -> Self {
Self {
field,
dim_id,
weight,
heap_factor: 1.0,
combiner: MultiValueCombiner::default(),
over_fetch_factor: 2.0,
}
}
pub fn with_heap_factor(mut self, heap_factor: f32) -> Self {
self.heap_factor = heap_factor;
self
}
pub fn with_combiner(mut self, combiner: MultiValueCombiner) -> Self {
self.combiner = combiner;
self
}
pub fn with_over_fetch_factor(mut self, factor: f32) -> Self {
self.over_fetch_factor = factor.max(1.0);
self
}
fn validate(&self, reader: &SegmentReader) -> crate::Result<()> {
let entry = reader
.schema()
.get_field_entry(self.field)
.ok_or_else(|| crate::Error::FieldNotFound(self.field.0.to_string()))?;
if entry.field_type != crate::dsl::FieldType::SparseVector {
return Err(crate::Error::InvalidFieldType {
expected: "sparse_vector".to_string(),
got: format!("{:?}", entry.field_type),
});
}
if !self.weight.is_finite() {
return Err(crate::Error::Query(
"sparse term query weight must be finite".to_string(),
));
}
if !self.heap_factor.is_finite() || !(0.0..=1.0).contains(&self.heap_factor) {
return Err(crate::Error::Query(format!(
"sparse heap_factor must be finite and in [0, 1], got {}",
self.heap_factor
)));
}
if !self.over_fetch_factor.is_finite() || self.over_fetch_factor < 1.0 {
return Err(crate::Error::Query(format!(
"sparse over_fetch_factor must be finite and at least 1, got {}",
self.over_fetch_factor
)));
}
self.combiner.validate().map_err(crate::Error::Query)
}
fn bmp_fallback_scorer<'a>(
&self,
reader: &'a SegmentReader,
limit: usize,
) -> crate::Result<Box<dyn Scorer + 'a>> {
if let Some(bmp) = reader.bmp_index(self.field) {
let executor_limit =
crate::query::planner::bounded_sparse_executor_limit(limit, self.over_fetch_factor)
.min(bmp.num_virtual_docs as usize);
let results = crate::query::bmp::execute_bmp(
bmp,
reader.schema().index_label(),
reader.schema().get_field_name(self.field).unwrap_or("?"),
&[(self.dim_id, self.weight)],
executor_limit,
self.heap_factor,
0,
)?;
let combined = crate::segment::combine_ordinal_results(
results.into_iter().map(|r| (r.doc_id, r.ordinal, r.score)),
self.combiner,
limit,
);
return Ok(Box::new(
crate::query::planner::VectorTopKResultScorer::new(combined, self.field.0),
));
}
Ok(Box::new(crate::query::EmptyScorer))
}
fn make_scorer<'a>(
&self,
reader: &'a SegmentReader,
) -> crate::Result<Option<SparseTermScorer<'a>>> {
let si = match reader.sparse_index(self.field) {
Some(si) => si,
None => return Ok(None),
};
let (skip_start, skip_count, global_max, block_data_offset) =
match si.get_skip_range_full(self.dim_id) {
Some(v) => v,
None => return Ok(None),
};
let cursor = crate::query::TermCursor::sparse(
si,
self.weight,
skip_start,
skip_count,
global_max,
block_data_offset,
);
Ok(Some(SparseTermScorer {
cursor,
field_id: self.field.0,
}))
}
}
impl Query for SparseTermQuery {
fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
let query = self.clone();
Box::pin(async move {
query.validate(reader)?;
let mut scorer = match query.make_scorer(reader)? {
Some(s) => s,
None => return query.bmp_fallback_scorer(reader, limit),
};
scorer.cursor.ensure_block_loaded().await.ok();
Ok(Box::new(scorer) as Box<dyn Scorer + 'a>)
})
}
#[cfg(feature = "sync")]
fn scorer_sync<'a>(
&self,
reader: &'a SegmentReader,
limit: usize,
) -> crate::Result<Box<dyn Scorer + 'a>> {
self.validate(reader)?;
let mut scorer = match self.make_scorer(reader)? {
Some(s) => s,
None => return self.bmp_fallback_scorer(reader, limit),
};
scorer.cursor.ensure_block_loaded_sync().ok();
Ok(Box::new(scorer) as Box<dyn Scorer + 'a>)
}
fn count_estimate<'a>(&self, reader: &'a SegmentReader) -> CountFuture<'a> {
let field = self.field;
let dim_id = self.dim_id;
Box::pin(async move {
let si = match reader.sparse_index(field) {
Some(si) => si,
None => return Ok(0),
};
match si.get_skip_range_full(dim_id) {
Some((_, skip_count, _, _)) => Ok((skip_count * 256) as u32),
None => Ok(0),
}
})
}
fn decompose(&self) -> crate::query::QueryDecomposition {
crate::query::QueryDecomposition::SparseTerms(vec![crate::query::SparseTermQueryInfo {
field: self.field,
dim_id: self.dim_id,
weight: self.weight,
heap_factor: self.heap_factor,
combiner: self.combiner,
over_fetch_factor: self.over_fetch_factor,
max_superblocks: 0,
}])
}
}
struct SparseTermScorer<'a> {
cursor: crate::query::TermCursor<'a>,
field_id: u32,
}
impl crate::query::docset::DocSet for SparseTermScorer<'_> {
fn doc(&self) -> DocId {
let d = self.cursor.doc();
if d == u32::MAX { TERMINATED } else { d }
}
fn advance(&mut self) -> DocId {
match self.cursor.advance_sync() {
Ok(d) if d == u32::MAX => TERMINATED,
Ok(d) => d,
Err(_) => TERMINATED,
}
}
fn seek(&mut self, target: DocId) -> DocId {
match self.cursor.seek_sync(target) {
Ok(d) if d == u32::MAX => TERMINATED,
Ok(d) => d,
Err(_) => TERMINATED,
}
}
fn size_hint(&self) -> u32 {
0
}
}
impl Scorer for SparseTermScorer<'_> {
fn score(&self) -> Score {
self.cursor.score()
}
fn matched_positions(&self) -> Option<MatchedPositions> {
let ordinal = self.cursor.ordinal();
let score = self.cursor.score();
if score == 0.0 {
return None;
}
Some(vec![(
self.field_id,
vec![ScoredPosition::new(ordinal as u32, score)],
)])
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::dsl::Field;
#[test]
fn test_sparse_vector_query_new() {
let sparse = vec![(1, 0.5), (5, 0.3), (10, 0.2)];
let query = SparseVectorQuery::new(Field(0), sparse.clone());
assert_eq!(query.field, Field(0));
assert_eq!(query.vector, sparse);
}
#[test]
fn test_sparse_vector_query_from_indices_weights() {
let query =
SparseVectorQuery::from_indices_weights(Field(0), vec![1, 5, 10], vec![0.5, 0.3, 0.2]);
assert_eq!(query.vector, vec![(1, 0.5), (5, 0.3), (10, 0.2)]);
}
#[test]
fn max_query_dims_cannot_exceed_executor_mask_width() {
let vector: Vec<(u32, f32)> = (0..100).map(|dim| (dim, dim as f32 + 1.0)).collect();
let query = SparseVectorQuery::new(Field(0), vector).with_max_query_dims(usize::MAX);
assert_eq!(query.pruned_dims().len(), crate::query::MAX_QUERY_TERMS);
assert!(query.pruned_dims().iter().all(|(dim, _)| *dim >= 36));
}
}