Skip to content

Commit 0c9a7ec

Browse files
author
Yinwei Li
committed
Add search_iterator() and example.
Signed-off-by: Yinwei Li <yinwei.li@zilliz.com>
1 parent 158cd82 commit 0c9a7ec

File tree

5 files changed

+811
-37
lines changed

5 files changed

+811
-37
lines changed

‎examples/iterator.rs‎

Lines changed: 134 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@ use milvus::client::Client;
22
use milvus::data::FieldColumn;
33
use milvus::error::Error;
44
use milvus::index::{IndexParams, IndexType, MetricType};
5-
use milvus::iterator::QueryIteratorOptions;
5+
use milvus::iterator::{QueryIteratorOptions, SearchIteratorOptions};
66
use milvus::options::LoadOptions;
77
use milvus::schema::{CollectionSchemaBuilder, FieldSchema};
8+
use milvus::value::Value;
89
use rand::Rng;
910
use std::collections::HashMap;
1011

@@ -16,6 +17,133 @@ const PICTURE: &str = "picture";
1617
const DIM: i64 = 8;
1718
const NUM_ENTITIES: i64 = 10000;
1819

20+
async fn test_search_iterator(client: &Client) -> Result<(), Error> {
21+
// Generate a random vector for search
22+
let mut rng = rand::thread_rng();
23+
let search_vector: Vec<f32> = (0..DIM).map(|_| rng.gen()).collect();
24+
25+
// Create search iterator options
26+
let search_options = SearchIteratorOptions::new()
27+
.batch_size(50)
28+
.anns_field(PICTURE.to_string())
29+
.add_search_param("metric_type".to_string(), "L2".to_string())
30+
.add_search_param("params".to_string(), "{\"nprobe\": 10}".to_string())
31+
.reduce_stop_for_best(true)
32+
.limit(10);
33+
34+
println!("Testing search iterator...");
35+
println!("Search vector: {:?}", search_vector);
36+
37+
// Create search iterator
38+
let mut search_iterator = client
39+
.search_iterator(
40+
COLLECTION_NAME,
41+
vec![Value::FloatArray(search_vector.into())],
42+
search_options,
43+
)
44+
.await
45+
.expect("Failed to create search iterator");
46+
47+
let mut page_idx = 0;
48+
loop {
49+
let res = search_iterator.next().await?;
50+
match res {
51+
Some(results) => {
52+
if results.is_empty() {
53+
println!("search iteration finished, close");
54+
search_iterator.close();
55+
break;
56+
}
57+
58+
// Print each search result in Python-like format
59+
for (i, result) in results.iter().enumerate() {
60+
println!("Search result {}: {{", i);
61+
println!(" size: {}", result.size);
62+
println!(" scores: {:?}", result.score);
63+
println!(" ids: {:?}", result.id);
64+
65+
// Print field data in entity format
66+
if !result.field.is_empty() {
67+
println!(" entity: {{");
68+
for field in &result.field {
69+
println!(" {}: {:?}", field.name, field.value);
70+
}
71+
println!(" }}");
72+
}
73+
println!("}}");
74+
}
75+
76+
page_idx += 1;
77+
println!("page{}-------------------------", page_idx);
78+
}
79+
None => {
80+
println!("search iteration finished, close");
81+
search_iterator.close();
82+
break;
83+
}
84+
}
85+
}
86+
87+
Ok(())
88+
}
89+
90+
async fn test_query_iterator(client: &Client) -> Result<(), Error> {
91+
// Create query iterator options
92+
let query_options = QueryIteratorOptions::new()
93+
.batch_size(50)
94+
.filter(format!("10 <= {} <= 25", AGE))
95+
.output_fields(vec![USER_ID.to_string(), AGE.to_string()])
96+
.limit(1000);
97+
98+
println!("Testing query iterator...");
99+
100+
// Create query iterator
101+
let mut query_iterator = client
102+
.query_iterator(COLLECTION_NAME, query_options)
103+
.await
104+
.expect("Failed to create query iterator");
105+
106+
let mut page_idx = 0;
107+
loop {
108+
let res = query_iterator.next().await?;
109+
match res {
110+
Some(fields) => {
111+
if fields.is_empty() {
112+
println!("query iteration finished, close");
113+
query_iterator.close();
114+
break;
115+
}
116+
117+
// Print each query result in Python-like format
118+
let num_records = if fields.is_empty() {
119+
0
120+
} else {
121+
fields[0].len()
122+
};
123+
for i in 0..num_records {
124+
println!("Query result {}: {{", i);
125+
for field in &fields {
126+
if let Some(value) = field.get(i) {
127+
println!(" {}: {:?}", field.name, value);
128+
}
129+
}
130+
println!("}}");
131+
}
132+
133+
page_idx += 1;
134+
println!("page{}-------------------------", page_idx);
135+
}
136+
None => {
137+
println!("query iteration finished, close");
138+
query_iterator.close();
139+
break;
140+
}
141+
}
142+
}
143+
144+
Ok(())
145+
}
146+
19147
#[tokio::main]
20148
async fn main() -> Result<(), Error> {
21149
let client = Client::new("http://localhost:19530").await?;
@@ -77,32 +205,13 @@ async fn main() -> Result<(), Error> {
77205
println!("Collection loaded and indexed.");
78206

79207
// Query Iterator 示例
80-
let expr = format!("10 <= {} <= 25", AGE);
81-
let output_fields = vec![USER_ID.to_string(), AGE.to_string()];
82-
let options = QueryIteratorOptions::new()
83-
.filter(expr)
84-
.batch_size(50)
85-
.output_fields(output_fields)
86-
.reduce_stop_for_best(true);
208+
println!("=== Testing Query Iterator ===");
209+
test_query_iterator(&client).await?;
87210

88-
let mut iterator = client
89-
.query_iterator(COLLECTION_NAME, options)
90-
.await
91-
.expect("Failed to create query iterator");
211+
// Search Iterator 示例
212+
println!("\n=== Testing Search Iterator ===");
213+
test_search_iterator(&client).await?;
92214

93-
let mut page_idx = 0;
94-
while let Some(batch) = iterator.next().await.expect("Query failed") {
95-
if batch.is_empty() {
96-
println!("query iteration finished, close");
97-
iterator.close();
98-
break;
99-
}
100-
for row in &batch {
101-
println!("{:?}", row);
102-
}
103-
page_idx += 1;
104-
println!("page{}-------------------------", page_idx);
105-
}
106215
client.drop_collection(COLLECTION_NAME).await?;
107216

108217
Ok(())

‎src/collection.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -839,6 +839,7 @@ pub type ParamValue = serde_json::Value;
839839
pub use serde_json::json as ParamValue;
840840

841841
// search result for a single vector
842+
#[derive(Clone)]
842843
pub struct SearchResult<'a> {
843844
pub size: i64,
844845
pub id: Vec<Value<'a>>,

0 commit comments

Comments
 (0)