diff --git a/jitexpr/src/ast/mod.rs b/jitexpr/src/ast/mod.rs index 14f66f503..de61743c1 100644 --- a/jitexpr/src/ast/mod.rs +++ b/jitexpr/src/ast/mod.rs @@ -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; diff --git a/jitexpr/src/ast/presence.rs b/jitexpr/src/ast/presence.rs index 279687a0c..b9bdcc395 100644 --- a/jitexpr/src/ast/presence.rs +++ b/jitexpr/src/ast/presence.rs @@ -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), /// Satisfied when all of the conditions are satisfied. - All(Vec), + All(ConditionSet), /// Satisfied when at least one of the conditions is satisfied. - Any(Vec), + Any(ConditionSet), } -impl PresenceCondition { - /// Builds the simplified conjunction of `conditions`. - pub fn all(conditions: impl IntoIterator) -> PresenceCondition { - let mut children: Vec = 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); + +impl ConditionSet { + /// Returns the children, in canonical order. + pub fn iter(&self) -> impl Iterator { + self.0.iter() + } +} + +impl VariablePresenceCondition { + /// Builds the canonical conjunction of `conditions`. + pub fn all( + conditions: impl IntoIterator, + ) -> VariablePresenceCondition { + let mut children: Vec = 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) -> PresenceCondition { - let mut children: Vec = Vec::new(); + /// Builds the canonical disjunction of `conditions`. + pub fn any( + conditions: impl IntoIterator, + ) -> VariablePresenceCondition { + let mut children: Vec = 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, 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::All(conditions) + fn all(conditions: Vec) -> VariablePresenceCondition { + VariablePresenceCondition::all(conditions) } - fn any(conditions: Vec) -> PresenceCondition { - PresenceCondition::Any(conditions) + fn any(conditions: Vec) -> 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![&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), + Any(Vec), + } + + 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 = + 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 { + 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 = 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( diff --git a/src/query/all_query.rs b/src/query/all_query.rs index 5431a3a1b..18d411b18 100644 --- a/src/query/all_query.rs +++ b/src/query/all_query.rs @@ -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 } } } diff --git a/src/query/doc_predicate_query/jitexpr_predicate.rs b/src/query/doc_predicate_query/jitexpr_predicate.rs index 94023ddcf..659940ac5 100644 --- a/src/query/doc_predicate_query/jitexpr_predicate.rs +++ b/src/query/doc_predicate_query/jitexpr_predicate.rs @@ -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> = 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), +} + +/// 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> = 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> = 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() { diff --git a/src/query/doc_predicate_query/mod.rs b/src/query/doc_predicate_query/mod.rs index 11e79dfd9..621f4627e 100644 --- a/src/query/doc_predicate_query/mod.rs +++ b/src/query/doc_predicate_query/mod.rs @@ -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 { doc_predicate: TSegmentDocPredicate, + necessary_condition: Box, doc: DocId, - max_doc: DocId, +} + +impl DocPredicateDocSet { + /// Creates a `DocPredicateDocSet` positioned on its first matching document. + fn new(doc_predicate: TSegmentDocPredicate, necessary_condition: Box) -> 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, + 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 DocSet @@ -80,27 +152,29 @@ impl 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 DocSet } fn size_hint(&self) -> u32 { - self.max_doc + self.necessary_condition.size_hint() } -} -impl DocPredicateDocSet { - 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 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) + 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 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; - 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 { /// 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>, + }, } impl From

for ConstOrVariableSegmentPredicate

{ 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, + } + + 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, + num_evals: Arc, + } + + impl EvenWithNecessaryCondition { + fn new(necessary_condition: Vec) -> Self { + EvenWithNecessaryCondition { + necessary_condition, + num_evals: Arc::default(), + } + } + } + + impl DocPredicate for EvenWithNecessaryCondition { + type SegmentDocPredicate = EvenDocIds; + + fn doc_predicate( + &self, + _segment_reader: &SegmentReader, + ) -> crate::Result> { + 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, 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::(), 0..30), + ) { + let candidates: Vec = candidates.into_iter().collect(); + let expected: Vec = 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 = 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); diff --git a/src/query/exist_query.rs b/src/query/exist_query.rs index fcda85fff..a0121fbe9 100644 --- a/src/query/exist_query.rs +++ b/src/query/exist_query.rs @@ -182,7 +182,7 @@ impl Weight for FastFieldExistsWeight { } } -enum ExistsColumnIndex { +pub(crate) enum ExistsColumnIndex { Optional(OptionalIndex), Multivalued(MultiValueIndex), }