Hotfix: prevent optimizer infinite loop with deferred points and multi vectors (#9285)

* Always optimize deferred points

* Add test (#9288)

* Hotfix: test indexing of deferred multivector under indeixing threshold (#9286)

---------

Co-authored-by: Arnaud Gourlay <arnaud.gourlay@gmail.com>

lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs

@@ -7,22 +7,30 @@ pub use shard::optimizers::indexing_optimizer::IndexingOptimizer;
 
 #[cfg(test)]
 mod tests {
-    use std::collections::BTreeMap;
+    use std::collections::{BTreeMap, HashMap};
+    use std::num::{NonZeroU64, NonZeroUsize};
     use std::path::PathBuf;
 
     use common::counter::hardware_counter::HardwareCounterCell;
     use fs_err as fs;
     use itertools::Itertools;
     use rand::rng;
-    use segment::data_types::vectors::DEFAULT_VECTOR_NAME;
+    use segment::data_types::named_vectors::NamedVectors;
+    use segment::data_types::vectors::{
+        DEFAULT_VECTOR_NAME, MultiDenseVectorInternal, VectorInternal,
+    };
     use segment::entry::ReadSegmentEntry;
+    use segment::entry::entry_point::SegmentEntry;
     use segment::fixtures::index_fixtures::random_vector;
+    use segment::fixtures::payload_fixtures::random_multi_vector;
     use segment::json_path::JsonPath;
     use segment::payload_json;
+    use segment::segment_constructor::build_segment;
     use segment::segment_constructor::simple_segment_constructor::{VECTOR1_NAME, VECTOR2_NAME};
     use segment::types::{
-        Distance, HnswConfig, HnswGlobalConfig, PayloadSchemaType, QuantizationConfig, SegmentType,
-        VectorNameBuf,
+        Distance, HnswConfig, HnswGlobalConfig, Indexes, MultiVectorComparator, MultiVectorConfig,
+        PayloadSchemaType, QuantizationConfig, SegmentConfig, SegmentType, VectorDataConfig,
+        VectorNameBuf, VectorStorageType,
     };
     use shard::operations::optimization::OptimizerThresholds;
     use shard::optimizers::segment_optimizer::SegmentOptimizer;
@@ -194,6 +202,192 @@ mod tests {
         }
     }
 
+    /// A multivector's deferred-point threshold is computed assuming a fixed inner-vector
+    /// count (`MULTIVECTOR_SIZE = 16`, see `CollectionParams::get_deferred_point_id`). This
+    /// makes points become "deferred" at a far smaller storage size than the actual data
+    /// occupies, so a segment can hold deferred points while staying well below the indexing
+    /// threshold.
+    ///
+    /// Before the fix, the indexing optimizer rebuilt such a segment as a plain (non-HNSW)
+    /// segment because it was below the indexing threshold. Plain segments don't promote
+    /// deferred points, so `has_deferred_points()` stayed `true` and the optimizer kept
+    /// re-selecting the segment forever (infinite loop).
+    ///
+    /// This test proves the fix: a below-threshold segment with deferred points is optimized
+    /// into an HNSW-indexed segment (which promotes the deferred points) and is no longer
+    /// selected for optimization afterwards.
+    #[test]
+    fn test_deferred_points_multivector_optimization() {
+        init();
+
+        // Multivector named vector. With float32 elements the per-point size used to derive
+        // the deferred-point threshold is `ELEMENT_BYTES * DIM * MULTIVECTOR_SIZE`.
+        const VECTOR_NAME: &str = "vector";
+        const DIM: usize = 16;
+        const MULTIVECTOR_SIZE: usize = 16; // mirrors CollectionParams::get_deferred_point_id
+        const ELEMENT_BYTES: usize = 4; // float32
+
+        // Deferred-point byte threshold, sized so points start deferring at internal offset 100.
+        let deferred_threshold_bytes =
+            NonZeroUsize::new(ELEMENT_BYTES * DIM * MULTIVECTOR_SIZE * 100).unwrap(); // 102_400
+
+        let collection_params = CollectionParams {
+            vectors: VectorsConfig::Multi(BTreeMap::from([(
+                VECTOR_NAME.to_owned(),
+                VectorParams {
+                    size: NonZeroU64::new(DIM as u64).unwrap(),
+                    distance: Distance::Dot,
+                    hnsw_config: None,
+                    quantization_config: None,
+                    on_disk: None,
+                    datatype: None,
+                    multivector_config: Some(MultiVectorConfig::default()),
+                },
+            )])),
+            ..CollectionParams::empty()
+        };
+
+        let hnsw_config = HnswConfig::default();
+
+        // Deferred-point offset, derived exactly as production does it. Because the multivector
+        // assumes 16 inner vectors per point, the threshold of 102_400 bytes maps to only 100
+        // points (102_400 / (4 * 16 * 16)), even though each point actually stores far less.
+        let deferred_internal_id =
+            collection_params.get_deferred_point_id(&hnsw_config, Some(deferred_threshold_bytes));
+        assert_eq!(
+            deferred_internal_id,
+            Some(100),
+            "points should start deferring at internal offset 100",
+        );
+
+        // Build a multivector segment holding deferred points: insert more points than the
+        // deferred offset, each with a single inner vector so the actual storage size stays
+        // tiny (DIM * ELEMENT_BYTES bytes/point) - far below the indexing threshold.
+        let segments_dir = Builder::new().prefix("segments_dir").tempdir().unwrap();
+        let segments_temp_dir = Builder::new()
+            .prefix("segments_temp_dir")
+            .tempdir()
+            .unwrap();
+
+        const NUM_POINTS: u64 = 120;
+        let segment_config = SegmentConfig {
+            vector_data: HashMap::from([(
+                VECTOR_NAME.to_owned(),
+                VectorDataConfig {
+                    size: DIM,
+                    distance: Distance::Dot,
+                    storage_type: VectorStorageType::default(),
+                    index: Indexes::Plain {},
+                    quantization_config: None,
+                    multivector_config: Some(MultiVectorConfig::default()),
+                    datatype: None,
+                },
+            )]),
+            sparse_vector_data: Default::default(),
+            payload_storage_type: Default::default(),
+        };
+
+        let mut segment = build_segment(
+            segments_dir.path(),
+            &segment_config,
+            deferred_internal_id,
+            true,
+        )
+        .unwrap();
+
+        let mut rnd = rng();
+        let hw_counter = HardwareCounterCell::new();
+        for n in 0..NUM_POINTS {
+            let multi_vec = random_multi_vector(&mut rnd, DIM, 1);
+            let mut named = NamedVectors::default();
+            named.insert(
+                VECTOR_NAME.to_owned(),
+                VectorInternal::MultiDense(multi_vec),
+            );
+            segment
+                .upsert_point(n, n.into(), named, &hw_counter)
+                .unwrap();
+        }
+
+        // The segment holds deferred points, yet its vectors stay below the indexing threshold.
+        assert!(
+            segment.has_deferred_points(),
+            "segment should hold deferred points",
+        );
+        let vectors_size_bytes = segment
+            .available_vectors_size_in_bytes(VECTOR_NAME)
+            .unwrap();
+
+        let mut holder = SegmentHolder::default();
+        let segment_id = holder.add_new(segment);
+        let locked_holder = LockedSegmentHolder::new(holder);
+
+        // Indexing threshold set BELOW the deferred byte threshold but ABOVE the actual segment
+        // size, so the segment does NOT exceed the indexing threshold by size.
+        let indexing_threshold_kb = 50; // 51_200 bytes
+        assert!(
+            vectors_size_bytes < indexing_threshold_kb * 1024,
+            "segment ({vectors_size_bytes} bytes) must stay below the indexing threshold",
+        );
+        assert!(
+            indexing_threshold_kb * 1024 < deferred_threshold_bytes.get(),
+            "indexing threshold must be under the deferred-point threshold",
+        );
+
+        let index_optimizer = new_indexing_optimizer(
+            1,
+            OptimizerThresholds {
+                max_segment_size_kb: 1000,
+                memmap_threshold_kb: 1000,
+                indexing_threshold_kb,
+                deferred_internal_id,
+            },
+            segments_dir.path().to_owned(),
+            segments_temp_dir.path().to_owned(),
+            collection_params,
+            hnsw_config,
+            HnswGlobalConfig::default(),
+            None,
+        );
+
+        // The segment is selected for optimization solely because it has deferred points.
+        let suggested_to_optimize = index_optimizer.plan_optimizations_for_test(&locked_holder);
+        let suggested_to_optimize = suggested_to_optimize.into_iter().exactly_one().unwrap();
+        assert!(suggested_to_optimize.contains(&segment_id));
+
+        index_optimizer.optimize_for_test(locked_holder.clone(), suggested_to_optimize);
+
+        // The fix: the optimized segment is HNSW-indexed even though it was below the indexing
+        // threshold, which promotes the deferred points.
+        let infos = locked_holder
+            .read()
+            .iter()
+            .map(|(_sid, segment)| segment.get().read().info())
+            .collect_vec();
+        assert!(
+            infos
+                .iter()
+                .any(|info| info.segment_type == SegmentType::Indexed),
+            "optimized segment must be HNSW-indexed to promote deferred points",
+        );
+
+        // No segment holds deferred points anymore, so the optimizer no longer loops on it.
+        let still_has_deferred = locked_holder
+            .read()
+            .iter()
+            .any(|(_sid, segment)| segment.get().read().has_deferred_points());
+        assert!(
+            !still_has_deferred,
+            "deferred points must be promoted after optimization",
+        );
+
+        let suggested_after = index_optimizer.plan_optimizations_for_test(&locked_holder);
+        assert!(
+            suggested_after.is_empty(),
+            "no further optimization should be required (no infinite loop)",
+        );
+    }
+
     #[test]
     fn test_indexing_optimizer() {
         init();
@@ -747,4 +941,171 @@ mod tests {
                 );
             });
     }
+
+    /// Multi vectors with deferred points below the indexing threshold must not cause an
+    /// infinite optimization loop.
+    ///
+    /// Deferred points are tracked with a static internal-id offset (`deferred_internal_id`).
+    /// For multi vectors this offset is computed assuming a fixed number of sub vectors per
+    /// point (see [`CollectionParams::get_deferred_point_id`] and its `MULTIVECTOR_SIZE`
+    /// constant). When a user uploads small multi vectors (e.g. a single sub vector per
+    /// point), the real storage size stays well below the indexing threshold, yet the segment
+    /// still ends up with deferred points because the offset was sized for much larger points.
+    ///
+    /// The indexing optimizer always triggers for segments that have deferred points, so it
+    /// can promote them quickly by building an HNSW index. Before commit "Always optimize
+    /// deferred points", the optimizer skipped building the HNSW index while the segment was
+    /// still below the indexing threshold. Because building the index is what promotes
+    /// deferred points, they stayed deferred, the optimizer kept being triggered, and the
+    /// same segment was re-optimized forever -- an infinite loop.
+    ///
+    /// This test reproduces that scenario. Before the fix it loops forever (and hits the
+    /// iteration limit, failing the test); after the fix the optimizer builds an HNSW index,
+    /// promotes the deferred points, and terminates.
+    #[test]
+    fn test_deferred_multivector_below_indexing_threshold_no_infinite_loop() {
+        init();
+
+        let dim = 4;
+        let indexing_threshold_kb = 10;
+        let indexing_threshold_bytes = indexing_threshold_kb * 1024;
+
+        let segments_dir = Builder::new().prefix("segments_dir").tempdir().unwrap();
+        let segments_temp_dir = Builder::new()
+            .prefix("segments_temp_dir")
+            .tempdir()
+            .unwrap();
+
+        // A single multi vector (one vector configured as a multivector).
+        let mut multi_params = VectorParamsBuilder::new(dim as u64, Distance::Dot).build();
+        multi_params.multivector_config = Some(MultiVectorConfig {
+            comparator: MultiVectorComparator::MaxSim,
+        });
+        let collection_params = CollectionParams {
+            vectors: VectorsConfig::Single(multi_params),
+            ..CollectionParams::empty()
+        };
+
+        let hnsw_config = HnswConfig::default();
+
+        // The deferred offset, computed exactly like production does. For multi vectors this
+        // assumes a fixed (large) number of sub vectors per point, so the offset is small
+        // relative to the number of points even when each point holds a single sub vector.
+        let deferred_internal_id = collection_params
+            .get_deferred_point_id(&hnsw_config, NonZeroUsize::new(indexing_threshold_bytes))
+            .expect("a multi vector with HNSW indexing should produce a deferred offset");
+
+        // Upload more points than the deferred offset (so we get deferred points), while each
+        // point stores only a single sub vector so the real storage size stays well below the
+        // indexing threshold.
+        let num_points = u64::from(deferred_internal_id) * 2;
+        assert!(
+            num_points > u64::from(deferred_internal_id),
+            "test must upload more points than the deferred offset to create deferred points",
+        );
+
+        // Build a plain, appendable segment for this collection with the deferred offset set,
+        // mimicking a live appendable segment that has accumulated deferred points.
+        let segment_optimizer_config =
+            build_segment_optimizer_config(&collection_params, &hnsw_config, &None);
+        let segment_config = segment_optimizer_config.plain_segment_config();
+
+        let hw_counter = HardwareCounterCell::new();
+        let mut segment = build_segment(
+            segments_dir.path(),
+            &segment_config,
+            Some(deferred_internal_id),
+            true,
+        )
+        .unwrap();
+
+        for i in 0..num_points {
+            let mut vectors = NamedVectors::default();
+            vectors.insert(
+                DEFAULT_VECTOR_NAME.into(),
+                VectorInternal::MultiDense(
+                    MultiDenseVectorInternal::try_from_matrix(vec![vec![0.5; dim]]).unwrap(),
+                ),
+            );
+            segment
+                .upsert_point(100, (i + 1).into(), vectors, &hw_counter)
+                .unwrap();
+        }
+
+        // Sanity: the segment has deferred points but stays below the indexing threshold.
+        assert!(
+            segment.has_deferred_points(),
+            "segment should have deferred points",
+        );
+        let stored_bytes = segment
+            .available_vectors_size_in_bytes(DEFAULT_VECTOR_NAME)
+            .unwrap();
+        assert!(
+            stored_bytes < indexing_threshold_bytes,
+            "segment must stay below the indexing threshold \
+             ({stored_bytes} >= {indexing_threshold_bytes})",
+        );
+
+        let mut holder = SegmentHolder::default();
+        holder.add_new(segment);
+        let locked_holder = LockedSegmentHolder::new(holder);
+
+        let index_optimizer = new_indexing_optimizer(
+            2,
+            OptimizerThresholds {
+                max_segment_size_kb: 1000,
+                memmap_threshold_kb: 1_000_000, // never put on disk
+                indexing_threshold_kb,          // real size stays below this
+                deferred_internal_id: Some(deferred_internal_id),
+            },
+            segments_dir.path().to_owned(),
+            segments_temp_dir.path().to_owned(),
+            collection_params.clone(),
+            hnsw_config,
+            HnswGlobalConfig::default(),
+            Default::default(),
+        );
+
+        // Drive the optimizer until there is nothing left to do. Before the fix the optimizer
+        // re-optimizes the same segment forever because the deferred points are never
+        // promoted; the iteration limit turns that infinite loop into a test failure.
+        const MAX_ITERATIONS: usize = 16;
+        let mut iterations = 0;
+        loop {
+            let suggested_to_optimize = index_optimizer.plan_optimizations_for_test(&locked_holder);
+            if suggested_to_optimize.is_empty() {
+                break;
+            }
+            assert!(
+                iterations < MAX_ITERATIONS,
+                "indexing optimizer is stuck in an infinite loop: deferred points below the \
+                 indexing threshold are never promoted",
+            );
+            let batch = suggested_to_optimize.into_iter().next().unwrap();
+            index_optimizer.optimize_for_test(locked_holder.clone(), batch);
+            iterations += 1;
+        }
+
+        // The deferred points must have been promoted: no segment has deferred points left,
+        // and an HNSW index was built to make them searchable.
+        let holder = locked_holder.read();
+        let any_deferred = holder
+            .iter()
+            .any(|(_, segment)| segment.get().read().has_deferred_points());
+        assert!(!any_deferred, "deferred points should have been promoted");
+
+        let any_hnsw = holder.iter().any(|(_, segment)| {
+            segment
+                .get()
+                .read()
+                .config()
+                .vector_data
+                .values()
+                .any(|vector| matches!(vector.index, Indexes::Hnsw(_)))
+        });
+        assert!(
+            any_hnsw,
+            "an HNSW index should have been created to promote the deferred points",
+        );
+    }
 }

