clt_database/index_method/
toy_vector_sparse_ivf.rs

1use std::{
2    collections::{BTreeSet, HashSet, VecDeque},
3    sync::atomic::Ordering,
4};
5
6use turso_parser::ast::{self, SortOrder};
7
8use crate::numeric::Numeric;
9use crate::util::quote_identifier;
10use crate::{
11    index_method::{
12        open_index_cursor, open_table_cursor, parse_patterns, IndexMethod, IndexMethodAttachment,
13        IndexMethodConfiguration, IndexMethodCursor, IndexMethodDefinition,
14        BACKING_BTREE_INDEX_METHOD_NAME, TOY_VECTOR_SPARSE_IVF_INDEX_METHOD_NAME,
15    },
16    return_if_io,
17    storage::btree::{BTreeCursor, BTreeKey, CursorTrait},
18    sync::Arc,
19    translate::collate::CollationSeq,
20    types::{IOResult, ImmutableRecord, KeyInfo, SeekKey, SeekOp, SeekResult},
21    vdbe::Register,
22    vector::{
23        operations,
24        vector_types::{Vector, VectorType},
25    },
26    Connection, LimboError, Result, Value, ValueRef,
27};
28
29/// Simple inverted index for sparse vectors
30/// > CREATE INDEX t_idx ON t USING toy_vector_sparse_ivf (embedding)
31///
32/// It accept single column which must contain vector encoded in sparse format (e.g. vector32_sparse(...))
33/// It can handle jaccard similarity scoring queries like the following:
34/// > SELECT vector_distance_jaccard(embedding, ?) as d FROM t ORDER BY d LIMIT ?
35#[derive(Debug)]
36pub struct VectorSparseInvertedIndexMethod;
37
38#[derive(Debug)]
39pub struct VectorSparseInvertedIndexMethodAttachment {
40    configuration: IndexMethodConfiguration,
41    patterns: Vec<ast::Select>,
42}
43
44#[derive(Debug)]
45pub enum VectorSparseInvertedIndexInsertState {
46    Init,
47    Prepare {
48        vector: Option<Vector<'static>>,
49        sum: f64,
50        rowid: i64,
51        idx: usize,
52    },
53    SeekInverted {
54        vector: Option<Vector<'static>>,
55        sum: f64,
56        key: Option<ImmutableRecord>,
57        rowid: i64,
58        idx: usize,
59    },
60    InsertInverted {
61        vector: Option<Vector<'static>>,
62        sum: f64,
63        key: Option<ImmutableRecord>,
64        rowid: i64,
65        idx: usize,
66    },
67    SeekStats {
68        vector: Option<Vector<'static>>,
69        sum: f64,
70        key: Option<ImmutableRecord>,
71        rowid: i64,
72        idx: usize,
73    },
74    ReadStats {
75        vector: Option<Vector<'static>>,
76        sum: f64,
77        rowid: i64,
78        idx: usize,
79    },
80    UpdateStats {
81        vector: Option<Vector<'static>>,
82        sum: f64,
83        key: Option<ImmutableRecord>,
84        rowid: i64,
85        idx: usize,
86    },
87}
88
89#[derive(Debug)]
90pub enum VectorSparseInvertedIndexDeleteState {
91    Init,
92    Prepare {
93        vector: Option<Vector<'static>>,
94        sum: f64,
95        rowid: i64,
96        idx: usize,
97    },
98    SeekInverted {
99        vector: Option<Vector<'static>>,
100        sum: f64,
101        key: Option<ImmutableRecord>,
102        rowid: i64,
103        idx: usize,
104    },
105    NextInverted {
106        vector: Option<Vector<'static>>,
107        sum: f64,
108        rowid: i64,
109        idx: usize,
110    },
111    DeleteInverted {
112        vector: Option<Vector<'static>>,
113        sum: f64,
114        rowid: i64,
115        idx: usize,
116    },
117    SeekStats {
118        vector: Option<Vector<'static>>,
119        sum: f64,
120        key: Option<ImmutableRecord>,
121        rowid: i64,
122        idx: usize,
123    },
124    ReadStats {
125        vector: Option<Vector<'static>>,
126        sum: f64,
127        rowid: i64,
128        idx: usize,
129    },
130    UpdateStats {
131        vector: Option<Vector<'static>>,
132        sum: f64,
133        key: Option<ImmutableRecord>,
134        rowid: i64,
135        idx: usize,
136    },
137}
138
139#[derive(Debug, PartialEq)]
140struct FloatOrd(f64);
141
142impl Eq for FloatOrd {}
143impl PartialOrd for FloatOrd {
144    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
145        Some(self.cmp(other))
146    }
147}
148impl Ord for FloatOrd {
149    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
150        self.0.total_cmp(&other.0)
151    }
152}
153
154#[derive(Debug)]
155struct ComponentStat {
156    position: u32,
157    cnt: i64,
158    min: f64,
159    max: f64,
160}
161
162fn parse_stat_row(record: Option<&ImmutableRecord>) -> Result<ComponentStat> {
163    let Some(record) = record else {
164        return Err(LimboError::Corrupt(
165            "stats index corrupted: expected row".to_string(),
166        ));
167    };
168    let ValueRef::Numeric(Numeric::Integer(position)) = record.get_value(0)? else {
169        return Err(LimboError::Corrupt(
170            "stats index corrupted: expected integer".to_string(),
171        ));
172    };
173    let ValueRef::Numeric(Numeric::Integer(cnt)) = record.get_value(1)? else {
174        return Err(LimboError::Corrupt(
175            "stats index corrupted: expected integer".to_string(),
176        ));
177    };
178    let ValueRef::Numeric(Numeric::Float(min)) = record.get_value(2)? else {
179        return Err(LimboError::Corrupt(
180            "stats index corrupted: expected float".to_string(),
181        ));
182    };
183    let ValueRef::Numeric(Numeric::Float(max)) = record.get_value(3)? else {
184        return Err(LimboError::Corrupt(
185            "stats index corrupted: expected float".to_string(),
186        ));
187    };
188    Ok(ComponentStat {
189        position: position as u32,
190        cnt,
191        min: f64::from(min),
192        max: f64::from(max),
193    })
194}
195#[derive(Debug)]
196struct ComponentRow {
197    position: u32,
198    sum: f64,
199    rowid: i64,
200}
201
202fn parse_inverted_index_row(record: Option<&ImmutableRecord>) -> Result<ComponentRow> {
203    let Some(record) = record else {
204        return Err(LimboError::Corrupt(
205            "inverted index corrupted: expected row".to_string(),
206        ));
207    };
208    let ValueRef::Numeric(Numeric::Integer(position)) = record.get_value(0)? else {
209        return Err(LimboError::Corrupt(
210            "inverted index corrupted: expected integer".to_string(),
211        ));
212    };
213    let ValueRef::Numeric(Numeric::Float(sum)) = record.get_value(1)? else {
214        return Err(LimboError::Corrupt(
215            "inverted index corrupted: expected float".to_string(),
216        ));
217    };
218    let ValueRef::Numeric(Numeric::Integer(rowid)) = record.get_value(2)? else {
219        return Err(LimboError::Corrupt(
220            "inverted index corrupted: expected integer".to_string(),
221        ));
222    };
223    Ok(ComponentRow {
224        position: position as u32,
225        sum: f64::from(sum),
226        rowid,
227    })
228}
229
230#[derive(Debug)]
231enum VectorSparseInvertedIndexSearchState {
232    Init,
233    CollectComponentsSeek {
234        sum: f64,
235        vector: Option<Vector<'static>>,
236        idx: usize,
237        components: Option<Vec<(ComponentStat, f32)>>,
238        limit: i64,
239        key: Option<ImmutableRecord>,
240    },
241    CollectComponentsRead {
242        sum: f64,
243        vector: Option<Vector<'static>>,
244        idx: usize,
245        components: Option<Vec<(ComponentStat, f32)>>,
246        limit: i64,
247    },
248    Seek {
249        sum: f64,
250        components: Option<VecDeque<ComponentStat>>,
251        collected: Option<HashSet<i64>>,
252        distances: Option<BTreeSet<(FloatOrd, i64)>>,
253        limit: i64,
254        key: Option<ImmutableRecord>,
255        sum_threshold: Option<f64>,
256        component: Option<u32>,
257    },
258    Read {
259        sum: f64,
260        components: Option<VecDeque<ComponentStat>>,
261        collected: Option<HashSet<i64>>,
262        distances: Option<BTreeSet<(FloatOrd, i64)>>,
263        limit: i64,
264        sum_threshold: Option<f64>,
265        component: u32,
266        current: Option<Vec<i64>>,
267    },
268    Next {
269        sum: f64,
270        components: Option<VecDeque<ComponentStat>>,
271        collected: Option<HashSet<i64>>,
272        distances: Option<BTreeSet<(FloatOrd, i64)>>,
273        limit: i64,
274        sum_threshold: Option<f64>,
275        component: u32,
276        current: Option<Vec<i64>>,
277    },
278    EvaluateSeek {
279        sum: f64,
280        components: Option<VecDeque<ComponentStat>>,
281        collected: Option<HashSet<i64>>,
282        distances: Option<BTreeSet<(FloatOrd, i64)>>,
283        limit: i64,
284        current: Option<VecDeque<i64>>,
285        rowid: Option<i64>,
286    },
287    EvaluateRead {
288        sum: f64,
289        components: Option<VecDeque<ComponentStat>>,
290        collected: Option<HashSet<i64>>,
291        distances: Option<BTreeSet<(FloatOrd, i64)>>,
292        limit: i64,
293        current: Option<VecDeque<i64>>,
294        rowid: i64,
295    },
296}
297
298#[derive(Debug, PartialEq)]
299pub enum ScanOrder {
300    DatasetFrequencyAsc,
301    QueryWeightDesc,
302}
303
304pub struct VectorSparseInvertedIndexMethodCursor {
305    configuration: IndexMethodConfiguration,
306    delta: f64,
307    scan_portion: f64,
308    scan_order: ScanOrder,
309    inverted_index_btree: String,
310    inverted_index_cursor: Option<BTreeCursor>,
311    stats_btree: String,
312    stats_cursor: Option<BTreeCursor>,
313    main_btree: Option<BTreeCursor>,
314    insert_state: VectorSparseInvertedIndexInsertState,
315    delete_state: VectorSparseInvertedIndexDeleteState,
316    search_state: VectorSparseInvertedIndexSearchState,
317    search_result: VecDeque<(i64, f64)>,
318}
319
320impl IndexMethod for VectorSparseInvertedIndexMethod {
321    fn attach(
322        &self,
323        configuration: &IndexMethodConfiguration,
324    ) -> Result<Arc<dyn IndexMethodAttachment>> {
325        let query_pattern1 = format!(
326            "SELECT vector_distance_jaccard({}, ?) as distance FROM {} ORDER BY distance LIMIT ?",
327            configuration.columns[0].name, configuration.table_name
328        );
329        let query_pattern2 = format!(
330            "SELECT vector_distance_jaccard(?, {}) as distance FROM {} ORDER BY distance LIMIT ?",
331            configuration.columns[0].name, configuration.table_name
332        );
333        Ok(Arc::new(VectorSparseInvertedIndexMethodAttachment {
334            configuration: configuration.clone(),
335            patterns: parse_patterns(&[&query_pattern1, &query_pattern2])?,
336        }))
337    }
338}
339
340impl IndexMethodAttachment for VectorSparseInvertedIndexMethodAttachment {
341    fn definition<'a>(&'a self) -> IndexMethodDefinition<'a> {
342        IndexMethodDefinition {
343            method_name: TOY_VECTOR_SPARSE_IVF_INDEX_METHOD_NAME,
344            index_name: &self.configuration.index_name,
345            patterns: self.patterns.as_slice(),
346            backing_btree: false,
347            results_materialized: true,
348        }
349    }
350    fn init(&self) -> Result<Box<dyn IndexMethodCursor>> {
351        Ok(Box::new(VectorSparseInvertedIndexMethodCursor::new(
352            self.configuration.clone(),
353        )))
354    }
355}
356
357impl VectorSparseInvertedIndexMethodCursor {
358    pub fn new(configuration: IndexMethodConfiguration) -> Self {
359        let inverted_index_btree = format!("{}_inverted_index", configuration.index_name);
360        let stats_btree = format!("{}_stats", configuration.index_name);
361        let delta = match configuration.parameters.get("delta") {
362            Some(&Value::Numeric(Numeric::Float(delta))) => f64::from(delta),
363            _ => 0.0,
364        };
365        let scan_portion = match configuration.parameters.get("scan_portion") {
366            Some(&Value::Numeric(Numeric::Float(scan_portion))) => f64::from(scan_portion),
367            _ => 1.0,
368        };
369        let scan_order = match configuration.parameters.get("scan_order") {
370            Some(Value::Text(scan_order)) if scan_order.as_str() == "dataset_frequency_asc" => {
371                ScanOrder::DatasetFrequencyAsc
372            }
373            Some(Value::Text(scan_order)) if scan_order.as_str() == "query_weight_desc" => {
374                ScanOrder::QueryWeightDesc
375            }
376            _ => ScanOrder::QueryWeightDesc,
377        };
378        Self {
379            configuration,
380            delta,
381            scan_portion,
382            scan_order,
383            inverted_index_btree,
384            inverted_index_cursor: None,
385            stats_btree,
386            stats_cursor: None,
387            main_btree: None,
388            search_result: VecDeque::new(),
389            insert_state: VectorSparseInvertedIndexInsertState::Init,
390            delete_state: VectorSparseInvertedIndexDeleteState::Init,
391            search_state: VectorSparseInvertedIndexSearchState::Init,
392        }
393    }
394}
395
396fn key_info() -> KeyInfo {
397    KeyInfo {
398        collation: CollationSeq::Binary,
399        sort_order: SortOrder::Asc,
400        nulls_order: None,
401    }
402}
403
404impl IndexMethodCursor for VectorSparseInvertedIndexMethodCursor {
405    fn create(&mut self, connection: &Arc<Connection>, database_id: usize) -> Result<IOResult<()>> {
406        // we need to properly track subprograms and propagate result to the root program to make this execution async
407
408        let columns = &self.configuration.columns;
409        let columns = columns.iter().map(|x| x.name.as_str()).collect::<Vec<_>>();
410        let db_prefix = connection
411            .get_database_name_by_index(database_id)
412            .filter(|name| name != "main")
413            .map(|name| format!("{}.", quote_identifier(&name)))
414            .unwrap_or_default();
415        let quoted_table = quote_identifier(&self.configuration.table_name);
416        let quoted_cols = columns
417            .iter()
418            .map(|c| quote_identifier(c))
419            .collect::<Vec<_>>()
420            .join(", ");
421        let inverted_index_create = format!(
422            "CREATE INDEX {db_prefix}{} ON {quoted_table} USING {BACKING_BTREE_INDEX_METHOD_NAME} ({quoted_cols})",
423            quote_identifier(&self.inverted_index_btree),
424        );
425        let stats_index_create = format!(
426            "CREATE INDEX {db_prefix}{} ON {quoted_table} USING {BACKING_BTREE_INDEX_METHOD_NAME} ({quoted_cols})",
427            quote_identifier(&self.stats_btree),
428        );
429        for sql in [inverted_index_create, stats_index_create] {
430            let mut stmt = connection.prepare(&sql)?;
431            // by default we set needs_stmt_subtransactions = true to all write transaction
432            // this will lead to Busy error here - because Transaction opcode will be unable to acquire ownership to the subjournal as it already owned by parent statement which is still active
433            //
434            // as we run nested statement - we actually don't need subjournal as it already started before in the parent statement
435            // so, this is hacky way to fix the situation for toy index for now, but we need to implement proper helpers in order to avoid similar errors in other code later
436            stmt.program
437                .prepared
438                .needs_stmt_subtransactions
439                .store(false, Ordering::Relaxed);
440            connection.start_nested();
441            let result = stmt.run_ignore_rows();
442            connection.end_nested();
443            result?;
444        }
445
446        Ok(IOResult::Done(()))
447    }
448
449    fn destroy(
450        &mut self,
451        connection: &Arc<Connection>,
452        database_id: usize,
453    ) -> Result<IOResult<()>> {
454        let db_prefix = connection
455            .get_database_name_by_index(database_id)
456            .filter(|name| name != "main")
457            .map(|name| format!("{}.", quote_identifier(&name)))
458            .unwrap_or_default();
459        let inverted_index_drop = format!(
460            "DROP INDEX {db_prefix}{}",
461            quote_identifier(&self.inverted_index_btree)
462        );
463        let stats_index_drop = format!(
464            "DROP INDEX {db_prefix}{}",
465            quote_identifier(&self.stats_btree)
466        );
467        for sql in [inverted_index_drop, stats_index_drop] {
468            let mut stmt = connection.prepare(&sql)?;
469            connection.start_nested();
470            let result = stmt.run_ignore_rows();
471            connection.end_nested();
472            result?;
473        }
474
475        Ok(IOResult::Done(()))
476    }
477
478    fn open_read(
479        &mut self,
480        connection: &Arc<Connection>,
481        database_id: usize,
482    ) -> Result<IOResult<()>> {
483        self.inverted_index_cursor = Some(open_index_cursor(
484            connection,
485            database_id,
486            &self.configuration.table_name,
487            &self.inverted_index_btree,
488            // component, length, rowid
489            [key_info(), key_info(), key_info()],
490        )?);
491        self.stats_cursor = Some(open_index_cursor(
492            connection,
493            database_id,
494            &self.configuration.table_name,
495            &self.stats_btree,
496            // component
497            [key_info()],
498        )?);
499        self.main_btree = Some(open_table_cursor(
500            connection,
501            database_id,
502            &self.configuration.table_name,
503        )?);
504        Ok(IOResult::Done(()))
505    }
506
507    fn open_write(
508        &mut self,
509        connection: &Arc<Connection>,
510        database_id: usize,
511    ) -> Result<IOResult<()>> {
512        self.inverted_index_cursor = Some(open_index_cursor(
513            connection,
514            database_id,
515            &self.configuration.table_name,
516            &self.inverted_index_btree,
517            // component, length, rowid
518            [key_info(), key_info(), key_info()],
519        )?);
520        self.stats_cursor = Some(open_index_cursor(
521            connection,
522            database_id,
523            &self.configuration.table_name,
524            &self.stats_btree,
525            // component
526            [key_info()],
527        )?);
528        Ok(IOResult::Done(()))
529    }
530
531    fn insert(&mut self, values: &[Register]) -> Result<IOResult<()>> {
532        let Some(inverted_cursor) = &mut self.inverted_index_cursor else {
533            return Err(LimboError::InternalError(
534                "inverted cursor must be opened".to_string(),
535            ));
536        };
537        let Some(stats_cursor) = &mut self.stats_cursor else {
538            return Err(LimboError::InternalError(
539                "stats cursor must be opened".to_string(),
540            ));
541        };
542        loop {
543            tracing::debug!("insert_state: {:?}", self.insert_state);
544            match &mut self.insert_state {
545                VectorSparseInvertedIndexInsertState::Init => {
546                    let Some(vector) = values[0].get_value().to_blob() else {
547                        return Err(LimboError::InternalError(
548                            "first value must be sparse vector".to_string(),
549                        ));
550                    };
551                    let vector = Vector::from_vec(vector.to_vec())?;
552                    if !matches!(vector.vector_type, VectorType::Float32Sparse) {
553                        return Err(LimboError::InternalError(
554                            "first value must be sparse vector".to_string(),
555                        ));
556                    }
557                    let Some(rowid) = values[1].get_value().as_int() else {
558                        return Err(LimboError::InternalError(
559                            "second value must be i64 rowid".to_string(),
560                        ));
561                    };
562                    let sum = vector.as_f32_sparse().values.iter().sum::<f32>() as f64;
563                    self.insert_state = VectorSparseInvertedIndexInsertState::Prepare {
564                        vector: Some(vector),
565                        sum,
566                        rowid,
567                        idx: 0,
568                    }
569                }
570                VectorSparseInvertedIndexInsertState::Prepare {
571                    vector,
572                    sum,
573                    rowid,
574                    idx,
575                } => {
576                    let Some(v) = vector.as_ref() else {
577                        return Err(LimboError::InternalError(
578                            "vector must be present in Prepare state".to_string(),
579                        ));
580                    };
581                    if *idx == v.as_f32_sparse().idx.len() {
582                        self.insert_state = VectorSparseInvertedIndexInsertState::Init;
583                        return Ok(IOResult::Done(()));
584                    }
585                    let position = v.as_f32_sparse().idx[*idx];
586                    let key = ImmutableRecord::from_values(
587                        &[
588                            Value::from_i64(position as i64),
589                            Value::from_f64(*sum),
590                            Value::from_i64(*rowid),
591                        ],
592                        3,
593                    )?;
594                    tracing::debug!(
595                        "insert_state: seek: component={}, sum={}, rowid={}",
596                        position,
597                        *sum,
598                        *rowid,
599                    );
600                    self.insert_state = VectorSparseInvertedIndexInsertState::SeekInverted {
601                        vector: vector.take(),
602                        sum: *sum,
603                        idx: *idx,
604                        rowid: *rowid,
605                        key: Some(key),
606                    };
607                }
608                VectorSparseInvertedIndexInsertState::SeekInverted {
609                    vector,
610                    sum,
611                    rowid,
612                    idx,
613                    key,
614                } => {
615                    let Some(k) = key.as_ref() else {
616                        return Err(LimboError::InternalError(
617                            "key must be present in SeekInverted state".to_string(),
618                        ));
619                    };
620                    let result =
621                        return_if_io!(inverted_cursor
622                            .seek(SeekKey::IndexKey(k), SeekOp::GE { eq_only: true }));
623                    tracing::debug!("insert_state: seek: result={:?}", result);
624                    self.insert_state = VectorSparseInvertedIndexInsertState::InsertInverted {
625                        vector: vector.take(),
626                        sum: *sum,
627                        idx: *idx,
628                        rowid: *rowid,
629                        key: key.take(),
630                    };
631                }
632                VectorSparseInvertedIndexInsertState::InsertInverted {
633                    vector,
634                    sum,
635                    rowid,
636                    idx,
637                    key,
638                } => {
639                    let Some(k) = key.as_ref() else {
640                        return Err(LimboError::InternalError(
641                            "key must be present in InsertInverted state".to_string(),
642                        ));
643                    };
644                    return_if_io!(inverted_cursor.insert(&BTreeKey::IndexKey(k)));
645
646                    let Some(v) = vector.as_ref() else {
647                        return Err(LimboError::InternalError(
648                            "vector must be present in InsertInverted state".to_string(),
649                        ));
650                    };
651                    let position = v.as_f32_sparse().idx[*idx];
652                    let key = ImmutableRecord::from_values(&[Value::from_i64(position as i64)], 1)?;
653                    self.insert_state = VectorSparseInvertedIndexInsertState::SeekStats {
654                        vector: vector.take(),
655                        sum: *sum,
656                        idx: *idx,
657                        rowid: *rowid,
658                        key: Some(key),
659                    };
660                }
661                VectorSparseInvertedIndexInsertState::SeekStats {
662                    vector,
663                    sum,
664                    key,
665                    rowid,
666                    idx,
667                } => {
668                    let Some(k) = key.as_ref() else {
669                        return Err(LimboError::InternalError(
670                            "key must be present in SeekStats state".to_string(),
671                        ));
672                    };
673                    let result = return_if_io!(
674                        stats_cursor.seek(SeekKey::IndexKey(k), SeekOp::GE { eq_only: true })
675                    );
676                    match result {
677                        SeekResult::Found => {
678                            self.insert_state = VectorSparseInvertedIndexInsertState::ReadStats {
679                                vector: vector.take(),
680                                sum: *sum,
681                                idx: *idx,
682                                rowid: *rowid,
683                            };
684                        }
685                        SeekResult::NotFound | SeekResult::TryAdvance => {
686                            let Some(v) = vector.as_ref() else {
687                                return Err(LimboError::InternalError(
688                                    "vector must be present in SeekStats state".to_string(),
689                                ));
690                            };
691                            let position = v.as_f32_sparse().idx[*idx];
692                            let value = v.as_f32_sparse().values[*idx] as f64;
693                            tracing::debug!(
694                                "update stats(insert): {} (cnt={}, min={}, max={})",
695                                position,
696                                1,
697                                value,
698                                value,
699                            );
700                            let key = ImmutableRecord::from_values(
701                                &[
702                                    Value::from_i64(position as i64),
703                                    Value::from_i64(1),
704                                    Value::from_f64(value),
705                                    Value::from_f64(value),
706                                ],
707                                4,
708                            )?;
709                            self.insert_state = VectorSparseInvertedIndexInsertState::UpdateStats {
710                                vector: vector.take(),
711                                sum: *sum,
712                                idx: *idx,
713                                rowid: *rowid,
714                                key: Some(key),
715                            };
716                        }
717                    }
718                }
719                VectorSparseInvertedIndexInsertState::ReadStats {
720                    vector,
721                    sum,
722                    rowid,
723                    idx,
724                } => {
725                    let record = return_if_io!(stats_cursor.record());
726                    let component = parse_stat_row(record)?;
727                    let Some(v) = vector.as_ref() else {
728                        return Err(LimboError::InternalError(
729                            "vector must be present in ReadStats state".to_string(),
730                        ));
731                    };
732                    let position = v.as_f32_sparse().idx[*idx];
733                    let value = v.as_f32_sparse().values[*idx] as f64;
734                    tracing::debug!(
735                        "update stats(insert): {} (cnt={}, min={}, max={})",
736                        position,
737                        component.cnt + 1,
738                        value.min(component.min),
739                        value.max(component.max),
740                    );
741                    let key = ImmutableRecord::from_values(
742                        &[
743                            Value::from_i64(position as i64),
744                            Value::from_i64(component.cnt + 1),
745                            Value::from_f64(value.min(component.min)),
746                            Value::from_f64(value.max(component.max)),
747                        ],
748                        4,
749                    )?;
750                    self.insert_state = VectorSparseInvertedIndexInsertState::UpdateStats {
751                        vector: vector.take(),
752                        sum: *sum,
753                        idx: *idx,
754                        rowid: *rowid,
755                        key: Some(key),
756                    };
757                }
758                VectorSparseInvertedIndexInsertState::UpdateStats {
759                    vector,
760                    sum,
761                    key,
762                    rowid,
763                    idx,
764                } => {
765                    let Some(k) = key.as_ref() else {
766                        return Err(LimboError::InternalError(
767                            "key must be present in UpdateStats state".to_string(),
768                        ));
769                    };
770                    return_if_io!(stats_cursor.insert(&BTreeKey::IndexKey(k)));
771
772                    self.insert_state = VectorSparseInvertedIndexInsertState::Prepare {
773                        vector: vector.take(),
774                        sum: *sum,
775                        idx: *idx + 1,
776                        rowid: *rowid,
777                    };
778                }
779            }
780        }
781    }
782
783    fn delete(&mut self, values: &[Register]) -> Result<IOResult<()>> {
784        let Some(cursor) = &mut self.inverted_index_cursor else {
785            return Err(LimboError::InternalError(
786                "cursor must be opened".to_string(),
787            ));
788        };
789        let Some(stats_cursor) = &mut self.stats_cursor else {
790            return Err(LimboError::InternalError(
791                "stats cursor must be opened".to_string(),
792            ));
793        };
794        loop {
795            tracing::debug!("delete_state: {:?}", self.delete_state);
796            match &mut self.delete_state {
797                VectorSparseInvertedIndexDeleteState::Init => {
798                    let Some(vector) = values[0].get_value().to_blob() else {
799                        return Err(LimboError::InternalError(
800                            "first value must be sparse vector".to_string(),
801                        ));
802                    };
803                    let vector = Vector::from_vec(vector.to_vec())?;
804                    if !matches!(vector.vector_type, VectorType::Float32Sparse) {
805                        return Err(LimboError::InternalError(
806                            "first value must be sparse vector".to_string(),
807                        ));
808                    }
809                    let Some(rowid) = values[1].get_value().as_int() else {
810                        return Err(LimboError::InternalError(
811                            "second value must be i64 rowid".to_string(),
812                        ));
813                    };
814                    let sum = vector.as_f32_sparse().values.iter().sum::<f32>() as f64;
815                    self.delete_state = VectorSparseInvertedIndexDeleteState::Prepare {
816                        vector: Some(vector),
817                        sum,
818                        rowid,
819                        idx: 0,
820                    }
821                }
822                VectorSparseInvertedIndexDeleteState::Prepare {
823                    vector,
824                    sum,
825                    rowid,
826                    idx,
827                } => {
828                    let Some(v) = vector.as_ref() else {
829                        return Err(LimboError::InternalError(
830                            "vector must be present in Prepare state".to_string(),
831                        ));
832                    };
833                    if *idx == v.as_f32_sparse().idx.len() {
834                        self.delete_state = VectorSparseInvertedIndexDeleteState::Init;
835                        return Ok(IOResult::Done(()));
836                    }
837                    let position = v.as_f32_sparse().idx[*idx];
838                    let key = ImmutableRecord::from_values(
839                        &[
840                            Value::from_i64(position as i64),
841                            Value::from_f64(*sum),
842                            Value::from_i64(*rowid),
843                        ],
844                        3,
845                    )?;
846                    self.delete_state = VectorSparseInvertedIndexDeleteState::SeekInverted {
847                        vector: vector.take(),
848                        idx: *idx,
849                        sum: *sum,
850                        rowid: *rowid,
851                        key: Some(key),
852                    };
853                }
854                VectorSparseInvertedIndexDeleteState::SeekInverted {
855                    vector,
856                    sum,
857                    rowid,
858                    idx,
859                    key,
860                } => {
861                    let component_idx = vector
862                        .as_ref()
863                        .and_then(|v| v.as_f32_sparse().idx.get(*idx).copied())
864                        .ok_or_else(|| {
865                            LimboError::InternalError(
866                                "vector must be present in SeekInverted state".to_string(),
867                            )
868                        })?;
869                    tracing::debug!(
870                        "delete_state: seek: component={}, sum={}, rowid={}",
871                        component_idx,
872                        *sum,
873                        *rowid,
874                    );
875                    let Some(k) = key.as_ref() else {
876                        return Err(LimboError::InternalError(
877                            "key must be present in SeekInverted state".to_string(),
878                        ));
879                    };
880                    let result = return_if_io!(
881                        cursor.seek(SeekKey::IndexKey(k), SeekOp::GE { eq_only: true })
882                    );
883                    match result {
884                        SeekResult::Found => {
885                            self.delete_state =
886                                VectorSparseInvertedIndexDeleteState::DeleteInverted {
887                                    vector: vector.take(),
888                                    sum: *sum,
889                                    idx: *idx,
890                                    rowid: *rowid,
891                                };
892                        }
893                        SeekResult::TryAdvance => {
894                            self.delete_state =
895                                VectorSparseInvertedIndexDeleteState::NextInverted {
896                                    vector: vector.take(),
897                                    sum: *sum,
898                                    idx: *idx,
899                                    rowid: *rowid,
900                                };
901                        }
902                        SeekResult::NotFound => {
903                            return Err(LimboError::Corrupt("inverted index corrupted".to_string()))
904                        }
905                    }
906                }
907                VectorSparseInvertedIndexDeleteState::NextInverted {
908                    vector,
909                    sum,
910                    rowid,
911                    idx,
912                } => {
913                    return_if_io!(cursor.next());
914                    if !cursor.has_record() {
915                        return Err(LimboError::Corrupt("inverted index corrupted".to_string()));
916                    }
917                    self.delete_state = VectorSparseInvertedIndexDeleteState::DeleteInverted {
918                        vector: vector.take(),
919                        sum: *sum,
920                        idx: *idx,
921                        rowid: *rowid,
922                    };
923                }
924                VectorSparseInvertedIndexDeleteState::DeleteInverted {
925                    vector,
926                    sum,
927                    rowid,
928                    idx,
929                } => {
930                    return_if_io!(cursor.delete());
931                    let Some(v) = vector.as_ref() else {
932                        return Err(LimboError::InternalError(
933                            "vector must be present in DeleteInverted state".to_string(),
934                        ));
935                    };
936                    let position = v.as_f32_sparse().idx[*idx];
937                    let key = ImmutableRecord::from_values(&[Value::from_i64(position as i64)], 1)?;
938                    self.delete_state = VectorSparseInvertedIndexDeleteState::SeekStats {
939                        vector: vector.take(),
940                        sum: *sum,
941                        idx: *idx,
942                        rowid: *rowid,
943                        key: Some(key),
944                    };
945                }
946                VectorSparseInvertedIndexDeleteState::SeekStats {
947                    vector,
948                    sum,
949                    key,
950                    rowid,
951                    idx,
952                } => {
953                    let Some(k) = key.as_ref() else {
954                        return Err(LimboError::InternalError(
955                            "key must be present in SeekStats state".to_string(),
956                        ));
957                    };
958                    let result = return_if_io!(
959                        stats_cursor.seek(SeekKey::IndexKey(k), SeekOp::GE { eq_only: true })
960                    );
961                    match result {
962                        SeekResult::Found => {
963                            self.delete_state = VectorSparseInvertedIndexDeleteState::ReadStats {
964                                vector: vector.take(),
965                                sum: *sum,
966                                idx: *idx,
967                                rowid: *rowid,
968                            };
969                        }
970                        SeekResult::NotFound | SeekResult::TryAdvance => {
971                            return Err(LimboError::Corrupt(
972                                "stats index corrupted: can't find component row".to_string(),
973                            ))
974                        }
975                    }
976                }
977                VectorSparseInvertedIndexDeleteState::ReadStats {
978                    vector,
979                    sum,
980                    rowid,
981                    idx,
982                } => {
983                    let record = return_if_io!(stats_cursor.record());
984                    let component = parse_stat_row(record)?;
985                    let Some(v) = vector.as_ref() else {
986                        return Err(LimboError::InternalError(
987                            "vector must be present in ReadStats state".to_string(),
988                        ));
989                    };
990                    let position = v.as_f32_sparse().idx[*idx];
991                    tracing::debug!(
992                        "update stats(delete): {} (cnt={}, min={}, max={})",
993                        position,
994                        component.cnt - 1,
995                        component.min,
996                        component.max,
997                    );
998                    let key = ImmutableRecord::from_values(
999                        &[
1000                            Value::from_i64(position as i64),
1001                            Value::from_i64(component.cnt - 1),
1002                            Value::from_f64(component.min),
1003                            Value::from_f64(component.max),
1004                        ],
1005                        4,
1006                    )?;
1007                    self.delete_state = VectorSparseInvertedIndexDeleteState::UpdateStats {
1008                        vector: vector.take(),
1009                        sum: *sum,
1010                        idx: *idx,
1011                        rowid: *rowid,
1012                        key: Some(key),
1013                    };
1014                }
1015                VectorSparseInvertedIndexDeleteState::UpdateStats {
1016                    vector,
1017                    sum,
1018                    key,
1019                    rowid,
1020                    idx,
1021                } => {
1022                    let Some(k) = key.as_ref() else {
1023                        return Err(LimboError::InternalError(
1024                            "key must be present in UpdateStats state".to_string(),
1025                        ));
1026                    };
1027                    return_if_io!(stats_cursor.insert(&BTreeKey::IndexKey(k)));
1028
1029                    self.delete_state = VectorSparseInvertedIndexDeleteState::Prepare {
1030                        vector: vector.take(),
1031                        sum: *sum,
1032                        idx: *idx + 1,
1033                        rowid: *rowid,
1034                    };
1035                }
1036            }
1037        }
1038    }
1039
1040    fn query_start(&mut self, values: &[Register]) -> Result<IOResult<bool>> {
1041        let Some(inverted) = &mut self.inverted_index_cursor else {
1042            return Err(LimboError::InternalError(
1043                "cursor must be opened".to_string(),
1044            ));
1045        };
1046        let Some(stats) = &mut self.stats_cursor else {
1047            return Err(LimboError::InternalError(
1048                "cursor must be opened".to_string(),
1049            ));
1050        };
1051        let Some(main) = &mut self.main_btree else {
1052            return Err(LimboError::InternalError(
1053                "cursor must be opened".to_string(),
1054            ));
1055        };
1056        loop {
1057            tracing::debug!("query_state: {:?}", self.search_state);
1058            match &mut self.search_state {
1059                VectorSparseInvertedIndexSearchState::Init => {
1060                    let Some(vector) = values[1].get_value().to_blob() else {
1061                        return Err(LimboError::InternalError(
1062                            "first value must be sparse vector".to_string(),
1063                        ));
1064                    };
1065                    let Some(limit) = values[2].get_value().as_int() else {
1066                        return Err(LimboError::InternalError(
1067                            "second value must be i64 limit parameter".to_string(),
1068                        ));
1069                    };
1070                    let vector = Vector::from_vec(vector.to_vec())?;
1071                    if !matches!(vector.vector_type, VectorType::Float32Sparse) {
1072                        return Err(LimboError::InternalError(
1073                            "first value must be sparse vector".to_string(),
1074                        ));
1075                    }
1076                    let sparse = vector.as_f32_sparse();
1077                    let sum = sparse.values.iter().sum::<f32>() as f64;
1078                    self.search_state =
1079                        VectorSparseInvertedIndexSearchState::CollectComponentsSeek {
1080                            sum,
1081                            vector: Some(vector),
1082                            idx: 0,
1083                            components: Some(Vec::new()),
1084                            key: None,
1085                            limit,
1086                        };
1087                }
1088                VectorSparseInvertedIndexSearchState::CollectComponentsSeek {
1089                    sum,
1090                    vector,
1091                    idx,
1092                    components,
1093                    limit,
1094                    key,
1095                } => {
1096                    let Some(v) = vector.as_ref() else {
1097                        return Err(LimboError::InternalError(
1098                            "vector must be present in CollectComponentsSeek state".to_string(),
1099                        ));
1100                    };
1101                    let p = &v.as_f32_sparse().idx[*idx..];
1102                    if p.is_empty() && key.is_none() {
1103                        let Some(mut components) = components.take() else {
1104                            return Err(LimboError::InternalError(
1105                                "components must be present in CollectComponentsSeek state"
1106                                    .to_string(),
1107                            ));
1108                        };
1109                        match self.scan_order {
1110                            ScanOrder::DatasetFrequencyAsc => {
1111                                // order by cnt ASC in order to check low-cardinality components first
1112                                components.sort_by_key(|(c, _)| c.cnt);
1113                            }
1114                            ScanOrder::QueryWeightDesc => {
1115                                // order by weight DESC in order to check high-impact components first
1116                                components
1117                                    .sort_by_key(|(_, w)| std::cmp::Reverse(FloatOrd(*w as f64)));
1118                            }
1119                        }
1120                        let take = (components.len() as f64 * self.scan_portion).ceil() as usize;
1121                        let components = components
1122                            .into_iter()
1123                            .take(take)
1124                            .map(|(c, _)| c)
1125                            .collect::<Vec<_>>();
1126
1127                        tracing::debug!(
1128                            "query_start: components: {:?}, delta: {}, scan_portion: {}, scan_order: {:?}",
1129                            components,
1130                            self.delta,
1131                            self.scan_portion,
1132                            self.scan_order,
1133                        );
1134                        self.search_state = VectorSparseInvertedIndexSearchState::Seek {
1135                            sum: *sum,
1136                            components: Some(components.into()),
1137                            collected: Some(HashSet::default()),
1138                            distances: Some(BTreeSet::new()),
1139                            limit: *limit,
1140                            key: None,
1141                            component: None,
1142                            sum_threshold: None,
1143                        };
1144                        continue;
1145                    }
1146                    if key.is_none() {
1147                        let Some(v) = vector.as_ref() else {
1148                            return Err(LimboError::InternalError(
1149                                "vector must be present in CollectComponentsSeek state".to_string(),
1150                            ));
1151                        };
1152                        let position = v.as_f32_sparse().idx[*idx];
1153                        *key = Some(ImmutableRecord::from_values(
1154                            &[Value::from_i64(position as i64)],
1155                            1,
1156                        )?);
1157                    }
1158                    let Some(k) = key.as_ref() else {
1159                        return Err(LimboError::InternalError(
1160                            "key must be present in CollectComponentsSeek state".to_string(),
1161                        ));
1162                    };
1163                    let result = return_if_io!(
1164                        stats.seek(SeekKey::IndexKey(k), SeekOp::GE { eq_only: true })
1165                    );
1166                    match result {
1167                        SeekResult::Found => {
1168                            self.search_state =
1169                                VectorSparseInvertedIndexSearchState::CollectComponentsRead {
1170                                    sum: *sum,
1171                                    vector: vector.take(),
1172                                    idx: *idx,
1173                                    components: components.take(),
1174                                    limit: *limit,
1175                                };
1176                        }
1177                        SeekResult::NotFound | SeekResult::TryAdvance => {
1178                            self.search_state =
1179                                VectorSparseInvertedIndexSearchState::CollectComponentsSeek {
1180                                    sum: *sum,
1181                                    components: components.take(),
1182                                    vector: vector.take(),
1183                                    idx: *idx + 1,
1184                                    limit: *limit,
1185                                    key: None,
1186                                };
1187                        }
1188                    }
1189                }
1190                VectorSparseInvertedIndexSearchState::CollectComponentsRead {
1191                    sum,
1192                    vector,
1193                    idx,
1194                    components,
1195                    limit,
1196                } => {
1197                    let record = return_if_io!(stats.record());
1198                    let Some(v) = vector.as_ref() else {
1199                        return Err(LimboError::InternalError(
1200                            "vector must be present in CollectComponentsRead state".to_string(),
1201                        ));
1202                    };
1203                    let value = v.as_f32_sparse().values[*idx];
1204                    let component = parse_stat_row(record)?;
1205                    let Some(comps) = components.as_mut() else {
1206                        return Err(LimboError::InternalError(
1207                            "components must be present in CollectComponentsRead state".to_string(),
1208                        ));
1209                    };
1210                    comps.push((component, value));
1211                    self.search_state =
1212                        VectorSparseInvertedIndexSearchState::CollectComponentsSeek {
1213                            sum: *sum,
1214                            components: components.take(),
1215                            vector: vector.take(),
1216                            idx: *idx + 1,
1217                            limit: *limit,
1218                            key: None,
1219                        };
1220                }
1221                VectorSparseInvertedIndexSearchState::Seek {
1222                    sum,
1223                    components,
1224                    collected,
1225                    distances,
1226                    limit,
1227                    key,
1228                    component,
1229                    sum_threshold,
1230                } => {
1231                    let Some(c) = components.as_ref() else {
1232                        return Err(LimboError::InternalError(
1233                            "components must be present in Seek state".to_string(),
1234                        ));
1235                    };
1236                    if c.is_empty() && key.is_none() {
1237                        let Some(distances) = distances.take() else {
1238                            return Err(LimboError::InternalError(
1239                                "distances must be present in Seek state".to_string(),
1240                            ));
1241                        };
1242                        self.search_result = distances.iter().map(|(d, i)| (*i, d.0)).collect();
1243                        return Ok(IOResult::Done(!self.search_result.is_empty()));
1244                    }
1245                    if key.is_none() {
1246                        // we estimate jaccard distance with the following approach:
1247                        // J = min(L, M1 + M2 + ... + Mr) / (Q + N - min(L, M1 + M2 + ... + Mr))
1248                        // so we want J > best + delta; define M1 + M2 + ... + Mr = M
1249                        // J = min(L, M) / (Q + L - min(L, M)) > best + delta
1250                        // we need to consider two cases:
1251                        // 1. L < M: J = L / (Q + L - L) > best + delta => L > (best + delta) * Q
1252                        // 2. L > M: J = M / (Q + L - M) > best + delta => M > (best + delta) * (Q + L - M) => L < M / (best + delta) - (Q - M)
1253                        // so we have two intervals: [(best + delta) * Q .. M] and [M .. M / (best + delta) - (Q - M)]
1254                        // to simplify code for now we will pick upper bound from second range if it is not degenerate, otherwise check first range
1255                        let m = c.iter().map(|c| c.max).sum::<f64>().min(*sum);
1256                        let Some(dists) = distances.as_ref() else {
1257                            return Err(LimboError::InternalError(
1258                                "distances must be present in Seek state".to_string(),
1259                            ));
1260                        };
1261                        if dists.len() >= *limit as usize {
1262                            if let Some((max_threshold, _)) = dists.last() {
1263                                let best = 1.0 - max_threshold.0;
1264                                let delta = self.delta;
1265                                let q = *sum;
1266
1267                                if best > 0.0 {
1268                                    let first_range_l = (best + delta) * q;
1269                                    let second_range_r = m / (best + delta) - (q - m);
1270                                    if m <= second_range_r {
1271                                        *sum_threshold = Some(second_range_r);
1272                                    } else if first_range_l <= m {
1273                                        *sum_threshold = Some(m);
1274                                    } else {
1275                                        *sum_threshold = Some(-1.0);
1276                                    }
1277                                    tracing::debug!(
1278                                        "sum_threshold={:?}, max_threshold={}, remained_sum={}, sum={}, components={:?}",
1279                                        sum_threshold,
1280                                        best,
1281                                        m,
1282                                        sum,
1283                                        c
1284                                    );
1285                                }
1286                            }
1287                        }
1288                        let Some(comps) = components.as_mut() else {
1289                            return Err(LimboError::InternalError(
1290                                "components must be present in Seek state".to_string(),
1291                            ));
1292                        };
1293                        let Some(c) = comps.pop_front() else {
1294                            return Err(LimboError::InternalError(
1295                                "components queue must not be empty in Seek state".to_string(),
1296                            ));
1297                        };
1298                        *key = Some(ImmutableRecord::from_values(
1299                            &[Value::from_i64(c.position as i64)],
1300                            1,
1301                        )?);
1302                        *component = Some(c.position);
1303                    }
1304                    let Some(k) = key.as_ref() else {
1305                        return Err(LimboError::InternalError(
1306                            "key must be present in Seek state".to_string(),
1307                        ));
1308                    };
1309                    let result = return_if_io!(
1310                        inverted.seek(SeekKey::IndexKey(k), SeekOp::GE { eq_only: false })
1311                    );
1312                    match result {
1313                        SeekResult::Found => {
1314                            let Some(comp) = component.take() else {
1315                                return Err(LimboError::InternalError(
1316                                    "component must be present in Seek state".to_string(),
1317                                ));
1318                            };
1319                            self.search_state = VectorSparseInvertedIndexSearchState::Read {
1320                                sum: *sum,
1321                                components: components.take(),
1322                                collected: collected.take(),
1323                                distances: distances.take(),
1324                                current: Some(Vec::new()),
1325                                limit: *limit,
1326                                sum_threshold: sum_threshold.take(),
1327                                component: comp,
1328                            };
1329                        }
1330                        SeekResult::TryAdvance | SeekResult::NotFound => {
1331                            let Some(comp) = component.take() else {
1332                                return Err(LimboError::InternalError(
1333                                    "component must be present in Seek state".to_string(),
1334                                ));
1335                            };
1336                            self.search_state = VectorSparseInvertedIndexSearchState::Next {
1337                                sum: *sum,
1338                                components: components.take(),
1339                                collected: collected.take(),
1340                                distances: distances.take(),
1341                                current: Some(Vec::new()),
1342                                limit: *limit,
1343                                sum_threshold: sum_threshold.take(),
1344                                component: comp,
1345                            };
1346                        }
1347                    }
1348                }
1349                VectorSparseInvertedIndexSearchState::Read {
1350                    sum,
1351                    components,
1352                    collected,
1353                    distances,
1354                    limit,
1355                    sum_threshold,
1356                    component,
1357                    current,
1358                } => {
1359                    let record = return_if_io!(inverted.record());
1360                    let row = parse_inverted_index_row(record)?;
1361                    if row.position != *component
1362                        || (sum_threshold.is_some()
1363                            && row.sum
1364                                > sum_threshold.ok_or_else(|| {
1365                                    LimboError::InternalError(
1366                                        "sum_threshold must be present when checked".to_string(),
1367                                    )
1368                                })?)
1369                    {
1370                        let Some(mut current) = current.take() else {
1371                            return Err(LimboError::InternalError(
1372                                "current must be present in Read state".to_string(),
1373                            ));
1374                        };
1375                        current.sort_unstable();
1376
1377                        self.search_state = VectorSparseInvertedIndexSearchState::EvaluateSeek {
1378                            sum: *sum,
1379                            components: components.take(),
1380                            collected: collected.take(),
1381                            distances: distances.take(),
1382                            limit: *limit,
1383                            current: Some(current.into()),
1384                            rowid: None,
1385                        };
1386                        continue;
1387                    }
1388                    let Some(coll) = collected.as_mut() else {
1389                        return Err(LimboError::InternalError(
1390                            "collected must be present in Read state".to_string(),
1391                        ));
1392                    };
1393                    if coll.insert(row.rowid) {
1394                        let Some(curr) = current.as_mut() else {
1395                            return Err(LimboError::InternalError(
1396                                "current must be present in Read state".to_string(),
1397                            ));
1398                        };
1399                        curr.push(row.rowid);
1400                    }
1401
1402                    self.search_state = VectorSparseInvertedIndexSearchState::Next {
1403                        sum: *sum,
1404                        components: components.take(),
1405                        collected: collected.take(),
1406                        distances: distances.take(),
1407                        limit: *limit,
1408                        sum_threshold: *sum_threshold,
1409                        component: *component,
1410                        current: current.take(),
1411                    };
1412                }
1413                VectorSparseInvertedIndexSearchState::Next {
1414                    sum,
1415                    components,
1416                    collected,
1417                    distances,
1418                    limit,
1419                    sum_threshold,
1420                    component,
1421                    current,
1422                } => {
1423                    return_if_io!(inverted.next());
1424                    if !inverted.has_record() {
1425                        let Some(mut current) = current.take() else {
1426                            return Err(LimboError::InternalError(
1427                                "current must be present in Next state".to_string(),
1428                            ));
1429                        };
1430                        current.sort_unstable();
1431
1432                        self.search_state = VectorSparseInvertedIndexSearchState::EvaluateSeek {
1433                            sum: *sum,
1434                            components: components.take(),
1435                            collected: collected.take(),
1436                            distances: distances.take(),
1437                            limit: *limit,
1438                            current: Some(current.into()),
1439                            rowid: None,
1440                        };
1441                    } else {
1442                        self.search_state = VectorSparseInvertedIndexSearchState::Read {
1443                            sum: *sum,
1444                            components: components.take(),
1445                            collected: collected.take(),
1446                            distances: distances.take(),
1447                            limit: *limit,
1448                            sum_threshold: *sum_threshold,
1449                            component: *component,
1450                            current: current.take(),
1451                        };
1452                    }
1453                }
1454                VectorSparseInvertedIndexSearchState::EvaluateSeek {
1455                    sum,
1456                    components,
1457                    collected,
1458                    distances,
1459                    limit,
1460                    current,
1461                    rowid,
1462                } => {
1463                    let Some(c) = current.as_ref() else {
1464                        return Err(LimboError::InternalError(
1465                            "current must be present in EvaluateSeek state".to_string(),
1466                        ));
1467                    };
1468                    if c.is_empty() && rowid.is_none() {
1469                        self.search_state = VectorSparseInvertedIndexSearchState::Seek {
1470                            sum: *sum,
1471                            components: components.take(),
1472                            collected: collected.take(),
1473                            distances: distances.take(),
1474                            limit: *limit,
1475                            component: None,
1476                            key: None,
1477                            sum_threshold: None,
1478                        };
1479                        continue;
1480                    }
1481                    if rowid.is_none() {
1482                        let Some(curr) = current.as_mut() else {
1483                            return Err(LimboError::InternalError(
1484                                "current must be present in EvaluateSeek state".to_string(),
1485                            ));
1486                        };
1487                        *rowid = Some(curr.pop_front().ok_or_else(|| {
1488                            LimboError::InternalError(
1489                                "current queue must not be empty in EvaluateSeek state".to_string(),
1490                            )
1491                        })?);
1492                    }
1493
1494                    let Some(rid) = rowid.as_ref() else {
1495                        return Err(LimboError::InternalError(
1496                            "rowid must be present in EvaluateSeek state".to_string(),
1497                        ));
1498                    };
1499                    let rowid = *rid;
1500                    let k = SeekKey::TableRowId(rowid);
1501                    let result = return_if_io!(main.seek(k, SeekOp::GE { eq_only: true }));
1502                    if !matches!(result, SeekResult::Found) {
1503                        return Err(LimboError::Corrupt(
1504                            "vector_sparse_ivf corrupted: unable to find rowid in main table"
1505                                .to_string(),
1506                        ));
1507                    };
1508                    self.search_state = VectorSparseInvertedIndexSearchState::EvaluateRead {
1509                        sum: *sum,
1510                        components: components.take(),
1511                        collected: collected.take(),
1512                        distances: distances.take(),
1513                        limit: *limit,
1514                        current: current.take(),
1515                        rowid,
1516                    };
1517                }
1518                VectorSparseInvertedIndexSearchState::EvaluateRead {
1519                    sum,
1520                    components,
1521                    collected,
1522                    distances,
1523                    limit,
1524                    current,
1525                    rowid,
1526                } => {
1527                    let record = return_if_io!(main.record());
1528                    if let Some(record) = record {
1529                        let column_idx = self.configuration.columns[0].pos_in_table;
1530                        let ValueRef::Blob(data) = record.get_value(column_idx)? else {
1531                            return Err(LimboError::InternalError(
1532                                "table column value must be sparse vector".to_string(),
1533                            ));
1534                        };
1535                        let data = Vector::from_vec(data.to_vec())?;
1536                        if !matches!(data.vector_type, VectorType::Float32Sparse) {
1537                            return Err(LimboError::InternalError(
1538                                "table column value must be sparse vector".to_string(),
1539                            ));
1540                        }
1541                        let Some(arg) = values[1].get_value().to_blob() else {
1542                            return Err(LimboError::InternalError(
1543                                "first value must be sparse vector".to_string(),
1544                            ));
1545                        };
1546                        let arg = Vector::from_vec(arg.to_vec())?;
1547                        if !matches!(arg.vector_type, VectorType::Float32Sparse) {
1548                            return Err(LimboError::InternalError(
1549                                "first value must be sparse vector".to_string(),
1550                            ));
1551                        }
1552                        tracing::debug!(
1553                            "vector: {:?}, query: {:?}",
1554                            data.as_f32_sparse(),
1555                            arg.as_f32_sparse()
1556                        );
1557                        let distance = operations::jaccard::vector_distance_jaccard(&data, &arg)?;
1558                        let Some(dists) = distances.as_mut() else {
1559                            return Err(LimboError::InternalError(
1560                                "distances must be present in EvaluateRead state".to_string(),
1561                            ));
1562                        };
1563                        dists.insert((FloatOrd(distance), *rowid));
1564                        if dists.len() > *limit as usize {
1565                            let _ = dists.pop_last();
1566                        }
1567                    }
1568
1569                    self.search_state = VectorSparseInvertedIndexSearchState::EvaluateSeek {
1570                        sum: *sum,
1571                        components: components.take(),
1572                        collected: collected.take(),
1573                        distances: distances.take(),
1574                        limit: *limit,
1575                        current: current.take(),
1576                        rowid: None,
1577                    };
1578                }
1579            }
1580        }
1581    }
1582
1583    fn query_rowid(&mut self) -> Result<IOResult<Option<i64>>> {
1584        let Some(result) = self.search_result.front() else {
1585            return Err(LimboError::InternalError(
1586                "search_result must not be empty when query_rowid is called".to_string(),
1587            ));
1588        };
1589        Ok(IOResult::Done(Some(result.0)))
1590    }
1591
1592    fn query_column(&mut self, _: usize) -> Result<IOResult<Value>> {
1593        let Some(result) = self.search_result.front() else {
1594            return Err(LimboError::InternalError(
1595                "search_result must not be empty when query_column is called".to_string(),
1596            ));
1597        };
1598        Ok(IOResult::Done(Value::from_f64(result.1)))
1599    }
1600
1601    fn query_next(&mut self) -> Result<IOResult<bool>> {
1602        let _ = self.search_result.pop_front();
1603        Ok(IOResult::Done(!self.search_result.is_empty()))
1604    }
1605}
clt_database/index_method/toy_vector_sparse_ivf.rs

clt_database/index_method/
toy_vector_sparse_ivf.rs
