lance_index/vector/
utils.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use arrow::{
5    array::AsArray,
6    compute::cast,
7    datatypes::{Float16Type, Float32Type, Float64Type},
8};
9use arrow_array::{Array, ArrayRef, BooleanArray, FixedSizeListArray};
10use arrow_schema::{DataType, Field};
11use lance_arrow::FixedSizeListArrayExt;
12use lance_core::{Error, Result};
13use lance_io::encodings::plain::bytes_to_array;
14use lance_linalg::distance::DistanceType;
15use prost::bytes;
16use std::sync::LazyLock;
17use std::{ops::Range, sync::Arc};
18
19use super::pb;
20use crate::pb::Tensor;
21use crate::vector::flat::storage::FlatFloatStorage;
22use crate::vector::hnsw::HNSW;
23use crate::vector::hnsw::builder::{HnswBuildParams, HnswQueryParams};
24use crate::vector::v3::subindex::IvfSubIndex;
25
26enum SimpleIndexStatus {
27    Auto,
28    Enabled,
29    Disabled,
30}
31
32static USE_HNSW_SPEEDUP_INDEXING: LazyLock<SimpleIndexStatus> = LazyLock::new(|| {
33    if let Ok(v) = std::env::var("LANCE_USE_HNSW_SPEEDUP_INDEXING") {
34        if v == "enabled" {
35            SimpleIndexStatus::Enabled
36        } else if v == "disabled" {
37            SimpleIndexStatus::Disabled
38        } else {
39            SimpleIndexStatus::Auto
40        }
41    } else {
42        SimpleIndexStatus::Auto
43    }
44});
45
46#[derive(Debug)]
47pub struct SimpleIndex {
48    store: FlatFloatStorage,
49    index: HNSW,
50}
51
52impl SimpleIndex {
53    pub fn try_new(store: FlatFloatStorage) -> Result<Self> {
54        let hnsw = HNSW::index_vectors(
55            &store,
56            HnswBuildParams::default().ef_construction(15).num_edges(12),
57        )?;
58        Ok(Self { store, index: hnsw })
59    }
60
61    // train HNSW over the centroids to speed up finding the nearest clusters,
62    // only train if all conditions are met:
63    //  - the centroids are float32s or uint8s
64    //  - `num_centroids * dimension >= 1_000_000`
65    //      we benchmarked that it's 2x faster in the case of 1024 centroids and 1024 dimensions,
66    //      so set the threshold to 1_000_000.
67    pub fn may_train_index(
68        centroids: ArrayRef,
69        dimension: usize,
70        distance_type: DistanceType,
71    ) -> Result<Option<Self>> {
72        match *USE_HNSW_SPEEDUP_INDEXING {
73            SimpleIndexStatus::Auto => {
74                if centroids.len() < 1_000_000 {
75                    return Ok(None);
76                }
77            }
78            SimpleIndexStatus::Disabled => return Ok(None),
79            _ => {}
80        }
81
82        let f32_centroids = match centroids.data_type() {
83            DataType::Float16 | DataType::Float32 => {
84                cast(&centroids, &DataType::Float32).map_err(|e| Error::index(e.to_string()))?
85            }
86            _ => return Ok(None),
87        };
88        let fsl = FixedSizeListArray::try_new_from_values(f32_centroids, dimension as i32)?;
89        let store = FlatFloatStorage::new(fsl, distance_type);
90        Self::try_new(store).map(Some)
91    }
92
93    pub(crate) fn search(&self, query: ArrayRef) -> Result<(u32, f32)> {
94        let query = cast(&query, &DataType::Float32).map_err(|e| Error::index(e.to_string()))?;
95        let res = self.index.search_basic(
96            query,
97            1,
98            &HnswQueryParams {
99                ef: 15,
100                lower_bound: None,
101                upper_bound: None,
102                dist_q_c: 0.0,
103            },
104            None,
105            &self.store,
106        )?;
107        Ok((res[0].id, res[0].dist.0))
108    }
109}
110
111#[inline]
112#[allow(dead_code)]
113pub(crate) fn prefetch_arrow_array(array: &dyn Array) -> Result<()> {
114    match array.data_type() {
115        DataType::FixedSizeList(_, _) => {
116            let array = array.as_fixed_size_list();
117            return prefetch_arrow_array(array.values());
118        }
119        DataType::Float16 => {
120            let array = array.as_primitive::<Float16Type>();
121            do_prefetch(array.values().as_ptr_range())
122        }
123        DataType::Float32 => {
124            let array = array.as_primitive::<Float32Type>();
125            do_prefetch(array.values().as_ptr_range())
126        }
127        DataType::Float64 => {
128            let array = array.as_primitive::<Float64Type>();
129            do_prefetch(array.values().as_ptr_range())
130        }
131        _ => {
132            return Err(Error::invalid_input(format!(
133                "Unsupported data type for prefetch: {}",
134                array.data_type()
135            )));
136        }
137    }
138
139    Ok(())
140}
141
142#[inline]
143pub(crate) fn do_prefetch<T>(ptrs: Range<*const T>) {
144    // TODO use rust intrinsics instead of x86 intrinsics
145    // TODO finish this
146    unsafe {
147        let (ptr, end_ptr) = (ptrs.start as *const i8, ptrs.end as *const i8);
148        let mut current_ptr = ptr;
149        while current_ptr < end_ptr {
150            const CACHE_LINE_SIZE: usize = 64;
151            #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
152            {
153                use core::arch::x86_64::{_MM_HINT_T0, _mm_prefetch};
154                _mm_prefetch(current_ptr, _MM_HINT_T0);
155            }
156            current_ptr = current_ptr.add(CACHE_LINE_SIZE);
157        }
158    }
159}
160
161impl From<pb::tensor::DataType> for DataType {
162    fn from(dt: pb::tensor::DataType) -> Self {
163        match dt {
164            pb::tensor::DataType::Uint8 => Self::UInt8,
165            pb::tensor::DataType::Uint16 => Self::UInt16,
166            pb::tensor::DataType::Uint32 => Self::UInt32,
167            pb::tensor::DataType::Uint64 => Self::UInt64,
168            pb::tensor::DataType::Float16 => Self::Float16,
169            pb::tensor::DataType::Float32 => Self::Float32,
170            pb::tensor::DataType::Float64 => Self::Float64,
171            pb::tensor::DataType::Bfloat16 => unimplemented!(),
172        }
173    }
174}
175
176impl TryFrom<&DataType> for pb::tensor::DataType {
177    type Error = Error;
178
179    fn try_from(dt: &DataType) -> Result<Self> {
180        match dt {
181            DataType::UInt8 => Ok(Self::Uint8),
182            DataType::UInt16 => Ok(Self::Uint16),
183            DataType::UInt32 => Ok(Self::Uint32),
184            DataType::UInt64 => Ok(Self::Uint64),
185            DataType::Float16 => Ok(Self::Float16),
186            DataType::Float32 => Ok(Self::Float32),
187            DataType::Float64 => Ok(Self::Float64),
188            _ => Err(Error::index(format!(
189                "pb tensor type not supported: {:?}",
190                dt
191            ))),
192        }
193    }
194}
195
196impl TryFrom<DataType> for pb::tensor::DataType {
197    type Error = Error;
198
199    fn try_from(dt: DataType) -> Result<Self> {
200        (&dt).try_into()
201    }
202}
203
204impl TryFrom<&FixedSizeListArray> for pb::Tensor {
205    type Error = Error;
206
207    fn try_from(array: &FixedSizeListArray) -> Result<Self> {
208        let mut tensor = Self::default();
209        tensor.data_type = pb::tensor::DataType::try_from(array.value_type())? as i32;
210        tensor.shape = vec![array.len() as u32, array.value_length() as u32];
211        let flat_array = array.values();
212        tensor.data = flat_array.into_data().buffers()[0].to_vec();
213        Ok(tensor)
214    }
215}
216
217impl TryFrom<&pb::Tensor> for FixedSizeListArray {
218    type Error = Error;
219
220    fn try_from(tensor: &Tensor) -> Result<Self> {
221        if tensor.shape.len() != 2 {
222            return Err(Error::index(format!(
223                "only accept 2-D tensor shape, got: {:?}",
224                tensor.shape
225            )));
226        }
227        let dim = tensor.shape[1] as usize;
228        let num_rows = tensor.shape[0] as usize;
229
230        let data = bytes::Bytes::from(tensor.data.clone());
231        let flat_array = bytes_to_array(
232            &DataType::from(pb::tensor::DataType::try_from(tensor.data_type).unwrap()),
233            data,
234            dim * num_rows,
235            0,
236        )?;
237
238        if flat_array.len() != dim * num_rows {
239            return Err(Error::index(format!(
240                "Tensor shape {:?} does not match to data len: {}",
241                tensor.shape,
242                flat_array.len()
243            )));
244        }
245
246        let field = Field::new("item", flat_array.data_type().clone(), true);
247        Ok(Self::try_new(
248            Arc::new(field),
249            dim as i32,
250            flat_array,
251            None,
252        )?)
253    }
254}
255
256/// Check if all vectors in the FixedSizeListArray are finite
257/// null values are considered as not finite
258/// returns a BooleanArray
259/// with the same length as the FixedSizeListArray
260/// with true for finite values and false for non-finite values
261pub fn is_finite(fsl: &FixedSizeListArray) -> BooleanArray {
262    let is_finite = fsl
263        .iter()
264        .map(|v| match v {
265            Some(v) => match v.data_type() {
266                DataType::Float16 => {
267                    let v = v.as_primitive::<Float16Type>();
268                    v.null_count() == 0 && v.values().iter().all(|v| v.is_finite())
269                }
270                DataType::Float32 => {
271                    let v = v.as_primitive::<Float32Type>();
272                    v.null_count() == 0 && v.values().iter().all(|v| v.is_finite())
273                }
274                DataType::Float64 => {
275                    let v = v.as_primitive::<Float64Type>();
276                    v.null_count() == 0 && v.values().iter().all(|v| v.is_finite())
277                }
278                _ => v.null_count() == 0,
279            },
280            None => false,
281        })
282        .collect::<Vec<_>>();
283    BooleanArray::from(is_finite)
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289
290    use arrow_array::{Float16Array, Float32Array, Float64Array};
291    use half::f16;
292    use lance_arrow::FixedSizeListArrayExt;
293    use num_traits::identities::Zero;
294
295    use arrow::compute::cast;
296    use rstest::rstest;
297
298    fn build_index(centroids: ArrayRef, dim: usize) -> SimpleIndex {
299        let f32_centroids = cast(&centroids, &DataType::Float32).unwrap();
300        let fsl = FixedSizeListArray::try_new_from_values(f32_centroids, dim as i32).unwrap();
301        let store = FlatFloatStorage::new(fsl, DistanceType::L2);
302        SimpleIndex::try_new(store).unwrap()
303    }
304
305    #[rstest]
306    #[case::f16(Arc::new(Float16Array::from(
307        (0..100).flat_map(|i| std::iter::repeat_n(f16::from_f32(i as f32), 16)).collect::<Vec<_>>(),
308    )) as ArrayRef)]
309    #[case::f32(Arc::new(Float32Array::from(
310        (0..100).flat_map(|i| std::iter::repeat_n(i as f32, 16)).collect::<Vec<_>>(),
311    )) as ArrayRef)]
312    fn test_simple_index_nearest_centroid(#[case] centroids: ArrayRef) {
313        let index = build_index(centroids, 16);
314        let query: ArrayRef = Arc::new(Float32Array::from(vec![42.1f32; 16]));
315        let (id, _) = index.search(query).unwrap();
316        assert_eq!(id, 42);
317    }
318
319    #[test]
320    fn test_simple_index_rejects_f64() {
321        let centroids: ArrayRef = Arc::new(Float64Array::from(vec![0.0; 1600]));
322        let result = SimpleIndex::may_train_index(centroids, 16, DistanceType::L2).unwrap();
323        assert!(result.is_none());
324    }
325
326    #[test]
327    fn test_fsl_to_tensor() {
328        let fsl =
329            FixedSizeListArray::try_new_from_values(Float16Array::from(vec![f16::zero(); 20]), 5)
330                .unwrap();
331        let tensor = pb::Tensor::try_from(&fsl).unwrap();
332        assert_eq!(tensor.data_type, pb::tensor::DataType::Float16 as i32);
333        assert_eq!(tensor.shape, vec![4, 5]);
334        assert_eq!(tensor.data.len(), 20 * 2);
335
336        let fsl =
337            FixedSizeListArray::try_new_from_values(Float32Array::from(vec![0.0; 20]), 5).unwrap();
338        let tensor = pb::Tensor::try_from(&fsl).unwrap();
339        assert_eq!(tensor.data_type, pb::tensor::DataType::Float32 as i32);
340        assert_eq!(tensor.shape, vec![4, 5]);
341        assert_eq!(tensor.data.len(), 20 * 4);
342
343        let fsl =
344            FixedSizeListArray::try_new_from_values(Float64Array::from(vec![0.0; 20]), 5).unwrap();
345        let tensor = pb::Tensor::try_from(&fsl).unwrap();
346        assert_eq!(tensor.data_type, pb::tensor::DataType::Float64 as i32);
347        assert_eq!(tensor.shape, vec![4, 5]);
348        assert_eq!(tensor.data.len(), 20 * 8);
349    }
350}
lance_index/vector/utils.rs

lance_index/vector/
utils.rs
