diff --git a/python/python/lance/util.py b/python/python/lance/util.py index 217f632031f..e99c9a6029a 100644 --- a/python/python/lance/util.py +++ b/python/python/lance/util.py @@ -78,7 +78,7 @@ class KMeans: >>> import numpy as np >>> import lance >>> data = np.random.randn(1000, 128).astype(np.float32) - >>> kmeans = lance.util.KMeans(8, metric_type="cosine") + >>> kmeans = lance.util.KMeans(8, metric_type="l2") >>> kmeans.fit(data) >>> centroids = np.stack(kmeans.centroids.to_numpy(zero_copy_only=False)) >>> clusters = kmeans.predict(data) diff --git a/python/python/tests/test_kmeans.py b/python/python/tests/test_kmeans.py index f9a1a1a240a..da2be6fdd86 100644 --- a/python/python/tests/test_kmeans.py +++ b/python/python/tests/test_kmeans.py @@ -16,30 +16,6 @@ import numpy as np import pyarrow as pa import pytest -from numpy.linalg import norm - - -def test_train_cosine(): - kmeans = lance.util.KMeans(32, metric_type="cosine") - data = np.random.randn(1000, 128).astype(np.float32) - - assert kmeans.centroids is None - kmeans.fit(data) - assert kmeans.centroids is not None - centroids = kmeans.centroids.to_numpy_ndarray() - assert centroids.shape == (32, 128) - - # test predict - pred = kmeans.predict(data) - - # compute predict using numpy brute-force - expected = [] - for row in data: - # Cosine distance - dist = 1 - np.dot(centroids, row) / (norm(centroids, axis=1) * norm(row)) - cluster_id = np.argmin(dist) - expected.append(cluster_id) - assert np.allclose(pred, expected) def test_invalid_inputs(): diff --git a/rust/lance-index/src/vector/ivf.rs b/rust/lance-index/src/vector/ivf.rs index f2c033be1f1..e6ab032a0c2 100644 --- a/rust/lance-index/src/vector/ivf.rs +++ b/rust/lance-index/src/vector/ivf.rs @@ -253,20 +253,17 @@ impl IvfImpl { ) -> Self { let mut transforms: Vec> = vec![]; - // Re-enable it after search path fixed. - // if metric_type == MetricType::Cosine { - // transforms.push(Arc::new(super::transform::NormalizeTransformer::new( - // vector_column, - // ))); - // }; + let mt = if metric_type == MetricType::Cosine { + transforms.push(Arc::new(super::transform::NormalizeTransformer::new( + vector_column, + ))); + MetricType::L2 + } else { + metric_type + }; - // TODO: add range filter - let ivf_transform = Arc::new(IvfTransformer::new( - centroids.clone(), - metric_type, - vector_column, - )); - transforms.push(ivf_transform.clone() as Arc); + let ivf_transform = Arc::new(IvfTransformer::new(centroids.clone(), mt, vector_column)); + transforms.push(ivf_transform.clone()); if let Some(range) = range { transforms.push(Arc::new(transform::PartitionFilter::new( @@ -349,16 +346,15 @@ impl Ivf for IvfImpl { self.compute_partitions(original).await? }; let dim = original.value_length() as usize; - let mut residual_arr: Vec<::Native> = - Vec::with_capacity(original.values().len()); - flatten_arr + let residual_arr = flatten_arr .as_slice() .chunks_exact(dim) .zip(part_ids.values()) - .for_each(|(vector, &part_id)| { + .flat_map(|(vector, &part_id)| { let centroid = self.centroids.row(part_id as usize).unwrap(); - residual_arr.extend(vector.iter().zip(centroid.iter()).map(|(&v, &c)| v - c)); - }); + vector.iter().zip(centroid.iter()).map(|(&v, &c)| v - c) + }) + .collect::>(); let arr = T::ArrayType::from(residual_arr); Ok(FixedSizeListArray::try_new_from_values(arr, dim as i32)?) } @@ -375,12 +371,13 @@ impl Ivf for IvfImpl { ), location: Default::default(), })?; - // todo: hold kmeans in this struct. - let kmeans = KMeans::::with_centroids( - self.centroids.data().clone(), - self.dimension(), - self.metric_type, - ); + let mt = if self.metric_type == MetricType::Cosine { + MetricType::L2 + } else { + self.metric_type + }; + let kmeans = + KMeans::::with_centroids(self.centroids.data().clone(), self.dimension(), mt); Ok(kmeans.find_partitions(query.as_slice(), nprobes)?) } } diff --git a/rust/lance-index/src/vector/ivf/transformer.rs b/rust/lance-index/src/vector/ivf/transformer.rs new file mode 100644 index 00000000000..a2d474de257 --- /dev/null +++ b/rust/lance-index/src/vector/ivf/transformer.rs @@ -0,0 +1,17 @@ +// Copyright 2024 Lance Developers. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! IVF Transformer +//! +//! It transforms a column of vectors into a column of IVF diff --git a/rust/lance-index/src/vector/kmeans.rs b/rust/lance-index/src/vector/kmeans.rs index 31c7e7c4332..07eea350303 100644 --- a/rust/lance-index/src/vector/kmeans.rs +++ b/rust/lance-index/src/vector/kmeans.rs @@ -21,13 +21,13 @@ use std::sync::Arc; use lance_core::{Error, Result}; use lance_linalg::{ - distance::{Cosine, Dot, MetricType, L2}, + distance::{Dot, MetricType, L2}, kmeans::{KMeans, KMeansParams}, }; /// Train KMeans model and returns the centroids of each cluster. #[allow(clippy::too_many_arguments)] -pub async fn train_kmeans( +pub async fn train_kmeans( array: &T::ArrayType, centroids: Option>, dimension: usize, diff --git a/rust/lance-index/src/vector/pq.rs b/rust/lance-index/src/vector/pq.rs index c0033e747f3..bc95be969cc 100644 --- a/rust/lance-index/src/vector/pq.rs +++ b/rust/lance-index/src/vector/pq.rs @@ -24,10 +24,8 @@ use async_trait::async_trait; use lance_arrow::floats::FloatArray; use lance_arrow::*; use lance_core::{Error, Result}; -use lance_linalg::distance::{ - cosine_distance_batch, dot_distance_batch, l2_distance_batch, Cosine, Dot, L2, -}; -use lance_linalg::kernels::{argmin, argmin_value_float, normalize}; +use lance_linalg::distance::{dot_distance_batch, l2_distance_batch, Dot, L2}; +use lance_linalg::kernels::{argmin, argmin_value_float}; use lance_linalg::{distance::MetricType, MatrixView}; use snafu::{location, Location}; pub mod builder; @@ -82,7 +80,7 @@ pub trait ProductQuantizer: Send + Sync + std::fmt::Debug { // // TODO: move this to be pub(crate) once we have a better way to test it. #[derive(Debug)] -pub struct ProductQuantizerImpl { +pub struct ProductQuantizerImpl { /// Number of bits for the centroids. /// /// Only support 8, as one of `u8` byte now. @@ -117,7 +115,7 @@ pub struct ProductQuantizerImpl { pub codebook: Arc, } -impl ProductQuantizerImpl { +impl ProductQuantizerImpl { /// Create a [`ProductQuantizer`] with pre-trained codebook. pub fn new( m: usize, @@ -126,6 +124,11 @@ impl ProductQuantizerImpl { codebook: Arc, metric_type: MetricType, ) -> Self { + assert_ne!( + metric_type, + MetricType::Cosine, + "Product quantization does not support cosine, use normalized L2 instead" + ); assert_eq!(nbits, 8, "nbits can only be 8"); Self { num_bits: nbits, @@ -201,12 +204,12 @@ impl ProductQuantizerImpl { lance_linalg::distance::DistanceType::L2 => { l2_distance_batch(sub_vec, centroids, sub_vector_width) } - lance_linalg::distance::DistanceType::Cosine => { - cosine_distance_batch(sub_vec, centroids, sub_vector_width) - } lance_linalg::distance::DistanceType::Dot => { dot_distance_batch(sub_vec, centroids, sub_vector_width) } + lance_linalg::distance::DistanceType::Cosine => { + panic!("There should not be cosine for PQ"); + } }; argmin_value_float(distances).map(|(_, v)| v).unwrap_or(0.0) }) @@ -321,7 +324,7 @@ impl ProductQuantizerImpl { } #[async_trait] -impl ProductQuantizer for ProductQuantizerImpl { +impl ProductQuantizer for ProductQuantizerImpl { fn as_any(&self) -> &dyn Any { self } @@ -331,36 +334,13 @@ impl ProductQuantizer for Produ .as_fixed_size_list_opt() .ok_or(Error::Index { message: format!( - "Expect to be a float vector array, got: {:?}", + "Expect to be a FixedSizeList vector array, got: {:?} array", data.data_type() ), location: location!(), })? .clone(); - let fsl = if self.metric_type == MetricType::Cosine { - // Normalize cosine vectors to unit length. - let values = fsl - .values() - .as_any() - .downcast_ref::() - .ok_or(Error::Index { - message: format!( - "Expect to be a float vector array, got: {:?}", - fsl.value_type() - ), - location: location!(), - })? - .as_slice() - .chunks(self.dimension) - .flat_map(normalize) - .collect::>(); - let data = T::ArrayType::from(values); - FixedSizeListArray::try_new_from_values(data, self.dimension as i32)? - } else { - fsl - }; - let num_sub_vectors = self.num_sub_vectors; let dim = self.dimension; let num_rows = fsl.len(); @@ -435,20 +415,10 @@ impl ProductQuantizer for Produ match self.metric_type { MetricType::L2 => self.l2_distances(query, code), MetricType::Cosine => { - let query: &T::ArrayType = query.as_any().downcast_ref().ok_or(Error::Index { - message: format!( - "Build cosine distance table, type mismatch: {}", - query.data_type() - ), - location: Default::default(), - })?; - - // Normalized query vector. - let query = T::ArrayType::from(normalize(query.as_slice()).collect::>()); // L2 over normalized vectors: ||x - y|| = x^2 + y^2 - 2 * xy = 1 + 1 - 2 * xy = 2 * (1 - xy) // Cosine distance: 1 - |xy| / (||x|| * ||y||) = 1 - xy / (x^2 * y^2) = 1 - xy / (1 * 1) = 1 - xy // Therefore, Cosine = L2 / 2 - let l2_dists = self.l2_distances(&query, code)?; + let l2_dists = self.l2_distances(query, code)?; Ok(l2_dists.values().iter().map(|v| *v / 2.0).collect()) } MetricType::Dot => self.dot_distances(query, code), @@ -506,7 +476,7 @@ mod tests { use approx::assert_relative_eq; use arrow_array::{ types::{Float16Type, Float32Type}, - Float16Array, Float32Array, + Float16Array, }; use half::f16; use lance_testing::datagen::generate_random_array; @@ -535,28 +505,6 @@ mod tests { assert_eq!(tensor.shape, vec![256, 16]); } - #[tokio::test] - async fn test_empty_dist_iter() { - let pq = ProductQuantizerImpl:: { - num_bits: 8, - num_sub_vectors: 4, - dimension: 16, - codebook: Arc::new(Float32Array::from_iter_values( - (0..256 * 16).map(|v| v as f32), - )), - metric_type: MetricType::Cosine, - }; - - let data = Float32Array::from_iter_values(repeat(0.0).take(16)); - let data = FixedSizeListArray::try_new_from_values(data, 16).unwrap(); - let rst = pq.transform(&data).await; - assert!(rst.is_err()); - assert!(rst - .unwrap_err() - .to_string() - .contains("it is likely that distance is NaN")); - } - #[tokio::test] async fn test_l2_distance() { const DIM: usize = 512; diff --git a/rust/lance-index/src/vector/pq/builder.rs b/rust/lance-index/src/vector/pq/builder.rs index b51772f9dd5..c2116ed3c55 100644 --- a/rust/lance-index/src/vector/pq/builder.rs +++ b/rust/lance-index/src/vector/pq/builder.rs @@ -27,7 +27,7 @@ use arrow_schema::DataType; use futures::{stream, StreamExt, TryStreamExt}; use lance_arrow::{ArrowFloatType, FloatArray}; use lance_core::{Error, Result}; -use lance_linalg::distance::{Cosine, Dot, L2}; +use lance_linalg::distance::{Dot, L2}; use lance_linalg::{distance::MetricType, MatrixView}; use rand::{self, SeedableRng}; use snafu::{location, Location}; @@ -95,23 +95,23 @@ impl PQBuildParams { } } - pub async fn build_from_matrix( + pub async fn build_from_matrix( &self, data: &MatrixView, metric_type: MetricType, ) -> Result> { - let (data, mt) = if metric_type == MetricType::Cosine { - // Use normalize L2 to train for cosine distance. - (data.normalize(), MetricType::L2) - } else { - (data.clone(), metric_type) - }; - - let sub_vectors = divide_to_subvectors(&data, self.num_sub_vectors); + assert_ne!( + metric_type, + MetricType::Cosine, + "PQ code does not support cosine" + ); + + const REDOS: usize = 1; + + let sub_vectors = divide_to_subvectors(data, self.num_sub_vectors); let num_centroids = 2_usize.pow(self.num_bits as u32); let dimension = data.num_columns(); let sub_vector_dimension = dimension / self.num_sub_vectors; - const REDOS: usize = 1; let d = stream::iter(sub_vectors.into_iter()) .map(|sub_vec| async move { @@ -124,7 +124,7 @@ impl PQBuildParams { self.max_iters as u32, REDOS, rng.clone(), - mt, + metric_type, self.sample_rate, ) .await @@ -149,6 +149,8 @@ impl PQBuildParams { } /// Build a [ProductQuantizer] from the given data. + /// + /// If the [MetricType] is [MetricType::Cosine], the input data will be normalized. pub async fn build( &self, data: &dyn Array, @@ -185,7 +187,7 @@ impl PQBuildParams { } fn create_typed_pq< - T: ArrowFloatType> + ArrowNumericType + L2 + Cosine + Dot, + T: ArrowFloatType> + ArrowNumericType + L2 + Dot, >( proto: &Pq, metric_type: MetricType, @@ -202,6 +204,12 @@ fn create_typed_pq< /// Load ProductQuantizer from Protobuf pub fn from_proto(proto: &Pq, metric_type: MetricType) -> Result> { + let mt = if metric_type == MetricType::Cosine { + MetricType::L2 + } else { + metric_type + }; + if let Some(tensor) = &proto.codebook_tensor { let fsl = FixedSizeListArray::try_from(tensor)?; @@ -209,21 +217,15 @@ pub fn from_proto(proto: &Pq, metric_type: MetricType) -> Result { unimplemented!() } - pb::tensor::DataType::Float16 => Ok(create_typed_pq::( - proto, - metric_type, - fsl.values(), - )), - pb::tensor::DataType::Float32 => Ok(create_typed_pq::( - proto, - metric_type, - fsl.values(), - )), - pb::tensor::DataType::Float64 => Ok(create_typed_pq::( - proto, - metric_type, - fsl.values(), - )), + pb::tensor::DataType::Float16 => { + Ok(create_typed_pq::(proto, mt, fsl.values())) + } + pb::tensor::DataType::Float32 => { + Ok(create_typed_pq::(proto, mt, fsl.values())) + } + pb::tensor::DataType::Float64 => { + Ok(create_typed_pq::(proto, mt, fsl.values())) + } _ => Err(Error::Index { message: format!("PQ builder: unsupported data type: {:?}", tensor.data_type), location: location!(), diff --git a/rust/lance-linalg/src/kernels.rs b/rust/lance-linalg/src/kernels.rs index 985ade4e15b..50949a32c7d 100644 --- a/rust/lance-linalg/src/kernels.rs +++ b/rust/lance-linalg/src/kernels.rs @@ -17,12 +17,11 @@ use std::iter::Sum; use std::sync::Arc; use std::{collections::hash_map::DefaultHasher, hash::Hash, hash::Hasher}; -use arrow_array::cast::AsArray; -use arrow_array::types::{Float16Type, Float32Type, Float64Type}; use arrow_array::{ - cast::{as_largestring_array, as_primitive_array, as_string_array}, + cast::{as_largestring_array, as_primitive_array, as_string_array, AsArray}, types::{ - Int16Type, Int32Type, Int64Type, Int8Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type, + Float16Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, + UInt16Type, UInt32Type, UInt64Type, UInt8Type, }, Array, ArrayRef, ArrowNumericType, ArrowPrimitiveType, FixedSizeListArray, GenericStringArray, OffsetSizeTrait, PrimitiveArray, UInt64Array, diff --git a/rust/lance-linalg/src/kmeans.rs b/rust/lance-linalg/src/kmeans.rs index 8d73ab703b6..18fffa48fe3 100644 --- a/rust/lance-linalg/src/kmeans.rs +++ b/rust/lance-linalg/src/kmeans.rs @@ -36,8 +36,8 @@ use rand::prelude::*; use rand::Rng; use tracing::instrument; -use crate::distance::{dot_distance_batch, norm_l2}; -use crate::kernels::{argmax, argmin_value_float, normalize}; +use crate::distance::dot_distance_batch; +use crate::kernels::{argmax, argmin_value_float}; use crate::{ distance::{ dot_distance, @@ -240,16 +240,6 @@ impl KMeanMembership { split_clusters(&mut cluster_cnts, &mut new_centroids, dimension); - if self.metric_type == MetricType::Cosine { - // Need normalize centroids - for i in 0..self.k { - let norm = norm_l2(&new_centroids[i * dimension..(i + 1) * dimension]); - new_centroids[i * dimension..(i + 1) * dimension] - .iter_mut() - .for_each(|v| *v /= T::Native::from_f32(norm).unwrap()); - } - } - Ok(KMeans { centroids: Arc::new(new_centroids.into()), dimension, @@ -365,6 +355,8 @@ where } /// Train a [`KMeans`] model with full parameters. + /// + /// If the MetricType is `Cosine`, the input vectors will be normalized with each iteration. pub async fn new_with_params( data: &FixedSizeListArray, k: usize, @@ -402,7 +394,8 @@ where let mut best_kmeans = Self::empty(k, dimension, params.metric_type); let mut best_stddev = f32::MAX; - let rng = rand::rngs::SmallRng::from_entropy(); + // TODO: use seed for Rng. + let rng = SmallRng::from_entropy(); for redo in 1..=params.redos { let mut kmeans = if let Some(centroids) = params.centroids.as_ref() { // Use existing centroids. @@ -475,8 +468,6 @@ where /// /// - *data*: a `N * dimension` floating array. Not necessarily normalized. /// - /// If the metric type is cosine, the vector will be normalized internally. - /// pub async fn compute_membership(&self, data: Arc) -> KMeanMembership { let dimension = self.dimension; let n = data.len() / self.dimension; @@ -498,14 +489,6 @@ where return compute_partitions_l2(centroids_array, values, dimension) .collect(); } - MetricType::Cosine => { - let normalized = values - .chunks(dimension) - .flat_map(normalize) - .collect::>(); - return compute_partitions_l2(centroids_array, &normalized, dimension) - .collect(); - } MetricType::Dot => values .chunks_exact(dimension) .map(|vector| { @@ -513,6 +496,9 @@ where argmin_value(centroid_stream.map(|cent| dot_distance(vector, cent))) }) .collect::>(), + MetricType::Cosine => { + panic!("KMeans: should not use cosine distance to train kmeans, use L2 instead."); + } } }) .await @@ -547,8 +533,9 @@ where l2_distance_batch(query, self.centroids.as_slice(), self.dimension).collect() } MetricType::Cosine => { - let normalized = normalize(query).collect::>(); - l2_distance_batch(&normalized, self.centroids.as_slice(), self.dimension).collect() + panic!( + "KMeans::find_partitions: cosine is not supported, use Normalized L2 instead" + ); } MetricType::Dot => { dot_distance_batch(query, self.centroids.as_slice(), self.dimension).collect() @@ -686,9 +673,7 @@ mod tests { use std::iter::repeat; use super::*; - use approx::{assert_relative_eq, relative_eq}; - use crate::distance::cosine_distance_batch; use arrow_array::types::Float32Type; use arrow_array::Float32Array; use lance_arrow::*; @@ -736,51 +721,6 @@ mod tests { assert_eq!(expected, actual); } - #[tokio::test] - async fn test_cosine_kmeans() { - const DIM: usize = 2; - let data = Float32Array::from(vec![1.0, 0.0, 0.0, 1.0, -1.0, -1.0, -2.0, -2.5_f32]); - let vectors = FixedSizeListArray::try_new_from_values(data, DIM as i32).unwrap(); - let params: KMeansParams = KMeansParams::cosine(); - - let expected_x = 0.5_f32.powf(0.5); - - let kmeans = KMeans::new_with_params(&vectors, 2, ¶ms).await.unwrap(); - assert!(kmeans - .centroids - .as_slice() - .iter() - .any(|&v| relative_eq!(expected_x, v) || relative_eq!(-0.7808688, v))); - - // All centroids are normalized - kmeans.centroids.as_slice().chunks(DIM).for_each(|cent| { - assert_relative_eq!(1.0, cent.iter().map(|&x| x.powi(2)).sum::()); - }) - } - - #[tokio::test] - async fn test_cosine_find_partitions() { - const DIM: usize = 8; - const K: usize = 16; - let data = generate_random_array(DIM * K * 100); - let vectors = FixedSizeListArray::try_new_from_values(data.clone(), DIM as i32).unwrap(); - let params = KMeansParams::::cosine(); - let kmeans = KMeans::new_with_params(&vectors, K, ¶ms).await.unwrap(); - - // Use cosine distance to brute-force to find the nearest neighbors. - // without normalization. - let query = &data.as_slice()[0..DIM]; - let dists = Float32Array::from_iter_values(cosine_distance_batch( - query, - kmeans.centroids.as_slice(), - DIM, - )); - let expected = sort_to_indices(&dists, None, Some(4)).unwrap(); - - let actual = kmeans.find_partitions(query, 4).unwrap(); - assert_eq!(expected, actual); - } - #[tokio::test] async fn test_l2_with_nans() { const DIM: usize = 8; @@ -810,22 +750,4 @@ mod tests { .iter() .for_each(|cd| assert!(cd.is_none())); } - - #[tokio::test] - async fn test_train_cosine_kmeans_with_invalid_values() { - const DIM: usize = 8; - const K: usize = 32; - const NUM_CENTROIDS: usize = 16 * 2048; - let centroids = Arc::new(generate_random_array(DIM * NUM_CENTROIDS)); - for val in [0.0, f32::NAN, f32::NEG_INFINITY, f32::INFINITY] { - let values = Float32Array::from_iter_values(repeat(val).take(DIM * K)); - let kmeans = - KMeans::::with_centroids(centroids.clone(), DIM, MetricType::Cosine); - let membership = kmeans.compute_membership(values.into()).await; - membership - .cluster_id_and_distances - .iter() - .for_each(|cd| assert!(cd.is_none())); - } - } } diff --git a/rust/lance-testing/src/datagen.rs b/rust/lance-testing/src/datagen.rs index 175cab621c0..1e2be98e160 100644 --- a/rust/lance-testing/src/datagen.rs +++ b/rust/lance-testing/src/datagen.rs @@ -216,6 +216,14 @@ pub fn generate_random_array(n: usize) -> Float32Array { Float32Array::from_iter_values(repeat_with(|| rng.gen::()).take(n)) } +/// Create a random float32 array where each element is uniformly distributed a +/// given range. +pub fn generate_random_array_with_range(n: usize, range: Range) -> Float32Array { + let mut rng = rand::thread_rng(); + let distribution = Uniform::new(range.start, range.end); + Float32Array::from_iter_values(repeat_with(|| distribution.sample(&mut rng)).take(n)) +} + /// Create a random float32 array where each element is uniformly /// distributed across the given range pub fn generate_scaled_random_array(n: usize, min: f32, max: f32) -> Float32Array { diff --git a/rust/lance/src/index/vector/fixture_test.rs b/rust/lance/src/index/vector/fixture_test.rs index f0c32deb10b..6801811747e 100644 --- a/rust/lance/src/index/vector/fixture_test.rs +++ b/rust/lance/src/index/vector/fixture_test.rs @@ -23,6 +23,7 @@ mod test { sync::{Arc, Mutex}, }; + use approx::assert_relative_eq; use arrow::array::AsArray; use arrow_array::{FixedSizeListArray, Float32Array, RecordBatch}; use arrow_schema::{DataType, Field, Schema}; @@ -87,7 +88,7 @@ mod test { let key: &Float32Array = query.key.as_primitive(); assert_eq!(key.len(), self.assert_query_value.len()); for (i, &v) in key.iter().zip(self.assert_query_value.iter()) { - assert_eq!(v, i.unwrap()); + assert_relative_eq!(v, i.unwrap()); } Ok(self.ret_val.clone()) } @@ -129,10 +130,9 @@ mod test { let centroids = FixedSizeListArray::try_new_from_values(centroids, 2).unwrap(); let mut ivf = Ivf::new(Arc::new(centroids)); // Add 4 partitions - ivf.add_partition(0, 0); - ivf.add_partition(0, 0); - ivf.add_partition(0, 0); - ivf.add_partition(0, 0); + for _ in 0..4 { + ivf.add_partition(0, 0); + } // hold on to this pointer, because the index only holds a weak reference let session = Arc::new(Session::default()); @@ -177,20 +177,20 @@ mod test { } in [ // L2 should residualize with the correct centroid TestCase { - query: vec![1.0, 1.0], + query: vec![1.0; 2], metric: MetricType::L2, - expected_query_at_subindex: vec![0.0, 0.0], + expected_query_at_subindex: vec![0.0; 2], }, // Cosine should normalize and residualize TestCase { - query: vec![1.0, 1.0], + query: vec![1.0; 2], metric: MetricType::Cosine, - expected_query_at_subindex: vec![0.0, 0.0], + expected_query_at_subindex: vec![1.0 / 2.0_f32.sqrt() - 1.0; 2], }, TestCase { - query: vec![2.0, 2.0], + query: vec![2.0; 2], metric: MetricType::Cosine, - expected_query_at_subindex: vec![0.0, 0.0], + expected_query_at_subindex: vec![2.0 / 8.0_f32.sqrt() - 1.0; 2], }, ] { let q = Query { diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 59a84956834..48452cd91f0 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -1,4 +1,4 @@ -// Copyright 2023 Lance Developers. +// Copyright 2024 Lance Developers. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -41,7 +41,7 @@ use lance_index::{ optimize::OptimizeOptions, vector::{ ivf::{builder::load_precomputed_partitions, shuffler::shuffle_dataset, IvfBuildParams}, - pq::{PQBuildParams, ProductQuantizer, ProductQuantizerImpl}, + pq::{PQBuildParams, ProductQuantizer}, Query, DIST_COL, }, Index, IndexType, @@ -55,18 +55,21 @@ use lance_io::{ traits::{Reader, WriteExt, Writer}, }; use lance_linalg::distance::{Cosine, Dot, MetricType, L2}; +use lance_linalg::kernels::{normalize_arrow, normalize_fsl}; use log::{debug, info}; use object_store::path::Path; use rand::{rngs::SmallRng, SeedableRng}; use roaring::RoaringBitmap; use serde::Serialize; use snafu::{location, Location}; -use tracing::{instrument, span, Level}; +use tracing::instrument; use uuid::Uuid; -#[cfg(feature = "opq")] -use super::opq::train_opq; -use super::{pq::PQIndex, utils::maybe_sample_training_data, VectorIndex}; +use super::{ + pq::{build_pq_model, PQIndex}, + utils::maybe_sample_training_data, + VectorIndex, +}; use crate::dataset::builder::DatasetBuilder; use crate::{ dataset::Dataset, @@ -407,13 +410,20 @@ impl Index for IVFIndex { impl VectorIndex for IVFIndex { #[instrument(level = "debug", skip_all, name = "IVFIndex::search")] async fn search(&self, query: &Query, pre_filter: Arc) -> Result { - let partition_ids = - self.ivf - .find_partitions(&query.key, query.nprobes, self.metric_type)?; + let mut query = query.clone(); + let mt = if self.metric_type == MetricType::Cosine { + let key = normalize_arrow(&query.key)?; + query.key = key; + MetricType::L2 + } else { + self.metric_type + }; + + let partition_ids = self.ivf.find_partitions(&query.key, query.nprobes, mt)?; assert!(partition_ids.len() <= query.nprobes); let part_ids = partition_ids.values().to_vec(); let batches = stream::iter(part_ids) - .map(|part_id| self.search_in_partition(part_id as usize, query, pre_filter.clone())) + .map(|part_id| self.search_in_partition(part_id as usize, &query, pre_filter.clone())) .buffer_unordered(num_cpus::get()) .try_collect::>() .await?; @@ -561,7 +571,7 @@ pub(crate) struct Ivf { /// /// It is a 2-D `(num_partitions * dimension)` of float32 array, 64-bit aligned via Arrow /// memory allocator. - centroids: Arc, + pub(crate) centroids: Arc, /// Offset of each partition in the file. offsets: Vec, @@ -580,7 +590,7 @@ impl Ivf { } /// Ivf model dimension. - fn dimension(&self) -> usize { + pub(super) fn dimension(&self) -> usize { self.centroids.value_length() as usize } @@ -725,6 +735,75 @@ fn sanity_check_params(ivf: &IvfBuildParams, pq: &PQBuildParams) -> Result<()> { Ok(()) } +/// Build IVF model from the dataset. +/// +/// Parameters +/// ---------- +/// - *dataset*: Dataset instance +/// - *column*: vector column. +/// - *dim*: vector dimension. +/// - *metric_type*: distance metric type. +/// - *params*: IVF build parameters. +/// +/// Returns +/// ------- +/// - IVF model. +/// +/// Visibility: pub(super) for testing +#[instrument(level = "debug", skip_all, name = "build_ivf_model")] +pub(super) async fn build_ivf_model( + dataset: &Dataset, + column: &str, + dim: usize, + metric_type: MetricType, + params: &IvfBuildParams, +) -> Result { + if let Some(centroids) = params.centroids.as_ref() { + info!("Pre-computed IVF centroids is provided, skip IVF training"); + if centroids.values().len() != params.num_partitions * dim { + return Err(Error::Index { + message: format!( + "IVF centroids length mismatch: {} != {}", + centroids.len(), + params.num_partitions * dim, + ), + location: location!(), + }); + } + return Ok(Ivf::new(centroids.clone())); + } + let sample_size_hint = params.num_partitions * params.sample_rate; + + let start = std::time::Instant::now(); + info!( + "Loading training data for IVF. Sample size: {}", + sample_size_hint + ); + let training_data = maybe_sample_training_data(dataset, column, sample_size_hint).await?; + info!( + "Finished loading training data in {:02} seconds", + start.elapsed().as_secs_f32() + ); + + // If metric type is cosine, normalize the training data, and after this point, + // treat the metric type as L2. + let (training_data, mt) = if metric_type == MetricType::Cosine { + let training_data = normalize_fsl(&training_data)?; + (training_data, MetricType::L2) + } else { + (training_data, metric_type) + }; + + info!("Start to train IVF model"); + let start = std::time::Instant::now(); + let ivf = train_ivf_model(&training_data, mt, params).await?; + info!( + "Trained IVF model in {:02} seconds", + start.elapsed().as_secs_f32() + ); + Ok(ivf) +} + /// Build IVF(PQ) index pub async fn build_ivf_pq_index( dataset: &Dataset, @@ -758,151 +837,14 @@ pub async fn build_ivf_pq_index( }); }; - // Maximum to train [IvfBuildParams::sample_size](default 256) vectors per centroid, see Faiss. - let sample_size_hint = std::cmp::max( - ivf_params.num_partitions, - lance_index::vector::pq::num_centroids(pq_params.num_bits as u32), - ) * ivf_params.sample_rate; - - let mut training_data = if ivf_params.centroids.is_none() { - let start = std::time::Instant::now(); - log::info!( - "Loading training data for IVF. Sample size: {}", - sample_size_hint - ); - let data = Some(maybe_sample_training_data(dataset, column, sample_size_hint).await?); - log::info!( - "Finished loading training data in {:02} seconds", - start.elapsed().as_secs_f32() - ); - data - } else { - None - }; + let ivf_model = build_ivf_model(dataset, column, dim, metric_type, ivf_params).await?; - #[cfg(feature = "opq")] - let mut transforms: Vec> = vec![]; - #[cfg(not(feature = "opq"))] - let transforms: Vec> = vec![]; - - let start = std::time::Instant::now(); - // Train IVF partitions. - let ivf_model = if let Some(centroids) = &ivf_params.centroids { - if centroids.values().len() != ivf_params.num_partitions * dim { - return Err(Error::Index { - message: format!( - "IVF centroids length mismatch: {} != {}", - centroids.len(), - ivf_params.num_partitions * dim, - ), - location: location!(), - }); - } - Ivf::new(centroids.clone()) + let ivf_residual = if matches!(metric_type, MetricType::Cosine | MetricType::L2) { + Some(&ivf_model) } else { - // Transform training data if necessary. - for transform in transforms.iter() { - if let Some(training_data) = &mut training_data { - *training_data = transform.transform(training_data).await?; - } - } - - info!("Start to train IVF model"); - train_ivf_model(training_data.as_ref().unwrap(), metric_type, ivf_params).await? - }; - info!( - "Trained IVF model in {:02} seconds", - start.elapsed().as_secs_f32() - ); - - let start = std::time::Instant::now(); - let pq: Arc = if let Some(codebook) = &pq_params.codebook { - match codebook.data_type() { - DataType::Float16 => Arc::new(ProductQuantizerImpl::::new( - pq_params.num_sub_vectors, - pq_params.num_bits as u32, - dim, - Arc::new(codebook.as_primitive().clone()), - metric_type, - )), - DataType::Float32 => Arc::new(ProductQuantizerImpl::::new( - pq_params.num_sub_vectors, - pq_params.num_bits as u32, - dim, - Arc::new(codebook.as_primitive().clone()), - metric_type, - )), - DataType::Float64 => Arc::new(ProductQuantizerImpl::::new( - pq_params.num_sub_vectors, - pq_params.num_bits as u32, - dim, - Arc::new(codebook.as_primitive().clone()), - metric_type, - )), - _ => { - return Err(Error::Index { - message: format!("Wrong codebook data type: {:?}", codebook.data_type()), - location: location!(), - }); - } - } - } else { - info!( - "Start to train PQ code: PQ{}, bits={}", - pq_params.num_sub_vectors, pq_params.num_bits - ); - let expected_sample_size = - lance_index::vector::pq::num_centroids(pq_params.num_bits as u32) - * pq_params.sample_rate; - let training_data = if let Some(training_data) = training_data { - if training_data.value_length() as usize > expected_sample_size { - training_data.sample(expected_sample_size)? - } else { - training_data - } - } else { - let start = std::time::Instant::now(); - log::info!( - "Loading training data for PQ. Sample size: {}", - expected_sample_size - ); - let data = maybe_sample_training_data(dataset, column, expected_sample_size).await?; - log::info!( - "Finished loading training data in {:02} seconds", - start.elapsed().as_secs_f32() - ); - data - }; - - // TODO: consolidate IVF models to `lance_index`. - let ivf2 = lance_index::vector::ivf::new_ivf( - ivf_model.centroids.values(), - ivf_model.dimension(), - metric_type, - vec![], - None, - )?; - - info!( - "starting to compute partitions for PQ training, sample size: {}", - training_data.value_length() - ); - // Compute the residual vector to train Product Quantizer. - // TODO: maybe use precomputed partitions here. since these are aggressively down sampled - // the time to compute them is not that bad. - let part_ids = ivf2.compute_partitions(&training_data).await?; - - let training_data = if ivf_params.use_residual { - span!(Level::INFO, "compute residual for PQ training") - .in_scope(|| ivf2.compute_residual(&training_data, Some(&part_ids))) - .await? - } else { - training_data - }; - info!("Start train PQ: params={:#?}", pq_params); - pq_params.build(&training_data, metric_type).await? + None }; - info!("Trained PQ in: {} seconds", start.elapsed().as_secs_f32()); + let pq = build_pq_model(dataset, column, dim, metric_type, pq_params, ivf_residual).await?; // Transform data, compute residuals and sort by partition ids. let mut scanner = dataset.scan(); @@ -930,7 +872,7 @@ pub async fn build_ivf_pq_index( column, index_name, uuid, - &transforms, + &[], ivf_model, pq, metric_type, @@ -1164,6 +1106,10 @@ async fn train_ivf_model( metric_type: MetricType, params: &IvfBuildParams, ) -> Result { + assert!( + metric_type != MetricType::Cosine, + "Cosine metric should be done by normalized L2 distance", + ); let values = data.values(); let dim = data.value_length() as usize; match values.data_type() { @@ -1189,6 +1135,7 @@ mod tests { use std::collections::HashMap; use std::iter::repeat; + use std::ops::Range; use arrow_array::types::UInt64Type; use arrow_array::{cast::AsArray, RecordBatchIterator, RecordBatchReader, UInt64Array}; @@ -1196,8 +1143,8 @@ mod tests { use lance_core::utils::address::RowAddress; use lance_linalg::distance::l2_distance_batch; use lance_testing::datagen::{ - generate_random_array, generate_random_array_with_seed, generate_scaled_random_array, - sample_without_replacement, + generate_random_array, generate_random_array_with_range, generate_random_array_with_seed, + generate_scaled_random_array, sample_without_replacement, }; use rand::{seq::SliceRandom, thread_rng}; use tempfile::tempdir; @@ -1416,8 +1363,11 @@ mod tests { } } - async fn generate_test_dataset(test_uri: &str) -> (Dataset, Arc) { - let vectors = generate_random_array(1000 * DIM); + async fn generate_test_dataset( + test_uri: &str, + range: Range, + ) -> (Dataset, Arc) { + let vectors = generate_random_array_with_range(1000 * DIM, range); let metadata: HashMap = vec![("test".to_string(), "ivf_pq".to_string())] .into_iter() .collect(); @@ -1446,7 +1396,7 @@ mod tests { let test_dir = tempdir().unwrap(); let test_uri = test_dir.path().to_str().unwrap(); - let (mut dataset, vector_array) = generate_test_dataset(test_uri).await; + let (mut dataset, vector_array) = generate_test_dataset(test_uri, 0.0..1.0).await; let centroids = generate_random_array(2 * DIM); let ivf_centroids = FixedSizeListArray::try_new_from_values(centroids, DIM as i32).unwrap(); @@ -1661,14 +1611,13 @@ mod tests { let test_dir = tempdir().unwrap(); let test_uri = test_dir.path().to_str().unwrap(); - let (mut dataset, vector_array) = generate_test_dataset(test_uri).await; + let (mut dataset, vector_array) = generate_test_dataset(test_uri, 0.0..1.0).await; let centroids = generate_random_array(2 * DIM); let ivf_centroids = FixedSizeListArray::try_new_from_values(centroids, DIM as i32).unwrap(); let ivf_params = IvfBuildParams::try_with_centroids(2, Arc::new(ivf_centroids)).unwrap(); - let codebook = Arc::new(generate_random_array(256 * DIM)); - let pq_params = PQBuildParams::with_codebook(4, 8, codebook); + let pq_params = PQBuildParams::new(4, 8); let params = VectorIndexParams::with_ivf_pq_params(MetricType::Cosine, ivf_params, pq_params); @@ -1694,20 +1643,83 @@ mod tests { assert_eq!(5, results[0].num_rows()); for batch in results.iter() { let dist = &batch["_distance"]; - assert!(dist - .as_primitive::() + dist.as_primitive::() .values() .iter() - .all(|v| (0.0..2.0).contains(v))); + .for_each(|v| { + assert!( + (0.0..2.0).contains(v), + "Expect cosine value in range [0.0, 2.0], got: {}", + v + ) + }); } } + #[tokio::test] + async fn test_build_ivf_model_l2() { + let test_dir = tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + + let (dataset, _) = generate_test_dataset(test_uri, 1000.0..1100.0).await; + + let ivf_params = IvfBuildParams::new(2); + let ivf_model = build_ivf_model(&dataset, "vector", DIM, MetricType::L2, &ivf_params) + .await + .unwrap(); + assert_eq!(2, ivf_model.centroids.len()); + assert_eq!(32, ivf_model.centroids.value_length()); + assert_eq!(2, ivf_model.num_partitions()); + + // All centroids values should be in the range [1000, 1100] + ivf_model + .centroids + .values() + .as_primitive::() + .values() + .iter() + .for_each(|v| { + assert!((1000.0..1100.0).contains(v)); + }); + } + + #[tokio::test] + async fn test_build_ivf_model_cosine() { + let test_dir = tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + + let (dataset, _) = generate_test_dataset(test_uri, 1000.0..1100.0).await; + + let ivf_params = IvfBuildParams::new(2); + let ivf_model = build_ivf_model(&dataset, "vector", DIM, MetricType::Cosine, &ivf_params) + .await + .unwrap(); + assert_eq!(2, ivf_model.centroids.len()); + assert_eq!(32, ivf_model.centroids.value_length()); + assert_eq!(2, ivf_model.num_partitions()); + + // All centroids values should be in the range [1000, 1100] + ivf_model + .centroids + .values() + .as_primitive::() + .values() + .iter() + .for_each(|v| { + assert!( + (-1.0..1.0).contains(v), + "Expect cosine value in range [-1.0, 1.0], got: {}", + v + ); + }); + } + #[tokio::test] async fn test_create_ivf_pq_dot() { let test_dir = tempdir().unwrap(); let test_uri = test_dir.path().to_str().unwrap(); - let (mut dataset, vector_array) = generate_test_dataset(test_uri).await; + let (mut dataset, vector_array) = generate_test_dataset(test_uri, 0.0..1.0).await; let centroids = generate_random_array(2 * DIM); let ivf_centroids = FixedSizeListArray::try_new_from_values(centroids, DIM as i32).unwrap(); @@ -1895,11 +1907,7 @@ mod tests { true, )])); - let arr = generate_random_array(1000 * DIM) - .values() - .iter() - .map(|&v| v + 1000.0) - .collect::(); + let arr = generate_random_array_with_range(1000 * DIM, 1000.0..1001.0); let fsl = FixedSizeListArray::try_new_from_values(arr.clone(), DIM as i32).unwrap(); let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(fsl)]).unwrap(); let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); @@ -1916,9 +1924,7 @@ mod tests { .await .unwrap(); let ivf_idx = idx.as_any().downcast_ref::().unwrap(); - // All centroids are normalized. - // - // If not normalized, the centroids should be on the mean of original vector space + assert!(ivf_idx .ivf .centroids @@ -1933,14 +1939,16 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); - assert!(pq_idx + + // PQ code is on residual space + pq_idx .pq .codebook_as_fsl() .values() .as_primitive::() .values() .iter() - .all(|v| (0.0..=1.0).contains(v))); + .for_each(|v| assert!((-1.0..=1.0).contains(v), "Got {}", v)); let dataset = Dataset::open(test_uri).await.unwrap(); @@ -1961,7 +1969,7 @@ mod tests { .unwrap() .as_primitive::() .value(0); - println!("Row id: {}", row_id); + println!("Row id: {} query_id: {}", row_id, query_id); if row_id == (query_id as u64) { correct_times += 1; } diff --git a/rust/lance/src/index/vector/pq.rs b/rust/lance/src/index/vector/pq.rs index 1168130d50b..1e071e5bf87 100644 --- a/rust/lance/src/index/vector/pq.rs +++ b/rust/lance/src/index/vector/pq.rs @@ -1,4 +1,4 @@ -// Copyright 2023 Lance Developers. +// Copyright 2024 Lance Developers. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -15,33 +15,37 @@ use std::sync::Arc; use std::{any::Any, collections::HashMap}; +use arrow_array::types::{Float16Type, Float32Type, Float64Type}; use arrow_array::{ cast::{as_primitive_array, AsArray}, - FixedSizeListArray, RecordBatch, UInt64Array, UInt8Array, + Array, FixedSizeListArray, RecordBatch, UInt64Array, UInt8Array, }; use arrow_ord::sort::sort_to_indices; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; use arrow_select::take::take; use async_trait::async_trait; -// Re-export -use lance_core::utils::address::RowAddress; -use lance_core::ROW_ID_FIELD; -pub use lance_index::vector::pq::{PQBuildParams, ProductQuantizerImpl}; +use lance_core::{utils::address::RowAddress, ROW_ID_FIELD}; use lance_index::{ vector::{pq::ProductQuantizer, Query, DIST_COL}, Index, IndexType, }; -use lance_io::traits::Reader; -use lance_io::utils::read_fixed_stride_array; +use lance_io::{traits::Reader, utils::read_fixed_stride_array}; use lance_linalg::distance::MetricType; +use log::info; use roaring::RoaringBitmap; use serde_json::json; use snafu::{location, Location}; -use tracing::instrument; +use tracing::{instrument, span, Level}; + +// Re-export +pub use lance_index::vector::pq::{PQBuildParams, ProductQuantizerImpl}; +use lance_linalg::kernels::normalize_fsl; +use super::ivf::Ivf; use super::VectorIndex; use crate::index::prefilter::PreFilter; -use crate::{arrow::*, utils::tokio::spawn_cpu}; +use crate::index::vector::utils::maybe_sample_training_data; +use crate::{arrow::*, utils::tokio::spawn_cpu, Dataset}; use crate::{Error, Result}; /// Product Quantization Index. @@ -263,5 +267,244 @@ impl VectorIndex for PQIndex { } } +/// Train Product Quantizer model. +/// +/// Parameters: +/// - `dataset`: The dataset to train the PQ model. +/// - `column`: The column name of the dataset. +/// - `dim`: The dimension of the vectors. +/// - `metric_type`: The metric type of the vectors. +/// - `params`: The parameters to train the PQ model. +/// - `ivf`: If provided, the IVF model to compute the residual for PQ training. +pub(super) async fn build_pq_model( + dataset: &Dataset, + column: &str, + dim: usize, + metric_type: MetricType, + params: &PQBuildParams, + ivf: Option<&Ivf>, +) -> Result> { + if let Some(codebook) = ¶ms.codebook { + let mt = if metric_type == MetricType::Cosine { + info!("Normalize training data for PQ training: Cosine"); + MetricType::L2 + } else { + metric_type + }; + + return match codebook.data_type() { + DataType::Float16 => Ok(Arc::new(ProductQuantizerImpl::::new( + params.num_sub_vectors, + params.num_bits as u32, + dim, + Arc::new(codebook.as_primitive().clone()), + mt, + ))), + DataType::Float32 => Ok(Arc::new(ProductQuantizerImpl::::new( + params.num_sub_vectors, + params.num_bits as u32, + dim, + Arc::new(codebook.as_primitive().clone()), + mt, + ))), + DataType::Float64 => Ok(Arc::new(ProductQuantizerImpl::::new( + params.num_sub_vectors, + params.num_bits as u32, + dim, + Arc::new(codebook.as_primitive().clone()), + mt, + ))), + _ => { + return Err(Error::Index { + message: format!("Wrong codebook data type: {:?}", codebook.data_type()), + location: location!(), + }); + } + }; + } + info!( + "Start to train PQ code: PQ{}, bits={}", + params.num_sub_vectors, params.num_bits + ); + let expected_sample_size = + lance_index::vector::pq::num_centroids(params.num_bits as u32) * params.sample_rate; + info!( + "Loading training data for PQ. Sample size: {}", + expected_sample_size + ); + let start = std::time::Instant::now(); + let mut training_data = + maybe_sample_training_data(dataset, column, expected_sample_size).await?; + info!( + "Finished loading training data in {:02} seconds", + start.elapsed().as_secs_f32() + ); + + info!( + "starting to compute partitions for PQ training, sample size: {}", + training_data.value_length() + ); + + if metric_type == MetricType::Cosine { + info!("Normalize training data for PQ training: Cosine"); + training_data = normalize_fsl(&training_data)?; + } + + let training_data = if let Some(ivf) = ivf { + // Compute residual for PQ training. + // + // TODO: consolidate IVF models to `lance_index`. + let ivf2 = lance_index::vector::ivf::new_ivf( + ivf.centroids.values(), + ivf.dimension(), + MetricType::L2, + vec![], + None, + )?; + span!(Level::INFO, "compute residual for PQ training") + .in_scope(|| ivf2.compute_residual(&training_data, None)) + .await? + } else { + training_data + }; + info!("Start train PQ: params={:#?}", params); + let pq = params.build(&training_data, MetricType::L2).await?; + info!("Trained PQ in: {} seconds", start.elapsed().as_secs_f32()); + Ok(pq) +} + #[cfg(test)] -mod tests {} +mod tests { + use super::*; + use crate::index::vector::ivf::build_ivf_model; + use arrow_array::RecordBatchIterator; + use arrow_schema::{Field, Schema}; + use lance_index::vector::ivf::IvfBuildParams; + use lance_testing::datagen::generate_random_array_with_range; + use std::ops::Range; + use tempfile::tempdir; + + const DIM: usize = 128; + async fn generate_dataset( + test_uri: &str, + range: Range, + ) -> (Dataset, Arc) { + let vectors = generate_random_array_with_range(1000 * DIM, range); + let metadata: HashMap = vec![("test".to_string(), "ivf_pq".to_string())] + .into_iter() + .collect(); + + let schema = Arc::new( + Schema::new(vec![Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIM as i32, + ), + true, + )]) + .with_metadata(metadata), + ); + let fsl = Arc::new(FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap()); + let batch = RecordBatch::try_new(schema.clone(), vec![fsl.clone()]).unwrap(); + + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); + (Dataset::write(batches, test_uri, None).await.unwrap(), fsl) + } + + #[tokio::test] + async fn test_build_pq_model_l2() { + let test_dir = tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + + let (dataset, _) = generate_dataset(test_uri, 100.0..120.0).await; + + let centroids = generate_random_array_with_range(4 * DIM, -1.0..1.0); + let fsl = FixedSizeListArray::try_new_from_values(centroids, DIM as i32).unwrap(); + let ivf = Ivf::new(fsl.into()); + let params = PQBuildParams::new(16, 8); + let pq = build_pq_model(&dataset, "vector", DIM, MetricType::L2, ¶ms, Some(&ivf)) + .await + .unwrap(); + + assert_eq!(pq.num_sub_vectors(), 16); + assert_eq!(pq.num_bits(), 8); + assert_eq!(pq.dimension(), DIM); + + let codebook = pq.codebook_as_fsl(); + assert_eq!(codebook.len(), 256); + codebook + .values() + .as_primitive::() + .values() + .iter() + .for_each(|v| { + assert!((99.0..121.0).contains(v)); + }); + } + + #[tokio::test] + async fn test_build_pq_model_cosine() { + let test_dir = tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + + let (dataset, vectors) = generate_dataset(test_uri, 100.0..120.0).await; + + let ivf_params = IvfBuildParams::new(4); + let ivf = build_ivf_model(&dataset, "vector", DIM, MetricType::Cosine, &ivf_params) + .await + .unwrap(); + let params = PQBuildParams::new(16, 8); + let pq = build_pq_model( + &dataset, + "vector", + DIM, + MetricType::Cosine, + ¶ms, + Some(&ivf), + ) + .await + .unwrap(); + + assert_eq!(pq.num_sub_vectors(), 16); + assert_eq!(pq.num_bits(), 8); + assert_eq!(pq.dimension(), DIM); + + let codebook = pq.codebook_as_fsl(); + assert_eq!(codebook.len(), 256); + codebook + .values() + .as_primitive::() + .values() + .iter() + .for_each(|v| { + assert!((-1.0..1.0).contains(v)); + }); + + let vectors = normalize_fsl(&vectors).unwrap(); + let row = vectors.slice(0, 1); + + let ivf2 = lance_index::vector::ivf::new_ivf( + ivf.centroids.values(), + ivf.dimension(), + MetricType::L2, + vec![], + None, + ) + .unwrap(); + + let residual_query = ivf2.compute_residual(&row, None).await.unwrap(); + let pq_code = pq.transform(&residual_query).await.unwrap(); + let distances = pq + .compute_distances( + &residual_query.value(0), + pq_code.as_fixed_size_list().values().as_primitive(), + ) + .unwrap(); + assert!( + distances.values().iter().all(|&d| d <= 0.001), + "distances: {:?}", + distances + ); + } +}