use crate::dsl::Field;
use crate::segment::SegmentReader;
use crate::structures::TERMINATED;
use crate::structures::fast_field::{FAST_FIELD_MISSING, f64_to_sortable_u64, zigzag_encode};
use crate::{DocId, Score};
use super::docset::DocSet;
use super::traits::{CountFuture, Query, Scorer, ScorerFuture};
#[derive(Debug, Clone)]
pub enum RangeBound {
U64 { min: Option<u64>, max: Option<u64> },
I64 { min: Option<i64>, max: Option<i64> },
F64 { min: Option<f64>, max: Option<f64> },
}
impl RangeBound {
fn compile(&self) -> (u64, u64) {
match self {
RangeBound::U64 { min, max } => {
let lo = min.unwrap_or(0);
let hi = max.unwrap_or(u64::MAX - 1);
(lo, hi)
}
RangeBound::I64 { min, max } => {
let lo = min.map(zigzag_encode).unwrap_or(0);
let hi = max.map(zigzag_encode).unwrap_or(u64::MAX - 1);
(lo, hi)
}
RangeBound::F64 { min, max } => {
let lo = min.map(f64_to_sortable_u64).unwrap_or(0);
let hi = max.map(f64_to_sortable_u64).unwrap_or(u64::MAX - 1);
(lo, hi)
}
}
}
fn is_i64(&self) -> bool {
matches!(self, RangeBound::I64 { .. })
}
fn i64_bounds(&self) -> (i64, i64) {
match self {
RangeBound::I64 { min, max } => (min.unwrap_or(i64::MIN), max.unwrap_or(i64::MAX)),
_ => (i64::MIN, i64::MAX),
}
}
}
#[derive(Debug, Clone)]
pub struct RangeQuery {
pub field: Field,
pub bound: RangeBound,
}
impl std::fmt::Display for RangeQuery {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match &self.bound {
RangeBound::U64 { min, max } => write!(
f,
"Range({}:[{} TO {}])",
self.field.0,
min.map_or("*".to_string(), |v| v.to_string()),
max.map_or("*".to_string(), |v| v.to_string()),
),
RangeBound::I64 { min, max } => write!(
f,
"Range({}:[{} TO {}])",
self.field.0,
min.map_or("*".to_string(), |v| v.to_string()),
max.map_or("*".to_string(), |v| v.to_string()),
),
RangeBound::F64 { min, max } => write!(
f,
"Range({}:[{} TO {}])",
self.field.0,
min.map_or("*".to_string(), |v| v.to_string()),
max.map_or("*".to_string(), |v| v.to_string()),
),
}
}
}
impl RangeQuery {
pub fn new(field: Field, bound: RangeBound) -> Self {
Self { field, bound }
}
pub fn u64(field: Field, min: Option<u64>, max: Option<u64>) -> Self {
Self::new(field, RangeBound::U64 { min, max })
}
pub fn i64(field: Field, min: Option<i64>, max: Option<i64>) -> Self {
Self::new(field, RangeBound::I64 { min, max })
}
pub fn f64(field: Field, min: Option<f64>, max: Option<f64>) -> Self {
Self::new(field, RangeBound::F64 { min, max })
}
}
impl Query for RangeQuery {
fn scorer<'a>(&self, reader: &'a SegmentReader, _limit: usize) -> ScorerFuture<'a> {
let field = self.field;
let bound = self.bound.clone();
Box::pin(async move {
match RangeScorer::new(reader, field, &bound) {
Ok(scorer) => Ok(Box::new(scorer) as Box<dyn Scorer>),
Err(_) => Ok(Box::new(EmptyRangeScorer) as Box<dyn Scorer>),
}
})
}
#[cfg(feature = "sync")]
fn scorer_sync<'a>(
&self,
reader: &'a SegmentReader,
_limit: usize,
) -> crate::Result<Box<dyn Scorer + 'a>> {
match RangeScorer::new(reader, self.field, &self.bound) {
Ok(scorer) => Ok(Box::new(scorer) as Box<dyn Scorer + 'a>),
Err(_) => Ok(Box::new(EmptyRangeScorer) as Box<dyn Scorer + 'a>),
}
}
fn count_estimate<'a>(&self, reader: &'a SegmentReader) -> CountFuture<'a> {
let num_docs = reader.num_docs();
Box::pin(async move { Ok(num_docs / 2) })
}
fn is_filter(&self) -> bool {
true
}
fn as_doc_predicate<'a>(&self, reader: &'a SegmentReader) -> Option<super::DocPredicate<'a>> {
let fast_field = reader.fast_field(self.field.0)?;
let (raw_lo, raw_hi) = self.bound.compile();
let use_i64 = self.bound.is_i64();
let (i64_lo, i64_hi) = self.bound.i64_bounds();
Some(Box::new(move |doc_id: DocId| -> bool {
let raw = fast_field.get_u64(doc_id);
if raw == FAST_FIELD_MISSING {
return false;
}
if use_i64 {
let val = crate::structures::fast_field::zigzag_decode(raw);
val >= i64_lo && val <= i64_hi
} else {
raw >= raw_lo && raw <= raw_hi
}
}))
}
fn as_doc_bitset(&self, reader: &SegmentReader) -> Option<super::DocBitset> {
let pred = self.as_doc_predicate(reader)?;
Some(super::DocBitset::from_predicate(reader.num_docs(), &*pred))
}
}
struct RangeScorer<'a> {
fast_field: &'a crate::structures::fast_field::FastFieldReader,
raw_lo: u64,
raw_hi: u64,
i64_lo: i64,
i64_hi: i64,
use_i64: bool,
current: u32,
num_docs: u32,
}
struct EmptyRangeScorer;
impl<'a> RangeScorer<'a> {
fn new(
reader: &'a SegmentReader,
field: Field,
bound: &RangeBound,
) -> Result<Self, EmptyRangeScorer> {
let fast_field = reader.fast_field(field.0).ok_or(EmptyRangeScorer)?;
let num_docs = reader.num_docs();
let (raw_lo, raw_hi) = bound.compile();
let use_i64 = bound.is_i64();
let (i64_lo, i64_hi) = bound.i64_bounds();
let mut scorer = Self {
fast_field,
raw_lo,
raw_hi,
i64_lo,
i64_hi,
use_i64,
current: 0,
num_docs,
};
if num_docs > 0 && !scorer.matches(0) {
scorer.scan_forward();
}
Ok(scorer)
}
#[inline]
fn matches(&self, doc_id: DocId) -> bool {
let raw = self.fast_field.get_u64(doc_id);
if raw == FAST_FIELD_MISSING {
return false;
}
if self.use_i64 {
let val = crate::structures::fast_field::zigzag_decode(raw);
val >= self.i64_lo && val <= self.i64_hi
} else {
raw >= self.raw_lo && raw <= self.raw_hi
}
}
fn scan_forward(&mut self) {
loop {
self.current += 1;
if self.current >= self.num_docs {
self.current = self.num_docs;
return;
}
if self.matches(self.current) {
return;
}
}
}
}
impl DocSet for RangeScorer<'_> {
fn doc(&self) -> DocId {
if self.current >= self.num_docs {
TERMINATED
} else {
self.current
}
}
fn advance(&mut self) -> DocId {
self.scan_forward();
self.doc()
}
fn seek(&mut self, target: DocId) -> DocId {
if self.current >= self.num_docs {
return TERMINATED;
}
if target <= self.current {
return self.current;
}
self.current = target - 1;
self.scan_forward();
self.doc()
}
fn size_hint(&self) -> u32 {
self.num_docs.saturating_sub(self.current)
}
}
impl Scorer for RangeScorer<'_> {
fn score(&self) -> Score {
1.0
}
}
impl DocSet for EmptyRangeScorer {
fn doc(&self) -> DocId {
TERMINATED
}
fn advance(&mut self) -> DocId {
TERMINATED
}
fn seek(&mut self, _target: DocId) -> DocId {
TERMINATED
}
fn size_hint(&self) -> u32 {
0
}
}
impl Scorer for EmptyRangeScorer {
fn score(&self) -> Score {
0.0
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_range_bound_u64_compile() {
let b = RangeBound::U64 {
min: Some(10),
max: Some(100),
};
let (lo, hi) = b.compile();
assert_eq!(lo, 10);
assert_eq!(hi, 100);
}
#[test]
fn test_range_bound_f64_compile_preserves_order() {
let b1 = RangeBound::F64 {
min: Some(-1.0),
max: Some(1.0),
};
let (lo1, hi1) = b1.compile();
assert!(lo1 < hi1);
let b2 = RangeBound::F64 {
min: Some(0.0),
max: Some(100.0),
};
let (lo2, hi2) = b2.compile();
assert!(lo2 < hi2);
}
#[test]
fn test_range_bound_open_bounds() {
let b = RangeBound::U64 {
min: None,
max: None,
};
let (lo, hi) = b.compile();
assert_eq!(lo, 0);
assert_eq!(hi, u64::MAX - 1);
}
#[test]
fn test_range_query_constructors() {
let q = RangeQuery::u64(Field(0), Some(10), Some(100));
assert_eq!(q.field, Field(0));
assert!(matches!(
q.bound,
RangeBound::U64 {
min: Some(10),
max: Some(100)
}
));
let q = RangeQuery::i64(Field(1), Some(-50), Some(50));
assert!(matches!(
q.bound,
RangeBound::I64 {
min: Some(-50),
max: Some(50)
}
));
let q = RangeQuery::f64(Field(2), Some(0.5), Some(9.5));
assert!(matches!(q.bound, RangeBound::F64 { .. }));
}
}