diff --git a/src/query.rs b/src/query.rs index c83b15a..4f3ee68 100644 --- a/src/query.rs +++ b/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> { 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); } diff --git a/src/scroll.rs b/src/scroll.rs index ce82044..ea85de2 100644 --- a/src/scroll.rs +++ b/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); } diff --git a/src/search/from_args.rs b/src/search/from_args.rs index fa751ff..40eb105 100644 --- a/src/search/from_args.rs +++ b/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 } diff --git a/src/search/from_config.rs b/src/search/from_config.rs index 293c279..4ea59d6 100644 --- a/src/search/from_config.rs +++ b/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 }