qdrant/lib/quantization/tests/integration/stop_condition.rs

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::sync::atomic::{AtomicBool, Ordering};

    use quantization::encoded_storage::TestEncodedStorageBuilder;
    use quantization::encoded_vectors::{DistanceType, VectorParameters};
    use quantization::encoded_vectors_u8::{self, EncodedVectorsU8, ScalarQuantizationMethod};
    use quantization::{EncodedVectorsPQ, EncodingError, encoded_vectors_pq};

    #[test]
    fn stop_condition_u8() {
        let stopped = Arc::new(AtomicBool::new(false));
        let stopped_clone = stopped.clone();
        let stopped_ref = stopped.as_ref();

        let stop_thread = std::thread::spawn(move || {
            std::thread::sleep(std::time::Duration::from_millis(100));
            stopped_clone.store(true, Ordering::Relaxed);
        });

        let vectors_count = 1_000_000;
        let vector_dim = 8;
        let vector_parameters = VectorParameters {
            dim: vector_dim,
            deprecated_count: None,
            distance_type: DistanceType::Dot,
            invert: false,
        };
        let zero_vector = vec![0.0; vector_dim];

        let quantized_vector_size =
            encoded_vectors_u8::get_quantized_vector_size(&vector_parameters);
        assert_eq!(
            EncodedVectorsU8::encode(
                (0..vectors_count).map(|_| &zero_vector),
                TestEncodedStorageBuilder::new(None, quantized_vector_size),
                &vector_parameters,
                vectors_count,
                None,
                ScalarQuantizationMethod::Int8,
                None,
                stopped_ref,
            )
            .err(),
            Some(EncodingError::Stopped)
        );

        stop_thread.join().unwrap();
    }

    #[test]
    fn stop_condition_pq() {
        let stopped = Arc::new(AtomicBool::new(false));
        let stopped_clone = stopped.clone();
        let stopped_ref = stopped.as_ref();

        let stop_thread = std::thread::spawn(move || {
            std::thread::sleep(std::time::Duration::from_millis(300));
            stopped_clone.store(true, Ordering::Relaxed);
        });

        let vectors_count = 1_000_000;
        let vector_dim = 8;
        let vector_parameters = VectorParameters {
            dim: vector_dim,
            deprecated_count: None,
            distance_type: DistanceType::Dot,
            invert: false,
        };
        let zero_vector = vec![0.0; vector_dim];

        let quantized_vector_size =
            encoded_vectors_pq::get_quantized_vector_size(&vector_parameters, 2);
        assert_eq!(
            EncodedVectorsPQ::encode(
                (0..vectors_count).map(|_| &zero_vector),
                TestEncodedStorageBuilder::new(None, quantized_vector_size),
                &vector_parameters,
                vectors_count,
                2,
                1,
                None,
                stopped_ref,
            )
            .err(),
            Some(EncodingError::Stopped)
        );

        stop_thread.join().unwrap();
    }
}