fix: filter out infinite vectors (#3648)

when training IVF/PQ, we sample the dataset and filter out nulls, but we
still need to filter out NaN/INFs

---------

Signed-off-by: BubbleCal <bubble-cal@outlook.com>

Author: BubbleCal

rust/lance-index/src/vector/ivf.rs

@@ -118,10 +118,8 @@ impl IvfTransformer {
         vector_column: &str,
         range: Option<Range<u32>>,
     ) -> Self {
-        let mut transforms: Vec<Arc<dyn Transformer>> = vec![
-            Arc::new(KeepFiniteVectors::new(vector_column)),
-            Arc::new(super::transform::Flatten::new(vector_column)),
-        ];
+        let mut transforms: Vec<Arc<dyn Transformer>> =
+            vec![Arc::new(super::transform::Flatten::new(vector_column))];
 
         let dt = if distance_type == DistanceType::Cosine {
             transforms.push(Arc::new(super::transform::NormalizeTransformer::new(
@@ -131,6 +129,7 @@ impl IvfTransformer {
         } else {
             distance_type
         };
+        transforms.push(Arc::new(KeepFiniteVectors::new(vector_column)));
 
         let ivf_transform = Arc::new(PartitionTransformer::new(
             centroids.clone(),
@@ -159,10 +158,8 @@ impl IvfTransformer {
         pq: ProductQuantizer,
         range: Option<Range<u32>>,
     ) -> Self {
-        let mut transforms: Vec<Arc<dyn Transformer>> = vec![
-            Arc::new(KeepFiniteVectors::new(vector_column)),
-            Arc::new(super::transform::Flatten::new(vector_column)),
-        ];
+        let mut transforms: Vec<Arc<dyn Transformer>> =
+            vec![Arc::new(super::transform::Flatten::new(vector_column))];
 
         let distance_type = if distance_type == MetricType::Cosine {
             transforms.push(Arc::new(super::transform::NormalizeTransformer::new(
@@ -172,6 +169,7 @@ impl IvfTransformer {
         } else {
             distance_type
         };
+        transforms.push(Arc::new(KeepFiniteVectors::new(vector_column)));
 
         let partition_transform = Arc::new(PartitionTransformer::new(
             centroids.clone(),
@@ -210,10 +208,8 @@ impl IvfTransformer {
         sq: ScalarQuantizer,
         range: Option<Range<u32>>,
     ) -> Self {
-        let mut transforms: Vec<Arc<dyn Transformer>> = vec![
-            Arc::new(KeepFiniteVectors::new(vector_column)),
-            Arc::new(super::transform::Flatten::new(vector_column)),
-        ];
+        let mut transforms: Vec<Arc<dyn Transformer>> =
+            vec![Arc::new(super::transform::Flatten::new(vector_column))];
 
         let distance_type = if metric_type == MetricType::Cosine {
             transforms.push(Arc::new(super::transform::NormalizeTransformer::new(
@@ -223,6 +219,7 @@ impl IvfTransformer {
         } else {
             metric_type
         };
+        transforms.push(Arc::new(KeepFiniteVectors::new(vector_column)));
 
         let partition_transformer = Arc::new(PartitionTransformer::new(
             centroids.clone(),

rust/lance-index/src/vector/utils.rs

@@ -5,7 +5,7 @@ use arrow::{
     array::AsArray,
     datatypes::{Float16Type, Float32Type, Float64Type},
 };
-use arrow_array::{Array, FixedSizeListArray};
+use arrow_array::{Array, BooleanArray, FixedSizeListArray};
 use arrow_schema::{DataType, Field};
 use lance_core::{Error, Result};
 use lance_io::encodings::plain::bytes_to_array;
@@ -164,6 +164,36 @@ impl TryFrom<&pb::Tensor> for FixedSizeListArray {
     }
 }
 
+/// Check if all vectors in the FixedSizeListArray are finite
+/// null values are considered as not finite
+/// returns a BooleanArray
+/// with the same length as the FixedSizeListArray
+/// with true for finite values and false for non-finite values
+pub fn is_finite(fsl: &FixedSizeListArray) -> BooleanArray {
+    let is_finite = fsl
+        .iter()
+        .map(|v| match v {
+            Some(v) => match v.data_type() {
+                DataType::Float16 => {
+                    let v = v.as_primitive::<Float16Type>();
+                    v.null_count() == 0 && v.values().iter().all(|v| v.is_finite())
+                }
+                DataType::Float32 => {
+                    let v = v.as_primitive::<Float32Type>();
+                    v.null_count() == 0 && v.values().iter().all(|v| v.is_finite())
+                }
+                DataType::Float64 => {
+                    let v = v.as_primitive::<Float64Type>();
+                    v.null_count() == 0 && v.values().iter().all(|v| v.is_finite())
+                }
+                _ => v.null_count() == 0,
+            },
+            None => false,
+        })
+        .collect::<Vec<_>>();
+    BooleanArray::from(is_finite)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

rust/lance/src/index/vector/builder.rs

@@ -24,6 +24,7 @@ use lance_index::vector::quantizer::{
     QuantizationMetadata, QuantizationType, QuantizerBuildParams,
 };
 use lance_index::vector::storage::STORAGE_METADATA_KEY;
+use lance_index::vector::utils::is_finite;
 use lance_index::vector::v3::shuffler::IvfShufflerReader;
 use lance_index::vector::v3::subindex::SubIndexType;
 use lance_index::vector::{VectorIndex, LOSS_METADATA_KEY, PART_ID_COLUMN, PQ_CODE_COLUMN};
@@ -370,6 +371,10 @@ impl<S: IvfSubIndex + 'static, Q: Quantization + 'static> IvfIndexBuilder<S, Q>
             training_data
         };
 
+        // we filtered out nulls when sampling, but we still need to filter out NaNs and INFs here
+        let training_data = arrow::compute::filter(&training_data, &is_finite(&training_data))?;
+        let training_data = training_data.as_fixed_size_list();
+
         let training_data = match (self.ivf.as_ref(), Q::use_residual(self.distance_type)) {
             (Some(ivf), true) => {
                 let ivf_transformer = lance_index::vector::ivf::new_ivf_transformer(
@@ -378,9 +383,9 @@ impl<S: IvfSubIndex + 'static, Q: Quantization + 'static> IvfIndexBuilder<S, Q>
                     vec![],
                 );
                 span!(Level::INFO, "compute residual for PQ training")
-                    .in_scope(|| ivf_transformer.compute_residual(&training_data))?
+                    .in_scope(|| ivf_transformer.compute_residual(training_data))?
             }
-            _ => training_data,
+            _ => training_data.clone(),
         };
 
         info!("Start to train quantizer");

rust/lance/src/index/vector/ivf.rs

@@ -46,6 +46,7 @@ use lance_index::vector::flat::index::{FlatBinQuantizer, FlatIndex, FlatQuantize
 use lance_index::vector::ivf::storage::IvfModel;
 use lance_index::vector::pq::storage::transpose;
 use lance_index::vector::quantizer::QuantizationType;
+use lance_index::vector::utils::is_finite;
 use lance_index::vector::v3::shuffler::IvfShuffler;
 use lance_index::vector::v3::subindex::{IvfSubIndex, SubIndexType};
 use lance_index::{
@@ -1222,9 +1223,13 @@ pub async fn build_ivf_model(
         (training_data, metric_type)
     };
 
+    // we filtered out nulls when sampling, but we still need to filter out NaNs and INFs here
+    let training_data = arrow::compute::filter(&training_data, &is_finite(&training_data))?;
+    let training_data = training_data.as_fixed_size_list();
+
     info!("Start to train IVF model");
     let start = std::time::Instant::now();
-    let ivf = train_ivf_model(centroids, &training_data, mt, params).await?;
+    let ivf = train_ivf_model(centroids, training_data, mt, params).await?;
     info!(
         "Trained IVF model in {:02} seconds",
         start.elapsed().as_secs_f32()

rust/lance/src/index/vector/ivf/v2.rs

@@ -642,8 +642,8 @@ mod tests {
     use arrow::datatypes::{UInt64Type, UInt8Type};
     use arrow::{array::AsArray, datatypes::Float32Type};
     use arrow_array::{
-        Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, FixedSizeListArray, ListArray,
-        RecordBatch, RecordBatchIterator, UInt64Array,
+        Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, FixedSizeListArray, Float32Array,
+        ListArray, RecordBatch, RecordBatchIterator, UInt64Array,
     };
     use arrow_buffer::OffsetBuffer;
     use arrow_schema::{DataType, Field, Schema, SchemaRef};
@@ -1700,4 +1700,44 @@ mod tests {
             assert_lt!(*d, dists[k - 1]);
         });
     }
+
+    #[tokio::test]
+    async fn test_index_with_zero_vectors() {
+        let test_dir = tempdir().unwrap();
+        let test_uri = test_dir.path().to_str().unwrap();
+        let (batch, schema) = generate_batch::<Float32Type>(256, None, 0.0..1.0, false);
+        let vector_field = schema.field(1).clone();
+        let zero_batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![
+                Arc::new(UInt64Array::from(vec![256])),
+                Arc::new(
+                    FixedSizeListArray::try_new_from_values(
+                        Float32Array::from(vec![0.0; DIM]),
+                        DIM as i32,
+                    )
+                    .unwrap(),
+                ),
+            ],
+        )
+        .unwrap();
+        let batches = RecordBatchIterator::new(vec![batch, zero_batch].into_iter().map(Ok), schema);
+        let mut dataset = Dataset::write(
+            batches,
+            test_uri,
+            Some(WriteParams {
+                mode: crate::dataset::WriteMode::Overwrite,
+                ..Default::default()
+            }),
+        )
+        .await
+        .unwrap();
+
+        let vector_column = vector_field.name();
+        let params = VectorIndexParams::ivf_pq(4, 8, DIM / 8, DistanceType::Cosine, 50);
+        dataset
+            .create_index(&[vector_column], IndexType::Vector, None, &params, true)
+            .await
+            .unwrap();
+    }
 }