From ee4538d6c2e92eeac39afddf7f5c5d7e1cdb7cfb Mon Sep 17 00:00:00 2001 From: Stu Hood Date: Sat, 27 Dec 2025 16:04:24 -0700 Subject: [PATCH] `test_order_by_u64_prop` --- src/collector/sort_key/mod.rs | 82 +++++++++++++++++++++++++++++++++++ 1 file changed, 82 insertions(+) diff --git a/src/collector/sort_key/mod.rs b/src/collector/sort_key/mod.rs index 7a478877d..b89059809 100644 --- a/src/collector/sort_key/mod.rs +++ b/src/collector/sort_key/mod.rs @@ -594,4 +594,86 @@ mod tests { prop_assert_eq!(actual_doc_ids, expected_doc_ids); } } + + proptest! { + #[test] + fn test_order_by_u64_prop( + order in prop_oneof!(Just(Order::Desc), Just(Order::Asc)), + limit in 1..20_usize, + offset in 0..20_usize, + segments_data in proptest::collection::vec( + proptest::collection::vec( + proptest::option::of(0..100u64), + 1..10_usize // segment size + ), + 1..4_usize // num segments + ) + ) { + use crate::collector::sort_key::ComparatorEnum; + use crate::TantivyDocument; + + let mut schema_builder = Schema::builder(); + let field = schema_builder.add_u64_field("field", FAST); + let schema = schema_builder.build(); + let index = Index::create_in_ram(schema); + let mut index_writer = index.writer_for_tests().unwrap(); + + for segment_data in segments_data.into_iter() { + for val in segment_data.into_iter() { + let mut doc = TantivyDocument::default(); + if let Some(v) = val { + doc.add_u64(field, v); + } + index_writer.add_document(doc).unwrap(); + } + index_writer.commit().unwrap(); + } + + let searcher = index.reader().unwrap().searcher(); + + let top_collector = TopDocs::with_limit(limit) + .and_offset(offset) + .order_by((SortByStaticFastValue::::for_field("field"), order)); + + let actual_results = searcher.search(&AllQuery, &top_collector).unwrap(); + let actual_doc_ids: Vec = + actual_results.into_iter().map(|(_, doc)| doc).collect(); + + // Verification logic + let all_docs_collector = DocSetCollector; + let all_docs = searcher.search(&AllQuery, &all_docs_collector).unwrap(); + + let docs_with_keys: Vec<(Option, DocAddress)> = all_docs + .into_iter() + .map(|doc_addr| { + let reader = searcher.segment_reader(doc_addr.segment_ord); + let val = if let Some((col, _)) = reader.fast_fields().u64_lenient("field").unwrap() { + col.first(doc_addr.doc_id) + } else { + None + }; + (val, doc_addr) + }) + .collect(); + + let comparator = ComparatorEnum::from(order); + let mut comparable_docs: Vec> = docs_with_keys + .into_iter() + .map(|(sort_key, doc)| ComparableDoc { sort_key, doc }) + .collect(); + + comparable_docs.sort_by(|l, r| comparator.compare_doc(l, r)); + + let expected_results = comparable_docs + .into_iter() + .skip(offset) + .take(limit) + .collect::>(); + + let expected_doc_ids: Vec = + expected_results.into_iter().map(|cd| cd.doc).collect(); + + prop_assert_eq!(actual_doc_ids, expected_doc_ids); + } + } }