use std::collections::BTreeMap;
use uqa_core::{DocId, Payload, PostingEntry, PostingList, Predicate, Value};
use uqa_storage::BTreeIndex;
mod comparison;
#[derive(Clone)]
pub struct ColumnValueIndex {
index: Option<BTreeIndex>,
values: BTreeMap<DocId, Value>,
nulls: Vec<DocId>,
has_temporal: bool,
has_fallible_comparison: bool,
}
fn value_is_temporal(value: &Value) -> bool {
matches!(value, Value::Temporal(_))
}
fn value_is_nan(value: &Value) -> bool {
matches!(value, Value::Float(f) if f.is_nan())
}
fn predicate_targets_are_index_safe(predicate: &Predicate) -> bool {
let safe = |v: &Value| !value_is_temporal(v) && !value_is_nan(v);
match predicate {
Predicate::Equals(v)
| Predicate::NotEquals(v)
| Predicate::GreaterThan(v)
| Predicate::GreaterThanOrEqual(v)
| Predicate::LessThan(v)
| Predicate::LessThanOrEqual(v) => safe(v),
Predicate::InSet(values) => values.iter().all(safe),
Predicate::Between { low, high } => safe(low) && safe(high),
Predicate::IsNull | Predicate::IsNotNull => true,
}
}
impl ColumnValueIndex {
pub fn stored_value(&self, doc_id: DocId) -> Option<&Value> {
self.values.get(&doc_id)
}
pub fn contains(&self, doc_id: DocId) -> bool {
self.values.contains_key(&doc_id)
}
pub fn is_carried(&self) -> bool {
self.index.is_none()
}
pub fn build(field: &str, values: impl Iterator<Item = (DocId, Value)>) -> Self {
let mut built = Self {
index: Some(BTreeIndex::new(field)),
values: BTreeMap::new(),
nulls: Vec::new(),
has_temporal: false,
has_fallible_comparison: false,
};
for (doc_id, value) in values {
built.index_value(doc_id, &value);
built.values.insert(doc_id, value);
}
built.nulls.sort_unstable();
built.nulls.dedup();
built
}
pub fn build_carried(values: impl Iterator<Item = (DocId, Value)>) -> Self {
Self {
index: None,
values: values.collect(),
nulls: Vec::new(),
has_temporal: false,
has_fallible_comparison: false,
}
}
#[must_use]
pub fn with_use(self, field: &str, carried: bool) -> Self {
match (carried, self.is_carried()) {
(false, true) => Self::build(field, self.values.into_iter()),
(true, false) => Self::build_carried(self.values.into_iter()),
_ => self,
}
}
fn index_value(&mut self, doc_id: DocId, value: &Value) {
let Some(index) = self.index.as_mut() else {
return;
};
match value {
Value::Null => self.nulls.push(doc_id),
value => {
self.has_temporal |= value_is_temporal(value);
self.has_fallible_comparison |= uqa_sql::expr::value_comparison_can_fail(value);
index.insert(doc_id, value.clone());
}
}
}
pub fn insert(&mut self, doc_id: DocId, value: &Value) {
self.values.insert(doc_id, value.clone());
let Some(index) = self.index.as_mut() else {
return;
};
match value {
Value::Null => {
if let Err(pos) = self.nulls.binary_search(&doc_id) {
self.nulls.insert(pos, doc_id);
}
}
value => {
self.has_temporal |= value_is_temporal(value);
self.has_fallible_comparison |= uqa_sql::expr::value_comparison_can_fail(value);
index.insert(doc_id, value.clone());
}
}
}
pub fn remove(&mut self, doc_id: DocId, value: &Value) {
let stored = self.values.remove(&doc_id);
let Some(index) = self.index.as_mut() else {
return;
};
let value = stored.as_ref().unwrap_or(value);
match value {
Value::Null => {
if let Ok(pos) = self.nulls.binary_search(&doc_id) {
self.nulls.remove(pos);
}
}
value => index.remove(doc_id, value),
}
}
pub fn clear(&mut self) {
if let Some(index) = self.index.as_mut() {
index.clear();
}
self.values.clear();
self.nulls.clear();
self.has_temporal = false;
self.has_fallible_comparison = false;
}
pub fn scan(&self, predicate: &Predicate) -> Option<PostingList> {
let index = self.index.as_ref()?;
if !self.supports(predicate) {
return None;
}
match predicate {
Predicate::IsNull => Some(posting_list_from_sorted_ids(self.nulls.iter().copied())),
Predicate::IsNotNull => Some(index.scan(&Predicate::IsNotNull)),
Predicate::NotEquals(_) => unreachable!("unsupported predicates return above"),
predicate => Some(index.scan(predicate)),
}
}
pub fn estimate_cardinality(&self, predicate: &Predicate) -> Option<usize> {
let index = self.index.as_ref()?;
if !self.supports(predicate) {
return None;
}
Some(match predicate {
Predicate::IsNull => self.nulls.len(),
Predicate::IsNotNull => index.estimate_cardinality(predicate),
Predicate::NotEquals(_) => unreachable!("unsupported predicates return above"),
predicate => index.estimate_cardinality(predicate),
})
}
pub fn scan_observing(
&self,
predicate: &Predicate,
observe: impl FnOnce() -> Result<(), uqa_sql::SQLError>,
) -> Result<Option<PostingList>, uqa_sql::SQLError> {
if self.is_carried() {
return Ok(None);
}
if comparison::needs_sql_comparison(predicate, self.has_fallible_comparison) {
observe()?;
let mut ids = Vec::new();
for (&id, value) in &self.values {
if comparison::matches(value, predicate)? {
ids.push(id);
}
}
return Ok(Some(posting_list_from_sorted_ids(ids.into_iter())));
}
if !self.supports(predicate) {
return Ok(None);
}
observe()?;
Ok(self.scan(predicate))
}
pub fn supports(&self, predicate: &Predicate) -> bool {
!self.is_carried()
&& predicate_targets_are_index_safe(predicate)
&& !comparison::needs_sql_comparison(predicate, self.has_fallible_comparison)
&& !matches!(predicate, Predicate::NotEquals(_))
&& (matches!(predicate, Predicate::IsNull | Predicate::IsNotNull) || !self.has_temporal)
}
}
fn posting_list_from_sorted_ids(ids: impl Iterator<Item = DocId>) -> PostingList {
let entries: Vec<PostingEntry> = ids
.map(|doc_id| PostingEntry::new(doc_id, Payload::default()))
.collect();
PostingList::from_sorted_unchecked(entries)
}
#[cfg(test)]
mod tests;