Skip to main content

zvec_rust/
schema.rs

1use std::ffi::CStr;
2
3use crate::error::{check_error, to_cstring, Error, ErrorCode, Result};
4use crate::types::{DataType, IndexType, MetricType, QuantizeType};
5
6/// Parameters for configuring an index on a field.
7pub struct IndexParams {
8    pub(crate) handle: *mut zvec_rust_sys::zvec_index_params_t,
9    owned: bool,
10}
11
12impl IndexParams {
13    /// Returns the raw FFI handle.
14    ///
15    /// # Safety
16    /// The caller must not use the handle after the `IndexParams` is dropped.
17    pub unsafe fn as_raw(&self) -> *mut zvec_rust_sys::zvec_index_params_t {
18        self.handle
19    }
20
21    /// Creates HNSW index parameters.
22    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    /// Creates HNSW index parameters with quantization.
48    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    /// Creates IVF index parameters.
83    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    /// Creates Flat index parameters.
107    pub fn flat(metric: MetricType) -> Result<Self> {
108        unsafe {
109            let handle = zvec_rust_sys::zvec_index_params_create(IndexType::Flat as u32);
110            if handle.is_null() {
111                return Err(Error {
112                    code: ErrorCode::InternalError,
113                    message: "failed to create Flat index params".into(),
114                });
115            }
116            check_error(zvec_rust_sys::zvec_index_params_set_metric_type(
117                handle,
118                metric as u32,
119            ))?;
120            Ok(IndexParams {
121                handle,
122                owned: true,
123            })
124        }
125    }
126
127    /// Creates inverted index parameters for scalar fields.
128    pub fn invert(enable_range_opt: bool, enable_wildcard: bool) -> Result<Self> {
129        unsafe {
130            let handle = zvec_rust_sys::zvec_index_params_create(IndexType::Invert as u32);
131            if handle.is_null() {
132                return Err(Error {
133                    code: ErrorCode::InternalError,
134                    message: "failed to create Invert index params".into(),
135                });
136            }
137            check_error(zvec_rust_sys::zvec_index_params_set_invert_params(
138                handle,
139                enable_range_opt,
140                enable_wildcard,
141            ))?;
142            Ok(IndexParams {
143                handle,
144                owned: true,
145            })
146        }
147    }
148
149    /// Creates FTS (Full-Text Search) index parameters.
150    ///
151    /// All parameters are optional — passing `None` keeps the library default.
152    pub fn fts(
153        tokenizer_name: Option<&str>,
154        filters: Option<&[&str]>,
155        extra_params: Option<&str>,
156    ) -> Result<Self> {
157        unsafe {
158            let handle = zvec_rust_sys::zvec_index_params_create(IndexType::Fts as u32);
159            if handle.is_null() {
160                return Err(Error {
161                    code: ErrorCode::InternalError,
162                    message: "failed to create FTS index params".into(),
163                });
164            }
165            let c_tokenizer = tokenizer_name.map(to_cstring).transpose()?;
166            let c_extra = extra_params.map(to_cstring).transpose()?;
167
168            let filter_array = if let Some(f) = filters {
169                let arr = zvec_rust_sys::zvec_string_array_create(f.len());
170                for (i, s) in f.iter().enumerate() {
171                    let cs = to_cstring(s)?;
172                    zvec_rust_sys::zvec_string_array_add(arr, i, cs.as_ptr());
173                }
174                arr
175            } else {
176                std::ptr::null_mut()
177            };
178
179            let result = check_error(zvec_rust_sys::zvec_index_params_set_fts_params(
180                handle,
181                c_tokenizer
182                    .as_ref()
183                    .map_or(std::ptr::null(), |c| c.as_ptr()),
184                filter_array as *const _,
185                c_extra.as_ref().map_or(std::ptr::null(), |c| c.as_ptr()),
186            ));
187
188            if !filter_array.is_null() {
189                zvec_rust_sys::zvec_string_array_destroy(filter_array);
190            }
191            result?;
192            Ok(IndexParams {
193                handle,
194                owned: true,
195            })
196        }
197    }
198
199    /// Returns the index type.
200    pub fn index_type(&self) -> IndexType {
201        IndexType::from(unsafe { zvec_rust_sys::zvec_index_params_get_type(self.handle) })
202    }
203
204    /// Returns the metric type.
205    pub fn metric_type(&self) -> MetricType {
206        MetricType::from(unsafe { zvec_rust_sys::zvec_index_params_get_metric_type(self.handle) })
207    }
208
209    /// Returns the quantize type.
210    pub fn quantize_type(&self) -> QuantizeType {
211        QuantizeType::from(unsafe {
212            zvec_rust_sys::zvec_index_params_get_quantize_type(self.handle)
213        })
214    }
215
216    /// Sets the metric type.
217    pub fn set_metric_type(&mut self, metric: MetricType) -> Result<()> {
218        check_error(unsafe {
219            zvec_rust_sys::zvec_index_params_set_metric_type(self.handle, metric as u32)
220        })
221    }
222
223    /// Sets the quantize type.
224    pub fn set_quantize_type(&mut self, quantize: QuantizeType) -> Result<()> {
225        check_error(unsafe {
226            zvec_rust_sys::zvec_index_params_set_quantize_type(self.handle, quantize as u32)
227        })
228    }
229}
230
231impl Drop for IndexParams {
232    fn drop(&mut self) {
233        if self.owned && !self.handle.is_null() {
234            unsafe { zvec_rust_sys::zvec_index_params_destroy(self.handle) };
235        }
236    }
237}
238
239/// Schema definition for a single field in a collection.
240pub struct FieldSchema {
241    pub(crate) handle: *mut zvec_rust_sys::zvec_field_schema_t,
242    owned: bool,
243}
244
245impl FieldSchema {
246    /// Creates a new field schema.
247    ///
248    /// - `name`: Field name
249    /// - `data_type`: Data type of the field
250    /// - `nullable`: Whether the field can be null
251    /// - `dimension`: Vector dimension (0 for non-vector fields)
252    pub fn new(name: &str, data_type: DataType, nullable: bool, dimension: u32) -> Result<Self> {
253        let c_name = to_cstring(name)?;
254        let handle = unsafe {
255            zvec_rust_sys::zvec_field_schema_create(
256                c_name.as_ptr(),
257                data_type as u32,
258                nullable,
259                dimension,
260            )
261        };
262        if handle.is_null() {
263            return Err(Error {
264                code: ErrorCode::InternalError,
265                message: "failed to create field schema".into(),
266            });
267        }
268        Ok(FieldSchema {
269            handle,
270            owned: true,
271        })
272    }
273
274    /// Creates a non-owning wrapper around an existing handle.
275    #[allow(dead_code)]
276    pub(crate) fn from_borrowed(handle: *mut zvec_rust_sys::zvec_field_schema_t) -> Self {
277        FieldSchema {
278            handle,
279            owned: false,
280        }
281    }
282
283    /// Sets the index parameters for this field.
284    pub fn set_index_params(&mut self, params: &IndexParams) -> Result<()> {
285        check_error(unsafe {
286            zvec_rust_sys::zvec_field_schema_set_index_params(self.handle, params.handle)
287        })
288    }
289
290    /// Returns the field name.
291    pub fn name(&self) -> &str {
292        unsafe {
293            let ptr = zvec_rust_sys::zvec_field_schema_get_name(self.handle);
294            if ptr.is_null() {
295                return "";
296            }
297            CStr::from_ptr(ptr).to_str().unwrap_or("")
298        }
299    }
300
301    /// Returns the data type.
302    pub fn data_type(&self) -> DataType {
303        DataType::from(unsafe { zvec_rust_sys::zvec_field_schema_get_data_type(self.handle) })
304    }
305
306    /// Returns the dimension (for vector fields).
307    pub fn dimension(&self) -> u32 {
308        unsafe { zvec_rust_sys::zvec_field_schema_get_dimension(self.handle) }
309    }
310
311    /// Returns whether the field is nullable.
312    pub fn is_nullable(&self) -> bool {
313        unsafe { zvec_rust_sys::zvec_field_schema_is_nullable(self.handle) }
314    }
315
316    /// Returns whether this is a vector field.
317    pub fn is_vector_field(&self) -> bool {
318        unsafe { zvec_rust_sys::zvec_field_schema_is_vector_field(self.handle) }
319    }
320
321    /// Returns whether this is a dense vector field.
322    pub fn is_dense_vector(&self) -> bool {
323        unsafe { zvec_rust_sys::zvec_field_schema_is_dense_vector(self.handle) }
324    }
325
326    /// Returns whether this is a sparse vector field.
327    pub fn is_sparse_vector(&self) -> bool {
328        unsafe { zvec_rust_sys::zvec_field_schema_is_sparse_vector(self.handle) }
329    }
330
331    /// Returns whether this field has an index.
332    pub fn has_index(&self) -> bool {
333        unsafe { zvec_rust_sys::zvec_field_schema_has_index(self.handle) }
334    }
335
336    /// Returns the index type.
337    pub fn index_type(&self) -> IndexType {
338        IndexType::from(unsafe { zvec_rust_sys::zvec_field_schema_get_index_type(self.handle) })
339    }
340
341    /// Returns whether this is an array type.
342    pub fn is_array_type(&self) -> bool {
343        unsafe { zvec_rust_sys::zvec_field_schema_is_array_type(self.handle) }
344    }
345}
346
347impl Drop for FieldSchema {
348    fn drop(&mut self) {
349        if self.owned && !self.handle.is_null() {
350            unsafe { zvec_rust_sys::zvec_field_schema_destroy(self.handle) };
351        }
352    }
353}
354
355/// Schema definition for a collection, containing field definitions.
356pub struct CollectionSchema {
357    pub(crate) handle: *mut zvec_rust_sys::zvec_collection_schema_t,
358    owned: bool,
359}
360
361impl CollectionSchema {
362    /// Creates a new collection schema with the given name.
363    pub fn new(name: &str) -> Result<Self> {
364        let c_name = to_cstring(name)?;
365        let handle = unsafe { zvec_rust_sys::zvec_collection_schema_create(c_name.as_ptr()) };
366        if handle.is_null() {
367            return Err(Error {
368                code: ErrorCode::InternalError,
369                message: "failed to create collection schema".into(),
370            });
371        }
372        Ok(CollectionSchema {
373            handle,
374            owned: true,
375        })
376    }
377
378    /// Returns a builder for constructing a collection schema.
379    pub fn builder(name: &str) -> CollectionSchemaBuilder {
380        CollectionSchemaBuilder::new(name)
381    }
382
383    /// Creates a non-owning wrapper around an existing handle.
384    pub(crate) fn from_owned(handle: *mut zvec_rust_sys::zvec_collection_schema_t) -> Self {
385        CollectionSchema {
386            handle,
387            owned: true,
388        }
389    }
390
391    /// Adds a field to the schema.
392    pub fn add_field(&mut self, field: &FieldSchema) -> Result<()> {
393        check_error(unsafe {
394            zvec_rust_sys::zvec_collection_schema_add_field(self.handle, field.handle)
395        })
396    }
397
398    /// Returns the collection name.
399    pub fn name(&self) -> &str {
400        unsafe {
401            let ptr = zvec_rust_sys::zvec_collection_schema_get_name(self.handle);
402            if ptr.is_null() {
403                return "";
404            }
405            CStr::from_ptr(ptr).to_str().unwrap_or("")
406        }
407    }
408
409    /// Checks if a field exists in the schema.
410    pub fn has_field(&self, name: &str) -> bool {
411        let c_name = match to_cstring(name) {
412            Ok(s) => s,
413            Err(_) => return false,
414        };
415        unsafe { zvec_rust_sys::zvec_collection_schema_has_field(self.handle, c_name.as_ptr()) }
416    }
417
418    /// Checks if a field has an index.
419    pub fn has_index(&self, field_name: &str) -> bool {
420        let c_name = match to_cstring(field_name) {
421            Ok(s) => s,
422            Err(_) => return false,
423        };
424        unsafe { zvec_rust_sys::zvec_collection_schema_has_index(self.handle, c_name.as_ptr()) }
425    }
426
427    /// Drops a field from the schema.
428    pub fn drop_field(&mut self, name: &str) -> Result<()> {
429        let c_name = to_cstring(name)?;
430        check_error(unsafe {
431            zvec_rust_sys::zvec_collection_schema_drop_field(self.handle, c_name.as_ptr())
432        })
433    }
434
435    /// Adds an index to a field.
436    pub fn add_index(&mut self, field_name: &str, params: &IndexParams) -> Result<()> {
437        let c_name = to_cstring(field_name)?;
438        check_error(unsafe {
439            zvec_rust_sys::zvec_collection_schema_add_index(
440                self.handle,
441                c_name.as_ptr(),
442                params.handle,
443            )
444        })
445    }
446
447    /// Drops an index from a field.
448    pub fn drop_index(&mut self, field_name: &str) -> Result<()> {
449        let c_name = to_cstring(field_name)?;
450        check_error(unsafe {
451            zvec_rust_sys::zvec_collection_schema_drop_index(self.handle, c_name.as_ptr())
452        })
453    }
454
455    /// Sets the maximum document count per segment.
456    pub fn set_max_doc_count_per_segment(&mut self, count: u64) -> Result<()> {
457        check_error(unsafe {
458            zvec_rust_sys::zvec_collection_schema_set_max_doc_count_per_segment(self.handle, count)
459        })
460    }
461
462    /// Returns the maximum document count per segment.
463    pub fn max_doc_count_per_segment(&self) -> u64 {
464        unsafe { zvec_rust_sys::zvec_collection_schema_get_max_doc_count_per_segment(self.handle) }
465    }
466}
467
468impl Drop for CollectionSchema {
469    fn drop(&mut self) {
470        if self.owned && !self.handle.is_null() {
471            unsafe { zvec_rust_sys::zvec_collection_schema_destroy(self.handle) };
472        }
473    }
474}
475
476/// Builder for constructing a [`CollectionSchema`] with a fluent API.
477pub struct CollectionSchemaBuilder {
478    name: String,
479    fields: Vec<(FieldSchema, Option<IndexParams>)>,
480    max_doc_count_per_segment: Option<u64>,
481    deferred_error: Option<Error>,
482}
483
484impl CollectionSchemaBuilder {
485    /// Creates a new builder with the given collection name.
486    pub fn new(name: &str) -> Self {
487        CollectionSchemaBuilder {
488            name: name.to_string(),
489            fields: Vec::new(),
490            max_doc_count_per_segment: None,
491            deferred_error: None,
492        }
493    }
494
495    /// Adds a field to the schema.
496    pub fn add_field(mut self, field: FieldSchema) -> Self {
497        self.fields.push((field, None));
498        self
499    }
500
501    /// Adds a vector field with index parameters.
502    pub fn add_vector_field(
503        self,
504        name: &str,
505        data_type: DataType,
506        dimension: u32,
507        index_params: IndexParams,
508    ) -> Self {
509        match FieldSchema::new(name, data_type, false, dimension) {
510            Ok(field) => {
511                let mut s = self;
512                s.fields.push((field, Some(index_params)));
513                s
514            }
515            Err(e) => {
516                let mut s = self;
517                s.deferred_error = Some(e);
518                s
519            }
520        }
521    }
522
523    /// Adds a scalar field with an inverted index.
524    pub fn add_indexed_field(
525        self,
526        name: &str,
527        data_type: DataType,
528        index_params: IndexParams,
529    ) -> Self {
530        match FieldSchema::new(name, data_type, false, 0) {
531            Ok(field) => {
532                let mut s = self;
533                s.fields.push((field, Some(index_params)));
534                s
535            }
536            Err(e) => {
537                let mut s = self;
538                s.deferred_error = Some(e);
539                s
540            }
541        }
542    }
543
544    /// Sets the maximum document count per segment.
545    pub fn max_doc_count_per_segment(mut self, count: u64) -> Self {
546        self.max_doc_count_per_segment = Some(count);
547        self
548    }
549
550    /// Builds the collection schema.
551    pub fn build(self) -> Result<CollectionSchema> {
552        if let Some(e) = self.deferred_error {
553            return Err(e);
554        }
555        let mut schema = CollectionSchema::new(&self.name)?;
556
557        for (mut field, index_params) in self.fields {
558            if let Some(params) = &index_params {
559                field.set_index_params(params)?;
560            }
561            schema.add_field(&field)?;
562        }
563
564        if let Some(count) = self.max_doc_count_per_segment {
565            schema.set_max_doc_count_per_segment(count)?;
566        }
567
568        Ok(schema)
569    }
570}
571
572#[cfg(test)]
573mod tests {
574    use super::*;
575    use crate::types::MetricType;
576
577    #[test]
578    fn test_collection_schema_builder_default_values() {
579        let builder = CollectionSchemaBuilder::new("test_collection");
580        assert_eq!(builder.name, "test_collection");
581        assert!(builder.fields.is_empty());
582        assert!(builder.max_doc_count_per_segment.is_none());
583    }
584
585    #[test]
586    fn test_collection_schema_builder_add_field() {
587        let field = FieldSchema::new("test_field", DataType::Int32, false, 0).unwrap();
588        let builder = CollectionSchemaBuilder::new("test_collection").add_field(field);
589
590        assert_eq!(builder.fields.len(), 1);
591        assert_eq!(builder.fields[0].0.name(), "test_field");
592        assert!(builder.fields[0].1.is_none());
593    }
594
595    #[test]
596    fn test_collection_schema_builder_add_vector_field() {
597        let index_params = IndexParams::hnsw(MetricType::L2, 16, 200).unwrap();
598        let builder = CollectionSchemaBuilder::new("test_collection").add_vector_field(
599            "vector",
600            DataType::VectorFp32,
601            128,
602            index_params,
603        );
604
605        assert_eq!(builder.fields.len(), 1);
606        assert_eq!(builder.fields[0].0.name(), "vector");
607        assert!(builder.fields[0].1.is_some());
608    }
609
610    #[test]
611    fn test_collection_schema_builder_add_indexed_field() {
612        let index_params = IndexParams::invert(true, false).unwrap();
613        let builder = CollectionSchemaBuilder::new("test_collection").add_indexed_field(
614            "age",
615            DataType::Int32,
616            index_params,
617        );
618
619        assert_eq!(builder.fields.len(), 1);
620        assert_eq!(builder.fields[0].0.name(), "age");
621        assert!(builder.fields[0].1.is_some());
622    }
623
624    #[test]
625    fn test_collection_schema_builder_max_doc_count_per_segment() {
626        let builder =
627            CollectionSchemaBuilder::new("test_collection").max_doc_count_per_segment(1000);
628
629        assert_eq!(builder.max_doc_count_per_segment, Some(1000));
630    }
631
632    #[test]
633    fn test_collection_schema_builder_multiple_fields() {
634        let field1 = FieldSchema::new("id", DataType::Int64, false, 0).unwrap();
635        let field2 = FieldSchema::new("name", DataType::String, false, 0).unwrap();
636        let index_params = IndexParams::hnsw(MetricType::L2, 16, 200).unwrap();
637
638        let builder = CollectionSchemaBuilder::new("test_collection")
639            .add_field(field1)
640            .add_field(field2)
641            .add_vector_field("vector", DataType::VectorFp32, 128, index_params);
642
643        assert_eq!(builder.fields.len(), 3);
644        assert_eq!(builder.fields[0].0.name(), "id");
645        assert_eq!(builder.fields[1].0.name(), "name");
646        assert_eq!(builder.fields[2].0.name(), "vector");
647    }
648
649    #[test]
650    fn test_index_params_hnsw() {
651        let params = IndexParams::hnsw(MetricType::L2, 16, 200).unwrap();
652        assert_eq!(params.index_type(), IndexType::Hnsw);
653        assert_eq!(params.metric_type(), MetricType::L2);
654    }
655
656    #[test]
657    fn test_index_params_hnsw_with_quantize() {
658        let params =
659            IndexParams::hnsw_with_quantize(MetricType::L2, 16, 200, QuantizeType::Int8).unwrap();
660        assert_eq!(params.index_type(), IndexType::Hnsw);
661        assert_eq!(params.metric_type(), MetricType::L2);
662        assert_eq!(params.quantize_type(), QuantizeType::Int8);
663    }
664
665    #[test]
666    fn test_index_params_ivf() {
667        let params = IndexParams::ivf(MetricType::Ip, 100, 10, true).unwrap();
668        assert_eq!(params.index_type(), IndexType::Ivf);
669        assert_eq!(params.metric_type(), MetricType::Ip);
670    }
671
672    #[test]
673    fn test_index_params_flat() {
674        let params = IndexParams::flat(MetricType::Cosine).unwrap();
675        assert_eq!(params.index_type(), IndexType::Flat);
676        assert_eq!(params.metric_type(), MetricType::Cosine);
677    }
678
679    #[test]
680    fn test_index_params_invert() {
681        let params = IndexParams::invert(true, true).unwrap();
682        assert_eq!(params.index_type(), IndexType::Invert);
683    }
684
685    #[test]
686    fn test_field_schema_properties() {
687        let field = FieldSchema::new("test_field", DataType::VectorFp32, true, 128).unwrap();
688        assert_eq!(field.name(), "test_field");
689        assert_eq!(field.data_type(), DataType::VectorFp32);
690        assert_eq!(field.dimension(), 128);
691        assert!(field.is_nullable());
692    }
693
694    #[test]
695    fn test_collection_schema_builder_method() {
696        let builder = CollectionSchema::builder("test_collection");
697        assert_eq!(builder.name, "test_collection");
698        assert!(builder.fields.is_empty());
699    }
700
701    #[test]
702    fn test_field_schema_non_nullable() {
703        let field = FieldSchema::new("test_field", DataType::Int32, false, 0).unwrap();
704        assert!(!field.is_nullable());
705    }
706
707    #[test]
708    fn test_field_schema_scalar_dimension_zero() {
709        let field = FieldSchema::new("scalar_field", DataType::String, false, 0).unwrap();
710        assert_eq!(field.dimension(), 0);
711        assert!(!field.is_vector_field());
712    }
713
714    #[test]
715    fn test_field_schema_new_with_null_byte() {
716        let result = FieldSchema::new("invalid\0field", DataType::String, false, 0);
717        assert!(result.is_err());
718    }
719
720    #[test]
721    fn test_field_schema_new_success() {
722        let result = FieldSchema::new("valid_field", DataType::Int64, true, 0);
723        assert!(result.is_ok());
724        let field = result.unwrap();
725        assert_eq!(field.name(), "valid_field");
726        assert_eq!(field.data_type(), DataType::Int64);
727    }
728
729    #[test]
730    fn test_index_params_set_metric_type() {
731        let mut params = IndexParams::hnsw(MetricType::L2, 16, 200).unwrap();
732        params.set_metric_type(MetricType::Cosine).unwrap();
733        assert_eq!(params.metric_type(), MetricType::Cosine);
734    }
735
736    #[test]
737    fn test_index_params_set_quantize_type() {
738        let mut params =
739            IndexParams::hnsw_with_quantize(MetricType::L2, 16, 200, QuantizeType::Undefined)
740                .unwrap();
741        params.set_quantize_type(QuantizeType::Int8).unwrap();
742        assert_eq!(params.quantize_type(), QuantizeType::Int8);
743    }
744
745    #[test]
746    fn test_index_params_ivf_cosine() {
747        let params = IndexParams::ivf(MetricType::Cosine, 100, 10, false).unwrap();
748        assert_eq!(params.index_type(), IndexType::Ivf);
749        assert_eq!(params.metric_type(), MetricType::Cosine);
750    }
751
752    #[test]
753    fn test_index_params_flat_ip() {
754        let params = IndexParams::flat(MetricType::Ip).unwrap();
755        assert_eq!(params.index_type(), IndexType::Flat);
756        assert_eq!(params.metric_type(), MetricType::Ip);
757    }
758
759    #[test]
760    fn test_collection_schema_builder_chaining() {
761        let vector_index = IndexParams::hnsw(MetricType::L2, 16, 200).unwrap();
762        let scalar_index = IndexParams::invert(true, false).unwrap();
763
764        let builder = CollectionSchemaBuilder::new("chained_collection")
765            .add_field(FieldSchema::new("id", DataType::Int64, false, 0).unwrap())
766            .add_vector_field("embedding", DataType::VectorFp32, 128, vector_index)
767            .add_indexed_field("category", DataType::String, scalar_index)
768            .max_doc_count_per_segment(5000);
769
770        assert_eq!(builder.fields.len(), 3);
771        assert_eq!(builder.max_doc_count_per_segment, Some(5000));
772        assert_eq!(builder.fields[0].0.name(), "id");
773        assert_eq!(builder.fields[1].0.name(), "embedding");
774        assert_eq!(builder.fields[2].0.name(), "category");
775    }
776}