velesdb_core/collection/
search.rs

1//! Search implementation for Collection.
2
3use super::types::Collection;
4use crate::error::{Error, Result};
5use crate::index::VectorIndex;
6use crate::point::{Point, SearchResult};
7use crate::storage::{PayloadStorage, VectorStorage};
8
9/// Wrapper for f32 to implement Ord for `BinaryHeap` in hybrid search.
10#[derive(Debug, Clone, Copy, PartialEq)]
11struct OrderedFloat(f32);
12
13impl Eq for OrderedFloat {}
14
15impl PartialOrd for OrderedFloat {
16    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
17        Some(self.cmp(other))
18    }
19}
20
21impl Ord for OrderedFloat {
22    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
23        self.0
24            .partial_cmp(&other.0)
25            .unwrap_or(std::cmp::Ordering::Equal)
26    }
27}
28
29impl Collection {
30    /// Searches for the k nearest neighbors of the query vector.
31    ///
32    /// Uses HNSW index for fast approximate nearest neighbor search.
33    ///
34    /// # Errors
35    ///
36    /// Returns an error if the query vector dimension doesn't match the collection,
37    /// or if this is a metadata-only collection (use `query()` instead).
38    pub fn search(&self, query: &[f32], k: usize) -> Result<Vec<SearchResult>> {
39        let config = self.config.read();
40
41        // Metadata-only collections don't support vector search
42        if config.metadata_only {
43            return Err(Error::SearchNotSupported(config.name.clone()));
44        }
45
46        if query.len() != config.dimension {
47            return Err(Error::DimensionMismatch {
48                expected: config.dimension,
49                actual: query.len(),
50            });
51        }
52        drop(config);
53
54        // Use HNSW index for fast ANN search
55        let index_results = self.index.search(query, k);
56
57        let vector_storage = self.vector_storage.read();
58        let payload_storage = self.payload_storage.read();
59
60        // Map index results to SearchResult with full point data
61        let results: Vec<SearchResult> = index_results
62            .into_iter()
63            .filter_map(|(id, score)| {
64                // We need to fetch vector and payload
65                let vector = vector_storage.retrieve(id).ok().flatten()?;
66                let payload = payload_storage.retrieve(id).ok().flatten();
67
68                let point = Point {
69                    id,
70                    vector,
71                    payload,
72                };
73
74                Some(SearchResult::new(point, score))
75            })
76            .collect();
77
78        Ok(results)
79    }
80
81    /// Performs vector similarity search with custom `ef_search` parameter.
82    ///
83    /// Higher `ef_search` = better recall, slower search.
84    /// Default `ef_search` is 128 (Balanced mode).
85    ///
86    /// # Errors
87    ///
88    /// Returns an error if the query vector dimension doesn't match the collection.
89    pub fn search_with_ef(
90        &self,
91        query: &[f32],
92        k: usize,
93        ef_search: usize,
94    ) -> Result<Vec<SearchResult>> {
95        let config = self.config.read();
96
97        if query.len() != config.dimension {
98            return Err(Error::DimensionMismatch {
99                expected: config.dimension,
100                actual: query.len(),
101            });
102        }
103        drop(config);
104
105        // Convert ef_search to SearchQuality
106        let quality = match ef_search {
107            0..=64 => crate::SearchQuality::Fast,
108            65..=128 => crate::SearchQuality::Balanced,
109            129..=256 => crate::SearchQuality::Accurate,
110            _ => crate::SearchQuality::Perfect,
111        };
112
113        let index_results = self.index.search_with_quality(query, k, quality);
114
115        let vector_storage = self.vector_storage.read();
116        let payload_storage = self.payload_storage.read();
117
118        let results: Vec<SearchResult> = index_results
119            .into_iter()
120            .filter_map(|(id, score)| {
121                let vector = vector_storage.retrieve(id).ok().flatten()?;
122                let payload = payload_storage.retrieve(id).ok().flatten();
123
124                let point = Point {
125                    id,
126                    vector,
127                    payload,
128                };
129
130                Some(SearchResult::new(point, score))
131            })
132            .collect();
133
134        Ok(results)
135    }
136
137    /// Performs fast vector similarity search returning only IDs and scores.
138    ///
139    /// Perf: This is ~3-5x faster than `search()` because it skips vector/payload retrieval.
140    /// Use this when you only need IDs and scores, not full point data.
141    ///
142    /// # Arguments
143    ///
144    /// * `query` - Query vector
145    /// * `k` - Maximum number of results to return
146    ///
147    /// # Returns
148    ///
149    /// Vector of (id, score) tuples sorted by similarity.
150    ///
151    /// # Errors
152    ///
153    /// Returns an error if the query vector dimension doesn't match the collection.
154    pub fn search_ids(&self, query: &[f32], k: usize) -> Result<Vec<(u64, f32)>> {
155        let config = self.config.read();
156
157        if query.len() != config.dimension {
158            return Err(Error::DimensionMismatch {
159                expected: config.dimension,
160                actual: query.len(),
161            });
162        }
163        drop(config);
164
165        // Perf: Direct HNSW search without vector/payload retrieval
166        let results = self.index.search(query, k);
167        Ok(results)
168    }
169
170    /// Performs batch search for multiple query vectors in parallel with metadata filtering.
171    /// Supports a different filter for each query in the batch.
172    ///
173    /// # Arguments
174    ///
175    /// * `queries` - List of query vector slices
176    /// * `k` - Maximum number of results per query
177    /// * `filters` - List of optional filters (must match queries length)
178    ///
179    /// # Returns
180    ///
181    /// Vector of search results for each query, matching its respective filter.
182    ///
183    /// # Errors
184    ///
185    /// Returns an error if queries and filters have different lengths or dimension mismatch.
186    pub fn search_batch_with_filters(
187        &self,
188        queries: &[&[f32]],
189        k: usize,
190        filters: &[Option<crate::filter::Filter>],
191    ) -> Result<Vec<Vec<SearchResult>>> {
192        use crate::index::SearchQuality;
193
194        if queries.len() != filters.len() {
195            return Err(Error::Config(format!(
196                "Queries count ({}) does not match filters count ({})",
197                queries.len(),
198                filters.len()
199            )));
200        }
201
202        let config = self.config.read();
203        let dimension = config.dimension;
204        drop(config);
205
206        // Validate all query dimensions
207        for query in queries {
208            if query.len() != dimension {
209                return Err(Error::DimensionMismatch {
210                    expected: dimension,
211                    actual: query.len(),
212                });
213            }
214        }
215
216        // We need to retrieve more candidates for post-filtering
217        let candidates_k = k.saturating_mul(4).max(k + 10);
218        let index_results =
219            self.index
220                .search_batch_parallel(queries, candidates_k, SearchQuality::Balanced);
221
222        let vector_storage = self.vector_storage.read();
223        let payload_storage = self.payload_storage.read();
224
225        let mut all_results = Vec::with_capacity(queries.len());
226
227        for (query_results, filter_opt) in index_results.into_iter().zip(filters) {
228            let mut filtered_results: Vec<SearchResult> = query_results
229                .into_iter()
230                .filter_map(|(id, score)| {
231                    let payload = payload_storage.retrieve(id).ok().flatten();
232
233                    // Apply filter if present
234                    if let Some(ref filter) = filter_opt {
235                        if let Some(ref p) = payload {
236                            if !filter.matches(p) {
237                                return None;
238                            }
239                        } else if !filter.matches(&serde_json::Value::Null) {
240                            return None;
241                        }
242                    }
243
244                    let vector = vector_storage.retrieve(id).ok().flatten()?;
245                    Some(SearchResult {
246                        point: Point {
247                            id,
248                            vector,
249                            payload,
250                        },
251                        score,
252                    })
253                })
254                .collect();
255
256            // Sort and truncate to k
257            let higher_is_better = self.config.read().metric.higher_is_better();
258            if higher_is_better {
259                filtered_results.sort_by(|a, b| {
260                    b.score
261                        .partial_cmp(&a.score)
262                        .unwrap_or(std::cmp::Ordering::Equal)
263                });
264            } else {
265                filtered_results.sort_by(|a, b| {
266                    a.score
267                        .partial_cmp(&b.score)
268                        .unwrap_or(std::cmp::Ordering::Equal)
269                });
270            }
271            filtered_results.truncate(k);
272
273            all_results.push(filtered_results);
274        }
275
276        Ok(all_results)
277    }
278
279    /// Performs batch search for multiple query vectors in parallel with a single metadata filter.
280    ///
281    /// # Arguments
282    ///
283    /// * `queries` - List of query vector slices
284    /// * `k` - Maximum number of results per query
285    /// * `filter` - Metadata filter to apply to all results
286    ///
287    /// # Errors
288    ///
289    /// Returns an error if any query has incorrect dimension.
290    pub fn search_batch_with_filter(
291        &self,
292        queries: &[&[f32]],
293        k: usize,
294        filter: &crate::filter::Filter,
295    ) -> Result<Vec<Vec<SearchResult>>> {
296        let filters: Vec<Option<crate::filter::Filter>> = vec![Some(filter.clone()); queries.len()];
297        self.search_batch_with_filters(queries, k, &filters)
298    }
299
300    /// Performs batch search for multiple query vectors in parallel.
301    ///
302    /// This method is optimized for high throughput using parallel index traversal.
303    ///
304    /// # Arguments
305    ///
306    /// * `queries` - List of query vector slices
307    /// * `k` - Maximum number of results per query
308    ///
309    /// # Returns
310    ///
311    /// Vector of search results for each query, with full point data.
312    ///
313    /// # Errors
314    ///
315    /// Returns an error if any query vector dimension doesn't match the collection.
316    pub fn search_batch_parallel(
317        &self,
318        queries: &[&[f32]],
319        k: usize,
320    ) -> Result<Vec<Vec<SearchResult>>> {
321        use crate::index::SearchQuality;
322
323        let config = self.config.read();
324        let dimension = config.dimension;
325        drop(config);
326
327        // Validate all query dimensions first
328        for query in queries {
329            if query.len() != dimension {
330                return Err(Error::DimensionMismatch {
331                    expected: dimension,
332                    actual: query.len(),
333                });
334            }
335        }
336
337        // Perf: Use parallel HNSW search (P0 optimization)
338        let index_results = self
339            .index
340            .search_batch_parallel(queries, k, SearchQuality::Balanced);
341
342        // Map results to SearchResult with full point data
343        let vector_storage = self.vector_storage.read();
344        let payload_storage = self.payload_storage.read();
345
346        let results: Vec<Vec<SearchResult>> = index_results
347            .into_iter()
348            .map(|query_results: Vec<(u64, f32)>| {
349                query_results
350                    .into_iter()
351                    .filter_map(|(id, score)| {
352                        let vector = vector_storage.retrieve(id).ok().flatten()?;
353                        let payload = payload_storage.retrieve(id).ok().flatten();
354                        Some(SearchResult {
355                            point: Point {
356                                id,
357                                vector,
358                                payload,
359                            },
360                            score,
361                        })
362                    })
363                    .collect()
364            })
365            .collect();
366
367        Ok(results)
368    }
369
370    /// Executes a `VelesQL` query on this collection.
371    ///
372    /// This method unifies vector search, text search, and metadata filtering
373    /// into a single interface.
374    ///
375    /// # Arguments
376    ///
377    /// * `query` - Parsed `VelesQL` query
378    /// * `params` - Query parameters for resolving placeholders (e.g., $v)
379    ///
380    /// # Errors
381    ///
382    /// Returns an error if the query cannot be executed (e.g., missing parameters).
383    pub fn execute_query(
384        &self,
385        query: &crate::velesql::Query,
386        params: &std::collections::HashMap<String, serde_json::Value>,
387    ) -> Result<Vec<SearchResult>> {
388        let stmt = &query.select;
389        let limit = usize::try_from(stmt.limit.unwrap_or(10)).unwrap_or(usize::MAX);
390
391        // 1. Extract vector search (NEAR) if present
392        let mut vector_search = None;
393        let mut filter_condition = None;
394
395        if let Some(ref cond) = stmt.where_clause {
396            let mut extracted_cond = cond.clone();
397            vector_search = self.extract_vector_search(&mut extracted_cond, params)?;
398            filter_condition = Some(extracted_cond);
399        }
400
401        // 2. Resolve WITH clause options
402        let mut ef_search = None;
403        if let Some(ref with) = stmt.with_clause {
404            ef_search = with.get_ef_search();
405        }
406
407        // 3. Execute query based on extracted components
408        let results = match (vector_search, filter_condition) {
409            (Some(vector), Some(ref cond)) => {
410                // Check if condition contains MATCH for hybrid search
411                if let Some(text_query) = Self::extract_match_query(cond) {
412                    // Hybrid search: NEAR + MATCH
413                    self.hybrid_search(&vector, &text_query, limit, None)?
414                } else {
415                    // Vector search with metadata filter
416                    let filter =
417                        crate::filter::Filter::new(crate::filter::Condition::from(cond.clone()));
418                    self.search_with_filter(&vector, limit, &filter)?
419                }
420            }
421            (Some(vector), None) => {
422                // Pure vector search
423                if let Some(ef) = ef_search {
424                    self.search_with_ef(&vector, limit, ef)?
425                } else {
426                    self.search(&vector, limit)?
427                }
428            }
429            (None, Some(cond)) => {
430                // Metadata-only filter (table scan + filter)
431                // If it's a MATCH condition, use text search
432                if let crate::velesql::Condition::Match(ref m) = cond {
433                    // Pure text search - no filter needed
434                    self.text_search(&m.query, limit)
435                } else {
436                    // Generic metadata filter: perform a scan (fallback)
437                    let filter = crate::filter::Filter::new(crate::filter::Condition::from(cond));
438                    self.execute_scan_query(&filter, limit)
439                }
440            }
441            (None, None) => {
442                // SELECT * FROM docs LIMIT N (no WHERE)
443                self.execute_scan_query(
444                    &crate::filter::Filter::new(crate::filter::Condition::And {
445                        conditions: vec![],
446                    }),
447                    limit,
448                )
449            }
450        };
451
452        Ok(results)
453    }
454
455    /// Helper to extract MATCH query from any nested condition.
456    fn extract_match_query(condition: &crate::velesql::Condition) -> Option<String> {
457        use crate::velesql::Condition;
458        match condition {
459            Condition::Match(m) => Some(m.query.clone()),
460            Condition::And(left, right) => {
461                Self::extract_match_query(left).or_else(|| Self::extract_match_query(right))
462            }
463            Condition::Group(inner) => Self::extract_match_query(inner),
464            _ => None,
465        }
466    }
467
468    /// Internal helper to extract vector search from WHERE clause.
469    #[allow(clippy::self_only_used_in_recursion)]
470    fn extract_vector_search(
471        &self,
472        condition: &mut crate::velesql::Condition,
473        params: &std::collections::HashMap<String, serde_json::Value>,
474    ) -> Result<Option<Vec<f32>>> {
475        use crate::velesql::{Condition, VectorExpr};
476
477        match condition {
478            Condition::VectorSearch(vs) => {
479                let vec = match &vs.vector {
480                    VectorExpr::Literal(v) => v.clone(),
481                    VectorExpr::Parameter(name) => {
482                        let val = params.get(name).ok_or_else(|| {
483                            Error::Config(format!("Missing query parameter: ${name}"))
484                        })?;
485                        if let serde_json::Value::Array(arr) = val {
486                            #[allow(clippy::cast_possible_truncation)]
487                            arr.iter()
488                                .map(|v| {
489                                    v.as_f64().map(|f| f as f32).ok_or_else(|| {
490                                        Error::Config(format!(
491                                            "Invalid vector parameter ${name}: expected numbers"
492                                        ))
493                                    })
494                                })
495                                .collect::<Result<Vec<f32>>>()?
496                        } else {
497                            return Err(Error::Config(format!(
498                                "Invalid vector parameter ${name}: expected array"
499                            )));
500                        }
501                    }
502                };
503                Ok(Some(vec))
504            }
505            Condition::And(left, right) => {
506                if let Some(v) = self.extract_vector_search(left, params)? {
507                    return Ok(Some(v));
508                }
509                self.extract_vector_search(right, params)
510            }
511            Condition::Group(inner) => self.extract_vector_search(inner, params),
512            _ => Ok(None),
513        }
514    }
515
516    /// Fallback method for metadata-only queries without vector search.
517    fn execute_scan_query(
518        &self,
519        filter: &crate::filter::Filter,
520        limit: usize,
521    ) -> Vec<SearchResult> {
522        let payload_storage = self.payload_storage.read();
523        let vector_storage = self.vector_storage.read();
524
525        // Scan all points (slow fallback)
526        // In production, this should use metadata indexes
527        let mut results = Vec::new();
528
529        // We need all IDs to scan
530        let ids = vector_storage.ids();
531
532        for id in ids {
533            let payload = payload_storage.retrieve(id).ok().flatten();
534            let matches = match payload {
535                Some(ref p) => filter.matches(p),
536                None => filter.matches(&serde_json::Value::Null),
537            };
538
539            if matches {
540                if let Ok(Some(vector)) = vector_storage.retrieve(id) {
541                    results.push(SearchResult::new(
542                        Point {
543                            id,
544                            vector,
545                            payload,
546                        },
547                        1.0, // Constant score for scans
548                    ));
549                }
550            }
551
552            if results.len() >= limit {
553                break;
554            }
555        }
556
557        results
558    }
559
560    /// Performs full-text search using BM25.
561    ///
562    /// # Arguments
563    ///
564    /// * `query` - Text query to search for
565    /// * `k` - Maximum number of results to return
566    ///
567    /// # Returns
568    ///
569    /// Vector of search results sorted by BM25 score (descending).
570    #[must_use]
571    pub fn text_search(&self, query: &str, k: usize) -> Vec<SearchResult> {
572        let bm25_results = self.text_index.search(query, k);
573
574        let vector_storage = self.vector_storage.read();
575        let payload_storage = self.payload_storage.read();
576
577        bm25_results
578            .into_iter()
579            .filter_map(|(id, score)| {
580                let vector = vector_storage.retrieve(id).ok().flatten()?;
581                let payload = payload_storage.retrieve(id).ok().flatten();
582
583                let point = Point {
584                    id,
585                    vector,
586                    payload,
587                };
588
589                Some(SearchResult::new(point, score))
590            })
591            .collect()
592    }
593
594    /// Performs hybrid search combining vector similarity and full-text search.
595    ///
596    /// Uses Reciprocal Rank Fusion (RRF) to combine results from both searches.
597    ///
598    /// # Arguments
599    ///
600    /// * `vector_query` - Query vector for similarity search
601    /// * `text_query` - Text query for BM25 search
602    /// * `k` - Maximum number of results to return
603    /// * `vector_weight` - Weight for vector results (0.0-1.0, default 0.5)
604    ///
605    /// # Performance (v0.9+)
606    ///
607    /// - **Streaming RRF**: `BinaryHeap` maintains top-k during fusion (O(n log k) vs O(n log n))
608    /// - **Vector-first gating**: Text search limited to 2k candidates for efficiency
609    /// - **`FxHashMap`**: Faster hashing for score aggregation
610    ///
611    /// # Errors
612    ///
613    /// Returns an error if the query vector dimension doesn't match.
614    pub fn hybrid_search(
615        &self,
616        vector_query: &[f32],
617        text_query: &str,
618        k: usize,
619        vector_weight: Option<f32>,
620    ) -> Result<Vec<SearchResult>> {
621        use std::cmp::Reverse;
622        use std::collections::BinaryHeap;
623
624        let config = self.config.read();
625        if vector_query.len() != config.dimension {
626            return Err(Error::DimensionMismatch {
627                expected: config.dimension,
628                actual: vector_query.len(),
629            });
630        }
631        drop(config);
632
633        let weight = vector_weight.unwrap_or(0.5).clamp(0.0, 1.0);
634        let text_weight = 1.0 - weight;
635
636        // Get vector search results (more than k to allow for fusion)
637        let vector_results = self.index.search(vector_query, k * 2);
638
639        // Get BM25 text search results
640        let text_results = self.text_index.search(text_query, k * 2);
641
642        // Perf: Apply RRF with FxHashMap for faster hashing
643        // RRF score = weight / (rank + 60) - the constant 60 is standard (Cormack et al.)
644        let mut fused_scores: rustc_hash::FxHashMap<u64, f32> =
645            rustc_hash::FxHashMap::with_capacity_and_hasher(
646                vector_results.len() + text_results.len(),
647                rustc_hash::FxBuildHasher,
648            );
649
650        // Add vector scores with RRF
651        #[allow(clippy::cast_precision_loss)]
652        for (rank, (id, _)) in vector_results.iter().enumerate() {
653            let rrf_score = weight / (rank as f32 + 60.0);
654            *fused_scores.entry(*id).or_insert(0.0) += rrf_score;
655        }
656
657        // Add text scores with RRF
658        #[allow(clippy::cast_precision_loss)]
659        for (rank, (id, _)) in text_results.iter().enumerate() {
660            let rrf_score = text_weight / (rank as f32 + 60.0);
661            *fused_scores.entry(*id).or_insert(0.0) += rrf_score;
662        }
663
664        // Perf: Streaming top-k with BinaryHeap (O(n log k) vs O(n log n) for full sort)
665        // Use min-heap of size k: always keep the k highest scores
666        let mut top_k: BinaryHeap<Reverse<(OrderedFloat, u64)>> = BinaryHeap::with_capacity(k + 1);
667
668        for (id, score) in fused_scores {
669            top_k.push(Reverse((OrderedFloat(score), id)));
670            if top_k.len() > k {
671                top_k.pop(); // Remove smallest
672            }
673        }
674
675        // Extract and sort descending
676        let mut scored_ids: Vec<(u64, f32)> = top_k
677            .into_iter()
678            .map(|Reverse((OrderedFloat(s), id))| (id, s))
679            .collect();
680        scored_ids.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
681
682        // Fetch full point data
683        let vector_storage = self.vector_storage.read();
684        let payload_storage = self.payload_storage.read();
685
686        let results: Vec<SearchResult> = scored_ids
687            .into_iter()
688            .filter_map(|(id, score)| {
689                let vector = vector_storage.retrieve(id).ok().flatten()?;
690                let payload = payload_storage.retrieve(id).ok().flatten();
691
692                let point = Point {
693                    id,
694                    vector,
695                    payload,
696                };
697
698                Some(SearchResult::new(point, score))
699            })
700            .collect();
701
702        Ok(results)
703    }
704
705    /// Searches for the k nearest neighbors with metadata filtering.
706    ///
707    /// Performs post-filtering: retrieves more candidates from HNSW,
708    /// then filters by metadata conditions.
709    ///
710    /// # Arguments
711    ///
712    /// * `query` - Query vector
713    /// * `k` - Maximum number of results to return
714    /// * `filter` - Metadata filter to apply
715    ///
716    /// # Errors
717    ///
718    /// Returns an error if the query vector dimension doesn't match the collection.
719    pub fn search_with_filter(
720        &self,
721        query: &[f32],
722        k: usize,
723        filter: &crate::filter::Filter,
724    ) -> Result<Vec<SearchResult>> {
725        let config = self.config.read();
726
727        if query.len() != config.dimension {
728            return Err(Error::DimensionMismatch {
729                expected: config.dimension,
730                actual: query.len(),
731            });
732        }
733        drop(config);
734
735        // Post-filtering strategy: retrieve more candidates than k, then filter
736        // Heuristic: retrieve 4x candidates to account for filtered-out results
737        let candidates_k = k.saturating_mul(4).max(k + 10);
738        let index_results = self.index.search(query, candidates_k);
739
740        let vector_storage = self.vector_storage.read();
741        let payload_storage = self.payload_storage.read();
742
743        // Map index results to SearchResult with full point data, applying filter
744        let mut results: Vec<SearchResult> = index_results
745            .into_iter()
746            .filter_map(|(id, score)| {
747                let vector = vector_storage.retrieve(id).ok().flatten()?;
748                let payload = payload_storage.retrieve(id).ok().flatten();
749
750                // Apply filter - if no payload, filter fails
751                let payload_ref = payload.as_ref()?;
752                if !filter.matches(payload_ref) {
753                    return None;
754                }
755
756                let point = Point {
757                    id,
758                    vector,
759                    payload,
760                };
761
762                Some(SearchResult::new(point, score))
763            })
764            .take(k)
765            .collect();
766
767        // Ensure results are sorted by score (should already be, but defensive)
768        results.sort_by(|a, b| {
769            b.score
770                .partial_cmp(&a.score)
771                .unwrap_or(std::cmp::Ordering::Equal)
772        });
773
774        Ok(results)
775    }
776
777    /// Performs full-text search with metadata filtering.
778    ///
779    /// # Arguments
780    ///
781    /// * `query` - Text query to search for
782    /// * `k` - Maximum number of results to return
783    /// * `filter` - Metadata filter to apply
784    ///
785    /// # Returns
786    ///
787    /// Vector of search results sorted by BM25 score (descending).
788    #[must_use]
789    pub fn text_search_with_filter(
790        &self,
791        query: &str,
792        k: usize,
793        filter: &crate::filter::Filter,
794    ) -> Vec<SearchResult> {
795        // Retrieve more candidates for filtering
796        let candidates_k = k.saturating_mul(4).max(k + 10);
797        let bm25_results = self.text_index.search(query, candidates_k);
798
799        let vector_storage = self.vector_storage.read();
800        let payload_storage = self.payload_storage.read();
801
802        bm25_results
803            .into_iter()
804            .filter_map(|(id, score)| {
805                let vector = vector_storage.retrieve(id).ok().flatten()?;
806                let payload = payload_storage.retrieve(id).ok().flatten();
807
808                // Apply filter - if no payload, filter fails
809                let payload_ref = payload.as_ref()?;
810                if !filter.matches(payload_ref) {
811                    return None;
812                }
813
814                let point = Point {
815                    id,
816                    vector,
817                    payload,
818                };
819
820                Some(SearchResult::new(point, score))
821            })
822            .take(k)
823            .collect()
824    }
825
826    /// Performs hybrid search (vector + text) with metadata filtering.
827    ///
828    /// Uses Reciprocal Rank Fusion (RRF) to combine results from both searches,
829    /// then applies metadata filter.
830    ///
831    /// # Arguments
832    ///
833    /// * `vector_query` - Query vector for similarity search
834    /// * `text_query` - Text query for BM25 search
835    /// * `k` - Maximum number of results to return
836    /// * `vector_weight` - Weight for vector results (0.0-1.0, default 0.5)
837    /// * `filter` - Metadata filter to apply
838    ///
839    /// # Errors
840    ///
841    /// Returns an error if the query vector dimension doesn't match.
842    pub fn hybrid_search_with_filter(
843        &self,
844        vector_query: &[f32],
845        text_query: &str,
846        k: usize,
847        vector_weight: Option<f32>,
848        filter: &crate::filter::Filter,
849    ) -> Result<Vec<SearchResult>> {
850        let config = self.config.read();
851        if vector_query.len() != config.dimension {
852            return Err(Error::DimensionMismatch {
853                expected: config.dimension,
854                actual: vector_query.len(),
855            });
856        }
857        drop(config);
858
859        let weight = vector_weight.unwrap_or(0.5).clamp(0.0, 1.0);
860        let text_weight = 1.0 - weight;
861
862        // Get more candidates for filtering
863        let candidates_k = k.saturating_mul(4).max(k + 10);
864
865        // Get vector search results
866        let vector_results = self.index.search(vector_query, candidates_k);
867
868        // Get BM25 text search results
869        let text_results = self.text_index.search(text_query, candidates_k);
870
871        // Apply RRF (Reciprocal Rank Fusion)
872        let mut fused_scores: rustc_hash::FxHashMap<u64, f32> = rustc_hash::FxHashMap::default();
873
874        #[allow(clippy::cast_precision_loss)]
875        for (rank, (id, _)) in vector_results.iter().enumerate() {
876            let rrf_score = weight / (rank as f32 + 60.0);
877            *fused_scores.entry(*id).or_insert(0.0) += rrf_score;
878        }
879
880        #[allow(clippy::cast_precision_loss)]
881        for (rank, (id, _)) in text_results.iter().enumerate() {
882            let rrf_score = text_weight / (rank as f32 + 60.0);
883            *fused_scores.entry(*id).or_insert(0.0) += rrf_score;
884        }
885
886        // Sort by fused score
887        let mut scored_ids: Vec<_> = fused_scores.into_iter().collect();
888        scored_ids.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
889
890        // Fetch full point data and apply filter
891        let vector_storage = self.vector_storage.read();
892        let payload_storage = self.payload_storage.read();
893
894        let results: Vec<SearchResult> = scored_ids
895            .into_iter()
896            .filter_map(|(id, score)| {
897                let vector = vector_storage.retrieve(id).ok().flatten()?;
898                let payload = payload_storage.retrieve(id).ok().flatten();
899
900                // Apply filter - if no payload, filter fails
901                let payload_ref = payload.as_ref()?;
902                if !filter.matches(payload_ref) {
903                    return None;
904                }
905
906                let point = Point {
907                    id,
908                    vector,
909                    payload,
910                };
911
912                Some(SearchResult::new(point, score))
913            })
914            .take(k)
915            .collect();
916
917        Ok(results)
918    }
919
920    /// Performs multi-query search with result fusion.
921    ///
922    /// This method executes parallel searches for multiple query vectors and fuses
923    /// the results using the specified fusion strategy. Ideal for Multiple Query
924    /// Generation (MQG) pipelines where multiple reformulations of a user query
925    /// are searched simultaneously.
926    ///
927    /// # Arguments
928    ///
929    /// * `vectors` - Slice of query vectors (all must have same dimension)
930    /// * `top_k` - Maximum number of results to return after fusion
931    /// * `fusion` - Strategy for combining results (Average, Maximum, RRF, Weighted)
932    /// * `filter` - Optional metadata filter to apply to all queries
933    ///
934    /// # Returns
935    ///
936    /// Vector of `SearchResult` sorted by fused score descending.
937    ///
938    /// # Errors
939    ///
940    /// Returns an error if:
941    /// - `vectors` is empty
942    /// - Any vector has incorrect dimension
943    /// - More than 10 vectors are provided (configurable limit)
944    ///
945    /// # Example
946    ///
947    /// ```rust,ignore
948    /// use velesdb_core::fusion::FusionStrategy;
949    ///
950    /// // Search with 4 query reformulations using weighted fusion
951    /// let results = collection.multi_query_search(
952    ///     &[&query1, &query2, &query3, &query4],
953    ///     10,
954    ///     FusionStrategy::Weighted {
955    ///         avg_weight: 0.6,
956    ///         max_weight: 0.3,
957    ///         hit_weight: 0.1,
958    ///     },
959    ///     None,
960    /// )?;
961    /// ```
962    #[allow(clippy::needless_pass_by_value)]
963    pub fn multi_query_search(
964        &self,
965        vectors: &[&[f32]],
966        top_k: usize,
967        fusion: crate::fusion::FusionStrategy,
968        filter: Option<&crate::filter::Filter>,
969    ) -> Result<Vec<SearchResult>> {
970        const MAX_VECTORS: usize = 10;
971
972        // Validation: non-empty
973        if vectors.is_empty() {
974            return Err(Error::Config(
975                "multi_query_search requires at least one vector".into(),
976            ));
977        }
978
979        // Validation: max vectors limit
980        if vectors.len() > MAX_VECTORS {
981            return Err(Error::Config(format!(
982                "multi_query_search supports at most {MAX_VECTORS} vectors, got {}",
983                vectors.len()
984            )));
985        }
986
987        // Validation: dimension consistency
988        let config = self.config.read();
989        let dimension = config.dimension;
990        drop(config);
991
992        for vector in vectors {
993            if vector.len() != dimension {
994                return Err(Error::DimensionMismatch {
995                    expected: dimension,
996                    actual: vector.len(),
997                });
998            }
999        }
1000
1001        // Calculate overfetch factor for better fusion quality
1002        let overfetch_k = match top_k {
1003            0..=10 => top_k * 20,
1004            11..=50 => top_k * 10,
1005            51..=100 => top_k * 5,
1006            _ => top_k * 2,
1007        };
1008
1009        // Execute parallel batch search
1010        let batch_results =
1011            self.index
1012                .search_batch_parallel(vectors, overfetch_k, crate::SearchQuality::Balanced);
1013
1014        // Apply filter if present (pre-fusion filtering)
1015        let filtered_results: Vec<Vec<(u64, f32)>> = if let Some(f) = filter {
1016            let payload_storage = self.payload_storage.read();
1017            batch_results
1018                .into_iter()
1019                .map(|query_results| {
1020                    query_results
1021                        .into_iter()
1022                        .filter(|(id, _score)| {
1023                            if let Ok(Some(payload)) = payload_storage.retrieve(*id) {
1024                                f.matches(&payload)
1025                            } else {
1026                                false
1027                            }
1028                        })
1029                        .collect()
1030                })
1031                .collect()
1032        } else {
1033            batch_results
1034        };
1035
1036        // Fuse results using the specified strategy
1037        let fused = fusion
1038            .fuse(filtered_results)
1039            .map_err(|e| Error::Config(format!("Fusion error: {e}")))?;
1040
1041        // Fetch full point data for top_k results
1042        let vector_storage = self.vector_storage.read();
1043        let payload_storage = self.payload_storage.read();
1044
1045        let results: Vec<SearchResult> = fused
1046            .into_iter()
1047            .take(top_k)
1048            .filter_map(|(id, score)| {
1049                let vector = vector_storage.retrieve(id).ok().flatten()?;
1050                let payload = payload_storage.retrieve(id).ok().flatten();
1051
1052                let point = Point {
1053                    id,
1054                    vector,
1055                    payload,
1056                };
1057
1058                Some(SearchResult::new(point, score))
1059            })
1060            .collect();
1061
1062        Ok(results)
1063    }
1064
1065    /// Performs multi-query search returning only IDs and fused scores.
1066    ///
1067    /// This is a faster variant of `multi_query_search` that skips fetching
1068    /// vector and payload data. Use when you only need document IDs.
1069    ///
1070    /// # Arguments
1071    ///
1072    /// * `vectors` - Slice of query vectors
1073    /// * `top_k` - Maximum number of results
1074    /// * `fusion` - Fusion strategy
1075    ///
1076    /// # Returns
1077    ///
1078    /// Vector of `(id, fused_score)` tuples sorted by score descending.
1079    ///
1080    /// # Errors
1081    ///
1082    /// Returns an error if vectors is empty, exceeds max limit, or has dimension mismatch.
1083    #[allow(clippy::needless_pass_by_value)]
1084    pub fn multi_query_search_ids(
1085        &self,
1086        vectors: &[&[f32]],
1087        top_k: usize,
1088        fusion: crate::fusion::FusionStrategy,
1089    ) -> Result<Vec<(u64, f32)>> {
1090        const MAX_VECTORS: usize = 10;
1091
1092        if vectors.is_empty() {
1093            return Err(Error::Config(
1094                "multi_query_search requires at least one vector".into(),
1095            ));
1096        }
1097
1098        if vectors.len() > MAX_VECTORS {
1099            return Err(Error::Config(format!(
1100                "multi_query_search supports at most {MAX_VECTORS} vectors, got {}",
1101                vectors.len()
1102            )));
1103        }
1104
1105        let config = self.config.read();
1106        let dimension = config.dimension;
1107        drop(config);
1108
1109        for vector in vectors {
1110            if vector.len() != dimension {
1111                return Err(Error::DimensionMismatch {
1112                    expected: dimension,
1113                    actual: vector.len(),
1114                });
1115            }
1116        }
1117
1118        let overfetch_k = match top_k {
1119            0..=10 => top_k * 20,
1120            11..=50 => top_k * 10,
1121            51..=100 => top_k * 5,
1122            _ => top_k * 2,
1123        };
1124
1125        let batch_results =
1126            self.index
1127                .search_batch_parallel(vectors, overfetch_k, crate::SearchQuality::Balanced);
1128
1129        let fused = fusion
1130            .fuse(batch_results)
1131            .map_err(|e| Error::Config(format!("Fusion error: {e}")))?;
1132
1133        Ok(fused.into_iter().take(top_k).collect())
1134    }
1135}
velesdb_core/collection/search.rs

velesdb_core/collection/
search.rs
