mito2/sst/index/inverted_index/applier/builder/
between.rsuse datafusion_expr::Between;
use index::inverted_index::search::predicate::{Bound, Predicate, Range, RangePredicate};
use crate::error::Result;
use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder;
impl InvertedIndexApplierBuilder<'_> {
pub(crate) fn collect_between(&mut self, between: &Between) -> Result<()> {
if between.negated {
return Ok(());
}
let Some(column_name) = Self::column_name(&between.expr) else {
return Ok(());
};
let Some((column_id, data_type)) = self.column_id_and_type(column_name)? else {
return Ok(());
};
let Some(low) = Self::nonnull_lit(&between.low) else {
return Ok(());
};
let Some(high) = Self::nonnull_lit(&between.high) else {
return Ok(());
};
let predicate = Predicate::Range(RangePredicate {
range: Range {
lower: Some(Bound {
inclusive: true,
value: Self::encode_lit(low, data_type.clone())?,
}),
upper: Some(Bound {
inclusive: true,
value: Self::encode_lit(high, data_type)?,
}),
},
});
self.add_predicate(column_id, predicate);
Ok(())
}
}
#[cfg(test)]
mod tests {
use std::collections::HashSet;
use super::*;
use crate::error::Error;
use crate::sst::index::inverted_index::applier::builder::tests::{
encoded_string, field_column, int64_lit, nonexistent_column, string_lit, tag_column,
test_object_store, test_region_metadata,
};
use crate::sst::index::puffin_manager::PuffinManagerFactory;
#[test]
fn test_collect_between_basic() {
let (_d, facotry) = PuffinManagerFactory::new_for_test_block("test_collect_between_basic_");
let metadata = test_region_metadata();
let mut builder = InvertedIndexApplierBuilder::new(
"test".to_string(),
test_object_store(),
&metadata,
HashSet::from_iter([1, 2, 3]),
facotry,
);
let between = Between {
negated: false,
expr: Box::new(tag_column()),
low: Box::new(string_lit("abc")),
high: Box::new(string_lit("def")),
};
builder.collect_between(&between).unwrap();
let predicates = builder.output.get(&1).unwrap();
assert_eq!(predicates.len(), 1);
assert_eq!(
predicates[0],
Predicate::Range(RangePredicate {
range: Range {
lower: Some(Bound {
inclusive: true,
value: encoded_string("abc"),
}),
upper: Some(Bound {
inclusive: true,
value: encoded_string("def"),
}),
}
})
);
}
#[test]
fn test_collect_between_negated() {
let (_d, facotry) =
PuffinManagerFactory::new_for_test_block("test_collect_between_negated_");
let metadata = test_region_metadata();
let mut builder = InvertedIndexApplierBuilder::new(
"test".to_string(),
test_object_store(),
&metadata,
HashSet::from_iter([1, 2, 3]),
facotry,
);
let between = Between {
negated: true,
expr: Box::new(tag_column()),
low: Box::new(string_lit("abc")),
high: Box::new(string_lit("def")),
};
builder.collect_between(&between).unwrap();
assert!(builder.output.is_empty());
}
#[test]
fn test_collect_between_field_column() {
let (_d, facotry) =
PuffinManagerFactory::new_for_test_block("test_collect_between_field_column_");
let metadata = test_region_metadata();
let mut builder = InvertedIndexApplierBuilder::new(
"test".to_string(),
test_object_store(),
&metadata,
HashSet::from_iter([1, 2, 3]),
facotry,
);
let between = Between {
negated: false,
expr: Box::new(field_column()),
low: Box::new(string_lit("abc")),
high: Box::new(string_lit("def")),
};
builder.collect_between(&between).unwrap();
let predicates = builder.output.get(&3).unwrap();
assert_eq!(predicates.len(), 1);
assert_eq!(
predicates[0],
Predicate::Range(RangePredicate {
range: Range {
lower: Some(Bound {
inclusive: true,
value: encoded_string("abc"),
}),
upper: Some(Bound {
inclusive: true,
value: encoded_string("def"),
}),
}
})
);
}
#[test]
fn test_collect_between_type_mismatch() {
let (_d, facotry) =
PuffinManagerFactory::new_for_test_block("test_collect_between_type_mismatch_");
let metadata = test_region_metadata();
let mut builder = InvertedIndexApplierBuilder::new(
"test".to_string(),
test_object_store(),
&metadata,
HashSet::from_iter([1, 2, 3]),
facotry,
);
let between = Between {
negated: false,
expr: Box::new(tag_column()),
low: Box::new(int64_lit(123)),
high: Box::new(int64_lit(456)),
};
let res = builder.collect_between(&between);
assert!(matches!(res, Err(Error::FieldTypeMismatch { .. })));
assert!(builder.output.is_empty());
}
#[test]
fn test_collect_between_nonexistent_column() {
let (_d, facotry) =
PuffinManagerFactory::new_for_test_block("test_collect_between_nonexistent_column_");
let metadata = test_region_metadata();
let mut builder = InvertedIndexApplierBuilder::new(
"test".to_string(),
test_object_store(),
&metadata,
HashSet::from_iter([1, 2, 3]),
facotry,
);
let between = Between {
negated: false,
expr: Box::new(nonexistent_column()),
low: Box::new(string_lit("abc")),
high: Box::new(string_lit("def")),
};
let res = builder.collect_between(&between);
assert!(matches!(res, Err(Error::ColumnNotFound { .. })));
assert!(builder.output.is_empty());
}
}