feat: use --shard-key in search, query, and scroll requests

src/query.rs

@@ -5,6 +5,7 @@ use std::sync::atomic::AtomicBool;
 use anyhow::Result;
 use qdrant_client::Qdrant;
 use qdrant_client::qdrant::ScrollPointsBuilder;
+use qdrant_client::qdrant::shard_key::Key;
 
 use crate::args::Args;
 use crate::client::create_clients;
@@ -71,6 +72,9 @@ async fn get_uuids(args: &Args, client: &Qdrant) -> Result<Vec<String>> {
     let mut scroll_builder = ScrollPointsBuilder::new(&args.collection_name)
         .with_payload(true)
         .limit(args.num_vectors_or_default() as u32);
+    if let Some(shard_key) = &args.shard_key {
+        scroll_builder = scroll_builder.shard_key_selector(vec![Key::Keyword(shard_key.clone())]);
+    }
     if let Some(timeout) = args.timeout {
         scroll_builder = scroll_builder.timeout(timeout as u64);
     }

src/scroll.rs

@@ -3,6 +3,7 @@ use std::sync::{Arc, Mutex};
 
 use indicatif::ProgressBar;
 use qdrant_client::Qdrant;
+use qdrant_client::qdrant::shard_key::Key;
 use qdrant_client::qdrant::{PointId, Query, QueryPointsBuilder, Sample, ScrollPointsBuilder};
 
 use rand::{Rng, RngExt};
@@ -174,6 +175,11 @@ impl ScrollProcessor {
             request_builder = request_builder.read_consistency(read_consistency);
         }
 
+        if let Some(shard_key) = &self.args.shard_key {
+            request_builder =
+                request_builder.shard_key_selector(vec![Key::Keyword(shard_key.clone())]);
+        }
+
         if let Some(timeout) = self.args.timeout {
             request_builder = request_builder.timeout(timeout as u64);
         }
@@ -205,6 +211,11 @@ impl ScrollProcessor {
             request_builder = request_builder.read_consistency(read_consistency);
         }
 
+        if let Some(shard_key) = &self.args.shard_key {
+            request_builder =
+                request_builder.shard_key_selector(vec![Key::Keyword(shard_key.clone())]);
+        }
+
         if let Some(timeout) = self.args.timeout {
             request_builder = request_builder.timeout(timeout as u64);
         }
@@ -236,6 +247,11 @@ impl ScrollProcessor {
             request_builder = request_builder.read_consistency(read_consistency);
         }
 
+        if let Some(shard_key) = &self.args.shard_key {
+            request_builder =
+                request_builder.shard_key_selector(vec![Key::Keyword(shard_key.clone())]);
+        }
+
         if let Some(timeout) = self.args.timeout {
             request_builder = request_builder.timeout(timeout as u64);
         }

src/search/from_args.rs

@@ -5,6 +5,7 @@ use std::sync::{Arc, Mutex};
 
 use indicatif::ProgressBar;
 use qdrant_client::Qdrant;
+use qdrant_client::qdrant::shard_key::Key;
 use qdrant_client::qdrant::{
     PrefetchQueryBuilder, QuantizationSearchParamsBuilder, Query, QueryBatchPointsBuilder,
     QueryPointsBuilder, SearchParamsBuilder, SparseIndices, VectorInput,
@@ -166,6 +167,11 @@ impl SearchProcessor {
             request_builder = request_builder.read_consistency(read_consistency);
         }
 
+        if let Some(shard_key) = &self.args.shard_key {
+            request_builder =
+                request_builder.shard_key_selector(vec![Key::Keyword(shard_key.clone())]);
+        }
+
         request_builder
     }
 

src/search/from_config.rs

@@ -5,6 +5,7 @@ use std::sync::{Arc, Mutex};
 
 use indicatif::ProgressBar;
 use qdrant_client::Qdrant;
+use qdrant_client::qdrant::shard_key::Key;
 use qdrant_client::qdrant::{
     PrefetchQueryBuilder, QuantizationSearchParamsBuilder, Query, QueryBatchPointsBuilder,
     QueryPointsBuilder, SearchParamsBuilder, SparseIndices, VectorInput,
@@ -104,6 +105,11 @@ impl ConfigSearchProcessor {
             request_builder = request_builder.read_consistency(read_consistency);
         }
 
+        if let Some(shard_key) = &self.args.shard_key {
+            request_builder =
+                request_builder.shard_key_selector(vec![Key::Keyword(shard_key.clone())]);
+        }
+
         request_builder
     }