test_order_by_u64_prop

src/collector/sort_key/mod.rs

@@ -609,4 +609,86 @@ pub(crate) 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::<u64>::for_field("field"), order));
+
+        let actual_results = searcher.search(&AllQuery, &top_collector).unwrap();
+        let actual_doc_ids: Vec<DocAddress> =
+            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<u64>, 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<ComparableDoc<_, _>> = 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::<Vec<_>>();
+
+        let expected_doc_ids: Vec<DocAddress> =
+            expected_results.into_iter().map(|cd| cd.doc).collect();
+
+        prop_assert_eq!(actual_doc_ids, expected_doc_ids);
+    }
+    }
 }