rust/tantivy
1
0
Fork 0
mirror of https://github.com/quickwit-oss/tantivy.git synced 2026-11-04 01:37:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added possible necessary conditions to jitexpr.

Browse Source 
A function makes it possible to infer a necessary query from an expression to match.
We can then accelerate queries involving a calculated field by not even evaluating the expression
on docs that trivially do not match.
This commit is contained in:
Paul Masurel
2026-09-25 18:50:18 +02:00
parent cc09b9deae
commit f01ce425c2
6 changed files with 912 additions and 189 deletions
Download Patch File Download Diff File Expand all files Collapse all files
jitexpr/src/ast/mod.rs
+3 -1
View File
@@ -7,7 +7,9 @@ mod untyped_expr;
pub use infer_types::{InferredTypeSet, TypeError, infer_types, infer_types_with_target};
pub(crate) use infer_types::{infer_type_with_variable_types, infer_types_aux};
pub use literal::{Literal, NonFiniteFloat};
pub use presence::{PresenceCondition, required_presence, required_presence_for_true};
pub use presence::{
ConditionSet, VariablePresenceCondition, required_presence, required_presence_for_true,
};
pub(crate) use serde::format_variable_name;
pub use serde::{DeserializeError, deserialize, serialize};
pub use untyped_expr::UntypedExpr;
jitexpr/src/ast/presence.rs
+333 -126
View File
@@ -16,11 +16,12 @@ use crate::ast::{Function, Literal, UntypedExpr};
/// It is meant to be used as a necessary condition: it is implied by some property of an
/// expression (producing a value, or evaluating to `true`), but it does not imply it.
///
/// Conditions built with [`PresenceCondition::all`] and [`PresenceCondition::any`] are simplified:
/// `All` and `Any` then have at least two distinct children, none of which is `Always`, `Never`,
/// or a node of the same kind.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum PresenceCondition {
/// Conditions are kept in a canonical form, so that `Eq` and `Hash` identify conditions that
/// only differ by the order, grouping, or repetition of their children. Two equivalent
/// conditions may still compare as different (e.g. `a ∧ (a ∨ b)` and `a`): deciding logical
/// equivalence in general is too expensive.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum VariablePresenceCondition {
/// Always satisfied. This condition carries no information.
Always,
/// Never satisfied.
@@ -28,110 +29,124 @@ pub enum PresenceCondition {
/// Satisfied when the variable is present.
Present(Arc<str>),
/// Satisfied when all of the conditions are satisfied.
All(Vec<PresenceCondition>),
All(ConditionSet),
/// Satisfied when at least one of the conditions is satisfied.
Any(Vec<PresenceCondition>),
Any(ConditionSet),
}
impl PresenceCondition {
/// Builds the simplified conjunction of `conditions`.
pub fn all(conditions: impl IntoIterator<Item = PresenceCondition>) -> PresenceCondition {
let mut children: Vec<PresenceCondition> = Vec::new();
/// The children of a [`PresenceCondition::All`] or [`PresenceCondition::Any`] node.
///
/// It can only be built through [`PresenceCondition::all`] and [`PresenceCondition::any`], which
/// uphold the following hidden contract, on which the derived `Eq` and `Hash` rely:
/// - children are sorted and distinct, and there are at least two of them;
/// - no child is `Always`, `Never`, or a node of the same kind as the parent.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ConditionSet(Vec<VariablePresenceCondition>);
impl ConditionSet {
/// Returns the children, in canonical order.
pub fn iter(&self) -> impl Iterator<Item = &VariablePresenceCondition> {
self.0.iter()
}
}
impl VariablePresenceCondition {
/// Builds the canonical conjunction of `conditions`.
pub fn all(
conditions: impl IntoIterator<Item = VariablePresenceCondition>,
) -> VariablePresenceCondition {
let mut children: Vec<VariablePresenceCondition> = Vec::new();
for condition in conditions {
match condition {
PresenceCondition::Always => {}
PresenceCondition::Never => return PresenceCondition::Never,
PresenceCondition::All(grand_children) => {
for grand_child in grand_children {
push_unique(&mut children, grand_child);
}
}
condition => push_unique(&mut children, condition),
VariablePresenceCondition::Always => {}
VariablePresenceCondition::Never => return VariablePresenceCondition::Never,
VariablePresenceCondition::All(grand_children) => children.extend(grand_children.0),
condition => children.push(condition),
}
}
children.sort();
children.dedup();
match children.len() {
0 => PresenceCondition::Always,
0 => VariablePresenceCondition::Always,
1 => children.pop().unwrap(),
_ => PresenceCondition::All(children),
_ => VariablePresenceCondition::All(ConditionSet(children)),
}
}
/// Builds the simplified disjunction of `conditions`.
pub fn any(conditions: impl IntoIterator<Item = PresenceCondition>) -> PresenceCondition {
let mut children: Vec<PresenceCondition> = Vec::new();
/// Builds the canonical disjunction of `conditions`.
pub fn any(
conditions: impl IntoIterator<Item = VariablePresenceCondition>,
) -> VariablePresenceCondition {
let mut children: Vec<VariablePresenceCondition> = Vec::new();
for condition in conditions {
match condition {
PresenceCondition::Never => {}
PresenceCondition::Always => return PresenceCondition::Always,
PresenceCondition::Any(grand_children) => {
for grand_child in grand_children {
push_unique(&mut children, grand_child);
}
}
condition => push_unique(&mut children, condition),
VariablePresenceCondition::Never => {}
VariablePresenceCondition::Always => return VariablePresenceCondition::Always,
VariablePresenceCondition::Any(grand_children) => children.extend(grand_children.0),
condition => children.push(condition),
}
}
children.sort();
children.dedup();
match children.len() {
0 => PresenceCondition::Never,
0 => VariablePresenceCondition::Never,
1 => children.pop().unwrap(),
_ => PresenceCondition::Any(children),
_ => VariablePresenceCondition::Any(ConditionSet(children)),
}
}
/// Evaluates the condition, given the presence of each variable.
#[cfg(test)]
pub fn eval(&self, is_present: &mut impl FnMut(&str) -> bool) -> bool {
match self {
PresenceCondition::Always => true,
PresenceCondition::Never => false,
PresenceCondition::Present(variable_name) => is_present(variable_name),
PresenceCondition::All(conditions) => conditions
VariablePresenceCondition::Always => true,
VariablePresenceCondition::Never => false,
VariablePresenceCondition::Present(variable_name) => is_present(variable_name),
VariablePresenceCondition::All(conditions) => conditions
.iter()
.all(|condition| condition.eval(&mut *is_present)),
PresenceCondition::Any(conditions) => conditions
VariablePresenceCondition::Any(conditions) => conditions
.iter()
.any(|condition| condition.eval(&mut *is_present)),
}
}
}
fn push_unique(conditions: &mut Vec<PresenceCondition>, condition: PresenceCondition) {
if !conditions.contains(&condition) {
conditions.push(condition);
}
}
/// Returns a necessary condition for `expr` to evaluate to a present value.
///
/// A variable is considered present when its value is not null. A variable missing from the
/// variable types given to the compiler is null for all evaluations, and should therefore be
/// considered absent.
pub fn required_presence(expr: &UntypedExpr) -> PresenceCondition {
/// Returns a necessary presence condition for `expr` to evaluate to a non-null value.
pub fn required_presence(expr: &UntypedExpr) -> VariablePresenceCondition {
match expr {
// Literals never constrain variables. `none` could be mapped to `Never`, but some
// functions take literal configuration arguments (delimiters, precision, ...), and we
// prefer not to depend on how each of them handles a `none` there.
UntypedExpr::Literal(_) => PresenceCondition::Always,
UntypedExpr::Variable(variable_name) => PresenceCondition::Present(variable_name.clone()),
UntypedExpr::Literal(_) => VariablePresenceCondition::Always,
UntypedExpr::Variable(variable_name) => {
VariablePresenceCondition::Present(variable_name.clone())
}
UntypedExpr::FnCall { function, args } => required_presence_for_fn_call(*function, args),
}
}
/// Returns a necessary condition for `expr` to evaluate to a present `true`.
/// Returns a necessary presence condition for `expr` to evaluate to a present `true`.
///
/// See [`required_presence`] for the definition of a present variable.
pub fn required_presence_for_true(expr: &UntypedExpr) -> PresenceCondition {
pub fn required_presence_for_true(expr: &UntypedExpr) -> VariablePresenceCondition {
match expr {
UntypedExpr::Literal(Literal::Bool(true)) => PresenceCondition::Always,
UntypedExpr::Literal(Literal::Bool(true)) => VariablePresenceCondition::Always,
// No other literal is `true`.
UntypedExpr::Literal(_) => PresenceCondition::Never,
UntypedExpr::Variable(variable_name) => PresenceCondition::Present(variable_name.clone()),
UntypedExpr::Literal(_) => VariablePresenceCondition::Never,
UntypedExpr::Variable(variable_name) => {
VariablePresenceCondition::Present(variable_name.clone())
}
UntypedExpr::FnCall { function, args } => {
required_presence_for_true_for_fn_call(*function, args)
}
}
}
fn required_presence_for_fn_call(function: Function, args: &[UntypedExpr]) -> PresenceCondition {
fn required_presence_for_fn_call(
function: Function,
args: &[UntypedExpr],
) -> VariablePresenceCondition {
// This match must remain exhaustive: classifying a function returning a present value for a
// null argument as "null in, null out" would make callers skip matching documents.
match function {
@@ -168,18 +183,18 @@ fn required_presence_for_fn_call(function: Function, args: &[UntypedExpr]) -> Pr
| Function::Subtract
| Function::TextJoin
| Function::Trim
| Function::Upper => PresenceCondition::all(args.iter().map(required_presence)),
| Function::Upper => VariablePresenceCondition::all(args.iter().map(required_presence)),
// OR is null only if all of its arguments are null.
Function::Or => PresenceCondition::any(args.iter().map(required_presence)),
Function::Or => VariablePresenceCondition::any(args.iter().map(required_presence)),
// IF is null if its condition is null. Otherwise it takes the presence of the selected
// branch.
Function::If => {
let [condition, when_true, when_false] = args else {
return PresenceCondition::Always;
return VariablePresenceCondition::Always;
};
PresenceCondition::all([
VariablePresenceCondition::all([
required_presence(condition),
PresenceCondition::any([
VariablePresenceCondition::any([
required_presence(when_true),
required_presence(when_false),
]),
@@ -192,24 +207,26 @@ fn required_presence_for_fn_call(function: Function, args: &[UntypedExpr]) -> Pr
| Function::IsNull
| Function::Neq
| Function::Not
| Function::RegexpLike => PresenceCondition::Always,
| Function::RegexpLike => VariablePresenceCondition::Always,
}
}
fn required_presence_for_true_for_fn_call(
function: Function,
args: &[UntypedExpr],
) -> PresenceCondition {
) -> VariablePresenceCondition {
match function {
Function::And => PresenceCondition::all(args.iter().map(required_presence_for_true)),
Function::Or => PresenceCondition::any(args.iter().map(required_presence_for_true)),
Function::And => {
VariablePresenceCondition::all(args.iter().map(required_presence_for_true))
}
Function::Or => VariablePresenceCondition::any(args.iter().map(required_presence_for_true)),
Function::If => {
let [condition, when_true, when_false] = args else {
return PresenceCondition::Always;
return VariablePresenceCondition::Always;
};
PresenceCondition::all([
VariablePresenceCondition::all([
required_presence(condition),
PresenceCondition::any([
VariablePresenceCondition::any([
required_presence_for_true(when_true),
required_presence_for_true(when_false),
]),
@@ -217,14 +234,14 @@ fn required_presence_for_true_for_fn_call(
}
Function::IsNotNull => {
let [arg] = args else {
return PresenceCondition::Always;
return VariablePresenceCondition::Always;
};
required_presence(arg)
}
// REGEXP_LIKE returns `false` for a null input.
Function::RegexpLike => {
let Some(input) = args.first() else {
return PresenceCondition::Always;
return VariablePresenceCondition::Always;
};
required_presence(input)
}
@@ -240,6 +257,8 @@ fn required_presence_for_true_for_fn_call(
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use proptest::prelude::*;
use proptest::strategy::BoxedStrategy;
@@ -249,39 +268,42 @@ mod tests {
use crate::compile::{StringArena, compile};
use crate::types::{VarType, VariableValue};
fn present(variable_name: &str) -> PresenceCondition {
PresenceCondition::Present(Arc::from(variable_name))
fn present(variable_name: &str) -> VariablePresenceCondition {
VariablePresenceCondition::Present(Arc::from(variable_name))
}
fn all(conditions: Vec<PresenceCondition>) -> PresenceCondition {
PresenceCondition::All(conditions)
fn all(conditions: Vec<VariablePresenceCondition>) -> VariablePresenceCondition {
VariablePresenceCondition::all(conditions)
}
fn any(conditions: Vec<PresenceCondition>) -> PresenceCondition {
PresenceCondition::Any(conditions)
fn any(conditions: Vec<VariablePresenceCondition>) -> VariablePresenceCondition {
VariablePresenceCondition::any(conditions)
}
fn for_true(expr: &str) -> PresenceCondition {
fn for_true(expr: &str) -> VariablePresenceCondition {
required_presence_for_true(&deserialize(expr).unwrap())
}
fn for_value(expr: &str) -> PresenceCondition {
fn for_value(expr: &str) -> VariablePresenceCondition {
required_presence(&deserialize(expr).unwrap())
}
#[test]
fn test_all_simplification() {
assert_eq!(PresenceCondition::all([]), PresenceCondition::Always);
assert_eq!(
PresenceCondition::all([PresenceCondition::Always, present("a")]),
VariablePresenceCondition::all([]),
VariablePresenceCondition::Always
);
assert_eq!(
VariablePresenceCondition::all([VariablePresenceCondition::Always, present("a")]),
present("a")
);
assert_eq!(
PresenceCondition::all([present("a"), PresenceCondition::Never]),
PresenceCondition::Never
VariablePresenceCondition::all([present("a"), VariablePresenceCondition::Never]),
VariablePresenceCondition::Never
);
assert_eq!(
PresenceCondition::all([
VariablePresenceCondition::all([
present("a"),
all(vec![present("b"), present("a")]),
present("c"),
@@ -289,28 +311,203 @@ mod tests {
all(vec![present("a"), present("b"), present("c")])
);
assert_eq!(
PresenceCondition::all([any(vec![present("a"), present("b")]), present("c")]),
VariablePresenceCondition::all([any(vec![present("a"), present("b")]), present("c")]),
all(vec![any(vec![present("a"), present("b")]), present("c")])
);
}
#[test]
fn test_any_simplification() {
assert_eq!(PresenceCondition::any([]), PresenceCondition::Never);
assert_eq!(
PresenceCondition::any([PresenceCondition::Never, present("a")]),
VariablePresenceCondition::any([]),
VariablePresenceCondition::Never
);
assert_eq!(
VariablePresenceCondition::any([VariablePresenceCondition::Never, present("a")]),
present("a")
);
assert_eq!(
PresenceCondition::any([present("a"), PresenceCondition::Always]),
PresenceCondition::Always
VariablePresenceCondition::any([present("a"), VariablePresenceCondition::Always]),
VariablePresenceCondition::Always
);
assert_eq!(
PresenceCondition::any([present("a"), any(vec![present("b"), present("a")])]),
VariablePresenceCondition::any([present("a"), any(vec![present("b"), present("a")])]),
any(vec![present("a"), present("b")])
);
}
fn hash_of(condition: &VariablePresenceCondition) -> u64 {
let mut hasher = DefaultHasher::new();
condition.hash(&mut hasher);
hasher.finish()
}
fn assert_same(left: VariablePresenceCondition, right: VariablePresenceCondition) {
assert_eq!(left, right);
assert_eq!(hash_of(&left), hash_of(&right));
}
#[test]
fn test_canonical_order() {
let (a, b, c) = (present("a"), present("b"), present("c"));
assert_same(
VariablePresenceCondition::all([a.clone(), b.clone()]),
VariablePresenceCondition::all([b.clone(), a.clone()]),
);
assert_same(
VariablePresenceCondition::any([c.clone(), a.clone(), b.clone()]),
VariablePresenceCondition::any([b.clone(), c.clone(), a.clone()]),
);
assert_ne!(
VariablePresenceCondition::all([a.clone(), b.clone()]),
VariablePresenceCondition::any([a.clone(), b.clone()])
);
let VariablePresenceCondition::All(children) =
VariablePresenceCondition::all([c.clone(), a.clone()])
else {
panic!("expected an All node");
};
assert_eq!(children.iter().collect::<Vec<_>>(), vec![&a, &c]);
}
#[test]
fn test_canonical_grouping_and_repetition() {
let (a, b, c) = (present("a"), present("b"), present("c"));
assert_same(
VariablePresenceCondition::all([
a.clone(),
VariablePresenceCondition::all([b.clone(), c.clone()]),
]),
VariablePresenceCondition::all([
VariablePresenceCondition::all([c.clone(), a.clone()]),
b.clone(),
]),
);
assert_same(
VariablePresenceCondition::all([a.clone(), a.clone()]),
a.clone(),
);
assert_same(
VariablePresenceCondition::any([
VariablePresenceCondition::all([a.clone(), b.clone()]),
VariablePresenceCondition::all([b.clone(), a.clone()]),
]),
VariablePresenceCondition::all([a.clone(), b.clone()]),
);
}
/// A condition tree built without any normalization.
#[derive(Clone, Debug)]
enum RawCondition {
Always,
Never,
Present(usize),
All(Vec<RawCondition>),
Any(Vec<RawCondition>),
}
const RAW_VARIABLES: [&str; 4] = ["a", "b", "c", "d"];
impl RawCondition {
fn eval(&self, present_mask: u32) -> bool {
match self {
RawCondition::Always => true,
RawCondition::Never => false,
RawCondition::Present(variable_ord) => present_mask & (1 << variable_ord) != 0,
RawCondition::All(children) => {
children.iter().all(|child| child.eval(present_mask))
}
RawCondition::Any(children) => {
children.iter().any(|child| child.eval(present_mask))
}
}
}
/// Builds the canonical condition, visiting children in reverse order if `reverse`.
fn build(&self, reverse: bool) -> VariablePresenceCondition {
let build_children = |children: &[RawCondition]| {
let mut built: Vec<VariablePresenceCondition> =
children.iter().map(|child| child.build(reverse)).collect();
if reverse {
built.reverse();
}
built
};
match self {
RawCondition::Always => VariablePresenceCondition::Always,
RawCondition::Never => VariablePresenceCondition::Never,
RawCondition::Present(variable_ord) => present(RAW_VARIABLES[*variable_ord]),
RawCondition::All(children) => {
VariablePresenceCondition::all(build_children(children))
}
RawCondition::Any(children) => {
VariablePresenceCondition::any(build_children(children))
}
}
}
}
fn raw_conditions() -> impl Strategy<Value = RawCondition> {
let leaf = prop_oneof![
1 => Just(RawCondition::Always),
1 => Just(RawCondition::Never),
6 => (0..RAW_VARIABLES.len()).prop_map(RawCondition::Present),
];
leaf.prop_recursive(4, 32, 4, |inner| {
prop_oneof![
prop::collection::vec(inner.clone(), 0..4).prop_map(RawCondition::All),
prop::collection::vec(inner, 0..4).prop_map(RawCondition::Any),
]
})
}
/// Checks the hidden contract of `ConditionSet`, recursively.
fn assert_canonical(condition: &VariablePresenceCondition) {
let (children, is_all) = match condition {
VariablePresenceCondition::All(children) => (children, true),
VariablePresenceCondition::Any(children) => (children, false),
_ => return,
};
let children: Vec<&VariablePresenceCondition> = children.iter().collect();
assert!(children.len() >= 2, "{condition:?}");
assert!(
children.windows(2).all(|pair| pair[0] < pair[1]),
"{condition:?}"
);
for child in &children {
assert_canonical(child);
match (child, is_all) {
(VariablePresenceCondition::Always | VariablePresenceCondition::Never, _) => {
panic!("neutral or absorbing child in {condition:?}")
}
(VariablePresenceCondition::All(_), true)
| (VariablePresenceCondition::Any(_), false) => {
panic!("same-kind child in {condition:?}")
}
_ => {}
}
}
}
proptest! {
#[test]
fn proptest_canonical_form(raw in raw_conditions()) {
let condition = raw.build(false);
assert_canonical(&condition);
let reversed = raw.build(true);
prop_assert_eq!(&condition, &reversed);
prop_assert_eq!(hash_of(&condition), hash_of(&reversed));
for present_mask in 0..(1u32 << RAW_VARIABLES.len()) {
let mut is_present = |variable_name: &str| {
let variable_ord =
RAW_VARIABLES.iter().position(|name| *name == variable_name).unwrap();
present_mask & (1 << variable_ord) != 0
};
prop_assert_eq!(condition.eval(&mut is_present), raw.eval(present_mask));
}
}
}
fn presence_of<'a>(present_names: &'a [&'a str]) -> impl FnMut(&str) -> bool + 'a {
move |variable_name: &str| present_names.contains(&variable_name)
}
@@ -321,17 +518,17 @@ mod tests {
assert!(condition.eval(&mut presence_of(&["a", "c"])));
assert!(!condition.eval(&mut presence_of(&["a"])));
assert!(!condition.eval(&mut presence_of(&["b", "c"])));
assert!(PresenceCondition::Always.eval(&mut presence_of(&[])));
assert!(!PresenceCondition::Never.eval(&mut presence_of(&["a"])));
assert!(VariablePresenceCondition::Always.eval(&mut presence_of(&[])));
assert!(!VariablePresenceCondition::Never.eval(&mut presence_of(&["a"])));
}
#[test]
fn test_literals() {
assert_eq!(for_true("true"), PresenceCondition::Always);
assert_eq!(for_true("false"), PresenceCondition::Never);
assert_eq!(for_true("none"), PresenceCondition::Never);
assert_eq!(for_value("none"), PresenceCondition::Always);
assert_eq!(for_value("1u64"), PresenceCondition::Always);
assert_eq!(for_true("true"), VariablePresenceCondition::Always);
assert_eq!(for_true("false"), VariablePresenceCondition::Never);
assert_eq!(for_true("none"), VariablePresenceCondition::Never);
assert_eq!(for_value("none"), VariablePresenceCondition::Always);
assert_eq!(for_value("1u64"), VariablePresenceCondition::Always);
// `none` in a strict function is conservatively ignored.
assert_eq!(for_true("(EQ a none)"), present("a"));
}
@@ -353,11 +550,8 @@ mod tests {
for_value(r#"(CONCAT "," "true" (UPPER a) (SUBSTRING b 0i64 2i64))"#),
all(vec![present("a"), present("b")])
);
assert_eq!(
for_value(r#"(REGEXP_EXTRACT a "(x+)" 1u64)"#),
present("a")
);
assert_eq!(for_value("(ADD)"), PresenceCondition::Always);
assert_eq!(for_value(r#"(REGEXP_EXTRACT a "(x+)" 1u64)"#), present("a"));
assert_eq!(for_value("(ADD)"), VariablePresenceCondition::Always);
}
#[test]
@@ -375,10 +569,7 @@ mod tests {
all(vec![any(vec![present("a"), present("b")]), present("c")])
);
// AND is null as soon as one of its arguments is null.
assert_eq!(
for_value("(AND (NOT a) b)"),
present("b")
);
assert_eq!(for_value("(AND (NOT a) b)"), present("b"));
assert_eq!(
for_value("(OR (EQ a 1i64) (EQ b 2i64))"),
any(vec![present("a"), present("b")])
@@ -387,21 +578,30 @@ mod tests {
#[test]
fn test_null_tolerant_functions() {
assert_eq!(for_true("(NOT (EQ a 1i64))"), PresenceCondition::Always);
assert_eq!(for_true("(NEQ a 1i64)"), PresenceCondition::Always);
assert_eq!(for_true("(IS_NULL a)"), PresenceCondition::Always);
assert_eq!(for_value("(IS_NOT_NULL a)"), PresenceCondition::Always);
assert_eq!(for_true("(IS_NOT_NULL (ADD a b))"), all(vec![present("a"), present("b")]));
assert_eq!(for_value(r#"(REGEXP_LIKE a "x")"#), PresenceCondition::Always);
assert_eq!(
for_true("(NOT (EQ a 1i64))"),
VariablePresenceCondition::Always
);
assert_eq!(for_true("(NEQ a 1i64)"), VariablePresenceCondition::Always);
assert_eq!(for_true("(IS_NULL a)"), VariablePresenceCondition::Always);
assert_eq!(
for_value("(IS_NOT_NULL a)"),
VariablePresenceCondition::Always
);
assert_eq!(
for_true("(IS_NOT_NULL (ADD a b))"),
all(vec![present("a"), present("b")])
);
assert_eq!(
for_value(r#"(REGEXP_LIKE a "x")"#),
VariablePresenceCondition::Always
);
assert_eq!(for_true(r#"(REGEXP_LIKE a "x")"#), present("a"));
assert_eq!(
for_true("(OR (EQ a 1i64) (IS_NULL b))"),
PresenceCondition::Always
);
assert_eq!(
for_true("(AND (EQ a 1i64) (IS_NULL b))"),
present("a")
VariablePresenceCondition::Always
);
assert_eq!(for_true("(AND (EQ a 1i64) (IS_NULL b))"), present("a"));
}
#[test]
@@ -415,7 +615,10 @@ mod tests {
all(vec![present("c"), any(vec![present("a"), present("b")])])
);
assert_eq!(for_true("(IF c true false)"), present("c"));
assert_eq!(for_true("(IF c false false)"), PresenceCondition::Never);
assert_eq!(
for_true("(IF c false false)"),
VariablePresenceCondition::Never
);
assert_eq!(for_value("(IF c 1i64 a)"), present("c"));
}
@@ -591,7 +794,7 @@ mod tests {
expr_str: &str,
target_type: InferredTypeSet,
assignments: &[Assignment],
required: fn(&UntypedExpr) -> PresenceCondition,
required: fn(&UntypedExpr) -> VariablePresenceCondition,
holds: fn(VarType, VariableValue) -> bool,
) -> Result<(), TestCaseError> {
let expr = deserialize(expr_str).unwrap();
@@ -629,12 +832,12 @@ mod tests {
let args: Vec<VariableValue> = compiled_fn
.inputs()
.iter()
.map(|input| {
match assignment[variable_ord(&input.variable_name)] {
.map(
|input| match assignment[variable_ord(&input.variable_name)] {
Some(value_ord) => variable_value(input.r#type, value_ord),
None => VariableValue::none(),
}
})
},
)
.collect();
// SAFETY: Each slot follows the compiled input order, and uses the input type.
let result = unsafe { compiled_fn.call(&args, &mut string_arena) };
@@ -673,7 +876,11 @@ mod tests {
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(512))]
// Compiler debug assertions reject a fraction of the generated expressions.
#![proptest_config(ProptestConfig {
max_global_rejects: 1 << 16,
..ProptestConfig::with_cases(512)
})]
#[test]
fn proptest_required_presence_for_true_is_necessary(
src/query/all_query.rs
+1
View File
@@ -47,6 +47,7 @@ pub struct AllScorer {
impl AllScorer {
/// Creates a new AllScorer with `max_doc` docs.
pub fn new(max_doc: DocId) -> AllScorer {
let doc = if max_doc == 0u32 { TERMINATED } else { 0 };
AllScorer { doc: 0u32, max_doc }
}
}
src/query/doc_predicate_query/jitexpr_predicate.rs
+235 -8
View File
@@ -1,15 +1,21 @@
use std::collections::HashMap;
use std::io;
use columnar::{ColumnType, DynamicColumn, StrColumn};
use jitexpr::ast::{infer_types_with_target, InferredTypeSet, TypeError, UntypedExpr};
use columnar::{ColumnIndex, ColumnType, DynamicColumn, StrColumn};
use jitexpr::ast::{
infer_types_with_target, required_presence_for_true, InferredTypeSet, TypeError, UntypedExpr,
VariablePresenceCondition,
};
use jitexpr::compile::{CompiledFnCtx, StringArena};
use jitexpr::types::{VarType, VariableValue};
use super::{DocPredicate, SegmentDocPredicate};
use crate::index::SegmentReader;
use crate::query::doc_predicate_query::ConstOrVariableSegmentPredicate;
use crate::{DocId, TantivyError};
use crate::query::exist_query::{ExistsColumnIndex, ExistsDocSet};
use crate::query::union::SimpleUnion;
use crate::query::Intersection;
use crate::{DocId, DocSet, TantivyError};
/// A [`DocPredicate`] that evaluates a boolean JIT expression against fast fields.
///
@@ -25,6 +31,10 @@ use crate::{DocId, TantivyError};
///
/// Only a present `true` result matches.
///
/// Documents missing the fields required for the expression to be `true` are skipped without
/// being evaluated. For instance, `(EQ (ADD price 1u64) 10u64)` is only evaluated on the
/// documents having a `price` value.
///
/// ```
/// use tantivy::jitexpr::ast::deserialize;
/// use tantivy::query::doc_predicate_query::{DocPredicateQuery, JitExprPredicate};
@@ -36,6 +46,8 @@ use crate::{DocId, TantivyError};
pub struct JitExprPredicate {
expression: UntypedExpr,
inferred_inputs: Vec<(String, InferredTypeSet)>,
// A necessary condition, on the presence of the variables, for the expression to be `true`.
required_presence: VariablePresenceCondition,
}
impl JitExprPredicate {
@@ -47,9 +59,11 @@ impl JitExprPredicate {
.into_iter()
.map(|(name, types)| (name.to_string(), types))
.collect();
let required_presence = required_presence_for_true(&expression);
Ok(Self {
expression,
inferred_inputs,
required_presence,
})
}
@@ -91,6 +105,18 @@ impl DocPredicate for JitExprPredicate {
opened_columns.insert(name.as_str(), column);
}
// The variables are bound to the columns opened above, and only to them: the presence
// of a variable is the presence of a value in its column.
let necessary_condition: Option<Box<dyn DocSet>> = match build_necessary_condition_docset(
&self.required_presence,
&opened_columns,
segment_reader.max_doc(),
) {
ResolvedCondition::NoDocs => return Ok(ConstOrVariableSegmentPredicate::Const(false)),
ResolvedCondition::AllDocs => None,
ResolvedCondition::DocSet(necessary_condition) => Some(necessary_condition),
};
let compiled_fn = segment_reader
.index()
.expr_compilation_cache()
@@ -157,13 +183,87 @@ impl DocPredicate for JitExprPredicate {
.filter(|column_opt| matches!(column_opt, Some(DynamicColumn::Str(_))))
.count();
let num_inputs = columns_opt.len();
Ok(JitExprEvalState {
let predicate = JitExprEvalState {
compiled: compiled_fn.context(),
columns_opt,
string_inputs: vec![String::new(); num_string_inputs],
input_values: Vec::with_capacity(num_inputs),
};
Ok(ConstOrVariableSegmentPredicate::Variable {
predicate,
necessary_condition,
})
}
}
/// A [`PresenceCondition`] resolved against the columns of a segment.
enum ResolvedCondition {
/// Satisfied by all of the documents of the segment.
AllDocs,
/// Satisfied by none of the documents of the segment.
NoDocs,
/// Satisfied by the documents of the `DocSet`, positioned on its first document.
DocSet(Box<dyn DocSet>),
}
/// Builds a docset off the variable presence condition.
///
/// Variables are bound to the columns of `columns`. A variable missing from `columns` is null
/// for all documents.
fn build_necessary_condition_docset(
condition: &VariablePresenceCondition,
columns: &HashMap<&str, DynamicColumn>,
max_doc: DocId,
) -> ResolvedCondition {
match condition {
VariablePresenceCondition::Always => ResolvedCondition::AllDocs,
VariablePresenceCondition::Never => ResolvedCondition::NoDocs,
VariablePresenceCondition::Present(variable_name) => {
let Some(column) = columns.get(variable_name.as_ref()) else {
return ResolvedCondition::NoDocs;
};
let exists_column_index = match column.column_index() {
ColumnIndex::Empty { .. } => return ResolvedCondition::NoDocs,
ColumnIndex::Full => return ResolvedCondition::AllDocs,
ColumnIndex::Optional(optional_index) => {
ExistsColumnIndex::Optional(optional_index.clone())
}
ColumnIndex::Multivalued(multivalued_index) => {
ExistsColumnIndex::Multivalued(multivalued_index.clone())
}
};
ResolvedCondition::DocSet(Box::new(ExistsDocSet::new(exists_column_index)))
}
VariablePresenceCondition::All(conditions) => {
let mut doc_sets: Vec<Box<dyn DocSet>> = Vec::new();
for condition in conditions.iter() {
match build_necessary_condition_docset(condition, columns, max_doc) {
ResolvedCondition::AllDocs => {}
ResolvedCondition::NoDocs => return ResolvedCondition::NoDocs,
ResolvedCondition::DocSet(doc_set) => doc_sets.push(doc_set),
}
}
match doc_sets.len() {
0 => ResolvedCondition::AllDocs,
1 => ResolvedCondition::DocSet(doc_sets.pop().unwrap()),
_ => ResolvedCondition::DocSet(Box::new(Intersection::new(doc_sets, max_doc))),
}
}
VariablePresenceCondition::Any(conditions) => {
let mut doc_sets: Vec<Box<dyn DocSet>> = Vec::new();
for condition in conditions.iter() {
match build_necessary_condition_docset(condition, columns, max_doc) {
ResolvedCondition::AllDocs => return ResolvedCondition::AllDocs,
ResolvedCondition::NoDocs => {}
ResolvedCondition::DocSet(doc_set) => doc_sets.push(doc_set),
}
}
match doc_sets.len() {
0 => ResolvedCondition::NoDocs,
1 => ResolvedCondition::DocSet(doc_sets.pop().unwrap()),
_ => ResolvedCondition::DocSet(Box::new(SimpleUnion::build(doc_sets))),
}
}
.into())
}
}
@@ -317,10 +417,11 @@ fn load_str_input<'buffer>(
#[cfg(test)]
mod tests {
use super::*;
use crate::collector::Count;
use crate::collector::{Count, DocSetCollector};
use crate::query::doc_predicate_query::DocPredicateQuery;
use crate::schema::{Schema, FAST, STORED, STRING};
use crate::Index;
use crate::query::{EnableScoring, Query};
use crate::schema::{Schema, FAST, INDEXED, STORED, STRING};
use crate::{Index, TantivyDocument, Term};
fn create_index() -> Index {
let mut schema_builder = Schema::builder();
@@ -589,6 +690,132 @@ mod tests {
);
}
/// Two segments with sparse, multivalued, and segment-dependent columns, and deleted docs.
///
/// `label` only has values in the second segment.
fn create_sparse_index() -> Index {
let mut schema_builder = Schema::builder();
let id = schema_builder.add_u64_field("id", FAST | INDEXED);
let number = schema_builder.add_u64_field("number", FAST);
let score = schema_builder.add_i64_field("score", FAST);
let flag = schema_builder.add_bool_field("flag", FAST);
let label = schema_builder.add_text_field("label", STRING | FAST);
let tags = schema_builder.add_text_field("tags", STRING | FAST);
let index = Index::create_in_ram(schema_builder.build());
let mut writer = index.writer_for_tests().unwrap();
for segment_ord in 0..2u64 {
for i in 0..300u64 {
let mut doc = TantivyDocument::default();
doc.add_u64(id, segment_ord * 1000 + i);
if i % 3 == 0 {
doc.add_u64(number, i);
}
if i % 5 == 0 {
doc.add_i64(score, (i % 7) as i64 - 3);
}
if i % 2 == 0 {
doc.add_bool(flag, i % 4 == 0);
}
if segment_ord == 1 && i % 4 == 0 {
doc.add_text(label, ["a", "b", "ab"][(i % 3) as usize]);
}
if i % 6 == 0 {
doc.add_text(tags, "x");
doc.add_text(tags, "y");
} else if i % 6 == 1 {
doc.add_text(tags, "y");
}
writer.add_document(doc).unwrap();
}
writer.commit().unwrap();
}
writer.delete_term(Term::from_field_u64(id, 30));
writer.delete_term(Term::from_field_u64(id, 1060));
writer.commit().unwrap();
index
}
/// Returns a query matching the same documents as `query(expression)`, but requiring the
/// presence of no field, so that all documents are evaluated.
///
/// `(NOT true)` is a present `false`: the disjunction is `true` if and only if `expression` is.
/// As `NOT` requires the presence of no field, neither does the disjunction.
fn query_without_required_presence(expression: &str) -> DocPredicateQuery {
query(&format!("(OR {expression} (NOT true))"))
}
#[test]
fn test_required_presence_does_not_change_results() {
let index = create_sparse_index();
let searcher = index.reader().unwrap().searcher();
assert_eq!(searcher.segment_readers().len(), 2);
let expressions = [
"true",
"false",
"flag",
"(EQ number 33u64)",
"(EQ number 30u64)",
"(IS_NOT_NULL number)",
"(IS_NULL score)",
"(NOT (EQ number 3u64))",
"(NEQ score 1i64)",
"(GT (ADD number score) 10i64)",
"(OR (EQ number 3u64) (EQ score 1i64))",
"(AND flag (IS_NOT_NULL label))",
r#"(EQ label "a")"#,
r#"(OR (EQ label "b") flag)"#,
r#"(REGEXP_LIKE label "b")"#,
r#"(EQ (UPPER tags) "X")"#,
"(IF flag (GT number 100u64) (LT score 0i64))",
"(IS_NOT_NULL (IF flag number score))",
"(AND (EQ missing 1i64) flag)",
"(OR (IS_NOT_NULL missing) (EQ score 2i64))",
"(OR (IS_NULL missing) (EQ score 2i64))",
];
for expression in expressions {
let expected = searcher
.search(
&query_without_required_presence(expression),
&DocSetCollector,
)
.unwrap();
let accelerated = searcher
.search(&query(expression), &DocSetCollector)
.unwrap();
assert_eq!(accelerated, expected, "{expression}");
}
// Sanity checks: the index does exercise the predicates.
let count = |expression: &str| searcher.search(&query(expression), &Count).unwrap();
assert_eq!(count("(EQ number 33u64)"), 2);
// Doc 30 of the first segment is deleted.
assert_eq!(count("(EQ number 30u64)"), 1);
// `label` is "a" on 25 docs of the second segment, one of which (1060) is deleted.
assert_eq!(count(r#"(EQ label "a")"#), 24);
}
#[test]
fn test_required_presence_restricts_evaluated_docs() {
let index = create_sparse_index();
let searcher = index.reader().unwrap().searcher();
let segment_reader = searcher.segment_reader(0);
let size_hint = |query: DocPredicateQuery| {
query
.weight(EnableScoring::disabled_from_searcher(&searcher))
.unwrap()
.scorer(segment_reader, 1.0)
.unwrap()
.size_hint()
};
// `number` has a value in one doc out of three.
assert_eq!(size_hint(query("(GT number 10u64)")), 100);
assert_eq!(
size_hint(query_without_required_presence("(GT number 10u64)")),
300
);
// Nothing to require: all docs are evaluated.
assert_eq!(size_hint(query("(NOT (GT number 10u64))")), 300);
}
// THIS FAILS! due to our pick best possible column approach policy.
// #[test]
// fn test_multi_typed_field_picks_one() {
src/query/doc_predicate_query/mod.rs
+339 -53
View File
@@ -1,3 +1,4 @@
use std::cmp::Ordering;
use std::sync::Arc;
mod function_predicate;
@@ -12,7 +13,8 @@ use crate::docset::{SeekDangerResult, TERMINATED};
use crate::index::SegmentReader;
use crate::query::explanation::does_not_match;
use crate::query::{
AllWeight, ConstScorer, EmptyWeight, EnableScoring, Explanation, Query, Scorer, Weight,
AllScorer, AllWeight, ConstScorer, EmptyWeight, EnableScoring, Explanation, Query, Scorer,
Weight,
};
use crate::{DocId, DocSet, Score};
@@ -65,12 +67,82 @@ impl Weight for DocPredicateQuery {
}
}
/// A [`DocSet`] that walks documents by repeatedly evaluating a
/// [`SegmentDocPredicate`], starting from doc `0`.
/// The cost of a [`DocPredicateDocSet`] is the cost of its necessary condition,
/// multiplied by this factor.
const PREDICATE_EVAL_COST_FACTOR: u64 = 100;
/// A [`DocSet`] that walks the documents of a necessary condition, and evaluates a
/// [`SegmentDocPredicate`] on each of them.
///
/// Hidden contract: every document matching the predicate belongs to the necessary condition.
/// Documents outside of it are never evaluated, and are considered as not matching.
///
/// Hidden contract: whenever the `DocPredicateDocSet` is in a valid state, the necessary condition
/// is in a valid state too, positioned on `self.doc`.
pub struct DocPredicateDocSet<TSegmentDocPredicate> {
doc_predicate: TSegmentDocPredicate,
necessary_condition: Box<dyn DocSet>,
doc: DocId,
max_doc: DocId,
}
impl<TSegmentDocPredicate: SegmentDocPredicate> DocPredicateDocSet<TSegmentDocPredicate> {
/// Creates a `DocPredicateDocSet` positioned on its first matching document.
fn new(doc_predicate: TSegmentDocPredicate, necessary_condition: Box<dyn DocSet>) -> Self {
let first_candidate = necessary_condition.doc();
let mut doc_set = DocPredicateDocSet {
doc_predicate,
necessary_condition,
doc: first_candidate,
};
doc_set.find_match(first_candidate);
doc_set
}
/// Creates a `DocPredicateDocSet`, and seeks it to `target`, following
/// [`Weight::scorer_danger`]'s contract.
///
/// Documents before `target` are not evaluated.
fn new_seeked_to(
doc_predicate: TSegmentDocPredicate,
necessary_condition: Box<dyn DocSet>,
target: DocId,
) -> (SeekDangerResult, Self) {
let first_candidate = necessary_condition.doc();
let mut doc_set = DocPredicateDocSet {
doc_predicate,
necessary_condition,
doc: first_candidate,
};
let seek_result = match first_candidate.cmp(&target) {
Ordering::Less => doc_set.seek_danger(target),
Ordering::Equal => doc_set.eval_candidate(target),
Ordering::Greater => SeekDangerResult::SeekLowerBound(first_candidate),
};
(seek_result, doc_set)
}
/// Evaluates the predicate on `candidate`.
///
/// Hidden contract: the necessary condition is positioned on `candidate`.
fn eval_candidate(&mut self, candidate: DocId) -> SeekDangerResult {
if self.doc_predicate.eval(candidate) {
self.doc = candidate;
SeekDangerResult::Found
} else {
SeekDangerResult::SeekLowerBound(candidate + 1)
}
}
/// Advances to the first matching document at or after `candidate`.
///
/// Hidden contract: the necessary condition is positioned on `candidate`.
fn find_match(&mut self, mut candidate: DocId) -> DocId {
while candidate != TERMINATED && !self.doc_predicate.eval(candidate) {
candidate = self.necessary_condition.advance();
}
self.doc = candidate;
candidate
}
}
impl<TSegmentDocPredicate: SegmentDocPredicate> DocSet
@@ -80,27 +152,29 @@ impl<TSegmentDocPredicate: SegmentDocPredicate> DocSet
if self.doc == TERMINATED {
return TERMINATED;
}
self.find_match(self.doc + 1)
let candidate = self.necessary_condition.advance();
self.find_match(candidate)
}
fn seek(&mut self, target: DocId) -> DocId {
debug_assert!(target >= self.doc);
if self.doc == TERMINATED {
return TERMINATED;
// In a valid state, `self.doc` is a match (or TERMINATED).
if self.doc >= target {
return self.doc;
}
self.find_match(target)
let candidate = self.necessary_condition.seek(target);
self.find_match(candidate)
}
fn seek_danger(&mut self, target: DocId) -> SeekDangerResult {
if target >= self.max_doc {
self.doc = TERMINATED;
return SeekDangerResult::SeekLowerBound(TERMINATED);
}
if self.doc_predicate.eval(target) {
self.doc = target;
SeekDangerResult::Found
} else {
SeekDangerResult::SeekLowerBound(target + 1)
match self.necessary_condition.seek_danger(target) {
SeekDangerResult::Found => self.eval_candidate(target),
SeekDangerResult::SeekLowerBound(lower_bound) => {
if lower_bound == TERMINATED {
self.doc = TERMINATED;
}
SeekDangerResult::SeekLowerBound(lower_bound)
}
}
}
@@ -109,23 +183,13 @@ impl<TSegmentDocPredicate: SegmentDocPredicate> DocSet
}
fn size_hint(&self) -> u32 {
self.max_doc
self.necessary_condition.size_hint()
}
}
impl<TSegmentDocPredicate: SegmentDocPredicate> DocPredicateDocSet<TSegmentDocPredicate> {
fn find_match(&mut self, mut target: DocId) -> DocId {
loop {
match self.seek_danger(target) {
SeekDangerResult::Found => return target,
SeekDangerResult::SeekLowerBound(next_target) => {
if next_target >= TERMINATED {
return TERMINATED;
}
target = next_target;
}
}
}
fn cost(&self) -> u64 {
self.necessary_condition
.cost()
.saturating_mul(PREDICATE_EVAL_COST_FACTOR)
}
}
@@ -156,14 +220,14 @@ impl<TDocPredicate: DocPredicate> DocPredicateBoxable for TDocPredicate {
EmptyWeight.scorer(segment_reader, boost)
}
}
ConstOrVariableSegmentPredicate::Variable(doc_predicate) => {
let mut doc_set = DocPredicateDocSet {
doc_predicate,
doc: 0u32,
max_doc: segment_reader.max_doc(),
};
doc_set.doc = doc_set.find_match(0);
Ok(Box::new(ConstScorer::new(doc_set, boost)) as Box<dyn Scorer>)
ConstOrVariableSegmentPredicate::Variable {
predicate,
necessary_condition,
} => {
let necessary_condition = necessary_condition
.unwrap_or_else(|| Box::new(AllScorer::new(segment_reader.max_doc())));
let doc_set = DocPredicateDocSet::new(predicate, necessary_condition);
Ok(Box::new(ConstScorer::new(doc_set, boost)))
}
}
}
@@ -183,15 +247,15 @@ impl<TDocPredicate: DocPredicate> DocPredicateBoxable for TDocPredicate {
EmptyWeight.scorer_danger(segment_reader, target, boost)
}
}
ConstOrVariableSegmentPredicate::Variable(doc_predicate) => {
let mut doc_set = DocPredicateDocSet {
doc_predicate,
doc: target,
max_doc: segment_reader.max_doc(),
};
let seek_result = doc_set.seek_danger(target);
let scorer = Box::new(ConstScorer::new(doc_set, boost)) as Box<dyn Scorer>;
Ok((seek_result, scorer))
ConstOrVariableSegmentPredicate::Variable {
predicate,
necessary_condition,
} => {
let necessary_condition = necessary_condition
.unwrap_or_else(|| Box::new(AllScorer::new(segment_reader.max_doc())));
let (seek_result, doc_set) =
DocPredicateDocSet::new_seeked_to(predicate, necessary_condition, target);
Ok((seek_result, Box::new(ConstScorer::new(doc_set, boost))))
}
}
}
@@ -202,13 +266,27 @@ pub enum ConstOrVariableSegmentPredicate<P: SegmentDocPredicate> {
/// Can be emitted to hint that a predicate will be always true or false on a segment.
/// Returning Const instead of a variable is an optimization.
Const(bool),
/// Just a regular SegmentDocPredicate.
Variable(P),
/// A regular SegmentDocPredicate, evaluated document by document.
Variable {
/// The predicate to evaluate.
predicate: P,
/// An optional [`DocSet`] restricting the documents on which `predicate` is evaluated.
///
/// Hidden contract: it must contain every document for which `predicate.eval` returns
/// true. Documents outside of it are never evaluated, and are considered as not
/// matching. `None` means that every document of the segment must be evaluated.
///
/// The `DocSet` must be positioned on its first document.
necessary_condition: Option<Box<dyn DocSet>>,
},
}
impl<P: SegmentDocPredicate> From<P> for ConstOrVariableSegmentPredicate<P> {
fn from(predicate: P) -> Self {
ConstOrVariableSegmentPredicate::Variable(predicate)
ConstOrVariableSegmentPredicate::Variable {
predicate,
necessary_condition: None,
}
}
}
@@ -240,8 +318,13 @@ pub trait SegmentDocPredicate: Send + 'static {
#[cfg(test)]
pub(crate) mod tests {
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
use proptest::prelude::*;
use super::*;
use crate::collector::Count;
use crate::collector::{Count, DocSetCollector};
use crate::query::VecDocSet;
pub(crate) fn create_index_for_test(num_docs: u32) -> crate::Index {
let schema_builder = crate::schema::Schema::builder();
@@ -314,6 +397,209 @@ pub(crate) mod tests {
assert_eq!(scorer.doc(), 2);
}
/// Matches even doc ids, and counts its evaluations.
struct EvenDocIds {
num_evals: Arc<AtomicUsize>,
}
impl SegmentDocPredicate for EvenDocIds {
fn eval(&mut self, doc_id: DocId) -> bool {
self.num_evals.fetch_add(1, AtomicOrdering::Relaxed);
doc_id.is_multiple_of(2)
}
}
/// `EvenDocIds`, with a fixed necessary condition.
#[derive(Debug)]
struct EvenWithNecessaryCondition {
necessary_condition: Vec<DocId>,
num_evals: Arc<AtomicUsize>,
}
impl EvenWithNecessaryCondition {
fn new(necessary_condition: Vec<DocId>) -> Self {
EvenWithNecessaryCondition {
necessary_condition,
num_evals: Arc::default(),
}
}
}
impl DocPredicate for EvenWithNecessaryCondition {
type SegmentDocPredicate = EvenDocIds;
fn doc_predicate(
&self,
_segment_reader: &SegmentReader,
) -> crate::Result<ConstOrVariableSegmentPredicate<EvenDocIds>> {
Ok(ConstOrVariableSegmentPredicate::Variable {
predicate: EvenDocIds {
num_evals: self.num_evals.clone(),
},
necessary_condition: Some(Box::new(VecDocSet::from(
self.necessary_condition.clone(),
))),
})
}
}
#[test]
fn test_necessary_condition_restricts_evaluations() {
let index = create_index_for_test(10);
let searcher = index.reader().unwrap().searcher();
let predicate = EvenWithNecessaryCondition::new(vec![1, 2, 3, 4, 6, 9]);
let num_evals = predicate.num_evals.clone();
let query: DocPredicateQuery = predicate.into();
assert_eq!(searcher.search(&query, &DocSetCollector).unwrap().len(), 3);
assert_eq!(num_evals.load(AtomicOrdering::Relaxed), 6);
}
#[test]
fn test_necessary_condition_size_hint_and_cost() {
let index = create_index_for_test(10);
let searcher = index.reader().unwrap().searcher();
let query: DocPredicateQuery =
EvenWithNecessaryCondition::new(vec![1, 2, 3, 4, 6, 9]).into();
let weight = query
.weight(EnableScoring::disabled_from_searcher(&searcher))
.unwrap();
let scorer = weight.scorer(searcher.segment_reader(0), 1.0).unwrap();
assert_eq!(scorer.size_hint(), 6);
assert_eq!(scorer.cost(), 6 * PREDICATE_EVAL_COST_FACTOR);
// Without a necessary condition, all docs are candidates.
let scorer = even_doc_id_query()
.scorer(searcher.segment_reader(0), 1.0)
.unwrap();
assert_eq!(scorer.size_hint(), 10);
assert_eq!(scorer.cost(), 10 * PREDICATE_EVAL_COST_FACTOR);
}
#[test]
fn test_necessary_condition_scorer_danger() {
let index = create_index_for_test(10);
let searcher = index.reader().unwrap().searcher();
let segment_reader = searcher.segment_reader(0);
let scorer_danger = |necessary_condition: Vec<DocId>, target: DocId| {
let predicate = EvenWithNecessaryCondition::new(necessary_condition);
let num_evals = predicate.num_evals.clone();
let query: DocPredicateQuery = predicate.into();
let (seek_result, scorer) = query.scorer_danger(segment_reader, target, 1.0).unwrap();
(seek_result, scorer, num_evals.load(AtomicOrdering::Relaxed))
};
// The necessary condition starts after the target.
let (seek_result, mut scorer, num_evals) = scorer_danger(vec![4, 6], 1);
assert_eq!(seek_result, SeekDangerResult::SeekLowerBound(4));
assert_eq!(num_evals, 0);
assert_eq!(scorer.seek_danger(4), SeekDangerResult::Found);
assert_eq!(scorer.doc(), 4);
// The target is the first candidate, and matches.
let (seek_result, scorer, _) = scorer_danger(vec![2, 6], 2);
assert_eq!(seek_result, SeekDangerResult::Found);
assert_eq!(scorer.doc(), 2);
// The target is a candidate, but does not match.
let (seek_result, mut scorer, num_evals) = scorer_danger(vec![1, 3, 4], 3);
assert_eq!(seek_result, SeekDangerResult::SeekLowerBound(4));
assert_eq!(num_evals, 1);
assert_eq!(scorer.seek_danger(4), SeekDangerResult::Found);
// The target is not a candidate: it is not evaluated.
let (seek_result, _, num_evals) = scorer_danger(vec![1, 6], 2);
assert_eq!(seek_result, SeekDangerResult::SeekLowerBound(6));
assert_eq!(num_evals, 0);
// No match after the target. The lower bound can stop on a non-matching candidate.
let (seek_result, mut scorer, _) = scorer_danger(vec![1, 3], 2);
assert_eq!(seek_result, SeekDangerResult::SeekLowerBound(3));
assert_eq!(scorer.seek_danger(3), SeekDangerResult::SeekLowerBound(4));
assert_eq!(
scorer.seek_danger(4),
SeekDangerResult::SeekLowerBound(TERMINATED)
);
}
proptest! {
#[test]
fn proptest_necessary_condition_doc_set(
candidates in prop::collection::btree_set(0u32..200, 0..60),
modulo in 1u32..5,
targets in prop::collection::vec(0u32..220, 0..30),
advances in prop::collection::vec(any::<bool>(), 0..30),
) {
let candidates: Vec<DocId> = candidates.into_iter().collect();
let expected: Vec<DocId> = candidates
.iter()
.copied()
.filter(|doc| doc.is_multiple_of(modulo))
.collect();
let new_doc_set = || {
DocPredicateDocSet::new(
move |doc: DocId| doc.is_multiple_of(modulo),
Box::new(VecDocSet::from(candidates.clone())),
)
};
let first_match = |target: DocId| {
expected
.iter()
.copied()
.find(|doc| *doc >= target)
.unwrap_or(TERMINATED)
};
// advance
let mut doc_set = new_doc_set();
let mut matches: Vec<DocId> = Vec::new();
while doc_set.doc() != TERMINATED {
matches.push(doc_set.doc());
doc_set.advance();
}
prop_assert_eq!(&matches, &expected);
// interleaved seek and advance
let mut doc_set = new_doc_set();
for (target, advance) in targets.iter().zip(advances.iter()) {
let target = (*target).max(doc_set.doc());
if *advance && doc_set.doc() != TERMINATED {
let current = doc_set.doc();
prop_assert_eq!(doc_set.advance(), first_match(current + 1));
} else {
prop_assert_eq!(doc_set.seek(target), first_match(target));
}
}
// seek_danger, following its contract: strictly increasing targets, respecting
// the returned lower bounds.
let mut sorted_targets = targets.clone();
sorted_targets.sort_unstable();
sorted_targets.dedup();
let mut doc_set = new_doc_set();
let mut lower_bound = doc_set.doc();
let mut previous_target = None;
for requested_target in sorted_targets {
let target = requested_target.max(lower_bound);
if previous_target.is_some_and(|previous| previous >= target) {
continue;
}
previous_target = Some(target);
let next_match = first_match(target);
match doc_set.seek_danger(target) {
SeekDangerResult::Found => {
prop_assert_eq!(next_match, target);
prop_assert_eq!(doc_set.doc(), target);
}
SeekDangerResult::SeekLowerBound(bound) => {
prop_assert!(next_match != target || target == TERMINATED);
prop_assert!(bound > target || target == TERMINATED);
prop_assert!(bound <= next_match);
lower_bound = bound;
}
}
}
}
}
#[test]
fn test_doc_predicate_query_scorer_danger_target_past_max_doc() {
let index = create_index_for_test(4);
src/query/exist_query.rs
+1 -1
View File
@@ -182,7 +182,7 @@ impl Weight for FastFieldExistsWeight {
}
}
enum ExistsColumnIndex {
pub(crate) enum ExistsColumnIndex {
Optional(OptionalIndex),
Multivalued(MultiValueIndex),
}