use crate::cursor::column_batch::ColumnValues;
use std::sync::Arc;
#[derive(Debug, Clone, PartialEq)]
pub enum IndexPredicate {
Equal(linkrs_core::Value),
Range {
lower: Option<linkrs_core::Value>,
upper: Option<linkrs_core::Value>,
include_lower: bool,
include_upper: bool,
},
Prefix(linkrs_core::Value),
All,
}
#[derive(Debug, Clone, PartialEq)]
pub enum ScanPredicate {
ColumnEqual {
column: Arc<str>,
value: linkrs_core::Value,
},
ColumnRange {
column: Arc<str>,
lower: Option<linkrs_core::Value>,
upper: Option<linkrs_core::Value>,
include_lower: bool,
include_upper: bool,
},
}
impl ScanPredicate {
pub fn matches(&self, props: &[(std::sync::Arc<str>, linkrs_core::Value)]) -> bool {
let Some(value) = props
.iter()
.find(|(name, _)| name.as_ref() == self.column())
.map(|(_, v)| v)
else {
return false;
};
match self {
ScanPredicate::ColumnEqual {
value: expected, ..
} => compare_scalar(value, expected) == std::cmp::Ordering::Equal,
ScanPredicate::ColumnRange {
lower,
upper,
include_lower,
include_upper,
..
} => {
if let Some(lower) = lower {
let ord = compare_scalar(value, lower);
let passes = if *include_lower {
ord != std::cmp::Ordering::Less
} else {
ord == std::cmp::Ordering::Greater
};
if !passes {
return false;
}
}
if let Some(upper) = upper {
let ord = compare_scalar(value, upper);
let passes = if *include_upper {
ord != std::cmp::Ordering::Greater
} else {
ord == std::cmp::Ordering::Less
};
if !passes {
return false;
}
}
true
}
}
}
pub fn column(&self) -> &str {
match self {
ScanPredicate::ColumnEqual { column, .. } => column,
ScanPredicate::ColumnRange { column, .. } => column,
}
}
pub fn column_arc(&self) -> Arc<str> {
match self {
ScanPredicate::ColumnEqual { column, .. } => column.clone(),
ScanPredicate::ColumnRange { column, .. } => column.clone(),
}
}
pub fn merged_ranges(predicates: &[ScanPredicate]) -> Vec<PredicateRange> {
let mut ranges: Vec<PredicateRange> = Vec::new();
for predicate in predicates {
let (lower, include_lower, upper, include_upper) = match predicate {
ScanPredicate::ColumnEqual { value, .. } => {
(Some(value.clone()), true, Some(value.clone()), true)
}
ScanPredicate::ColumnRange {
lower,
include_lower,
upper,
include_upper,
..
} => (lower.clone(), *include_lower, upper.clone(), *include_upper),
};
let column = predicate.column_arc();
match ranges.iter_mut().find(|r| r.column == column) {
Some(range) => range.intersect(lower, include_lower, upper, include_upper),
None => ranges.push(PredicateRange {
column,
lower,
include_lower,
upper,
include_upper,
}),
}
}
ranges
}
}
#[derive(Debug, Clone)]
pub struct PredicateRange {
pub column: Arc<str>,
pub lower: Option<linkrs_core::Value>,
pub include_lower: bool,
pub upper: Option<linkrs_core::Value>,
pub include_upper: bool,
}
impl PredicateRange {
fn intersect(
&mut self,
lower: Option<linkrs_core::Value>,
include_lower: bool,
upper: Option<linkrs_core::Value>,
include_upper: bool,
) {
if let (_, Some(new_lower)) = (&self.lower, &lower) {
let replace = match &self.lower {
Some(old) => match compare_scalar(old, new_lower) {
std::cmp::Ordering::Equal => self.include_lower && !include_lower,
std::cmp::Ordering::Less => true,
std::cmp::Ordering::Greater => false,
},
None => true,
};
if replace {
self.lower = Some(new_lower.clone());
self.include_lower = include_lower;
}
}
if let (_, Some(new_upper)) = (&self.upper, &upper) {
let replace = match &self.upper {
Some(old) => match compare_scalar(old, new_upper) {
std::cmp::Ordering::Equal => self.include_upper && !include_upper,
std::cmp::Ordering::Greater => true,
std::cmp::Ordering::Less => false,
},
None => true,
};
if replace {
self.upper = Some(new_upper.clone());
self.include_upper = include_upper;
}
}
}
pub fn overlaps(&self, min: &linkrs_core::Value, max: &linkrs_core::Value) -> bool {
if let Some(ref lower) = self.lower {
let ord = compare_scalar(max, lower);
let reaches = if self.include_lower {
ord != std::cmp::Ordering::Less
} else {
ord == std::cmp::Ordering::Greater
};
if !reaches {
return false;
}
}
if let Some(ref upper) = self.upper {
let ord = compare_scalar(min, upper);
let under = if self.include_upper {
ord != std::cmp::Ordering::Greater
} else {
ord == std::cmp::Ordering::Less
};
if !under {
return false;
}
}
true
}
pub fn equality_len(&self) -> Option<usize> {
if !self.include_lower || !self.include_upper {
return None;
}
let lower = self.lower.as_ref()?;
let upper = self.upper.as_ref()?;
if compare_scalar(lower, upper) != std::cmp::Ordering::Equal {
return None;
}
crate::vertex::column_store::complex_len(lower)
}
pub fn equality_leaf_range(&self) -> Option<(linkrs_core::Value, linkrs_core::Value)> {
if !self.include_lower || !self.include_upper {
return None;
}
let lower = self.lower.as_ref()?;
let upper = self.upper.as_ref()?;
if compare_scalar(lower, upper) != std::cmp::Ordering::Equal {
return None;
}
crate::vertex::column_store::complex_leaf_range(lower)
}
pub fn equality_key_fp(&self) -> Option<u64> {
if !self.include_lower || !self.include_upper {
return None;
}
let lower = self.lower.as_ref()?;
let upper = self.upper.as_ref()?;
if compare_scalar(lower, upper) != std::cmp::Ordering::Equal {
return None;
}
let fp = crate::vertex::column_store::complex_key_fp(lower);
(fp != 0).then_some(fp)
}
}
impl ScanPredicate {
pub fn matches_column(&self, column: &ColumnValues, idx: usize) -> bool {
let Some(value) = column.value_at(idx) else {
return false;
};
self.matches_scalar(&value)
}
pub fn matches_value(&self, value: &linkrs_core::Value) -> bool {
self.matches_scalar(value)
}
fn matches_scalar(&self, value: &linkrs_core::Value) -> bool {
match self {
ScanPredicate::ColumnEqual {
value: expected, ..
} => compare_scalar(value, expected) == std::cmp::Ordering::Equal,
ScanPredicate::ColumnRange {
lower,
upper,
include_lower,
include_upper,
..
} => {
if let Some(lower) = lower {
let ord = compare_scalar(value, lower);
let passes = if *include_lower {
ord != std::cmp::Ordering::Less
} else {
ord == std::cmp::Ordering::Greater
};
if !passes {
return false;
}
}
if let Some(upper) = upper {
let ord = compare_scalar(value, upper);
let passes = if *include_upper {
ord != std::cmp::Ordering::Greater
} else {
ord == std::cmp::Ordering::Less
};
if !passes {
return false;
}
}
true
}
}
}
}
pub fn compare_stat_values(a: &linkrs_core::Value, b: &linkrs_core::Value) -> std::cmp::Ordering {
compare_scalar(a, b)
}
fn compare_scalar(a: &linkrs_core::Value, b: &linkrs_core::Value) -> std::cmp::Ordering {
match (as_i64(a), as_i64(b)) {
(Some(x), Some(y)) => x.cmp(&y),
_ => match (as_f64(a), as_f64(b)) {
(Some(x), Some(y)) => x.partial_cmp(&y).unwrap_or(std::cmp::Ordering::Equal),
_ => linkrs_core::Value::cmp(a, b),
},
}
}
fn as_i64(value: &linkrs_core::Value) -> Option<i64> {
match value {
linkrs_core::Value::SmallInt(v) => Some(*v as i64),
linkrs_core::Value::Int(v) => Some(*v as i64),
linkrs_core::Value::BigInt(v) => Some(*v),
_ => None,
}
}
fn as_f64(value: &linkrs_core::Value) -> Option<f64> {
match value {
linkrs_core::Value::Float(v) => Some(*v as f64),
linkrs_core::Value::Double(v) => Some(*v),
_ => None,
}
}