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 IVF RaBitQ index parameters.
107    ///
108    /// IVF RaBitQ combines inverted-file clustering with RaBitQ quantization for
109    /// memory-efficient approximate search.
110    ///
111    /// - `nlist`: number of cluster centers
112    /// - `total_bits`: total bits for RaBitQ quantization
113    /// - `sample_count`: sample count for training; `0` uses all vectors
114    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    /// Creates Flat index parameters.
146    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    /// Creates DiskANN index parameters.
167    ///
168    /// DiskANN is a disk-based graph index that keeps the bulk of the index on
169    /// disk, drastically reducing memory usage for large datasets.
170    ///
171    /// - `max_degree`: graph connectivity, i.e. max out-degree of the graph (typical: 64)
172    /// - `list_size`: build-time candidate list size during construction (typical: 100)
173    /// - `pq_chunk_num`: PQ chunk count; `0` disables product quantization
174    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    /// Creates inverted index parameters for scalar fields.
206    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    /// Creates FTS (Full-Text Search) index parameters.
228    ///
229    /// All parameters are optional — passing `None` keeps the library default.
230    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    /// Returns the index type.
278    pub fn index_type(&self) -> IndexType {
279        IndexType::from(unsafe { zvec_rust_sys::zvec_index_params_get_type(self.handle) })
280    }
281
282    /// Returns the metric type.
283    pub fn metric_type(&self) -> MetricType {
284        MetricType::from(unsafe { zvec_rust_sys::zvec_index_params_get_metric_type(self.handle) })
285    }
286
287    /// Returns the quantize type.
288    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    /// Returns whether vectors are randomly rotated before quantization.
295    pub fn quantizer_enable_rotate(&self) -> bool {
296        unsafe { zvec_rust_sys::zvec_index_params_get_quantizer_enable_rotate(self.handle) }
297    }
298
299    /// Sets the metric type.
300    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    /// Sets the quantize type.
307    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    /// Enables or disables random rotation before quantization.
314    ///
315    /// Rotating vectors before `Int8` / `Int4` quantization reduces the
316    /// quantization error. The rotation matrix is stored with the index and
317    /// applied to query vectors at search time, so it has no effect with other
318    /// quantize types.
319    ///
320    /// Only vector index params accept this option; scalar indexes (`invert`,
321    /// `fts`) return [`ErrorCode::InvalidArgument`].
322    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
337/// Schema definition for a single field in a collection.
338pub struct FieldSchema {
339    pub(crate) handle: *mut zvec_rust_sys::zvec_field_schema_t,
340    owned: bool,
341}
342
343impl FieldSchema {
344    /// Creates a new field schema.
345    ///
346    /// - `name`: Field name
347    /// - `data_type`: Data type of the field
348    /// - `nullable`: Whether the field can be null
349    /// - `dimension`: Vector dimension (0 for non-vector fields)
350    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    /// Creates a non-owning wrapper around an existing handle.
373    #[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    /// Sets the index parameters for this field.
382    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    /// Returns the field name.
389    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    /// Returns the data type.
400    pub fn data_type(&self) -> DataType {
401        DataType::from(unsafe { zvec_rust_sys::zvec_field_schema_get_data_type(self.handle) })
402    }
403
404    /// Returns the dimension (for vector fields).
405    pub fn dimension(&self) -> u32 {
406        unsafe { zvec_rust_sys::zvec_field_schema_get_dimension(self.handle) }
407    }
408
409    /// Returns whether the field is nullable.
410    pub fn is_nullable(&self) -> bool {
411        unsafe { zvec_rust_sys::zvec_field_schema_is_nullable(self.handle) }
412    }
413
414    /// Returns whether this is a vector field.
415    pub fn is_vector_field(&self) -> bool {
416        unsafe { zvec_rust_sys::zvec_field_schema_is_vector_field(self.handle) }
417    }
418
419    /// Returns whether this is a dense vector field.
420    pub fn is_dense_vector(&self) -> bool {
421        unsafe { zvec_rust_sys::zvec_field_schema_is_dense_vector(self.handle) }
422    }
423
424    /// Returns whether this is a sparse vector field.
425    pub fn is_sparse_vector(&self) -> bool {
426        unsafe { zvec_rust_sys::zvec_field_schema_is_sparse_vector(self.handle) }
427    }
428
429    /// Returns whether this field has an index.
430    pub fn has_index(&self) -> bool {
431        unsafe { zvec_rust_sys::zvec_field_schema_has_index(self.handle) }
432    }
433
434    /// Returns the index type.
435    pub fn index_type(&self) -> IndexType {
436        IndexType::from(unsafe { zvec_rust_sys::zvec_field_schema_get_index_type(self.handle) })
437    }
438
439    /// Returns whether this is an array type.
440    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
453/// Schema definition for a collection, containing field definitions.
454pub struct CollectionSchema {
455    pub(crate) handle: *mut zvec_rust_sys::zvec_collection_schema_t,
456    owned: bool,
457}
458
459impl CollectionSchema {
460    /// Creates a new collection schema with the given name.
461    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    /// Returns a builder for constructing a collection schema.
477    pub fn builder(name: &str) -> CollectionSchemaBuilder {
478        CollectionSchemaBuilder::new(name)
479    }
480
481    /// Creates a non-owning wrapper around an existing handle.
482    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    /// Adds a field to the schema.
490    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    /// Returns the collection name.
497    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    /// Checks if a field exists in the schema.
508    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    /// Checks if a field has an index.
517    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    /// Drops a field from the schema.
526    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    /// Adds an index to a field.
534    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    /// Drops an index from a field.
546    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    /// Sets the maximum document count per segment.
554    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    /// Returns the maximum document count per segment.
561    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
574/// Builder for constructing a [`CollectionSchema`] with a fluent API.
575pub 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    /// Creates a new builder with the given collection name.
584    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    /// Adds a field to the schema.
594    pub fn add_field(mut self, field: FieldSchema) -> Self {
595        self.fields.push((field, None));
596        self
597    }
598
599    /// Adds a vector field with index parameters.
600    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    /// Adds a scalar field with an inverted index.
622    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    /// Sets the maximum document count per segment.
643    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    /// Builds the collection schema.
649    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}