lance_index/vector/sq/
storage.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::ops::Range;
5
6use arrow::datatypes::Float64Type;
7use arrow::{compute::concat_batches, datatypes::Float16Type};
8use arrow_array::{
9    ArrayRef, RecordBatch, UInt8Array, UInt64Array,
10    cast::AsArray,
11    types::{Float32Type, UInt8Type, UInt64Type},
12};
13use arrow_schema::{DataType, SchemaRef};
14use async_trait::async_trait;
15use lance_arrow::ArrowFloatType;
16use lance_core::deepsize::DeepSizeOf;
17use lance_core::{Error, ROW_ID, Result};
18use lance_file::previous::reader::FileReader as PreviousFileReader;
19use lance_io::object_store::ObjectStore;
20use lance_linalg::distance::{DistanceType, dot_u8::dot_u8, l2_u8::l2_u8};
21use lance_table::format::SelfDescribingFileReader;
22use num_traits::AsPrimitive;
23use object_store::path::Path;
24use serde::{Deserialize, Serialize};
25use std::sync::Arc;
26
27use super::{ScalarQuantizer, scale_to_u8};
28use crate::frag_reuse::FragReuseIndex;
29use crate::{
30    INDEX_METADATA_SCHEMA_KEY, IndexMetadata,
31    vector::{
32        SQ_CODE_COLUMN,
33        quantizer::{QuantizerMetadata, QuantizerStorage},
34        storage::{DistCalculator, DistanceCalculatorOptions, QueryResidual, VectorStore},
35        transform::Transformer,
36    },
37};
38
39pub const SQ_METADATA_KEY: &str = "lance:sq";
40
41#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
42pub struct ScalarQuantizationMetadata {
43    pub dim: usize,
44    pub num_bits: u16,
45    pub bounds: Range<f64>,
46}
47
48impl DeepSizeOf for ScalarQuantizationMetadata {
49    fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize {
50        0
51    }
52}
53
54#[async_trait]
55impl QuantizerMetadata for ScalarQuantizationMetadata {
56    async fn load(reader: &PreviousFileReader) -> Result<Self> {
57        let metadata_str = reader
58            .schema()
59            .metadata
60            .get(SQ_METADATA_KEY)
61            .ok_or(Error::index(format!(
62                "Reading SQ metadata: metadata key {} not found",
63                SQ_METADATA_KEY
64            )))?;
65        serde_json::from_str(metadata_str)
66            .map_err(|_| Error::index(format!("Failed to parse index metadata: {}", metadata_str)))
67    }
68}
69
70/// An immutable chunk of ScalarQuantizationStorage.
71#[derive(Debug, Clone)]
72struct SQStorageChunk {
73    batch: RecordBatch,
74
75    dim: usize,
76
77    // Helper fields, references to the batch
78    // These fields share the `Arc` pointer to the columns in batch,
79    // so it does not take more memory.
80    row_ids: UInt64Array,
81    sq_codes: UInt8Array,
82}
83
84impl SQStorageChunk {
85    // Create a new chunk from a RecordBatch.
86    fn new(batch: RecordBatch) -> Result<Self> {
87        let row_ids = batch
88            .column_by_name(ROW_ID)
89            .ok_or(Error::index(
90                "Row ID column not found in the batch".to_owned(),
91            ))?
92            .as_primitive::<UInt64Type>()
93            .clone();
94        let fsl = batch
95            .column_by_name(SQ_CODE_COLUMN)
96            .ok_or(Error::index(
97                "SQ code column not found in the batch".to_owned(),
98            ))?
99            .as_fixed_size_list();
100        let dim = fsl.value_length() as usize;
101        let sq_codes = fsl
102            .values()
103            .as_primitive_opt::<UInt8Type>()
104            .ok_or(Error::index(
105                "SQ code column is not FixedSizeList<u8>".to_owned(),
106            ))?
107            .clone();
108        Ok(Self {
109            batch,
110            dim,
111            row_ids,
112            sq_codes,
113        })
114    }
115
116    /// Returns vector dimension
117    fn dim(&self) -> usize {
118        self.dim
119    }
120
121    fn len(&self) -> usize {
122        self.row_ids.len()
123    }
124
125    fn schema(&self) -> &SchemaRef {
126        self.batch.schema_ref()
127    }
128
129    #[inline]
130    fn row_id(&self, id: u32) -> u64 {
131        self.row_ids.value(id as usize)
132    }
133
134    /// Get a slice of SQ code for id
135    #[inline]
136    fn sq_code_slice(&self, id: u32) -> &[u8] {
137        // assert!(id < self.len() as u32);
138        &self.sq_codes.values()[id as usize * self.dim..(id + 1) as usize * self.dim]
139    }
140}
141
142impl DeepSizeOf for SQStorageChunk {
143    fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
144        self.batch.deep_size_of_children(context)
145    }
146}
147
148#[derive(Debug, Clone)]
149pub struct ScalarQuantizationStorage {
150    quantizer: ScalarQuantizer,
151
152    distance_type: DistanceType,
153
154    /// Chunks of storage
155    offsets: Vec<u32>,
156    chunks: Vec<SQStorageChunk>,
157}
158
159impl DeepSizeOf for ScalarQuantizationStorage {
160    fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
161        self.chunks
162            .iter()
163            .map(|c| c.deep_size_of_children(context))
164            .sum()
165    }
166}
167
168const SQ_CHUNK_CAPACITY: usize = 1024;
169
170impl ScalarQuantizationStorage {
171    pub fn try_new(
172        num_bits: u16,
173        distance_type: DistanceType,
174        bounds: Range<f64>,
175        batches: impl IntoIterator<Item = RecordBatch>,
176        frag_reuse_index: Option<Arc<FragReuseIndex>>,
177    ) -> Result<Self> {
178        let mut chunks = Vec::with_capacity(SQ_CHUNK_CAPACITY);
179        let mut offsets = Vec::with_capacity(SQ_CHUNK_CAPACITY + 1);
180        offsets.push(0);
181        for mut batch in batches.into_iter() {
182            if let Some(frag_reuse_index_ref) = frag_reuse_index.as_ref() {
183                batch = frag_reuse_index_ref.remap_row_ids_record_batch(batch, 0)?
184            }
185            offsets.push(offsets.last().unwrap() + batch.num_rows() as u32);
186            let chunk = SQStorageChunk::new(batch)?;
187            chunks.push(chunk);
188        }
189        let quantizer = ScalarQuantizer::with_bounds(num_bits, chunks[0].dim(), bounds);
190
191        Ok(Self {
192            quantizer,
193            distance_type,
194            offsets,
195            chunks,
196        })
197    }
198
199    /// Get the chunk that covers the id.
200    ///
201    /// Returns:
202    /// `(offset, chunk)`
203    ///
204    /// We did not check out of range in this call. But the out of range will
205    /// panic once you access the data in the last [SQStorageChunk].
206    fn chunk(&self, id: u32) -> (u32, &SQStorageChunk) {
207        match self.offsets.binary_search(&id) {
208            Ok(o) => (self.offsets[o], &self.chunks[o]),
209            Err(o) => (self.offsets[o - 1], &self.chunks[o - 1]),
210        }
211    }
212
213    pub async fn load(
214        object_store: &ObjectStore,
215        path: &Path,
216        frag_reuse_index: Option<Arc<FragReuseIndex>>,
217    ) -> Result<Self> {
218        let reader = PreviousFileReader::try_new_self_described(object_store, path, None).await?;
219        let schema = reader.schema();
220
221        let metadata_str = schema
222            .metadata
223            .get(INDEX_METADATA_SCHEMA_KEY)
224            .ok_or(Error::index(format!(
225                "Reading SQ storage: index key {} not found",
226                INDEX_METADATA_SCHEMA_KEY
227            )))?;
228        let index_metadata: IndexMetadata = serde_json::from_str(metadata_str).map_err(|_| {
229            Error::index(format!("Failed to parse index metadata: {}", metadata_str))
230        })?;
231        let distance_type = DistanceType::try_from(index_metadata.distance_type.as_str())?;
232        let metadata = ScalarQuantizationMetadata::load(&reader).await?;
233
234        Self::load_partition(
235            &reader,
236            0..reader.len(),
237            distance_type,
238            &metadata,
239            frag_reuse_index,
240        )
241        .await
242    }
243
244    fn optimize(self) -> Result<Self> {
245        if self.len() <= SQ_CHUNK_CAPACITY {
246            Ok(self)
247        } else {
248            let mut new = self.clone();
249            let batch = concat_batches(
250                self.chunks[0].schema(),
251                self.chunks.iter().map(|c| &c.batch),
252            )?;
253            new.offsets = vec![0, batch.num_rows() as u32];
254            new.chunks = vec![SQStorageChunk::new(batch)?];
255            Ok(new)
256        }
257    }
258}
259
260#[async_trait]
261impl QuantizerStorage for ScalarQuantizationStorage {
262    type Metadata = ScalarQuantizationMetadata;
263
264    fn try_from_batch(
265        batch: RecordBatch,
266        metadata: &Self::Metadata,
267        distance_type: DistanceType,
268        frag_reuse_index: Option<Arc<FragReuseIndex>>,
269    ) -> Result<Self>
270    where
271        Self: Sized,
272    {
273        Self::try_new(
274            metadata.num_bits,
275            distance_type,
276            metadata.bounds.clone(),
277            [batch],
278            frag_reuse_index,
279        )
280    }
281
282    fn metadata(&self) -> &Self::Metadata {
283        &self.quantizer.metadata
284    }
285
286    /// Load a partition of SQ storage from disk.
287    ///
288    /// Parameters
289    /// ----------
290    /// - *reader: file reader
291    /// - *range: row range of the partition
292    /// - *metric_type: metric type of the vectors
293    /// - *metadata: scalar quantization metadata
294    async fn load_partition(
295        reader: &PreviousFileReader,
296        range: std::ops::Range<usize>,
297        distance_type: DistanceType,
298        metadata: &Self::Metadata,
299        frag_reuse_index: Option<Arc<FragReuseIndex>>,
300    ) -> Result<Self> {
301        let schema = reader.schema();
302        let batch = reader.read_range(range, schema).await?;
303
304        Self::try_new(
305            metadata.num_bits,
306            distance_type,
307            metadata.bounds.clone(),
308            [batch],
309            frag_reuse_index,
310        )
311    }
312}
313
314impl VectorStore for ScalarQuantizationStorage {
315    type DistanceCalculator<'a> = SQDistCalculator<'a>;
316
317    fn to_batches(&self) -> Result<impl Iterator<Item = RecordBatch>> {
318        Ok(self.chunks.iter().map(|c| c.batch.clone()))
319    }
320
321    fn append_batch(&self, batch: RecordBatch, vector_column: &str) -> Result<Self> {
322        // TODO: use chunked storage
323        let transformer = super::transform::SQTransformer::new(
324            self.quantizer.clone(),
325            vector_column.to_string(),
326            SQ_CODE_COLUMN.to_string(),
327        );
328
329        let new_batch = transformer.transform(&batch)?;
330
331        // self.quantizer.transform(data)
332        let mut storage = self.clone();
333        let offset = self.len() as u32;
334        let new_chunk = SQStorageChunk::new(new_batch)?;
335        storage.offsets.push(offset + new_chunk.len() as u32);
336        storage.chunks.push(new_chunk);
337
338        storage.optimize()
339    }
340
341    fn schema(&self) -> &SchemaRef {
342        self.chunks[0].schema()
343    }
344
345    fn as_any(&self) -> &dyn std::any::Any {
346        self
347    }
348
349    fn len(&self) -> usize {
350        *self.offsets.last().unwrap() as usize
351    }
352
353    /// Return the [DistanceType] of the vectors.
354    fn distance_type(&self) -> DistanceType {
355        self.distance_type
356    }
357
358    fn row_id(&self, id: u32) -> u64 {
359        let (offset, chunk) = self.chunk(id);
360        chunk.row_id(id - offset)
361    }
362
363    fn row_ids(&self) -> impl Iterator<Item = &u64> {
364        self.chunks.iter().flat_map(|c| c.row_ids.values())
365    }
366
367    /// Create a [DistCalculator] to compute the distance between the query.
368    ///
369    /// Using dist calculator can be more efficient as it can pre-compute some
370    /// values.
371    fn dist_calculator(&self, query: ArrayRef, _dist_q_c: f32) -> Self::DistanceCalculator<'_> {
372        SQDistCalculator::new(query, self, self.quantizer.bounds())
373    }
374
375    fn dist_calculator_with_scratch<'a>(
376        &'a self,
377        query: ArrayRef,
378        _dist_q_c: f32,
379        _residual: Option<QueryResidual<'a>>,
380        f32_scratch: &'a mut Vec<f32>,
381        _options: DistanceCalculatorOptions,
382    ) -> Self::DistanceCalculator<'a> {
383        SQDistCalculator::new_with_scratch(query, self, self.quantizer.bounds(), f32_scratch)
384    }
385
386    fn dist_calculator_from_id(&self, id: u32) -> Self::DistanceCalculator<'_> {
387        let (offset, chunk) = self.chunk(id);
388        let query_sq_code = chunk.sq_code_slice(id - offset);
389        let bounds = self.quantizer.bounds();
390        let lower_bound = bounds.start as f32;
391        let value_scale = sq_value_scale(&bounds);
392        let query_dot = match self.distance_type {
393            DistanceType::Dot => Some(SQDotQuery::from_sq_code(query_sq_code)),
394            _ => None,
395        };
396        SQDistCalculator {
397            query_sq_code: SQQueryCode::Borrowed(query_sq_code),
398            query_dot,
399            scale: sq_distance_scale(&bounds),
400            lower_bound,
401            value_scale,
402            storage: self,
403        }
404    }
405}
406
407#[inline]
408fn sq_value_scale(bounds: &Range<f64>) -> f32 {
409    (bounds.end - bounds.start) as f32 / 255.0_f32
410}
411
412#[inline]
413fn sq_distance_scale(bounds: &Range<f64>) -> f32 {
414    let scale = sq_value_scale(bounds);
415    scale * scale
416}
417
418pub struct SQDistCalculator<'a> {
419    query_sq_code: SQQueryCode<'a>,
420    query_dot: Option<SQDotQuery<'a>>,
421    scale: f32,
422    lower_bound: f32,
423    value_scale: f32,
424    storage: &'a ScalarQuantizationStorage,
425}
426
427enum SQDotQuery<'a> {
428    Values { values: SQFloatQuery<'a>, sum: f32 },
429    SqCode { code: &'a [u8], sum: f32 },
430}
431
432enum SQFloatQuery<'a> {
433    Borrowed(&'a [f32]),
434    Owned(Vec<f32>),
435}
436
437impl SQFloatQuery<'_> {
438    fn as_slice(&self) -> &[f32] {
439        match self {
440            Self::Borrowed(values) => values,
441            Self::Owned(values) => values,
442        }
443    }
444}
445
446impl<'a> SQDotQuery<'a> {
447    fn from_values<T: ArrowFloatType>(values: &[T::Native]) -> Self
448    where
449        T::Native: AsPrimitive<f32>,
450    {
451        let values: Vec<_> = values.iter().map(|v| v.as_()).collect();
452        let sum = values.iter().sum();
453        Self::Values {
454            values: SQFloatQuery::Owned(values),
455            sum,
456        }
457    }
458
459    fn from_values_with_scratch<T: ArrowFloatType>(
460        values: &[T::Native],
461        scratch: &'a mut Vec<f32>,
462    ) -> Self
463    where
464        T::Native: AsPrimitive<f32>,
465    {
466        scratch.clear();
467        scratch.extend(values.iter().map(|v| v.as_()));
468        let sum = scratch.iter().sum();
469        Self::Values {
470            values: SQFloatQuery::Borrowed(scratch.as_slice()),
471            sum,
472        }
473    }
474
475    fn from_sq_code(sq_code: &'a [u8]) -> Self {
476        Self::SqCode {
477            code: sq_code,
478            sum: sq_code_sum(sq_code),
479        }
480    }
481}
482
483fn sq_code_sum(sq_code: &[u8]) -> f32 {
484    sq_code.iter().map(|code| *code as u32).sum::<u32>() as f32
485}
486
487enum SQQueryCode<'a> {
488    Borrowed(&'a [u8]),
489    Owned(Vec<u8>),
490}
491
492impl SQQueryCode<'_> {
493    #[inline]
494    fn as_slice(&self) -> &[u8] {
495        match self {
496            Self::Borrowed(query) => query,
497            Self::Owned(query) => query,
498        }
499    }
500}
501
502impl<'a> SQDistCalculator<'a> {
503    fn new(query: ArrayRef, storage: &'a ScalarQuantizationStorage, bounds: Range<f64>) -> Self {
504        // This is okay-ish to use hand-rolled dynamic dispatch here
505        // since we search 10s-100s of partitions, we can afford the overhead
506        // this could be annoying at indexing time for HNSW, which requires constructing the
507        // dist calculator frequently. However, HNSW isn't first-class citizen in Lance yet. so be it.
508        let (query_sq_code, query_dot) = match storage.distance_type {
509            DistanceType::Dot => {
510                let query_dot = match query.data_type() {
511                    DataType::Float16 => SQDotQuery::from_values::<Float16Type>(
512                        query.as_primitive::<Float16Type>().values(),
513                    ),
514                    DataType::Float32 => SQDotQuery::from_values::<Float32Type>(
515                        query.as_primitive::<Float32Type>().values(),
516                    ),
517                    DataType::Float64 => SQDotQuery::from_values::<Float64Type>(
518                        query.as_primitive::<Float64Type>().values(),
519                    ),
520                    _ => {
521                        panic!("Unsupported data type for ScalarQuantizationStorage");
522                    }
523                };
524                (SQQueryCode::Owned(Vec::new()), Some(query_dot))
525            }
526            DistanceType::L2 | DistanceType::Cosine => {
527                let query_sq_code = match query.data_type() {
528                    DataType::Float16 => scale_to_u8::<Float16Type>(
529                        query.as_primitive::<Float16Type>().values(),
530                        &bounds,
531                    ),
532                    DataType::Float32 => scale_to_u8::<Float32Type>(
533                        query.as_primitive::<Float32Type>().values(),
534                        &bounds,
535                    ),
536                    DataType::Float64 => scale_to_u8::<Float64Type>(
537                        query.as_primitive::<Float64Type>().values(),
538                        &bounds,
539                    ),
540                    _ => {
541                        panic!("Unsupported data type for ScalarQuantizationStorage");
542                    }
543                };
544                (SQQueryCode::Owned(query_sq_code), None)
545            }
546            _ => panic!("We should not reach here: sq distance can only be L2 or Dot"),
547        };
548        let lower_bound = bounds.start as f32;
549        let value_scale = sq_value_scale(&bounds);
550        Self {
551            query_sq_code,
552            query_dot,
553            scale: sq_distance_scale(&bounds),
554            lower_bound,
555            value_scale,
556            storage,
557        }
558    }
559
560    fn new_with_scratch(
561        query: ArrayRef,
562        storage: &'a ScalarQuantizationStorage,
563        bounds: Range<f64>,
564        f32_scratch: &'a mut Vec<f32>,
565    ) -> Self {
566        if storage.distance_type != DistanceType::Dot {
567            return Self::new(query, storage, bounds);
568        }
569
570        let query_dot = match query.data_type() {
571            DataType::Float16 => SQDotQuery::from_values_with_scratch::<Float16Type>(
572                query.as_primitive::<Float16Type>().values(),
573                f32_scratch,
574            ),
575            DataType::Float32 => SQDotQuery::from_values_with_scratch::<Float32Type>(
576                query.as_primitive::<Float32Type>().values(),
577                f32_scratch,
578            ),
579            DataType::Float64 => SQDotQuery::from_values_with_scratch::<Float64Type>(
580                query.as_primitive::<Float64Type>().values(),
581                f32_scratch,
582            ),
583            _ => {
584                panic!("Unsupported data type for ScalarQuantizationStorage");
585            }
586        };
587        let lower_bound = bounds.start as f32;
588        let value_scale = sq_value_scale(&bounds);
589        Self {
590            query_sq_code: SQQueryCode::Owned(Vec::new()),
591            query_dot: Some(query_dot),
592            scale: sq_distance_scale(&bounds),
593            lower_bound,
594            value_scale,
595            storage,
596        }
597    }
598
599    fn dot_distance(&self, sq_code: &[u8]) -> f32 {
600        let query = self
601            .query_dot
602            .as_ref()
603            .expect("SQ dot distance requires a dot query");
604        let dot = match query {
605            SQDotQuery::Values { values, sum } => {
606                let values = values.as_slice();
607                self.lower_bound * *sum
608                    + self.value_scale
609                        * sq_code
610                            .iter()
611                            .zip(values.iter())
612                            .map(|(code, query_value)| *code as f32 * *query_value)
613                            .sum::<f32>()
614            }
615            SQDotQuery::SqCode {
616                code: query_sq_code,
617                sum: query_code_sum,
618            } => {
619                let dim = sq_code.len() as f32;
620                let code_dot = dot_u8(sq_code, query_sq_code) as f32;
621                let code_sum = sq_code_sum(sq_code);
622                dim * self.lower_bound * self.lower_bound
623                    + self.lower_bound * self.value_scale * (code_sum + *query_code_sum)
624                    + self.scale * code_dot
625            }
626        };
627        1.0 - dot
628    }
629}
630
631impl DistCalculator for SQDistCalculator<'_> {
632    fn distance(&self, id: u32) -> f32 {
633        let (offset, chunk) = self.storage.chunk(id);
634        let sq_code = chunk.sq_code_slice(id - offset);
635        let query_sq_code = self.query_sq_code.as_slice();
636        match self.storage.distance_type {
637            DistanceType::L2 | DistanceType::Cosine => {
638                l2_u8(sq_code, query_sq_code) as f32 * self.scale
639            }
640            DistanceType::Dot => self.dot_distance(sq_code),
641            _ => panic!("We should not reach here: sq distance can only be L2 or Dot"),
642        }
643    }
644
645    fn distance_all(&self, _k_hint: usize) -> Vec<f32> {
646        let query_sq_code = self.query_sq_code.as_slice();
647        match self.storage.distance_type {
648            DistanceType::L2 | DistanceType::Cosine => self
649                .storage
650                .chunks
651                .iter()
652                .flat_map(|c| {
653                    c.sq_codes
654                        .values()
655                        .chunks_exact(c.dim())
656                        .map(|sq_codes| l2_u8(sq_codes, query_sq_code) as f32)
657                })
658                .map(|dist| dist * self.scale)
659                .collect(),
660            DistanceType::Dot => self
661                .storage
662                .chunks
663                .iter()
664                .flat_map(|c| {
665                    c.sq_codes
666                        .values()
667                        .chunks_exact(c.dim())
668                        .map(|sq_codes| self.dot_distance(sq_codes))
669                })
670                .collect(),
671            _ => panic!("We should not reach here: sq distance can only be L2 or Dot"),
672        }
673    }
674
675    #[allow(unused_variables)]
676    fn prefetch(&self, id: u32) {
677        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
678        {
679            const CACHE_LINE_SIZE: usize = 64;
680
681            let (offset, chunk) = self.storage.chunk(id);
682            let dim = chunk.dim();
683            let base_ptr = chunk.sq_code_slice(id - offset).as_ptr();
684
685            unsafe {
686                // Loop over the sq_code to prefetch each cache line
687                for offset in (0..dim).step_by(CACHE_LINE_SIZE) {
688                    {
689                        use core::arch::x86_64::{_MM_HINT_T0, _mm_prefetch};
690                        _mm_prefetch(base_ptr.add(offset) as *const i8, _MM_HINT_T0);
691                    }
692                }
693            }
694        }
695    }
696}
697
698#[cfg(test)]
699mod tests {
700    use super::*;
701
702    use std::iter::repeat_with;
703    use std::sync::Arc;
704
705    use arrow_array::{FixedSizeListArray, Float32Array};
706    use arrow_schema::{DataType, Field, Schema};
707    use lance_arrow::FixedSizeListArrayExt;
708    use lance_testing::datagen::generate_random_array;
709    use rand::prelude::*;
710
711    fn create_record_batch(row_ids: Range<u64>) -> RecordBatch {
712        const DIM: usize = 64;
713
714        let mut rng = rand::rng();
715        let row_ids = UInt64Array::from_iter_values(row_ids);
716        let sq_code = UInt8Array::from_iter_values(
717            repeat_with(|| rng.random::<u8>()).take(row_ids.len() * DIM),
718        );
719        let code_arr = FixedSizeListArray::try_new_from_values(sq_code, DIM as i32).unwrap();
720
721        let schema = Arc::new(Schema::new(vec![
722            Field::new(ROW_ID, DataType::UInt64, false),
723            Field::new(
724                SQ_CODE_COLUMN,
725                DataType::FixedSizeList(
726                    Arc::new(Field::new("item", DataType::UInt8, true)),
727                    DIM as i32,
728                ),
729                false,
730            ),
731        ]));
732        RecordBatch::try_new(schema, vec![Arc::new(row_ids), Arc::new(code_arr)]).unwrap()
733    }
734
735    fn create_record_batch_with_sq_codes(
736        row_ids: Vec<u64>,
737        sq_codes: Vec<u8>,
738        dim: usize,
739    ) -> RecordBatch {
740        assert_eq!(sq_codes.len(), row_ids.len() * dim);
741
742        let row_ids = UInt64Array::from_iter_values(row_ids);
743        let sq_code = UInt8Array::from_iter_values(sq_codes);
744        let code_arr = FixedSizeListArray::try_new_from_values(sq_code, dim as i32).unwrap();
745
746        let schema = Arc::new(Schema::new(vec![
747            Field::new(ROW_ID, DataType::UInt64, false),
748            Field::new(
749                SQ_CODE_COLUMN,
750                DataType::FixedSizeList(
751                    Arc::new(Field::new("item", DataType::UInt8, true)),
752                    dim as i32,
753                ),
754                false,
755            ),
756        ]));
757        RecordBatch::try_new(schema, vec![Arc::new(row_ids), Arc::new(code_arr)]).unwrap()
758    }
759
760    #[test]
761    fn test_get_chunks() {
762        const DIM: usize = 64;
763
764        let storage = ScalarQuantizationStorage::try_new(
765            8,
766            DistanceType::L2,
767            -0.7..0.7,
768            (0..4).map(|start| create_record_batch(start * 100..(start + 1) * 100)),
769            None,
770        )
771        .unwrap();
772
773        assert_eq!(storage.len(), 400);
774
775        let (offset, chunk) = storage.chunk(0);
776        assert_eq!(offset, 0);
777        assert_eq!(chunk.row_id(20), 20);
778
779        let (offset, _) = storage.chunk(50);
780        assert_eq!(offset, 0);
781
782        let row_ids = UInt64Array::from_iter_values(100..250);
783        let vector_data = generate_random_array(row_ids.len() * DIM);
784        let fsl = FixedSizeListArray::try_new_from_values(vector_data, DIM as i32).unwrap();
785
786        let schema = Arc::new(Schema::new(vec![
787            Field::new(ROW_ID, DataType::UInt64, false),
788            Field::new(
789                "vector",
790                DataType::FixedSizeList(
791                    Arc::new(Field::new("item", DataType::Float32, true)),
792                    DIM as i32,
793                ),
794                false,
795            ),
796        ]));
797
798        let second_batch =
799            RecordBatch::try_new(schema, vec![Arc::new(row_ids), Arc::new(fsl)]).unwrap();
800        let storage = storage.append_batch(second_batch, "vector").unwrap();
801
802        assert_eq!(storage.len(), 550);
803        let (offset, chunk) = storage.chunk(112);
804        assert_eq!(offset, 100);
805        assert_eq!(chunk.row_id(10), 110);
806
807        let (offset, chunk) = storage.chunk(432);
808        assert_eq!(offset, 400);
809        assert_eq!(chunk.row_id(5), 105);
810    }
811
812    #[test]
813    fn test_dot_distance_accounts_for_sq_offset() {
814        const DIM: usize = 2;
815
816        let storage = ScalarQuantizationStorage::try_new(
817            8,
818            DistanceType::Dot,
819            -1.0..1.0,
820            [create_record_batch_with_sq_codes(
821                vec![0, 1],
822                vec![
823                    255, 255, // [1.0, 1.0]
824                    0, 191, // [-1.0, 0.498]
825                ],
826                DIM,
827            )],
828            None,
829        )
830        .unwrap();
831
832        let query = Arc::new(Float32Array::from(vec![-1.0, 1.0])) as ArrayRef;
833        let calculator = storage.dist_calculator(query.clone(), 0.0);
834        let distances = calculator.distance_all(2);
835
836        assert!(
837            distances[1] < distances[0],
838            "expected [-1.0, 0.498] to rank before [1.0, 1.0], got {distances:?}"
839        );
840        assert!((calculator.distance(0) - 1.0).abs() < 1e-6);
841        assert!((calculator.distance(1) - -0.49803925).abs() < 1e-6);
842
843        let mut scratch = Vec::new();
844        let scratch_distances = {
845            let scratch_calculator = storage.dist_calculator_with_scratch(
846                query,
847                0.0,
848                None,
849                &mut scratch,
850                DistanceCalculatorOptions::default(),
851            );
852            scratch_calculator.distance_all(2)
853        };
854        assert_eq!(scratch_distances, distances);
855        assert_eq!(scratch.len(), DIM);
856
857        let stored_query_calculator = storage.dist_calculator_from_id(0);
858        assert!((stored_query_calculator.distance(1) - 1.5019608).abs() < 1e-6);
859    }
860}
lance_index/vector/sq/storage.rs

lance_index/vector/sq/
storage.rs
