zvec_rust/
query.rs

1use std::os::raw::c_void;
2
3use crate::error::{check_error, to_cstring, Error, ErrorCode, Result};
4
5/// HNSW-specific query parameters.
6pub struct HnswQueryParams {
7    pub(crate) handle: *mut zvec_rust_sys::zvec_hnsw_query_params_t,
8}
9
10impl HnswQueryParams {
11    /// Creates new HNSW query parameters.
12    pub fn new(ef: i32, radius: f32, is_linear: bool, is_using_refiner: bool) -> Self {
13        let handle = unsafe {
14            zvec_rust_sys::zvec_query_params_hnsw_create(ef, radius, is_linear, is_using_refiner)
15        };
16        HnswQueryParams { handle }
17    }
18
19    /// Sets the exploration factor.
20    pub fn set_ef(&mut self, ef: i32) -> Result<()> {
21        check_error(unsafe { zvec_rust_sys::zvec_query_params_hnsw_set_ef(self.handle, ef) })
22    }
23
24    /// Returns the exploration factor.
25    pub fn ef(&self) -> i32 {
26        unsafe { zvec_rust_sys::zvec_query_params_hnsw_get_ef(self.handle) }
27    }
28}
29
30impl Drop for HnswQueryParams {
31    fn drop(&mut self) {
32        if !self.handle.is_null() {
33            unsafe { zvec_rust_sys::zvec_query_params_hnsw_destroy(self.handle) };
34        }
35    }
36}
37
38/// IVF-specific query parameters.
39pub struct IvfQueryParams {
40    pub(crate) handle: *mut zvec_rust_sys::zvec_ivf_query_params_t,
41}
42
43impl IvfQueryParams {
44    /// Creates new IVF query parameters.
45    pub fn new(nprobe: i32, is_using_refiner: bool, scale_factor: f32) -> Self {
46        let handle = unsafe {
47            zvec_rust_sys::zvec_query_params_ivf_create(nprobe, is_using_refiner, scale_factor)
48        };
49        IvfQueryParams { handle }
50    }
51
52    /// Sets the number of probe clusters.
53    pub fn set_nprobe(&mut self, nprobe: i32) -> Result<()> {
54        check_error(unsafe { zvec_rust_sys::zvec_query_params_ivf_set_nprobe(self.handle, nprobe) })
55    }
56
57    /// Returns the number of probe clusters.
58    pub fn nprobe(&self) -> i32 {
59        unsafe { zvec_rust_sys::zvec_query_params_ivf_get_nprobe(self.handle) }
60    }
61}
62
63impl Drop for IvfQueryParams {
64    fn drop(&mut self) {
65        if !self.handle.is_null() {
66            unsafe { zvec_rust_sys::zvec_query_params_ivf_destroy(self.handle) };
67        }
68    }
69}
70
71/// IVF RaBitQ-specific query parameters.
72pub struct IvfRabitqQueryParams {
73    pub(crate) handle: *mut zvec_rust_sys::zvec_ivf_rabitq_query_params_t,
74}
75
76impl IvfRabitqQueryParams {
77    /// Creates new IVF RaBitQ query parameters.
78    pub fn new(nprobe: i32, radius: f32, is_linear: bool, is_using_refiner: bool) -> Self {
79        let handle = unsafe {
80            zvec_rust_sys::zvec_query_params_ivf_rabitq_create(
81                nprobe,
82                radius,
83                is_linear,
84                is_using_refiner,
85            )
86        };
87        IvfRabitqQueryParams { handle }
88    }
89
90    /// Sets the number of probe clusters.
91    pub fn set_nprobe(&mut self, nprobe: i32) -> Result<()> {
92        check_error(unsafe {
93            zvec_rust_sys::zvec_query_params_ivf_rabitq_set_nprobe(self.handle, nprobe)
94        })
95    }
96
97    /// Returns the number of probe clusters.
98    pub fn nprobe(&self) -> i32 {
99        unsafe { zvec_rust_sys::zvec_query_params_ivf_rabitq_get_nprobe(self.handle) }
100    }
101
102    /// Sets the candidate expansion factor used by the refiner.
103    pub fn set_scale_factor(&mut self, scale_factor: f32) -> Result<()> {
104        check_error(unsafe {
105            zvec_rust_sys::zvec_query_params_ivf_rabitq_set_scale_factor(self.handle, scale_factor)
106        })
107    }
108
109    /// Returns the candidate expansion factor used by the refiner.
110    pub fn scale_factor(&self) -> f32 {
111        unsafe { zvec_rust_sys::zvec_query_params_ivf_rabitq_get_scale_factor(self.handle) }
112    }
113}
114
115impl Drop for IvfRabitqQueryParams {
116    fn drop(&mut self) {
117        if !self.handle.is_null() {
118            unsafe { zvec_rust_sys::zvec_query_params_ivf_rabitq_destroy(self.handle) };
119        }
120    }
121}
122
123/// Flat-specific query parameters.
124pub struct FlatQueryParams {
125    pub(crate) handle: *mut zvec_rust_sys::zvec_flat_query_params_t,
126}
127
128impl FlatQueryParams {
129    /// Creates new Flat query parameters.
130    pub fn new(is_using_refiner: bool, scale_factor: f32) -> Self {
131        let handle =
132            unsafe { zvec_rust_sys::zvec_query_params_flat_create(is_using_refiner, scale_factor) };
133        FlatQueryParams { handle }
134    }
135}
136
137impl Drop for FlatQueryParams {
138    fn drop(&mut self) {
139        if !self.handle.is_null() {
140            unsafe { zvec_rust_sys::zvec_query_params_flat_destroy(self.handle) };
141        }
142    }
143}
144
145/// DiskANN-specific query parameters.
146pub struct DiskannQueryParams {
147    pub(crate) handle: *mut zvec_rust_sys::zvec_diskann_query_params_t,
148}
149
150impl DiskannQueryParams {
151    /// Creates new DiskANN query parameters.
152    ///
153    /// - `list_size`: search frontier size (default in the C library: 300)
154    pub fn new(list_size: i32) -> Self {
155        let handle = unsafe { zvec_rust_sys::zvec_query_params_diskann_create(list_size) };
156        DiskannQueryParams { handle }
157    }
158
159    /// Sets the search-time frontier size.
160    pub fn set_list_size(&mut self, list_size: i32) -> Result<()> {
161        check_error(unsafe {
162            zvec_rust_sys::zvec_query_params_diskann_set_list_size(self.handle, list_size)
163        })
164    }
165
166    /// Returns the search-time frontier size.
167    pub fn list_size(&self) -> i32 {
168        unsafe { zvec_rust_sys::zvec_query_params_diskann_get_list_size(self.handle) }
169    }
170}
171
172impl Drop for DiskannQueryParams {
173    fn drop(&mut self) {
174        if !self.handle.is_null() {
175            unsafe { zvec_rust_sys::zvec_query_params_diskann_destroy(self.handle) };
176        }
177    }
178}
179
180/// FTS-specific query parameters controlling the default boolean operator.
181pub struct FtsQueryParams {
182    pub(crate) handle: *mut zvec_rust_sys::zvec_fts_query_params_t,
183}
184
185impl FtsQueryParams {
186    /// Creates new FTS query parameters.
187    ///
188    /// `default_operator` sets the boolean operator for adjacent bare terms
189    /// ("OR" or "AND", case-insensitive). Pass `None` to use the library default.
190    pub fn new(default_operator: Option<&str>) -> Result<Self> {
191        let c_op = default_operator.map(to_cstring).transpose()?;
192        let handle = unsafe {
193            zvec_rust_sys::zvec_query_params_fts_create(
194                c_op.as_ref().map_or(std::ptr::null(), |c| c.as_ptr()),
195            )
196        };
197        if handle.is_null() {
198            return Err(Error {
199                code: ErrorCode::InternalError,
200                message: "failed to create FTS query params".into(),
201            });
202        }
203        Ok(FtsQueryParams { handle })
204    }
205
206    /// Sets the default boolean operator.
207    pub fn set_default_operator(&mut self, op: &str) -> Result<()> {
208        let c_op = to_cstring(op)?;
209        check_error(unsafe {
210            zvec_rust_sys::zvec_query_params_fts_set_default_operator(self.handle, c_op.as_ptr())
211        })
212    }
213
214    /// Returns the default boolean operator.
215    pub fn default_operator(&self) -> Option<String> {
216        unsafe {
217            let ptr = zvec_rust_sys::zvec_query_params_fts_get_default_operator(self.handle);
218            if ptr.is_null() {
219                return None;
220            }
221            Some(std::ffi::CStr::from_ptr(ptr).to_string_lossy().into_owned())
222        }
223    }
224}
225
226impl Drop for FtsQueryParams {
227    fn drop(&mut self) {
228        if !self.handle.is_null() {
229            unsafe { zvec_rust_sys::zvec_query_params_fts_destroy(self.handle) };
230        }
231    }
232}
233
234/// FTS query payload holding the query expression and match string.
235///
236/// - `query_string`: a boolean / advanced query expression
237/// - `match_string`: a natural-language match string
238pub struct Fts {
239    pub(crate) handle: *mut zvec_rust_sys::zvec_fts_t,
240}
241
242impl Fts {
243    /// Creates a new FTS query payload.
244    pub fn new() -> Result<Self> {
245        let handle = unsafe { zvec_rust_sys::zvec_fts_create() };
246        if handle.is_null() {
247            return Err(Error {
248                code: ErrorCode::InternalError,
249                message: "failed to create FTS payload".into(),
250            });
251        }
252        Ok(Fts { handle })
253    }
254
255    /// Sets the boolean / advanced query expression.
256    pub fn set_query_string(&mut self, query: &str) -> Result<()> {
257        let c = to_cstring(query)?;
258        check_error(unsafe { zvec_rust_sys::zvec_fts_set_query_string(self.handle, c.as_ptr()) })
259    }
260
261    /// Sets the natural-language match string.
262    pub fn set_match_string(&mut self, match_str: &str) -> Result<()> {
263        let c = to_cstring(match_str)?;
264        check_error(unsafe { zvec_rust_sys::zvec_fts_set_match_string(self.handle, c.as_ptr()) })
265    }
266
267    /// Returns the query expression, or `None` if not set.
268    pub fn query_string(&self) -> Option<String> {
269        unsafe {
270            let ptr = zvec_rust_sys::zvec_fts_get_query_string(self.handle);
271            if ptr.is_null() {
272                return None;
273            }
274            Some(std::ffi::CStr::from_ptr(ptr).to_string_lossy().into_owned())
275        }
276    }
277
278    /// Returns the match string, or `None` if not set.
279    pub fn match_string(&self) -> Option<String> {
280        unsafe {
281            let ptr = zvec_rust_sys::zvec_fts_get_match_string(self.handle);
282            if ptr.is_null() {
283                return None;
284            }
285            Some(std::ffi::CStr::from_ptr(ptr).to_string_lossy().into_owned())
286        }
287    }
288}
289
290impl Drop for Fts {
291    fn drop(&mut self) {
292        if !self.handle.is_null() {
293            unsafe { zvec_rust_sys::zvec_fts_destroy(self.handle) };
294        }
295    }
296}
297
298/// A scalar filter, full-text, or vector similarity search query.
299pub struct SearchQuery {
300    pub(crate) handle: *mut zvec_rust_sys::zvec_vector_query_t,
301}
302
303impl SearchQuery {
304    /// Creates a scalar query without a vector or full-text field.
305    ///
306    /// Use [`SearchQuery::set_filter`] to select documents and
307    /// [`SearchQuery::set_output_fields`] to restrict the returned fields.
308    /// Results are limited to `topk` and have no guaranteed ordering; callers
309    /// enumerating all matches must account for that limit.
310    pub fn scalar(topk: i32) -> Result<Self> {
311        let handle = unsafe { zvec_rust_sys::zvec_vector_query_create() };
312        if handle.is_null() {
313            return Err(Error {
314                code: ErrorCode::InternalError,
315                message: "failed to create scalar query".into(),
316            });
317        }
318
319        // Own the handle before fallible setup so an error cannot leak it.
320        let query = SearchQuery { handle };
321        check_error(unsafe { zvec_rust_sys::zvec_vector_query_set_topk(query.handle, topk) })?;
322        Ok(query)
323    }
324
325    /// Returns the raw FFI handle.
326    ///
327    /// # Safety
328    /// The caller must not use the handle after the `SearchQuery` is dropped.
329    pub unsafe fn as_raw(&self) -> *mut zvec_rust_sys::zvec_vector_query_t {
330        self.handle
331    }
332
333    /// Creates a `SearchQuery` from a raw FFI handle.
334    ///
335    /// # Safety
336    /// The caller must ensure the handle is valid and was created by the zvec C API.
337    /// The `SearchQuery` takes ownership and will call `zvec_vector_query_destroy` on drop.
338    pub unsafe fn from_raw(handle: *mut zvec_rust_sys::zvec_vector_query_t) -> Self {
339        SearchQuery { handle }
340    }
341
342    /// Creates a new vector query with the given field name, query vector, and topk.
343    pub fn new(field_name: &str, vector: &[f32], topk: i32) -> Result<Self> {
344        let handle = unsafe { zvec_rust_sys::zvec_vector_query_create() };
345        if handle.is_null() {
346            return Err(Error {
347                code: ErrorCode::InternalError,
348                message: "failed to create vector query".into(),
349            });
350        }
351
352        // Own the handle before fallible setup so an error cannot leak it.
353        let query = SearchQuery { handle };
354        let c_field = to_cstring(field_name)?;
355
356        check_error(unsafe {
357            zvec_rust_sys::zvec_vector_query_set_field_name(query.handle, c_field.as_ptr())
358        })?;
359        check_error(unsafe {
360            zvec_rust_sys::zvec_vector_query_set_query_vector(
361                query.handle,
362                vector.as_ptr() as *const c_void,
363                std::mem::size_of_val(vector),
364            )
365        })?;
366        check_error(unsafe { zvec_rust_sys::zvec_vector_query_set_topk(query.handle, topk) })?;
367
368        Ok(query)
369    }
370
371    /// Creates a pure full-text (keyword-only) search query with no dense vector.
372    ///
373    /// [`SearchQuery::new`] and [`SearchQueryBuilder::build`] both require a query
374    /// vector, so an FTS clause can only ever be layered on top of a dense vector
375    /// (hybrid search). This constructor reaches the C API's vector-less FTS path:
376    /// once an FTS payload is attached, the underlying `zvec_vector_query` does not
377    /// require a query vector, enabling keyword-only retrieval.
378    ///
379    /// `field_name` must be the FTS-indexed field.
380    pub fn fts(field_name: &str, fts: &Fts, topk: i32) -> Result<Self> {
381        let handle = unsafe { zvec_rust_sys::zvec_vector_query_create() };
382        if handle.is_null() {
383            return Err(Error {
384                code: ErrorCode::InternalError,
385                message: "failed to create vector query".into(),
386            });
387        }
388
389        // Wrap the handle before any fallible call so `Drop` frees it on early return.
390        let mut query = SearchQuery { handle };
391        let c_field = to_cstring(field_name)?;
392
393        check_error(unsafe {
394            zvec_rust_sys::zvec_vector_query_set_field_name(query.handle, c_field.as_ptr())
395        })?;
396        check_error(unsafe { zvec_rust_sys::zvec_vector_query_set_topk(query.handle, topk) })?;
397        query.set_fts(fts)?;
398
399        Ok(query)
400    }
401
402    /// Returns a builder for constructing a search query.
403    pub fn builder() -> SearchQueryBuilder {
404        SearchQueryBuilder::new()
405    }
406
407    /// Sets the filter expression.
408    pub fn set_filter(&mut self, filter: &str) -> Result<()> {
409        let c_filter = to_cstring(filter)?;
410        check_error(unsafe {
411            zvec_rust_sys::zvec_vector_query_set_filter(self.handle, c_filter.as_ptr())
412        })
413    }
414
415    /// Sets whether to include vector data in results.
416    pub fn set_include_vector(&mut self, include: bool) -> Result<()> {
417        check_error(unsafe {
418            zvec_rust_sys::zvec_vector_query_set_include_vector(self.handle, include)
419        })
420    }
421
422    /// Sets whether to include doc ID in results.
423    pub fn set_include_doc_id(&mut self, include: bool) -> Result<()> {
424        check_error(unsafe {
425            zvec_rust_sys::zvec_vector_query_set_include_doc_id(self.handle, include)
426        })
427    }
428
429    /// Sets the output fields to include in results.
430    pub fn set_output_fields(&mut self, fields: &[&str]) -> Result<()> {
431        let c_fields: Vec<_> = fields
432            .iter()
433            .map(|f| to_cstring(f))
434            .collect::<Result<Vec<_>>>()?;
435        let c_ptrs: Vec<_> = c_fields.iter().map(|f| f.as_ptr()).collect();
436        check_error(unsafe {
437            zvec_rust_sys::zvec_vector_query_set_output_fields(
438                self.handle,
439                c_ptrs.as_ptr(),
440                c_ptrs.len(),
441            )
442        })
443    }
444
445    /// Sets HNSW query parameters (takes ownership on success).
446    pub fn set_hnsw_params(&mut self, mut params: HnswQueryParams) -> Result<()> {
447        check_error(unsafe {
448            zvec_rust_sys::zvec_vector_query_set_hnsw_params(self.handle, params.handle)
449        })?;
450        // Ownership transferred to query only on success; prevent double-free
451        params.handle = std::ptr::null_mut();
452        Ok(())
453    }
454
455    /// Sets IVF query parameters (takes ownership on success).
456    pub fn set_ivf_params(&mut self, mut params: IvfQueryParams) -> Result<()> {
457        check_error(unsafe {
458            zvec_rust_sys::zvec_vector_query_set_ivf_params(self.handle, params.handle)
459        })?;
460        params.handle = std::ptr::null_mut();
461        Ok(())
462    }
463
464    /// Sets IVF RaBitQ query parameters (takes ownership on success).
465    pub fn set_ivf_rabitq_params(&mut self, mut params: IvfRabitqQueryParams) -> Result<()> {
466        check_error(unsafe {
467            zvec_rust_sys::zvec_vector_query_set_ivf_rabitq_params(self.handle, params.handle)
468        })?;
469        params.handle = std::ptr::null_mut();
470        Ok(())
471    }
472
473    /// Sets Flat query parameters (takes ownership on success).
474    pub fn set_flat_params(&mut self, mut params: FlatQueryParams) -> Result<()> {
475        check_error(unsafe {
476            zvec_rust_sys::zvec_vector_query_set_flat_params(self.handle, params.handle)
477        })?;
478        params.handle = std::ptr::null_mut();
479        Ok(())
480    }
481
482    /// Sets DiskANN query parameters (takes ownership on success).
483    pub fn set_diskann_params(&mut self, mut params: DiskannQueryParams) -> Result<()> {
484        check_error(unsafe {
485            zvec_rust_sys::zvec_vector_query_set_diskann_params(self.handle, params.handle)
486        })?;
487        params.handle = std::ptr::null_mut();
488        Ok(())
489    }
490
491    /// Sets FTS query parameters (takes ownership on success).
492    pub fn set_fts_params(&mut self, mut params: FtsQueryParams) -> Result<()> {
493        check_error(unsafe {
494            zvec_rust_sys::zvec_vector_query_set_fts_params(self.handle, params.handle)
495        })?;
496        params.handle = std::ptr::null_mut();
497        Ok(())
498    }
499
500    /// Sets FTS payload (payload is copied, caller retains ownership).
501    pub fn set_fts(&mut self, fts: &Fts) -> Result<()> {
502        check_error(unsafe { zvec_rust_sys::zvec_vector_query_set_fts(self.handle, fts.handle) })
503    }
504}
505
506impl Drop for SearchQuery {
507    fn drop(&mut self) {
508        if !self.handle.is_null() {
509            unsafe { zvec_rust_sys::zvec_vector_query_destroy(self.handle) };
510        }
511    }
512}
513
514/// Builder for constructing a [`SearchQuery`].
515pub struct SearchQueryBuilder {
516    field_name: Option<String>,
517    vector: Option<Vec<f32>>,
518    topk: i32,
519    filter: Option<String>,
520    include_vector: Option<bool>,
521    include_doc_id: Option<bool>,
522    output_fields: Option<Vec<String>>,
523    fts_query_string: Option<String>,
524    fts_match_string: Option<String>,
525}
526
527impl SearchQueryBuilder {
528    fn new() -> Self {
529        SearchQueryBuilder {
530            field_name: None,
531            vector: None,
532            topk: 10,
533            filter: None,
534            include_vector: None,
535            include_doc_id: None,
536            output_fields: None,
537            fts_query_string: None,
538            fts_match_string: None,
539        }
540    }
541
542    /// Sets the field name to query.
543    pub fn field_name(mut self, name: &str) -> Self {
544        self.field_name = Some(name.to_string());
545        self
546    }
547
548    /// Sets the query vector.
549    pub fn vector(mut self, vector: &[f32]) -> Self {
550        self.vector = Some(vector.to_vec());
551        self
552    }
553
554    /// Sets the number of results to return.
555    pub fn topk(mut self, topk: i32) -> Self {
556        self.topk = topk;
557        self
558    }
559
560    /// Sets the filter expression.
561    pub fn filter(mut self, filter: &str) -> Self {
562        self.filter = Some(filter.to_string());
563        self
564    }
565
566    /// Sets whether to include vector data in results.
567    pub fn include_vector(mut self, include: bool) -> Self {
568        self.include_vector = Some(include);
569        self
570    }
571
572    /// Sets whether to include doc ID in results.
573    pub fn include_doc_id(mut self, include: bool) -> Self {
574        self.include_doc_id = Some(include);
575        self
576    }
577
578    /// Sets the output fields.
579    pub fn output_fields(mut self, fields: &[&str]) -> Self {
580        self.output_fields = Some(fields.iter().map(|s| s.to_string()).collect());
581        self
582    }
583
584    /// Sets the FTS boolean / advanced query expression.
585    pub fn fts_query_string(mut self, query: &str) -> Self {
586        self.fts_query_string = Some(query.to_string());
587        self
588    }
589
590    /// Sets the FTS natural-language match string.
591    pub fn fts_match_string(mut self, match_str: &str) -> Self {
592        self.fts_match_string = Some(match_str.to_string());
593        self
594    }
595
596    /// Builds the search query.
597    ///
598    /// The query kind follows from what was set on the builder:
599    ///
600    /// - `vector(..)` plus `field_name(..)` builds a dense vector query; an FTS
601    ///   clause set alongside it is layered on top, producing a hybrid search.
602    /// - `field_name(..)` plus an FTS clause and no vector builds a keyword-only
603    ///   query, equivalent to [`SearchQuery::fts`].
604    /// - Neither a vector nor an FTS clause builds a scalar query, equivalent to
605    ///   [`SearchQuery::scalar`]; `field_name(..)` must then be omitted.
606    pub fn build(self) -> Result<SearchQuery> {
607        let SearchQueryBuilder {
608            field_name,
609            vector,
610            topk,
611            filter,
612            include_vector,
613            include_doc_id,
614            output_fields,
615            fts_query_string,
616            fts_match_string,
617        } = self;
618
619        let fts = if fts_query_string.is_some() || fts_match_string.is_some() {
620            let mut payload = Fts::new()?;
621            if let Some(query_string) = &fts_query_string {
622                payload.set_query_string(query_string)?;
623            }
624            if let Some(match_string) = &fts_match_string {
625                payload.set_match_string(match_string)?;
626            }
627            Some(payload)
628        } else {
629            None
630        };
631
632        let mut query = match (vector, fts) {
633            (Some(vector), fts) => {
634                let field_name = field_name.ok_or_else(|| Error {
635                    code: ErrorCode::InvalidArgument,
636                    message: "field_name is required".into(),
637                })?;
638                let mut query = SearchQuery::new(&field_name, &vector, topk)?;
639                if let Some(payload) = &fts {
640                    query.set_fts(payload)?;
641                }
642                query
643            }
644            (None, Some(payload)) => {
645                let field_name = field_name.ok_or_else(|| Error {
646                    code: ErrorCode::InvalidArgument,
647                    message: "field_name is required".into(),
648                })?;
649                SearchQuery::fts(&field_name, &payload, topk)?
650            }
651            (None, None) => {
652                if field_name.is_some() {
653                    return Err(Error {
654                        code: ErrorCode::InvalidArgument,
655                        message: "field_name requires a vector or an FTS clause; omit it to build a scalar query".into(),
656                    });
657                }
658                SearchQuery::scalar(topk)?
659            }
660        };
661
662        if let Some(filter) = &filter {
663            query.set_filter(filter)?;
664        }
665        if let Some(include) = include_vector {
666            query.set_include_vector(include)?;
667        }
668        if let Some(include) = include_doc_id {
669            query.set_include_doc_id(include)?;
670        }
671        if let Some(fields) = &output_fields {
672            let field_refs: Vec<&str> = fields.iter().map(|s| s.as_str()).collect();
673            query.set_output_fields(&field_refs)?;
674        }
675
676        Ok(query)
677    }
678}
679
680/// A grouped vector similarity search query.
681pub struct GroupBySearchQuery {
682    pub(crate) handle: *mut zvec_rust_sys::zvec_group_by_vector_query_t,
683}
684
685impl GroupBySearchQuery {
686    /// Creates a new group-by search query.
687    pub fn new(
688        field_name: &str,
689        group_by_field: &str,
690        vector: &[f32],
691        group_count: u32,
692        group_topk: u32,
693    ) -> Result<Self> {
694        let handle = unsafe { zvec_rust_sys::zvec_group_by_vector_query_create() };
695        if handle.is_null() {
696            return Err(Error {
697                code: ErrorCode::InternalError,
698                message: "failed to create group by vector query".into(),
699            });
700        }
701
702        // Own the handle before fallible setup so an error cannot leak it.
703        let query = GroupBySearchQuery { handle };
704        let c_field = to_cstring(field_name)?;
705        let c_group_field = to_cstring(group_by_field)?;
706
707        check_error(unsafe {
708            zvec_rust_sys::zvec_group_by_vector_query_set_field_name(query.handle, c_field.as_ptr())
709        })?;
710        check_error(unsafe {
711            zvec_rust_sys::zvec_group_by_vector_query_set_group_by_field_name(
712                query.handle,
713                c_group_field.as_ptr(),
714            )
715        })?;
716        check_error(unsafe {
717            zvec_rust_sys::zvec_group_by_vector_query_set_query_vector(
718                query.handle,
719                vector.as_ptr() as *const c_void,
720                std::mem::size_of_val(vector),
721            )
722        })?;
723        check_error(unsafe {
724            zvec_rust_sys::zvec_group_by_vector_query_set_group_count(query.handle, group_count)
725        })?;
726        check_error(unsafe {
727            zvec_rust_sys::zvec_group_by_vector_query_set_topk_per_group(query.handle, group_topk)
728        })?;
729
730        Ok(query)
731    }
732
733    /// Sets the filter expression.
734    pub fn set_filter(&mut self, filter: &str) -> Result<()> {
735        let c_filter = to_cstring(filter)?;
736        check_error(unsafe {
737            zvec_rust_sys::zvec_group_by_vector_query_set_filter(self.handle, c_filter.as_ptr())
738        })
739    }
740
741    /// Sets whether to include vector data in results.
742    pub fn set_include_vector(&mut self, include: bool) -> Result<()> {
743        check_error(unsafe {
744            zvec_rust_sys::zvec_group_by_vector_query_set_include_vector(self.handle, include)
745        })
746    }
747
748    /// Sets the output fields.
749    pub fn set_output_fields(&mut self, fields: &[&str]) -> Result<()> {
750        let c_fields: Vec<_> = fields
751            .iter()
752            .map(|f| to_cstring(f))
753            .collect::<Result<Vec<_>>>()?;
754        let c_ptrs: Vec<_> = c_fields.iter().map(|f| f.as_ptr()).collect();
755        check_error(unsafe {
756            zvec_rust_sys::zvec_group_by_vector_query_set_output_fields(
757                self.handle,
758                c_ptrs.as_ptr(),
759                c_ptrs.len(),
760            )
761        })
762    }
763
764    /// Sets HNSW query parameters (takes ownership on success).
765    pub fn set_hnsw_params(&mut self, mut params: HnswQueryParams) -> Result<()> {
766        check_error(unsafe {
767            zvec_rust_sys::zvec_group_by_vector_query_set_hnsw_params(self.handle, params.handle)
768        })?;
769        // Ownership transferred to query only on success; prevent double-free
770        params.handle = std::ptr::null_mut();
771        Ok(())
772    }
773
774    /// Sets IVF query parameters (takes ownership on success).
775    pub fn set_ivf_params(&mut self, mut params: IvfQueryParams) -> Result<()> {
776        check_error(unsafe {
777            zvec_rust_sys::zvec_group_by_vector_query_set_ivf_params(self.handle, params.handle)
778        })?;
779        // Ownership transferred to query only on success; prevent double-free
780        params.handle = std::ptr::null_mut();
781        Ok(())
782    }
783
784    /// Sets IVF RaBitQ query parameters (takes ownership on success).
785    pub fn set_ivf_rabitq_params(&mut self, mut params: IvfRabitqQueryParams) -> Result<()> {
786        check_error(unsafe {
787            zvec_rust_sys::zvec_group_by_vector_query_set_ivf_rabitq_params(
788                self.handle,
789                params.handle,
790            )
791        })?;
792        // Ownership transferred to query only on success; prevent double-free
793        params.handle = std::ptr::null_mut();
794        Ok(())
795    }
796
797    /// Sets Flat query parameters (takes ownership on success).
798    pub fn set_flat_params(&mut self, mut params: FlatQueryParams) -> Result<()> {
799        check_error(unsafe {
800            zvec_rust_sys::zvec_group_by_vector_query_set_flat_params(self.handle, params.handle)
801        })?;
802        // Ownership transferred to query only on success; prevent double-free
803        params.handle = std::ptr::null_mut();
804        Ok(())
805    }
806
807    /// Sets DiskANN query parameters (takes ownership on success).
808    pub fn set_diskann_params(&mut self, mut params: DiskannQueryParams) -> Result<()> {
809        check_error(unsafe {
810            zvec_rust_sys::zvec_group_by_vector_query_set_diskann_params(self.handle, params.handle)
811        })?;
812        // Ownership transferred to query only on success; prevent double-free
813        params.handle = std::ptr::null_mut();
814        Ok(())
815    }
816}
817
818impl Drop for GroupBySearchQuery {
819    fn drop(&mut self) {
820        if !self.handle.is_null() {
821            unsafe { zvec_rust_sys::zvec_group_by_vector_query_destroy(self.handle) };
822        }
823    }
824}
825
826#[cfg(test)]
827mod tests {
828    use super::*;
829
830    #[test]
831    fn test_vector_query_builder_default_values() {
832        let builder = SearchQueryBuilder::new();
833        assert!(builder.field_name.is_none());
834        assert!(builder.vector.is_none());
835        assert_eq!(builder.topk, 10);
836        assert!(builder.filter.is_none());
837        assert!(builder.include_vector.is_none());
838        assert!(builder.include_doc_id.is_none());
839        assert!(builder.output_fields.is_none());
840    }
841
842    #[test]
843    fn test_vector_query_builder_setters() {
844        let builder = SearchQueryBuilder::new()
845            .field_name("test_field")
846            .vector(&[1.0, 2.0, 3.0])
847            .topk(5)
848            .filter("age > 18")
849            .include_vector(true)
850            .include_doc_id(false)
851            .output_fields(&["name", "age"]);
852
853        assert_eq!(builder.field_name, Some("test_field".to_string()));
854        assert_eq!(builder.vector, Some(vec![1.0, 2.0, 3.0]));
855        assert_eq!(builder.topk, 5);
856        assert_eq!(builder.filter, Some("age > 18".to_string()));
857        assert_eq!(builder.include_vector, Some(true));
858        assert_eq!(builder.include_doc_id, Some(false));
859        assert_eq!(
860            builder.output_fields,
861            Some(vec!["name".to_string(), "age".to_string()])
862        );
863    }
864
865    #[test]
866    fn test_vector_query_builder_build_missing_field_name() {
867        let builder = SearchQueryBuilder::new().vector(&[1.0, 2.0, 3.0]);
868
869        let result = builder.build();
870        assert!(result.is_err());
871        if let Err(e) = result {
872            assert_eq!(e.code, ErrorCode::InvalidArgument);
873            assert!(e.message.contains("field_name is required"));
874        }
875    }
876
877    #[test]
878    fn test_vector_query_builder_build_field_name_without_vector_or_fts() {
879        let builder = SearchQueryBuilder::new().field_name("test_field");
880
881        let result = builder.build();
882        assert!(result.is_err());
883        if let Err(e) = result {
884            assert_eq!(e.code, ErrorCode::InvalidArgument);
885            assert!(e.message.contains("omit it to build a scalar query"));
886        }
887    }
888
889    #[test]
890    fn test_vector_query_builder_build_scalar_query() {
891        let query = SearchQueryBuilder::new()
892            .topk(3)
893            .filter("count > 1")
894            .include_doc_id(true)
895            .build()
896            .expect("scalar query without a field name or vector");
897
898        let handle = query.handle;
899        assert_eq!(
900            unsafe { zvec_rust_sys::zvec_vector_query_get_topk(handle) },
901            3
902        );
903        assert!(unsafe { zvec_rust_sys::zvec_vector_query_get_include_doc_id(handle) });
904        let filter = unsafe { zvec_rust_sys::zvec_vector_query_get_filter(handle) };
905        assert!(!filter.is_null());
906        assert_eq!(
907            unsafe { std::ffi::CStr::from_ptr(filter) }.to_str(),
908            Ok("count > 1")
909        );
910        assert!(unsafe { zvec_rust_sys::zvec_vector_query_get_fts(handle) }.is_null());
911    }
912
913    #[test]
914    fn test_vector_query_builder_build_fts_only_query() {
915        let query = SearchQueryBuilder::new()
916            .field_name("content")
917            .fts_match_string("hello world")
918            .topk(7)
919            .build()
920            .expect("keyword-only query without a vector");
921
922        let handle = query.handle;
923        assert_eq!(
924            unsafe { zvec_rust_sys::zvec_vector_query_get_topk(handle) },
925            7
926        );
927        let field_name = unsafe { zvec_rust_sys::zvec_vector_query_get_field_name(handle) };
928        assert!(!field_name.is_null());
929        assert_eq!(
930            unsafe { std::ffi::CStr::from_ptr(field_name) }.to_str(),
931            Ok("content")
932        );
933        assert!(!unsafe { zvec_rust_sys::zvec_vector_query_get_fts(handle) }.is_null());
934    }
935
936    #[test]
937    fn test_vector_query_builder_build_missing_field_name_with_fts() {
938        let result = SearchQueryBuilder::new().fts_match_string("hello").build();
939        assert!(result.is_err());
940        if let Err(e) = result {
941            assert_eq!(e.code, ErrorCode::InvalidArgument);
942            assert!(e.message.contains("field_name is required"));
943        }
944    }
945
946    #[test]
947    fn test_vector_query_builder_build_hybrid_query() {
948        let query = SearchQueryBuilder::new()
949            .field_name("embedding")
950            .vector(&[1.0, 2.0])
951            .fts_match_string("hello")
952            .topk(4)
953            .build()
954            .expect("hybrid query keeps the FTS clause on top of the vector");
955
956        let handle = query.handle;
957        assert_eq!(
958            unsafe { zvec_rust_sys::zvec_vector_query_get_topk(handle) },
959            4
960        );
961        assert!(!unsafe { zvec_rust_sys::zvec_vector_query_get_fts(handle) }.is_null());
962    }
963
964    #[test]
965    fn test_vector_query_builder_builder_method() {
966        let builder = SearchQuery::builder();
967        assert!(builder.field_name.is_none());
968        assert!(builder.vector.is_none());
969        assert_eq!(builder.topk, 10);
970    }
971
972    #[test]
973    fn test_vector_query_builder_partial_setters() {
974        let builder = SearchQueryBuilder::new()
975            .field_name("test_field")
976            .vector(&[1.0, 2.0])
977            .filter("status = 'active'");
978        assert_eq!(builder.field_name, Some("test_field".to_string()));
979        assert_eq!(builder.vector, Some(vec![1.0, 2.0]));
980        assert_eq!(builder.filter, Some("status = 'active'".to_string()));
981        assert!(builder.include_vector.is_none());
982        assert!(builder.output_fields.is_none());
983    }
984
985    #[test]
986    fn test_vector_query_builder_empty_vector() {
987        let builder = SearchQueryBuilder::new()
988            .field_name("test_field")
989            .vector(&[]);
990        assert_eq!(builder.vector, Some(vec![]));
991    }
992
993    #[test]
994    fn test_vector_query_builder_empty_output_fields() {
995        let builder = SearchQueryBuilder::new()
996            .field_name("test_field")
997            .vector(&[1.0])
998            .output_fields(&[]);
999        assert_eq!(builder.output_fields, Some(vec![]));
1000    }
1001
1002    #[test]
1003    fn test_vector_query_builder_topk_zero() {
1004        let builder = SearchQueryBuilder::new()
1005            .field_name("test_field")
1006            .vector(&[1.0])
1007            .topk(0);
1008        assert_eq!(builder.topk, 0);
1009    }
1010
1011    #[test]
1012    fn test_vector_query_builder_topk_negative() {
1013        let builder = SearchQueryBuilder::new()
1014            .field_name("test_field")
1015            .vector(&[1.0])
1016            .topk(-1);
1017        assert_eq!(builder.topk, -1);
1018    }
1019
1020    #[test]
1021    fn test_vector_query_builder_overwrite_field_name() {
1022        let builder = SearchQueryBuilder::new()
1023            .field_name("first_field")
1024            .field_name("second_field");
1025        assert_eq!(builder.field_name, Some("second_field".to_string()));
1026    }
1027
1028    #[test]
1029    fn test_vector_query_builder_large_vector() {
1030        let large_vector: Vec<f32> = (0..1024).map(|i| i as f32).collect();
1031        let builder = SearchQueryBuilder::new()
1032            .field_name("test_field")
1033            .vector(&large_vector);
1034        assert_eq!(builder.vector.as_ref().unwrap().len(), 1024);
1035    }
1036
1037    #[test]
1038    fn test_fts_query_no_vector() {
1039        // A pure FTS query must build without a query vector.
1040        // `SearchQueryBuilder::build` reaches the same path when an FTS clause is
1041        // set without a vector.
1042        let mut fts = Fts::new().expect("create fts payload");
1043        fts.set_match_string("hello world")
1044            .expect("set match string");
1045
1046        let query = SearchQuery::fts("content", &fts, 10);
1047        assert!(
1048            query.is_ok(),
1049            "pure FTS query should build without a vector"
1050        );
1051        assert!(!unsafe { query.unwrap().as_raw() }.is_null());
1052    }
1053
1054    #[test]
1055    fn test_diskann_query_params_create_and_getters() {
1056        let mut params = DiskannQueryParams::new(200);
1057        assert_eq!(params.list_size(), 200);
1058        params.set_list_size(300).expect("set list_size");
1059        assert_eq!(params.list_size(), 300);
1060    }
1061
1062    #[test]
1063    fn test_ivf_rabitq_query_params_create_and_getters() {
1064        let mut params = IvfRabitqQueryParams::new(16, 0.0, false, true);
1065        assert_eq!(params.nprobe(), 16);
1066        params.set_nprobe(32).expect("set nprobe");
1067        assert_eq!(params.nprobe(), 32);
1068        params.set_scale_factor(2.5).expect("set scale_factor");
1069        assert!((params.scale_factor() - 2.5).abs() < f32::EPSILON);
1070    }
1071
1072    #[test]
1073    fn test_group_by_query_new_rejects_nul_field_name() {
1074        // The native handle is already created when the field name conversion
1075        // fails, so it must be owned by the wrapper to avoid leaking it.
1076        let result = GroupBySearchQuery::new("embedding\0", "category", &[0.1, 0.2], 2, 2);
1077        assert!(result.is_err());
1078        if let Err(e) = result {
1079            assert_eq!(e.code, ErrorCode::InvalidArgument);
1080        }
1081    }
1082
1083    #[test]
1084    fn test_vector_query_new_rejects_nul_field_name() {
1085        let result = SearchQuery::new("embedding\0", &[0.1, 0.2], 2);
1086        assert!(result.is_err());
1087        if let Err(e) = result {
1088            assert_eq!(e.code, ErrorCode::InvalidArgument);
1089        }
1090    }
1091}
zvec_rust/query.rs

zvec_rust/
query.rs
