From b60cc4b7ead3cd20252be38249f5414e61743f3a Mon Sep 17 00:00:00 2001
From: trinity-1686a <trinity@quickwit.io>
Date: Mon, 6 May 2024 09:59:23 +0200
Subject: [PATCH] fix timestamp rewrite on boolean should (#4936)

* fix rewriting of timestamp in presence of onnly should sub-ast

* add test
---
 quickwit/quickwit-search/src/leaf.rs | 70 ++++++++++++++++++++++++++--
 1 file changed, 66 insertions(+), 4 deletions(-)

diff --git a/quickwit/quickwit-search/src/leaf.rs b/quickwit/quickwit-search/src/leaf.rs
index 87a4484151e..f839e206bb5 100644
--- a/quickwit/quickwit-search/src/leaf.rs
+++ b/quickwit/quickwit-search/src/leaf.rs
@@ -560,15 +560,30 @@ fn remove_redundant_timestamp_range(
                 bound.into_timestamp_nanos().into()
             }),
         };
-        if let QueryAst::Bool(bool_query) = &mut new_ast {
-            bool_query.filter.push(range.into());
+        new_ast = if let QueryAst::Bool(mut bool_query) = new_ast {
+            if bool_query.must.is_empty()
+                && bool_query.filter.is_empty()
+                && !bool_query.should.is_empty()
+            {
+                // we can't simply add a filter if we have some should but no must/filter. We must
+                // add a new layer of bool query
+                BoolQuery {
+                    must: vec![bool_query.into()],
+                    filter: vec![range.into()],
+                    ..Default::default()
+                }
+                .into()
+            } else {
+                bool_query.filter.push(range.into());
+                QueryAst::Bool(bool_query)
+            }
         } else {
-            new_ast = BoolQuery {
+            BoolQuery {
                 must: vec![new_ast],
                 filter: vec![range.into()],
                 ..Default::default()
             }
-            .into();
+            .into()
         }
     }
 
@@ -1199,4 +1214,51 @@ mod tests {
         remove_redundant_timestamp_range(&mut search_request, &split, &timestamp_field);
         assert_ast_eq(&search_request, &QueryAst::MatchAll);
     }
+
+    // regression test for #4935
+    #[test]
+    fn test_remove_timestamp_range_keep_should() {
+        let time1 = 1700001000;
+        let time2 = 1700002000;
+        let time3 = 1700003000;
+
+        let timestamp_field = "timestamp".to_string();
+
+        // cases where the bounds are larger than the split: no bound is emited
+        let split = SplitIdAndFooterOffsets {
+            timestamp_start: Some(time1),
+            timestamp_end: Some(time3),
+            ..SplitIdAndFooterOffsets::default()
+        };
+
+        let mut search_request = SearchRequest {
+            query_ast: serde_json::to_string(&QueryAst::Bool(BoolQuery {
+                should: vec![QueryAst::MatchAll],
+                ..BoolQuery::default()
+            }))
+            .unwrap(),
+            start_timestamp: Some(time2),
+            end_timestamp: None,
+            ..SearchRequest::default()
+        };
+        remove_redundant_timestamp_range(&mut search_request, &split, &timestamp_field);
+        assert_ast_eq(
+            &search_request,
+            &QueryAst::Bool(BoolQuery {
+                // original request
+                must: vec![QueryAst::Bool(BoolQuery {
+                    should: vec![QueryAst::MatchAll],
+                    ..BoolQuery::default()
+                })],
+                // time bound
+                filter: vec![RangeQuery {
+                    field: "timestamp".to_string(),
+                    lower_bound: Bound::Included(1_700_002_000_000_000_000u64.into()),
+                    upper_bound: Bound::Unbounded,
+                }
+                .into()],
+                ..BoolQuery::default()
+            }),
+        );
+    }
 }