add timeout to individual group by requests, non-decreasing

qdrant · Nov 2, 2023 · 10102c9 · 10102c9
1 parent 10a2eb1
commit 10102c9
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 9 deletions.
diff --git a/lib/collection/src/grouping/builder.rs b/lib/collection/src/grouping/builder.rs
@@ -81,6 +81,7 @@ where
             self.collection_by_name.clone(),
             self.read_consistency,
             self.shard_selection,
+            self.timeout,
         )
         .await?;
 

diff --git a/lib/collection/src/grouping/group_by.rs b/lib/collection/src/grouping/group_by.rs
@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 use std::future::Future;
+use std::time::Duration;
 
 use itertools::Itertools;
 use segment::data_types::vectors::{Named, DEFAULT_VECTOR_NAME};
@@ -132,6 +133,7 @@ impl GroupRequest {
         collection_by_name: F,
         read_consistency: Option<ReadConsistency>,
         shard_selection: Option<ShardId>,
+        timeout: Option<Duration>,
     ) -> CollectionResult<Vec<ScoredPoint>>
     where
         F: Fn(String) -> Fut,
@@ -154,13 +156,7 @@ impl GroupRequest {
                 request.with_vector = None;
 
                 collection
-                    .search(
-                        request,
-                        read_consistency,
-                        shard_selection,
-                        // Timeout is handled by dropping this Future
-                        None,
-                    )
+                    .search(request, read_consistency, shard_selection, timeout)
                     .await
             }
             SourceRequest::Recommend(mut request) => {
@@ -177,8 +173,7 @@ impl GroupRequest {
                     collection,
                     collection_by_name,
                     read_consistency,
-                    // Timeout is handled by dropping this Future
-                    None,
+                    timeout,
                 )
                 .await
             }
@@ -280,6 +275,7 @@ pub async fn group_by<'a, F, Fut>(
     collection_by_name: F,
     read_consistency: Option<ReadConsistency>,
     shard_selection: Option<ShardId>,
+    timeout: Option<Duration>,
 ) -> CollectionResult<Vec<PointGroup>>
 where
     F: Fn(String) -> Fut + Clone,
@@ -333,6 +329,7 @@ where
                 collection_by_name.clone(),
                 read_consistency,
                 shard_selection,
+                timeout,
             )
             .await?;
 
@@ -381,6 +378,7 @@ where
                     collection_by_name.clone(),
                     read_consistency,
                     shard_selection,
+                    timeout,
                 )
                 .await?;
 

diff --git a/lib/collection/tests/integration/grouping_test.rs b/lib/collection/tests/integration/grouping_test.rs
@@ -99,6 +99,7 @@ mod group_by {
             |_name| async { unreachable!() },
             resources.read_consistency,
             resources.shard_selection,
+            None,
         )
         .await;
 
@@ -154,6 +155,7 @@ mod group_by {
             |_name| async { unreachable!() },
             resources.read_consistency,
             resources.shard_selection,
+            None,
         )
         .await;
 
@@ -217,6 +219,7 @@ mod group_by {
             |_name| async { unreachable!() },
             resources.read_consistency,
             resources.shard_selection,
+            None,
         )
         .await;
 
@@ -252,6 +255,7 @@ mod group_by {
             |_name| async { unreachable!() },
             resources.read_consistency,
             resources.shard_selection,
+            None,
         )
         .await;
 
@@ -298,6 +302,7 @@ mod group_by {
             |_name| async { unreachable!() },
             read_consistency,
             shard_selection,
+            None,
         )
         .await;
 
@@ -342,6 +347,7 @@ mod group_by {
             |_name| async { unreachable!() },
             read_consistency,
             shard_selection,
+            None,
         )
         .await;
 
@@ -382,6 +388,7 @@ mod group_by {
             |_name| async { unreachable!() },
             read_consistency,
             shard_selection,
+            None,
         )
         .await;
 
@@ -422,6 +429,7 @@ mod group_by {
             |_name| async { unreachable!() },
             read_consistency,
             shard_selection,
+            None,
         )
         .await;
 
@@ -466,6 +474,7 @@ mod group_by {
             |_name| async { unreachable!() },
             read_consistency,
             shard_selection,
+            None,
         )
         .await;