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,12 @@ 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 +282,14 @@ 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,89 @@ mod tests {
         assert_eq!(search_response.failed_splits.len(), 1);
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_search_partial_hits_phase_filters_by_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: qast_json_helper("test", &["body"]),
+            max_hits: 10,
+            ..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().times(1).returning(|_req| {
+            Ok(quickwit_proto::search::LeafSearchResponse {
+                num_hits: 1,
+                partial_hits: vec![mock_partial_hit("split_inside", 1, 1)],
+                failed_splits: Vec::new(),
+                num_attempted_splits: 1,
+                ..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, 1);
+        assert_eq!(resp.hits.len(), 1);
+
+        Ok(())
+    }
 }