lance_index/vector/
storage.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Vector Storage, holding (quantized) vectors and providing distance calculation.
5
6use crate::vector::quantizer::QuantizerStorage;
7use arrow::compute::concat_batches;
8use arrow_array::{ArrayRef, RecordBatch};
9use arrow_schema::SchemaRef;
10use futures::prelude::stream::TryStreamExt;
11use lance_arrow::RecordBatchExt;
12use lance_core::deepsize::DeepSizeOf;
13use lance_core::{Error, ROW_ID, Result};
14use lance_encoding::decoder::FilterExpression;
15use lance_file::reader::FileReader;
16use lance_io::ReadBatchParams;
17use lance_io::scheduler::IoStats;
18use lance_linalg::distance::DistanceType;
19use prost::Message;
20use std::{
21    any::Any,
22    borrow::Cow,
23    collections::BinaryHeap,
24    mem::size_of,
25    ops::{Deref, DerefMut},
26    sync::Arc,
27};
28
29use crossbeam_queue::ArrayQueue;
30
31use crate::frag_reuse::FragReuseIndex;
32use crate::{
33    pb,
34    vector::{
35        ivf::storage::{IVF_METADATA_KEY, IvfModel},
36        quantizer::Quantization,
37    },
38};
39
40use super::graph::OrderedFloat;
41use super::graph::OrderedNode;
42use super::quantizer::{Quantizer, QuantizerMetadata};
43use super::{ApproxMode, DISTANCE_TYPE_KEY};
44
45/// <section class="warning">
46///  Internal API
47///
48///  API stability is not guaranteed
49/// </section>
50pub trait DistCalculator {
51    fn distance(&self, id: u32) -> f32;
52
53    // return the distances of all rows
54    // k_hint is a hint that can be used for optimization
55    fn distance_all(&self, k_hint: usize) -> Vec<f32>;
56
57    // Write the distances of all rows into caller-owned scratch buffers.
58    fn distance_all_with_scratch(
59        &self,
60        k_hint: usize,
61        dists: &mut Vec<f32>,
62        _u16_scratch: &mut Vec<u16>,
63        _u8_scratch: &mut Vec<u8>,
64        _u32_scratch: &mut Vec<u32>,
65    ) {
66        *dists = self.distance_all(k_hint);
67    }
68
69    fn prefetch(&self, _id: u32) {}
70
71    #[allow(clippy::too_many_arguments)]
72    fn accumulate_topk_with_scratch(
73        &self,
74        k: usize,
75        lower_bound: Option<f32>,
76        upper_bound: Option<f32>,
77        row_id: impl Fn(u32) -> u64,
78        res: &mut BinaryHeap<OrderedNode<u64>>,
79        dists: &mut Vec<f32>,
80        u16_scratch: &mut Vec<u16>,
81        u8_scratch: &mut Vec<u8>,
82        u32_scratch: &mut Vec<u32>,
83    ) {
84        if k == 0 {
85            return;
86        }
87
88        self.distance_all_with_scratch(k, dists, u16_scratch, u8_scratch, u32_scratch);
89        let lower_bound = lower_bound.unwrap_or(f32::MIN).into();
90        let upper_bound = upper_bound.unwrap_or(f32::MAX).into();
91        let mut max_dist = res.peek().map(|node| node.dist);
92
93        for (id, dist) in dists.iter().copied().enumerate() {
94            let dist = OrderedFloat(dist);
95            if dist < lower_bound || dist >= upper_bound {
96                continue;
97            }
98            if res.len() < k {
99                res.push(OrderedNode::new(row_id(id as u32), dist));
100                if res.len() == k {
101                    max_dist = res.peek().map(|node| node.dist);
102                }
103            } else if max_dist.is_some_and(|max_dist| max_dist > dist) {
104                res.pop();
105                res.push(OrderedNode::new(row_id(id as u32), dist));
106                max_dist = res.peek().map(|node| node.dist);
107            }
108        }
109    }
110
111    #[allow(clippy::too_many_arguments)]
112    fn accumulate_filtered_topk_with_scratch(
113        &self,
114        k: usize,
115        lower_bound: Option<f32>,
116        upper_bound: Option<f32>,
117        row_ids: impl Iterator<Item = (u32, u64)>,
118        accept_row: impl Fn(u64) -> bool,
119        res: &mut BinaryHeap<OrderedNode<u64>>,
120        _dists: &mut Vec<f32>,
121        _u16_scratch: &mut Vec<u16>,
122        _u8_scratch: &mut Vec<u8>,
123        _u32_scratch: &mut Vec<u32>,
124    ) {
125        if k == 0 {
126            return;
127        }
128
129        let lower_bound = lower_bound.unwrap_or(f32::MIN).into();
130        let upper_bound = upper_bound.unwrap_or(f32::MAX).into();
131        let mut max_dist = res.peek().map(|node| node.dist);
132
133        for (id, row_id) in row_ids {
134            if !accept_row(row_id) {
135                continue;
136            }
137            let dist = OrderedFloat(self.distance(id));
138            if dist < lower_bound || dist >= upper_bound {
139                continue;
140            }
141            if res.len() < k {
142                res.push(OrderedNode::new(row_id, dist));
143                if res.len() == k {
144                    max_dist = res.peek().map(|node| node.dist);
145                }
146            } else if max_dist.is_some_and(|max_dist| max_dist > dist) {
147                res.pop();
148                res.push(OrderedNode::new(row_id, dist));
149                max_dist = res.peek().map(|node| node.dist);
150            }
151        }
152    }
153}
154
155pub const STORAGE_METADATA_KEY: &str = "storage_metadata";
156
157#[derive(Debug)]
158pub struct QueryScratch {
159    pub distances: Vec<f32>,
160    pub query_f32: Vec<f32>,
161    pub u16: Vec<u16>,
162    pub u8: Vec<u8>,
163    pub u32: Vec<u32>,
164}
165
166impl QueryScratch {
167    pub const fn new() -> Self {
168        Self {
169            distances: Vec::new(),
170            query_f32: Vec::new(),
171            u16: Vec::new(),
172            u8: Vec::new(),
173            u32: Vec::new(),
174        }
175    }
176
177    pub fn with_capacity(capacity: QueryScratchCapacity) -> Self {
178        Self {
179            distances: vec![0.0; capacity.distances],
180            query_f32: vec![0.0; capacity.query_f32],
181            u16: vec![0; capacity.u16],
182            u8: vec![0; capacity.u8],
183            u32: vec![0; capacity.u32],
184        }
185    }
186}
187
188impl Default for QueryScratch {
189    fn default() -> Self {
190        Self::new()
191    }
192}
193
194impl DeepSizeOf for QueryScratch {
195    fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize {
196        self.distances.capacity() * size_of::<f32>()
197            + self.query_f32.capacity() * size_of::<f32>()
198            + self.u16.capacity() * size_of::<u16>()
199            + self.u8.capacity() * size_of::<u8>()
200            + self.u32.capacity() * size_of::<u32>()
201    }
202}
203
204#[derive(Clone, Copy, Debug, Default)]
205pub struct QueryScratchCapacity {
206    pub distances: usize,
207    pub query_f32: usize,
208    pub u16: usize,
209    pub u8: usize,
210    pub u32: usize,
211}
212
213impl QueryScratchCapacity {
214    pub const fn new(distances: usize, query_f32: usize, u16: usize, u8: usize) -> Self {
215        Self::new_with_u32(distances, query_f32, u16, u8, 0)
216    }
217
218    pub const fn new_with_u32(
219        distances: usize,
220        query_f32: usize,
221        u16: usize,
222        u8: usize,
223        u32: usize,
224    ) -> Self {
225        Self {
226            distances,
227            query_f32,
228            u16,
229            u8,
230            u32,
231        }
232    }
233
234    fn deep_size_bytes(&self) -> usize {
235        self.distances * size_of::<f32>()
236            + self.query_f32 * size_of::<f32>()
237            + self.u16 * size_of::<u16>()
238            + self.u8 * size_of::<u8>()
239            + self.u32 * size_of::<u32>()
240    }
241}
242
243#[derive(Clone, Copy, Debug, Default)]
244pub struct DistanceCalculatorOptions {
245    pub approx_mode: ApproxMode,
246}
247
248#[derive(Debug)]
249pub struct RabitRawQueryContext {
250    pub code_dim: usize,
251    pub ex_bits: u8,
252    pub rotated_query: Vec<f32>,
253    pub dist_table: Vec<f32>,
254    /// The rotated query zero-padded to a 64-dim multiple for the ex-dot
255    /// kernels; empty when `code_dim` is already aligned (the kernels then
256    /// read `rotated_query` directly).
257    pub ex_query: Vec<f32>,
258    pub sum_q: f32,
259}
260
261#[derive(Clone, Copy)]
262pub enum QueryResidual<'a> {
263    Centroid(&'a dyn arrow_array::Array),
264    RabitRawQuery {
265        rotated_centroid: Option<&'a [f32]>,
266        query: Option<&'a RabitRawQueryContext>,
267    },
268}
269
270#[derive(Debug)]
271pub struct QueryScratchPool {
272    scratches: ArrayQueue<QueryScratch>,
273    scratch_capacity: QueryScratchCapacity,
274}
275
276impl QueryScratchPool {
277    pub fn new(size: usize) -> Self {
278        Self::with_capacity(size, QueryScratchCapacity::default())
279    }
280
281    pub fn with_capacity(size: usize, capacity: QueryScratchCapacity) -> Self {
282        let size = size.max(1);
283        let scratches = ArrayQueue::new(size);
284        for _ in 0..size {
285            scratches
286                .push(QueryScratch::with_capacity(capacity))
287                .expect("query scratch pool should have spare capacity during initialization");
288        }
289        Self {
290            scratches,
291            scratch_capacity: capacity,
292        }
293    }
294
295    pub fn scratch(&self) -> QueryScratchGuard<'_> {
296        let (scratch, pooled) = if let Some(scratch) = self.scratches.pop() {
297            (scratch, true)
298        } else {
299            (QueryScratch::with_capacity(self.scratch_capacity), false)
300        };
301        QueryScratchGuard {
302            pool: self,
303            scratch: Some(scratch),
304            pooled,
305        }
306    }
307
308    pub fn with_scratch<T>(&self, f: impl FnOnce(&mut QueryScratch) -> T) -> T {
309        let mut scratch = self.scratch();
310        f(&mut scratch)
311    }
312}
313
314pub struct QueryScratchGuard<'a> {
315    pool: &'a QueryScratchPool,
316    scratch: Option<QueryScratch>,
317    pooled: bool,
318}
319
320impl Deref for QueryScratchGuard<'_> {
321    type Target = QueryScratch;
322
323    fn deref(&self) -> &Self::Target {
324        self.scratch
325            .as_ref()
326            .expect("query scratch guard should hold scratch")
327    }
328}
329
330impl DerefMut for QueryScratchGuard<'_> {
331    fn deref_mut(&mut self) -> &mut Self::Target {
332        self.scratch
333            .as_mut()
334            .expect("query scratch guard should hold scratch")
335    }
336}
337
338impl Drop for QueryScratchGuard<'_> {
339    fn drop(&mut self) {
340        if !self.pooled {
341            return;
342        }
343        if let Some(scratch) = self.scratch.take() {
344            match self.pool.scratches.push(scratch) {
345                Ok(()) => {}
346                Err(_) => unreachable!("query scratch pool should not exceed its capacity"),
347            }
348        }
349    }
350}
351
352impl DeepSizeOf for QueryScratchPool {
353    fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
354        let mut total = self.scratches.capacity() * size_of::<QueryScratch>();
355        let mut scratches = Vec::new();
356        while let Some(scratch) = self.scratches.pop() {
357            total += scratch.deep_size_of_children(context);
358            scratches.push(scratch);
359        }
360        let checked_out = self.scratches.capacity().saturating_sub(scratches.len());
361        total += checked_out * self.scratch_capacity.deep_size_bytes();
362        for scratch in scratches {
363            let _ = self.scratches.push(scratch);
364        }
365        total
366    }
367}
368
369/// Vector Storage is the abstraction to store the vectors.
370///
371/// It can be in-memory or on-disk, raw vector or quantized vectors.
372///
373/// It abstracts away the logic to compute the distance between vectors.
374///
375/// TODO: should we rename this to "VectorDistance"?;
376///
377/// <section class="warning">
378///  Internal API
379///
380///  API stability is not guaranteed
381/// </section>
382pub trait VectorStore: Send + Sync + Sized + Clone {
383    type DistanceCalculator<'a>: DistCalculator
384    where
385        Self: 'a;
386
387    fn as_any(&self) -> &dyn Any;
388
389    fn schema(&self) -> &SchemaRef;
390
391    fn to_batches(&self) -> Result<impl Iterator<Item = RecordBatch> + Send>;
392
393    fn len(&self) -> usize;
394
395    /// Returns true if this graph is empty.
396    fn is_empty(&self) -> bool {
397        self.len() == 0
398    }
399
400    /// Return [DistanceType].
401    fn distance_type(&self) -> DistanceType;
402
403    /// Get the lance ROW ID from one vector.
404    fn row_id(&self, id: u32) -> u64;
405
406    fn row_ids(&self) -> impl Iterator<Item = &u64>;
407
408    /// Append Raw [RecordBatch] into the Storage.
409    /// The storage implement will perform quantization if necessary.
410    fn append_batch(&self, batch: RecordBatch, vector_column: &str) -> Result<Self>;
411
412    /// Create a [DistCalculator] to compute the distance between the query.
413    ///
414    /// Using dist calculator can be more efficient as it can pre-compute some
415    /// values.
416    fn dist_calculator(&self, query: ArrayRef, dist_q_c: f32) -> Self::DistanceCalculator<'_>;
417
418    /// Create a [DistCalculator], reusing caller-owned scratch for query-time
419    /// precomputed state when the storage supports it.
420    fn dist_calculator_with_scratch<'a>(
421        &'a self,
422        query: ArrayRef,
423        dist_q_c: f32,
424        _residual: Option<QueryResidual<'a>>,
425        _f32_scratch: &'a mut Vec<f32>,
426        _options: DistanceCalculatorOptions,
427    ) -> Self::DistanceCalculator<'a> {
428        self.dist_calculator(query, dist_q_c)
429    }
430
431    fn dist_calculator_from_id(&self, id: u32) -> Self::DistanceCalculator<'_>;
432
433    fn dist_between(&self, u: u32, v: u32) -> f32 {
434        let dist_cal_u = self.dist_calculator_from_id(u);
435        dist_cal_u.distance(v)
436    }
437
438    fn prefers_candidate(&self, candidate: &OrderedNode, selected: &[OrderedNode]) -> bool {
439        let dist_cal_candidate = self.dist_calculator_from_id(candidate.id);
440        selected
441            .iter()
442            .all(|other| candidate.dist < OrderedFloat(dist_cal_candidate.distance(other.id)))
443    }
444}
445
446pub struct StorageBuilder<Q: Quantization> {
447    vector_column: String,
448    distance_type: DistanceType,
449    quantizer: Q,
450
451    frag_reuse_index: Option<Arc<FragReuseIndex>>,
452}
453
454impl<Q: Quantization> StorageBuilder<Q> {
455    pub fn new(
456        vector_column: String,
457        distance_type: DistanceType,
458        quantizer: Q,
459        frag_reuse_index: Option<Arc<FragReuseIndex>>,
460    ) -> Result<Self> {
461        Ok(Self {
462            vector_column,
463            distance_type,
464            quantizer,
465            frag_reuse_index,
466        })
467    }
468
469    pub fn build(&self, batches: Vec<RecordBatch>) -> Result<Q::Storage> {
470        let mut batch = concat_batches(batches[0].schema_ref(), batches.iter())?;
471
472        if batch.column_by_name(self.quantizer.column()).is_none() {
473            let vectors = batch
474                .column_by_name(&self.vector_column)
475                .ok_or(Error::index(format!(
476                    "Vector column {} not found in batch",
477                    self.vector_column
478                )))?;
479            let codes = self.quantizer.quantize(vectors)?;
480            batch = batch.drop_column(&self.vector_column)?.try_with_column(
481                arrow_schema::Field::new(self.quantizer.column(), codes.data_type().clone(), true),
482                codes,
483            )?;
484        }
485
486        debug_assert!(batch.column_by_name(ROW_ID).is_some());
487        debug_assert!(batch.column_by_name(self.quantizer.column()).is_some());
488
489        Q::Storage::try_from_batch(
490            batch,
491            &self.quantizer.metadata(None),
492            self.distance_type,
493            self.frag_reuse_index.clone(),
494        )
495    }
496}
497
498/// Loader to load partitioned PQ storage from disk.
499#[derive(Debug)]
500pub struct IvfQuantizationStorage<Q: Quantization> {
501    reader: FileReader,
502
503    distance_type: DistanceType,
504    metadata: Q::Metadata,
505
506    ivf: IvfModel,
507    frag_reuse_index: Option<Arc<FragReuseIndex>>,
508}
509
510impl<Q: Quantization> DeepSizeOf for IvfQuantizationStorage<Q> {
511    fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
512        self.metadata.deep_size_of_children(context) + self.ivf.deep_size_of_children(context)
513    }
514}
515
516impl<Q: Quantization> IvfQuantizationStorage<Q> {
517    /// Open a Loader.
518    ///
519    ///
520    pub async fn try_new(
521        reader: FileReader,
522        frag_reuse_index: Option<Arc<FragReuseIndex>>,
523    ) -> Result<Self> {
524        let schema = reader.schema();
525
526        let distance_type = DistanceType::try_from(
527            schema
528                .metadata
529                .get(DISTANCE_TYPE_KEY)
530                .ok_or(Error::index(format!("{} not found", DISTANCE_TYPE_KEY)))?
531                .as_str(),
532        )?;
533
534        let ivf_pos = schema
535            .metadata
536            .get(IVF_METADATA_KEY)
537            .ok_or(Error::index(format!("{} not found", IVF_METADATA_KEY)))?
538            .parse()
539            .map_err(|e| Error::index(format!("Failed to decode IVF metadata: {}", e)))?;
540        let ivf_bytes = reader.read_global_buffer(ivf_pos).await?;
541        let ivf = IvfModel::try_from(pb::Ivf::decode(ivf_bytes)?)?;
542
543        let mut metadata: Vec<String> = serde_json::from_str(
544            schema
545                .metadata
546                .get(STORAGE_METADATA_KEY)
547                .ok_or(Error::index(format!("{} not found", STORAGE_METADATA_KEY)))?
548                .as_str(),
549        )?;
550        debug_assert_eq!(metadata.len(), 1);
551        // for now the metadata is the same for all partitions, so we just store one
552        let metadata = metadata
553            .pop()
554            .ok_or(Error::index("metadata is empty".to_string()))?;
555        let mut metadata: Q::Metadata = serde_json::from_str(&metadata)?;
556        // we store large metadata (e.g. PQ codebook) in global buffer,
557        // and the schema metadata just contains a pointer to the buffer
558        if let Some(pos) = metadata.buffer_index() {
559            let bytes = reader.read_global_buffer(pos).await?;
560            metadata.parse_buffer(bytes)?;
561        }
562
563        Ok(Self {
564            reader,
565            distance_type,
566            metadata,
567            ivf,
568            frag_reuse_index,
569        })
570    }
571
572    /// Construct from pre-parsed metadata, skipping global buffer reads.
573    /// Used when reconstructing from a disk cache.
574    pub fn from_cached(
575        reader: FileReader,
576        ivf: IvfModel,
577        metadata: Q::Metadata,
578        distance_type: DistanceType,
579        frag_reuse_index: Option<Arc<FragReuseIndex>>,
580    ) -> Self {
581        Self {
582            reader,
583            distance_type,
584            metadata,
585            ivf,
586            frag_reuse_index,
587        }
588    }
589
590    pub fn reader(&self) -> &FileReader {
591        &self.reader
592    }
593
594    pub fn ivf(&self) -> &IvfModel {
595        &self.ivf
596    }
597
598    pub fn num_rows(&self) -> u64 {
599        self.reader.num_rows()
600    }
601
602    pub fn partition_size(&self, part_id: usize) -> usize {
603        self.ivf.partition_size(part_id)
604    }
605
606    pub fn quantizer(&self) -> Result<Quantizer> {
607        let metadata = self.metadata();
608        Q::from_metadata(metadata, self.distance_type)
609    }
610
611    pub fn metadata(&self) -> &Q::Metadata {
612        &self.metadata
613    }
614
615    pub fn distance_type(&self) -> DistanceType {
616        self.distance_type
617    }
618
619    pub fn schema(&self) -> SchemaRef {
620        Arc::new(self.reader.schema().as_ref().into())
621    }
622
623    /// Get the number of partitions in the storage.
624    pub fn num_partitions(&self) -> usize {
625        self.ivf.num_partitions()
626    }
627
628    /// Load a partition's quantization storage, optionally measuring the exact
629    /// I/O it performs into `io_stats`.
630    ///
631    /// When `io_stats` is `Some`, the partition is read through a reader whose
632    /// scheduler also records into the sink (a cheap clone that shares all
633    /// cached metadata, so no file is re-opened).  When `None`, the normal
634    /// uninstrumented reader is used.
635    pub async fn load_partition(
636        &self,
637        part_id: usize,
638        io_stats: Option<IoStats>,
639    ) -> Result<Q::Storage> {
640        let range = self.ivf.row_range(part_id);
641        let batch = if range.is_empty() {
642            let schema = self.reader.schema();
643            let arrow_schema = arrow_schema::Schema::from(schema.as_ref());
644            RecordBatch::new_empty(Arc::new(arrow_schema))
645        } else {
646            let reader = match &io_stats {
647                Some(io_stats) => Cow::Owned(self.reader.with_io_stats(io_stats.recorder())),
648                None => Cow::Borrowed(&self.reader),
649            };
650            let batches = reader
651                .read_stream(
652                    ReadBatchParams::Range(range),
653                    u32::MAX,
654                    1,
655                    FilterExpression::no_filter(),
656                )
657                .await?
658                .try_collect::<Vec<_>>()
659                .await?;
660            let schema = Arc::new(self.reader.schema().as_ref().into());
661            concat_batches(&schema, batches.iter())?
662        };
663        Q::Storage::try_from_batch(
664            batch,
665            self.metadata(),
666            self.distance_type,
667            self.frag_reuse_index.clone(),
668        )
669    }
670}
671
672#[cfg(test)]
673mod tests {
674    use super::{QueryScratchCapacity, QueryScratchPool};
675    use lance_core::deepsize::DeepSizeOf;
676
677    #[test]
678    fn test_query_scratch_pool_reuses_buffers() {
679        let pool = QueryScratchPool::new(1);
680        let first_ptrs = pool.with_scratch(|scratch| {
681            scratch.query_f32.clear();
682            scratch.query_f32.resize(16, 1.0);
683            scratch.distances.clear();
684            scratch.distances.resize(8, 2.0);
685            scratch.u16.clear();
686            scratch.u16.resize(4, 3);
687            scratch.u8.clear();
688            scratch.u8.resize(2, 4);
689            scratch.u32.clear();
690            scratch.u32.resize(3, 5);
691            (
692                scratch.query_f32.as_ptr(),
693                scratch.distances.as_ptr(),
694                scratch.u16.as_ptr(),
695                scratch.u8.as_ptr(),
696                scratch.u32.as_ptr(),
697            )
698        });
699
700        let second_ptrs = pool.with_scratch(|scratch| {
701            assert_eq!(scratch.query_f32.len(), 16);
702            assert!(scratch.query_f32.iter().all(|value| *value == 1.0));
703            assert_eq!(scratch.distances.len(), 8);
704            assert!(scratch.distances.iter().all(|value| *value == 2.0));
705            assert_eq!(scratch.u16.len(), 4);
706            assert!(scratch.u16.iter().all(|value| *value == 3));
707            assert_eq!(scratch.u8.len(), 2);
708            assert!(scratch.u8.iter().all(|value| *value == 4));
709            assert_eq!(scratch.u32.len(), 3);
710            assert!(scratch.u32.iter().all(|value| *value == 5));
711            (
712                scratch.query_f32.as_ptr(),
713                scratch.distances.as_ptr(),
714                scratch.u16.as_ptr(),
715                scratch.u8.as_ptr(),
716                scratch.u32.as_ptr(),
717            )
718        });
719
720        assert_eq!(first_ptrs, second_ptrs);
721    }
722
723    #[test]
724    fn test_query_scratch_pool_is_pool_owned() {
725        let first_pool = QueryScratchPool::new(1);
726        let second_pool = QueryScratchPool::new(1);
727
728        let first_ptr = first_pool.with_scratch(|scratch| {
729            scratch.query_f32.resize(16, 1.0);
730            scratch.query_f32.as_ptr()
731        });
732        let second_ptr = second_pool.with_scratch(|scratch| {
733            scratch.query_f32.resize(16, 1.0);
734            scratch.query_f32.as_ptr()
735        });
736
737        assert_ne!(first_ptr, second_ptr);
738    }
739
740    #[test]
741    fn test_query_scratch_pool_uses_temporary_scratch_when_empty() {
742        let pool =
743            QueryScratchPool::with_capacity(1, QueryScratchCapacity::new_with_u32(8, 16, 4, 2, 3));
744        let pooled = pool.scratch();
745        assert!(pooled.pooled);
746
747        let temporary = pool.scratch();
748        assert!(!temporary.pooled);
749        assert_eq!(temporary.distances.len(), 8);
750        assert_eq!(temporary.query_f32.len(), 16);
751        assert_eq!(temporary.u16.len(), 4);
752        assert_eq!(temporary.u8.len(), 2);
753        assert_eq!(temporary.u32.len(), 3);
754    }
755
756    #[test]
757    fn test_query_scratch_pool_deep_size_includes_buffer_capacity() {
758        let empty_size = QueryScratchPool::new(1).deep_size_of();
759        let pool =
760            QueryScratchPool::with_capacity(1, QueryScratchCapacity::new_with_u32(8, 16, 4, 2, 3));
761
762        assert!(pool.deep_size_of() > empty_size);
763
764        let idle_size = pool.deep_size_of();
765        let _checked_out = pool.scratch();
766
767        assert_eq!(pool.deep_size_of(), idle_size);
768    }
769
770    #[test]
771    fn test_query_scratch_pool_initializes_buffer_capacity() {
772        let pool =
773            QueryScratchPool::with_capacity(1, QueryScratchCapacity::new_with_u32(8, 16, 4, 2, 3));
774
775        pool.with_scratch(|scratch| {
776            assert_eq!(scratch.distances.len(), 8);
777            assert_eq!(scratch.distances.capacity(), 8);
778            assert_eq!(scratch.query_f32.len(), 16);
779            assert_eq!(scratch.query_f32.capacity(), 16);
780            assert_eq!(scratch.u16.len(), 4);
781            assert_eq!(scratch.u16.capacity(), 4);
782            assert_eq!(scratch.u8.len(), 2);
783            assert_eq!(scratch.u8.capacity(), 2);
784            assert_eq!(scratch.u32.len(), 3);
785            assert_eq!(scratch.u32.capacity(), 3);
786        });
787    }
788}
lance_index/vector/storage.rs

lance_index/vector/
storage.rs