lib/shard/src/optimizers/segment_optimizer.rs

@@ -169,6 +169,8 @@ pub trait SegmentOptimizer: Sync {
         // }
         let mut bytes_count_by_vector_name = HashMap::new();
 
+        let mut any_has_deferred = false;
+
         for segment in optimizing_segments {
             let segment = match segment {
                 LockedSegment::Original(segment) => segment,
@@ -180,6 +182,8 @@ pub trait SegmentOptimizer: Sync {
             };
             let locked_segment = segment.read();
 
+            any_has_deferred |= locked_segment.has_deferred_points();
+
             for vector_name in locked_segment.vector_names() {
                 let vector_size = locked_segment.available_vectors_size_in_bytes(&vector_name)?;
                 let size = bytes_count_by_vector_name.entry(vector_name).or_insert(0);
@@ -207,7 +211,13 @@ pub trait SegmentOptimizer: Sync {
         let mut sparse_vector_data = segment_optimizer_config.plain_sparse_vector_config.clone();
 
         // If indexing, change to HNSW index and quantization
-        if threshold_is_indexed {
+        // We must always create an HNSW index if we have deferred points to be able to promote them
+        if threshold_is_indexed || any_has_deferred {
+            if !threshold_is_indexed {
+                log::info!(
+                    "Segment has deferred points, but doesn't exceed indexing threshold. It will be optimized with HNSW index and quantization."
+                );
+            }
             vector_data.iter_mut().for_each(|(vector_name, config)| {
                 if let Some(vector_cfg) = segment_optimizer_config.dense_vector.get(vector_name) {
                     // Assign HNSW index