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

lance_index/vector/
utils.rs
