Add test to search splits contained in time range

Author: tontinton

quickwit/quickwit-indexing/src/test_utils.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::num::NonZeroUsize;
+use std::ops::RangeInclusive;
 use std::str::FromStr;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -251,8 +252,16 @@ pub struct MockSplitBuilder {
 
 impl MockSplitBuilder {
     pub fn new(split_id: &str) -> Self {
+        Self::new_with_time_range(split_id, Some(121000..=130198))
+    }
+
+    pub fn new_with_time_range(split_id: &str, time_range: Option<RangeInclusive<i64>>) -> Self {
         Self {
-            split_metadata: mock_split_meta(split_id, &IndexUid::for_test("test-index", 0)),
+            split_metadata: mock_split_meta(
+                split_id,
+                &IndexUid::for_test("test-index", 0),
+                time_range,
+            ),
         }
     }
 
@@ -277,14 +286,18 @@ pub fn mock_split(split_id: &str) -> Split {
 }
 
 /// Mock split meta helper.
-pub fn mock_split_meta(split_id: &str, index_uid: &IndexUid) -> SplitMetadata {
+pub fn mock_split_meta(
+    split_id: &str,
+    index_uid: &IndexUid,
+    time_range: Option<RangeInclusive<i64>>,
+) -> SplitMetadata {
     SplitMetadata {
         index_uid: index_uid.clone(),
         split_id: split_id.to_string(),
         partition_id: 13u64,
         num_docs: if split_id == "split1" { 1_000_000 } else { 10 },
         uncompressed_docs_size_in_bytes: 256,
-        time_range: Some(121000..=130198),
+        time_range,
         create_timestamp: 0,
         footer_offsets: 700..800,
         ..Default::default()

quickwit/quickwit-search/src/root.rs

@@ -5076,4 +5076,112 @@ mod tests {
         assert_eq!(search_response.failed_splits.len(), 1);
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_count_from_metastore_in_contained_time_range() -> anyhow::Result<()> {
+        let search_request = quickwit_proto::search::SearchRequest {
+            start_timestamp: Some(122_000),
+            end_timestamp: Some(129_000),
+            index_id_patterns: vec!["test-index".to_string()],
+            query_ast: serde_json::to_string(&QueryAst::MatchAll)
+                .expect("MatchAll should be JSON serializable."),
+            max_hits: 0,
+            ..Default::default()
+        };
+
+        let index_metadata = IndexMetadata::for_test("test-index", "ram:///test-index");
+        let index_uid = index_metadata.index_uid.clone();
+
+        let mut mock_metastore = MockMetastoreService::new();
+        mock_metastore
+            .expect_list_indexes_metadata()
+            .returning(move |_q| {
+                Ok(ListIndexesMetadataResponse::for_test(vec![
+                    index_metadata.clone(),
+                ]))
+            });
+        mock_metastore.expect_list_splits().returning(move |_req| {
+            let splits = vec![
+                MockSplitBuilder::new_with_time_range("split_before", Some(100_000..=110_000))
+                    .with_index_uid(&index_uid)
+                    .build(),
+                MockSplitBuilder::new_with_time_range(
+                    "split_overlap_start",
+                    Some(120_000..=123_000),
+                )
+                .with_index_uid(&index_uid)
+                .build(),
+                MockSplitBuilder::new_with_time_range("split_overlap_end", Some(128_000..=140_000))
+                    .with_index_uid(&index_uid)
+                    .build(),
+                MockSplitBuilder::new_with_time_range(
+                    "split_covering_whole",
+                    Some(100_000..=200_000),
+                )
+                .with_index_uid(&index_uid)
+                .build(),
+                MockSplitBuilder::new_with_time_range("split_inside", Some(124_000..=126_000))
+                    .with_index_uid(&index_uid)
+                    .build(),
+            ];
+            let resp = ListSplitsResponse::try_from_splits(splits).unwrap();
+            Ok(ServiceStream::from(vec![Ok(resp)]))
+        });
+
+        let mut mock_search = MockSearchService::new();
+        mock_search
+            .expect_leaf_search()
+            .withf(|leaf_search_req| {
+                let mut expected = HashSet::new();
+                expected.insert("split_before");
+                expected.insert("split_covering_whole");
+                expected.insert("split_overlap_end");
+                expected.insert("split_overlap_start");
+
+                leaf_search_req.leaf_requests.len() == 1
+                    && leaf_search_req.leaf_requests[0]
+                        .split_offsets
+                        .iter()
+                        .map(|s| s.split_id.as_str())
+                        .collect::<HashSet<&str>>()
+                        == expected
+            })
+            .times(1)
+            .returning(|req| {
+                Ok(quickwit_proto::search::LeafSearchResponse {
+                    num_hits: req.leaf_requests[0]
+                        .split_offsets
+                        .iter()
+                        .map(|s| s.num_docs)
+                        .sum(),
+                    partial_hits: vec![],
+                    failed_splits: Vec::new(),
+                    num_attempted_splits: 0,
+                    ..Default::default()
+                })
+            });
+        mock_search.expect_fetch_docs().returning(|fetch_req| {
+            Ok(quickwit_proto::search::FetchDocsResponse {
+                hits: get_doc_for_fetch_req(fetch_req),
+            })
+        });
+
+        let searcher_pool = searcher_pool_for_test([("127.0.0.1:1001", mock_search)]);
+        let search_job_placer = SearchJobPlacer::new(searcher_pool);
+        let cluster_client = ClusterClient::new(search_job_placer);
+
+        let ctx = SearcherContext::for_test();
+        let resp = root_search(
+            &ctx,
+            search_request,
+            MetastoreServiceClient::from_mock(mock_metastore),
+            &cluster_client,
+        )
+        .await?;
+
+        assert_eq!(resp.num_hits, 50);
+        assert_eq!(resp.hits.len(), 0);
+
+        Ok(())
+    }
 }