1use std::ffi::CStr;
2
3use crate::error::{check_error, to_cstring, Error, ErrorCode, Result};
4use crate::types::{DataType, IndexType, MetricType, QuantizeType};
5
6pub struct IndexParams {
8 pub(crate) handle: *mut zvec_rust_sys::zvec_index_params_t,
9 owned: bool,
10}
11
12impl IndexParams {
13 pub unsafe fn as_raw(&self) -> *mut zvec_rust_sys::zvec_index_params_t {
18 self.handle
19 }
20
21 pub fn hnsw(metric: MetricType, m: i32, ef_construction: i32) -> Result<Self> {
23 unsafe {
24 let handle = zvec_rust_sys::zvec_index_params_create(IndexType::Hnsw as u32);
25 if handle.is_null() {
26 return Err(Error {
27 code: ErrorCode::InternalError,
28 message: "failed to create HNSW index params".into(),
29 });
30 }
31 check_error(zvec_rust_sys::zvec_index_params_set_metric_type(
32 handle,
33 metric as u32,
34 ))?;
35 check_error(zvec_rust_sys::zvec_index_params_set_hnsw_params(
36 handle,
37 m,
38 ef_construction,
39 ))?;
40 Ok(IndexParams {
41 handle,
42 owned: true,
43 })
44 }
45 }
46
47 pub fn hnsw_with_quantize(
49 metric: MetricType,
50 m: i32,
51 ef_construction: i32,
52 quantize: QuantizeType,
53 ) -> Result<Self> {
54 unsafe {
55 let handle = zvec_rust_sys::zvec_index_params_create(IndexType::Hnsw as u32);
56 if handle.is_null() {
57 return Err(Error {
58 code: ErrorCode::InternalError,
59 message: "failed to create HNSW index params".into(),
60 });
61 }
62 check_error(zvec_rust_sys::zvec_index_params_set_metric_type(
63 handle,
64 metric as u32,
65 ))?;
66 check_error(zvec_rust_sys::zvec_index_params_set_hnsw_params(
67 handle,
68 m,
69 ef_construction,
70 ))?;
71 check_error(zvec_rust_sys::zvec_index_params_set_quantize_type(
72 handle,
73 quantize as u32,
74 ))?;
75 Ok(IndexParams {
76 handle,
77 owned: true,
78 })
79 }
80 }
81
82 pub fn ivf(metric: MetricType, n_list: i32, n_iters: i32, use_soar: bool) -> Result<Self> {
84 unsafe {
85 let handle = zvec_rust_sys::zvec_index_params_create(IndexType::Ivf as u32);
86 if handle.is_null() {
87 return Err(Error {
88 code: ErrorCode::InternalError,
89 message: "failed to create IVF index params".into(),
90 });
91 }
92 check_error(zvec_rust_sys::zvec_index_params_set_metric_type(
93 handle,
94 metric as u32,
95 ))?;
96 check_error(zvec_rust_sys::zvec_index_params_set_ivf_params(
97 handle, n_list, n_iters, use_soar,
98 ))?;
99 Ok(IndexParams {
100 handle,
101 owned: true,
102 })
103 }
104 }
105
106 pub fn ivf_rabitq(
115 metric: MetricType,
116 nlist: i32,
117 total_bits: i32,
118 sample_count: i32,
119 ) -> Result<Self> {
120 unsafe {
121 let handle = zvec_rust_sys::zvec_index_params_create(IndexType::IvfRabitq as u32);
122 if handle.is_null() {
123 return Err(Error {
124 code: ErrorCode::InternalError,
125 message: "failed to create IVF RaBitQ index params".into(),
126 });
127 }
128 check_error(zvec_rust_sys::zvec_index_params_set_metric_type(
129 handle,
130 metric as u32,
131 ))?;
132 check_error(zvec_rust_sys::zvec_index_params_set_ivf_rabitq_params(
133 handle,
134 nlist,
135 total_bits,
136 sample_count,
137 ))?;
138 Ok(IndexParams {
139 handle,
140 owned: true,
141 })
142 }
143 }
144
145 pub fn flat(metric: MetricType) -> Result<Self> {
147 unsafe {
148 let handle = zvec_rust_sys::zvec_index_params_create(IndexType::Flat as u32);
149 if handle.is_null() {
150 return Err(Error {
151 code: ErrorCode::InternalError,
152 message: "failed to create Flat index params".into(),
153 });
154 }
155 check_error(zvec_rust_sys::zvec_index_params_set_metric_type(
156 handle,
157 metric as u32,
158 ))?;
159 Ok(IndexParams {
160 handle,
161 owned: true,
162 })
163 }
164 }
165
166 pub fn diskann(
175 metric: MetricType,
176 max_degree: i32,
177 list_size: i32,
178 pq_chunk_num: i32,
179 ) -> Result<Self> {
180 unsafe {
181 let handle = zvec_rust_sys::zvec_index_params_create(IndexType::Diskann as u32);
182 if handle.is_null() {
183 return Err(Error {
184 code: ErrorCode::InternalError,
185 message: "failed to create DiskANN index params".into(),
186 });
187 }
188 check_error(zvec_rust_sys::zvec_index_params_set_metric_type(
189 handle,
190 metric as u32,
191 ))?;
192 check_error(zvec_rust_sys::zvec_index_params_set_diskann_params(
193 handle,
194 max_degree,
195 list_size,
196 pq_chunk_num,
197 ))?;
198 Ok(IndexParams {
199 handle,
200 owned: true,
201 })
202 }
203 }
204
205 pub fn invert(enable_range_opt: bool, enable_wildcard: bool) -> Result<Self> {
207 unsafe {
208 let handle = zvec_rust_sys::zvec_index_params_create(IndexType::Invert as u32);
209 if handle.is_null() {
210 return Err(Error {
211 code: ErrorCode::InternalError,
212 message: "failed to create Invert index params".into(),
213 });
214 }
215 check_error(zvec_rust_sys::zvec_index_params_set_invert_params(
216 handle,
217 enable_range_opt,
218 enable_wildcard,
219 ))?;
220 Ok(IndexParams {
221 handle,
222 owned: true,
223 })
224 }
225 }
226
227 pub fn fts(
231 tokenizer_name: Option<&str>,
232 filters: Option<&[&str]>,
233 extra_params: Option<&str>,
234 ) -> Result<Self> {
235 unsafe {
236 let handle = zvec_rust_sys::zvec_index_params_create(IndexType::Fts as u32);
237 if handle.is_null() {
238 return Err(Error {
239 code: ErrorCode::InternalError,
240 message: "failed to create FTS index params".into(),
241 });
242 }
243 let c_tokenizer = tokenizer_name.map(to_cstring).transpose()?;
244 let c_extra = extra_params.map(to_cstring).transpose()?;
245
246 let filter_array = if let Some(f) = filters {
247 let arr = zvec_rust_sys::zvec_string_array_create(f.len());
248 for (i, s) in f.iter().enumerate() {
249 let cs = to_cstring(s)?;
250 zvec_rust_sys::zvec_string_array_add(arr, i, cs.as_ptr());
251 }
252 arr
253 } else {
254 std::ptr::null_mut()
255 };
256
257 let result = check_error(zvec_rust_sys::zvec_index_params_set_fts_params(
258 handle,
259 c_tokenizer
260 .as_ref()
261 .map_or(std::ptr::null(), |c| c.as_ptr()),
262 filter_array as *const _,
263 c_extra.as_ref().map_or(std::ptr::null(), |c| c.as_ptr()),
264 ));
265
266 if !filter_array.is_null() {
267 zvec_rust_sys::zvec_string_array_destroy(filter_array);
268 }
269 result?;
270 Ok(IndexParams {
271 handle,
272 owned: true,
273 })
274 }
275 }
276
277 pub fn index_type(&self) -> IndexType {
279 IndexType::from(unsafe { zvec_rust_sys::zvec_index_params_get_type(self.handle) })
280 }
281
282 pub fn metric_type(&self) -> MetricType {
284 MetricType::from(unsafe { zvec_rust_sys::zvec_index_params_get_metric_type(self.handle) })
285 }
286
287 pub fn quantize_type(&self) -> QuantizeType {
289 QuantizeType::from(unsafe {
290 zvec_rust_sys::zvec_index_params_get_quantize_type(self.handle)
291 })
292 }
293
294 pub fn quantizer_enable_rotate(&self) -> bool {
296 unsafe { zvec_rust_sys::zvec_index_params_get_quantizer_enable_rotate(self.handle) }
297 }
298
299 pub fn set_metric_type(&mut self, metric: MetricType) -> Result<()> {
301 check_error(unsafe {
302 zvec_rust_sys::zvec_index_params_set_metric_type(self.handle, metric as u32)
303 })
304 }
305
306 pub fn set_quantize_type(&mut self, quantize: QuantizeType) -> Result<()> {
308 check_error(unsafe {
309 zvec_rust_sys::zvec_index_params_set_quantize_type(self.handle, quantize as u32)
310 })
311 }
312
313 pub fn set_quantizer_enable_rotate(&mut self, enable_rotate: bool) -> Result<()> {
323 check_error(unsafe {
324 zvec_rust_sys::zvec_index_params_set_quantizer_enable_rotate(self.handle, enable_rotate)
325 })
326 }
327}
328
329impl Drop for IndexParams {
330 fn drop(&mut self) {
331 if self.owned && !self.handle.is_null() {
332 unsafe { zvec_rust_sys::zvec_index_params_destroy(self.handle) };
333 }
334 }
335}
336
337pub struct FieldSchema {
339 pub(crate) handle: *mut zvec_rust_sys::zvec_field_schema_t,
340 owned: bool,
341}
342
343impl FieldSchema {
344 pub fn new(name: &str, data_type: DataType, nullable: bool, dimension: u32) -> Result<Self> {
351 let c_name = to_cstring(name)?;
352 let handle = unsafe {
353 zvec_rust_sys::zvec_field_schema_create(
354 c_name.as_ptr(),
355 data_type as u32,
356 nullable,
357 dimension,
358 )
359 };
360 if handle.is_null() {
361 return Err(Error {
362 code: ErrorCode::InternalError,
363 message: "failed to create field schema".into(),
364 });
365 }
366 Ok(FieldSchema {
367 handle,
368 owned: true,
369 })
370 }
371
372 #[allow(dead_code)]
374 pub(crate) fn from_borrowed(handle: *mut zvec_rust_sys::zvec_field_schema_t) -> Self {
375 FieldSchema {
376 handle,
377 owned: false,
378 }
379 }
380
381 pub fn set_index_params(&mut self, params: &IndexParams) -> Result<()> {
383 check_error(unsafe {
384 zvec_rust_sys::zvec_field_schema_set_index_params(self.handle, params.handle)
385 })
386 }
387
388 pub fn name(&self) -> &str {
390 unsafe {
391 let ptr = zvec_rust_sys::zvec_field_schema_get_name(self.handle);
392 if ptr.is_null() {
393 return "";
394 }
395 CStr::from_ptr(ptr).to_str().unwrap_or("")
396 }
397 }
398
399 pub fn data_type(&self) -> DataType {
401 DataType::from(unsafe { zvec_rust_sys::zvec_field_schema_get_data_type(self.handle) })
402 }
403
404 pub fn dimension(&self) -> u32 {
406 unsafe { zvec_rust_sys::zvec_field_schema_get_dimension(self.handle) }
407 }
408
409 pub fn is_nullable(&self) -> bool {
411 unsafe { zvec_rust_sys::zvec_field_schema_is_nullable(self.handle) }
412 }
413
414 pub fn is_vector_field(&self) -> bool {
416 unsafe { zvec_rust_sys::zvec_field_schema_is_vector_field(self.handle) }
417 }
418
419 pub fn is_dense_vector(&self) -> bool {
421 unsafe { zvec_rust_sys::zvec_field_schema_is_dense_vector(self.handle) }
422 }
423
424 pub fn is_sparse_vector(&self) -> bool {
426 unsafe { zvec_rust_sys::zvec_field_schema_is_sparse_vector(self.handle) }
427 }
428
429 pub fn has_index(&self) -> bool {
431 unsafe { zvec_rust_sys::zvec_field_schema_has_index(self.handle) }
432 }
433
434 pub fn index_type(&self) -> IndexType {
436 IndexType::from(unsafe { zvec_rust_sys::zvec_field_schema_get_index_type(self.handle) })
437 }
438
439 pub fn is_array_type(&self) -> bool {
441 unsafe { zvec_rust_sys::zvec_field_schema_is_array_type(self.handle) }
442 }
443}
444
445impl Drop for FieldSchema {
446 fn drop(&mut self) {
447 if self.owned && !self.handle.is_null() {
448 unsafe { zvec_rust_sys::zvec_field_schema_destroy(self.handle) };
449 }
450 }
451}
452
453pub struct CollectionSchema {
455 pub(crate) handle: *mut zvec_rust_sys::zvec_collection_schema_t,
456 owned: bool,
457}
458
459impl CollectionSchema {
460 pub fn new(name: &str) -> Result<Self> {
462 let c_name = to_cstring(name)?;
463 let handle = unsafe { zvec_rust_sys::zvec_collection_schema_create(c_name.as_ptr()) };
464 if handle.is_null() {
465 return Err(Error {
466 code: ErrorCode::InternalError,
467 message: "failed to create collection schema".into(),
468 });
469 }
470 Ok(CollectionSchema {
471 handle,
472 owned: true,
473 })
474 }
475
476 pub fn builder(name: &str) -> CollectionSchemaBuilder {
478 CollectionSchemaBuilder::new(name)
479 }
480
481 pub(crate) fn from_owned(handle: *mut zvec_rust_sys::zvec_collection_schema_t) -> Self {
483 CollectionSchema {
484 handle,
485 owned: true,
486 }
487 }
488
489 pub fn add_field(&mut self, field: &FieldSchema) -> Result<()> {
491 check_error(unsafe {
492 zvec_rust_sys::zvec_collection_schema_add_field(self.handle, field.handle)
493 })
494 }
495
496 pub fn name(&self) -> &str {
498 unsafe {
499 let ptr = zvec_rust_sys::zvec_collection_schema_get_name(self.handle);
500 if ptr.is_null() {
501 return "";
502 }
503 CStr::from_ptr(ptr).to_str().unwrap_or("")
504 }
505 }
506
507 pub fn has_field(&self, name: &str) -> bool {
509 let c_name = match to_cstring(name) {
510 Ok(s) => s,
511 Err(_) => return false,
512 };
513 unsafe { zvec_rust_sys::zvec_collection_schema_has_field(self.handle, c_name.as_ptr()) }
514 }
515
516 pub fn has_index(&self, field_name: &str) -> bool {
518 let c_name = match to_cstring(field_name) {
519 Ok(s) => s,
520 Err(_) => return false,
521 };
522 unsafe { zvec_rust_sys::zvec_collection_schema_has_index(self.handle, c_name.as_ptr()) }
523 }
524
525 pub fn drop_field(&mut self, name: &str) -> Result<()> {
527 let c_name = to_cstring(name)?;
528 check_error(unsafe {
529 zvec_rust_sys::zvec_collection_schema_drop_field(self.handle, c_name.as_ptr())
530 })
531 }
532
533 pub fn add_index(&mut self, field_name: &str, params: &IndexParams) -> Result<()> {
535 let c_name = to_cstring(field_name)?;
536 check_error(unsafe {
537 zvec_rust_sys::zvec_collection_schema_add_index(
538 self.handle,
539 c_name.as_ptr(),
540 params.handle,
541 )
542 })
543 }
544
545 pub fn drop_index(&mut self, field_name: &str) -> Result<()> {
547 let c_name = to_cstring(field_name)?;
548 check_error(unsafe {
549 zvec_rust_sys::zvec_collection_schema_drop_index(self.handle, c_name.as_ptr())
550 })
551 }
552
553 pub fn set_max_doc_count_per_segment(&mut self, count: u64) -> Result<()> {
555 check_error(unsafe {
556 zvec_rust_sys::zvec_collection_schema_set_max_doc_count_per_segment(self.handle, count)
557 })
558 }
559
560 pub fn max_doc_count_per_segment(&self) -> u64 {
562 unsafe { zvec_rust_sys::zvec_collection_schema_get_max_doc_count_per_segment(self.handle) }
563 }
564}
565
566impl Drop for CollectionSchema {
567 fn drop(&mut self) {
568 if self.owned && !self.handle.is_null() {
569 unsafe { zvec_rust_sys::zvec_collection_schema_destroy(self.handle) };
570 }
571 }
572}
573
574pub struct CollectionSchemaBuilder {
576 name: String,
577 fields: Vec<(FieldSchema, Option<IndexParams>)>,
578 max_doc_count_per_segment: Option<u64>,
579 deferred_error: Option<Error>,
580}
581
582impl CollectionSchemaBuilder {
583 pub fn new(name: &str) -> Self {
585 CollectionSchemaBuilder {
586 name: name.to_string(),
587 fields: Vec::new(),
588 max_doc_count_per_segment: None,
589 deferred_error: None,
590 }
591 }
592
593 pub fn add_field(mut self, field: FieldSchema) -> Self {
595 self.fields.push((field, None));
596 self
597 }
598
599 pub fn add_vector_field(
601 self,
602 name: &str,
603 data_type: DataType,
604 dimension: u32,
605 index_params: IndexParams,
606 ) -> Self {
607 match FieldSchema::new(name, data_type, false, dimension) {
608 Ok(field) => {
609 let mut s = self;
610 s.fields.push((field, Some(index_params)));
611 s
612 }
613 Err(e) => {
614 let mut s = self;
615 s.deferred_error = Some(e);
616 s
617 }
618 }
619 }
620
621 pub fn add_indexed_field(
623 self,
624 name: &str,
625 data_type: DataType,
626 index_params: IndexParams,
627 ) -> Self {
628 match FieldSchema::new(name, data_type, false, 0) {
629 Ok(field) => {
630 let mut s = self;
631 s.fields.push((field, Some(index_params)));
632 s
633 }
634 Err(e) => {
635 let mut s = self;
636 s.deferred_error = Some(e);
637 s
638 }
639 }
640 }
641
642 pub fn max_doc_count_per_segment(mut self, count: u64) -> Self {
644 self.max_doc_count_per_segment = Some(count);
645 self
646 }
647
648 pub fn build(self) -> Result<CollectionSchema> {
650 if let Some(e) = self.deferred_error {
651 return Err(e);
652 }
653 let mut schema = CollectionSchema::new(&self.name)?;
654
655 for (mut field, index_params) in self.fields {
656 if let Some(params) = &index_params {
657 field.set_index_params(params)?;
658 }
659 schema.add_field(&field)?;
660 }
661
662 if let Some(count) = self.max_doc_count_per_segment {
663 schema.set_max_doc_count_per_segment(count)?;
664 }
665
666 Ok(schema)
667 }
668}
669
670#[cfg(test)]
671mod tests {
672 use super::*;
673 use crate::types::MetricType;
674
675 #[test]
676 fn test_collection_schema_builder_default_values() {
677 let builder = CollectionSchemaBuilder::new("test_collection");
678 assert_eq!(builder.name, "test_collection");
679 assert!(builder.fields.is_empty());
680 assert!(builder.max_doc_count_per_segment.is_none());
681 }
682
683 #[test]
684 fn test_collection_schema_builder_add_field() {
685 let field = FieldSchema::new("test_field", DataType::Int32, false, 0).unwrap();
686 let builder = CollectionSchemaBuilder::new("test_collection").add_field(field);
687
688 assert_eq!(builder.fields.len(), 1);
689 assert_eq!(builder.fields[0].0.name(), "test_field");
690 assert!(builder.fields[0].1.is_none());
691 }
692
693 #[test]
694 fn test_collection_schema_builder_add_vector_field() {
695 let index_params = IndexParams::hnsw(MetricType::L2, 16, 200).unwrap();
696 let builder = CollectionSchemaBuilder::new("test_collection").add_vector_field(
697 "vector",
698 DataType::VectorFp32,
699 128,
700 index_params,
701 );
702
703 assert_eq!(builder.fields.len(), 1);
704 assert_eq!(builder.fields[0].0.name(), "vector");
705 assert!(builder.fields[0].1.is_some());
706 }
707
708 #[test]
709 fn test_collection_schema_builder_add_indexed_field() {
710 let index_params = IndexParams::invert(true, false).unwrap();
711 let builder = CollectionSchemaBuilder::new("test_collection").add_indexed_field(
712 "age",
713 DataType::Int32,
714 index_params,
715 );
716
717 assert_eq!(builder.fields.len(), 1);
718 assert_eq!(builder.fields[0].0.name(), "age");
719 assert!(builder.fields[0].1.is_some());
720 }
721
722 #[test]
723 fn test_collection_schema_builder_max_doc_count_per_segment() {
724 let builder =
725 CollectionSchemaBuilder::new("test_collection").max_doc_count_per_segment(1000);
726
727 assert_eq!(builder.max_doc_count_per_segment, Some(1000));
728 }
729
730 #[test]
731 fn test_collection_schema_builder_multiple_fields() {
732 let field1 = FieldSchema::new("id", DataType::Int64, false, 0).unwrap();
733 let field2 = FieldSchema::new("name", DataType::String, false, 0).unwrap();
734 let index_params = IndexParams::hnsw(MetricType::L2, 16, 200).unwrap();
735
736 let builder = CollectionSchemaBuilder::new("test_collection")
737 .add_field(field1)
738 .add_field(field2)
739 .add_vector_field("vector", DataType::VectorFp32, 128, index_params);
740
741 assert_eq!(builder.fields.len(), 3);
742 assert_eq!(builder.fields[0].0.name(), "id");
743 assert_eq!(builder.fields[1].0.name(), "name");
744 assert_eq!(builder.fields[2].0.name(), "vector");
745 }
746
747 #[test]
748 fn test_index_params_hnsw() {
749 let params = IndexParams::hnsw(MetricType::L2, 16, 200).unwrap();
750 assert_eq!(params.index_type(), IndexType::Hnsw);
751 assert_eq!(params.metric_type(), MetricType::L2);
752 }
753
754 #[test]
755 fn test_index_params_hnsw_with_quantize() {
756 let params =
757 IndexParams::hnsw_with_quantize(MetricType::L2, 16, 200, QuantizeType::Int8).unwrap();
758 assert_eq!(params.index_type(), IndexType::Hnsw);
759 assert_eq!(params.metric_type(), MetricType::L2);
760 assert_eq!(params.quantize_type(), QuantizeType::Int8);
761 }
762
763 #[test]
764 fn test_index_params_ivf() {
765 let params = IndexParams::ivf(MetricType::Ip, 100, 10, true).unwrap();
766 assert_eq!(params.index_type(), IndexType::Ivf);
767 assert_eq!(params.metric_type(), MetricType::Ip);
768 }
769
770 #[test]
771 fn test_index_params_flat() {
772 let params = IndexParams::flat(MetricType::Cosine).unwrap();
773 assert_eq!(params.index_type(), IndexType::Flat);
774 assert_eq!(params.metric_type(), MetricType::Cosine);
775 }
776
777 #[test]
778 fn test_index_params_diskann() {
779 let params = IndexParams::diskann(MetricType::L2, 32, 100, 0).unwrap();
780 assert_eq!(params.index_type(), IndexType::Diskann);
781 assert_eq!(params.metric_type(), MetricType::L2);
782 }
783
784 #[test]
785 fn test_index_params_invert() {
786 let params = IndexParams::invert(true, true).unwrap();
787 assert_eq!(params.index_type(), IndexType::Invert);
788 }
789
790 #[test]
791 fn test_field_schema_properties() {
792 let field = FieldSchema::new("test_field", DataType::VectorFp32, true, 128).unwrap();
793 assert_eq!(field.name(), "test_field");
794 assert_eq!(field.data_type(), DataType::VectorFp32);
795 assert_eq!(field.dimension(), 128);
796 assert!(field.is_nullable());
797 }
798
799 #[test]
800 fn test_collection_schema_builder_method() {
801 let builder = CollectionSchema::builder("test_collection");
802 assert_eq!(builder.name, "test_collection");
803 assert!(builder.fields.is_empty());
804 }
805
806 #[test]
807 fn test_field_schema_non_nullable() {
808 let field = FieldSchema::new("test_field", DataType::Int32, false, 0).unwrap();
809 assert!(!field.is_nullable());
810 }
811
812 #[test]
813 fn test_field_schema_scalar_dimension_zero() {
814 let field = FieldSchema::new("scalar_field", DataType::String, false, 0).unwrap();
815 assert_eq!(field.dimension(), 0);
816 assert!(!field.is_vector_field());
817 }
818
819 #[test]
820 fn test_field_schema_new_with_null_byte() {
821 let result = FieldSchema::new("invalid\0field", DataType::String, false, 0);
822 assert!(result.is_err());
823 }
824
825 #[test]
826 fn test_field_schema_new_success() {
827 let result = FieldSchema::new("valid_field", DataType::Int64, true, 0);
828 assert!(result.is_ok());
829 let field = result.unwrap();
830 assert_eq!(field.name(), "valid_field");
831 assert_eq!(field.data_type(), DataType::Int64);
832 }
833
834 #[test]
835 fn test_index_params_set_metric_type() {
836 let mut params = IndexParams::hnsw(MetricType::L2, 16, 200).unwrap();
837 params.set_metric_type(MetricType::Cosine).unwrap();
838 assert_eq!(params.metric_type(), MetricType::Cosine);
839 }
840
841 #[test]
842 fn test_index_params_set_quantize_type() {
843 let mut params =
844 IndexParams::hnsw_with_quantize(MetricType::L2, 16, 200, QuantizeType::Undefined)
845 .unwrap();
846 params.set_quantize_type(QuantizeType::Int8).unwrap();
847 assert_eq!(params.quantize_type(), QuantizeType::Int8);
848 }
849
850 #[test]
851 fn test_index_params_quantizer_enable_rotate() {
852 let mut params =
853 IndexParams::hnsw_with_quantize(MetricType::L2, 16, 200, QuantizeType::Int8).unwrap();
854 assert!(!params.quantizer_enable_rotate());
855 params.set_quantizer_enable_rotate(true).unwrap();
856 assert!(params.quantizer_enable_rotate());
857 params.set_quantizer_enable_rotate(false).unwrap();
858 assert!(!params.quantizer_enable_rotate());
859 }
860
861 #[test]
862 fn test_index_params_quantizer_enable_rotate_rejects_scalar_index() {
863 let mut params = IndexParams::invert(false, false).unwrap();
864 let err = params
865 .set_quantizer_enable_rotate(true)
866 .expect_err("scalar index params must reject the quantizer option");
867 assert_eq!(err.code, ErrorCode::InvalidArgument);
868 }
869
870 #[test]
871 fn test_index_params_ivf_cosine() {
872 let params = IndexParams::ivf(MetricType::Cosine, 100, 10, false).unwrap();
873 assert_eq!(params.index_type(), IndexType::Ivf);
874 assert_eq!(params.metric_type(), MetricType::Cosine);
875 }
876
877 #[test]
878 fn test_index_params_flat_ip() {
879 let params = IndexParams::flat(MetricType::Ip).unwrap();
880 assert_eq!(params.index_type(), IndexType::Flat);
881 assert_eq!(params.metric_type(), MetricType::Ip);
882 }
883
884 #[test]
885 fn test_collection_schema_builder_chaining() {
886 let vector_index = IndexParams::hnsw(MetricType::L2, 16, 200).unwrap();
887 let scalar_index = IndexParams::invert(true, false).unwrap();
888
889 let builder = CollectionSchemaBuilder::new("chained_collection")
890 .add_field(FieldSchema::new("id", DataType::Int64, false, 0).unwrap())
891 .add_vector_field("embedding", DataType::VectorFp32, 128, vector_index)
892 .add_indexed_field("category", DataType::String, scalar_index)
893 .max_doc_count_per_segment(5000);
894
895 assert_eq!(builder.fields.len(), 3);
896 assert_eq!(builder.max_doc_count_per_segment, Some(5000));
897 assert_eq!(builder.fields[0].0.name(), "id");
898 assert_eq!(builder.fields[1].0.name(), "embedding");
899 assert_eq!(builder.fields[2].0.name(), "category");
900 }
901}