diff --git a/src/query/boolean_query/mod.rs b/src/query/boolean_query/mod.rs index f03ca054f..516a46ec7 100644 --- a/src/query/boolean_query/mod.rs +++ b/src/query/boolean_query/mod.rs @@ -375,7 +375,16 @@ mod proptest_boolean_query { } } - fn create_index_with_boolean_permutations(num_fields: usize) -> (Index, Vec) { + fn doc_ids(num_docs: usize, num_fields: usize) -> impl Iterator { + let permutations = 1 << num_fields; + let copies = (num_docs as f32 / permutations as f32).ceil() as u32; + (0..(permutations * copies)).into_iter() + } + + fn create_index_with_boolean_permutations( + num_docs: usize, + num_fields: usize, + ) -> (Index, Vec) { let mut schema_builder = Schema::builder(); let fields: Vec = (0..num_fields) .map(|i| schema_builder.add_text_field(&format!("field_{}", i), TEXT)) @@ -384,10 +393,10 @@ mod proptest_boolean_query { let index = Index::create_in_ram(schema); let mut writer = index.writer_for_tests().unwrap(); - for doc_id in 0..(1 << num_fields) { + for doc_id in doc_ids(num_docs, num_fields) { let mut doc: BTreeMap<_, OwnedValue> = BTreeMap::default(); for (field_idx, &field) in fields.iter().enumerate() { - if (doc_id >> field_idx) & 1 == 1 { + if BooleanQueryAST::matches_field(doc_id, field_idx) { doc.insert(field, "true".into()); } } @@ -414,14 +423,15 @@ mod proptest_boolean_query { #[test] fn proptest_boolean_query() { + let num_docs = 10000; let num_fields = 8; - let (index, fields) = create_index_with_boolean_permutations(num_fields); + let (index, fields) = create_index_with_boolean_permutations(num_docs, num_fields); let searcher = index.reader().unwrap().searcher(); proptest!(|(ast in arb_boolean_query_ast(num_fields))| { let query = ast.to_query(&fields); let mut matching_docs = HashSet::new(); - for doc_id in 0..(1 << num_fields) { + for doc_id in doc_ids(num_docs, num_fields) { if ast.matches(doc_id as DocId) { matching_docs.insert(doc_id as DocId); }