| 317 | |
| 318 | #[test] |
| 319 | fn train_and_search() { |
| 320 | let vecs = make_vectors(1000, 16); |
| 321 | let refs: Vec<&[f32]> = vecs.iter().map(|v| v.as_slice()).collect(); |
| 322 | |
| 323 | let mut idx = IvfPqIndex::new( |
| 324 | 16, |
| 325 | IvfPqParams { |
| 326 | n_cells: 32, |
| 327 | pq_m: 4, |
| 328 | pq_k: 32, |
| 329 | nprobe: 8, |
| 330 | metric: DistanceMetric::L2, |
| 331 | }, |
| 332 | ); |
| 333 | idx.train(&refs); |
| 334 | idx.add_batch(&refs); |
| 335 | |
| 336 | assert_eq!(idx.len(), 1000); |
| 337 | |
| 338 | let query = &vecs[500]; |
| 339 | let results = idx.search(query, 5); |
| 340 | assert_eq!(results.len(), 5); |
| 341 | assert!( |
| 342 | results.iter().any(|r| r.id == 500), |
| 343 | "exact match not found in top-5" |
| 344 | ); |
| 345 | } |
| 346 | |
| 347 | #[test] |
| 348 | fn empty_index() { |