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 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 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 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 pub fn index_type(&self) -> IndexType {
201 IndexType::from(unsafe { zvec_rust_sys::zvec_index_params_get_type(self.handle) })
202 }
203
204 pub fn metric_type(&self) -> MetricType {
206 MetricType::from(unsafe { zvec_rust_sys::zvec_index_params_get_metric_type(self.handle) })
207 }
208
209 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 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 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
239pub struct FieldSchema {
241 pub(crate) handle: *mut zvec_rust_sys::zvec_field_schema_t,
242 owned: bool,
243}
244
245impl FieldSchema {
246 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 #[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 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 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 pub fn data_type(&self) -> DataType {
303 DataType::from(unsafe { zvec_rust_sys::zvec_field_schema_get_data_type(self.handle) })
304 }
305
306 pub fn dimension(&self) -> u32 {
308 unsafe { zvec_rust_sys::zvec_field_schema_get_dimension(self.handle) }
309 }
310
311 pub fn is_nullable(&self) -> bool {
313 unsafe { zvec_rust_sys::zvec_field_schema_is_nullable(self.handle) }
314 }
315
316 pub fn is_vector_field(&self) -> bool {
318 unsafe { zvec_rust_sys::zvec_field_schema_is_vector_field(self.handle) }
319 }
320
321 pub fn is_dense_vector(&self) -> bool {
323 unsafe { zvec_rust_sys::zvec_field_schema_is_dense_vector(self.handle) }
324 }
325
326 pub fn is_sparse_vector(&self) -> bool {
328 unsafe { zvec_rust_sys::zvec_field_schema_is_sparse_vector(self.handle) }
329 }
330
331 pub fn has_index(&self) -> bool {
333 unsafe { zvec_rust_sys::zvec_field_schema_has_index(self.handle) }
334 }
335
336 pub fn index_type(&self) -> IndexType {
338 IndexType::from(unsafe { zvec_rust_sys::zvec_field_schema_get_index_type(self.handle) })
339 }
340
341 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
355pub struct CollectionSchema {
357 pub(crate) handle: *mut zvec_rust_sys::zvec_collection_schema_t,
358 owned: bool,
359}
360
361impl CollectionSchema {
362 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 pub fn builder(name: &str) -> CollectionSchemaBuilder {
380 CollectionSchemaBuilder::new(name)
381 }
382
383 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 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 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 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 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 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 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 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 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 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
476pub 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 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 pub fn add_field(mut self, field: FieldSchema) -> Self {
497 self.fields.push((field, None));
498 self
499 }
500
501 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 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 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 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}