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

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicBool, AtomicUsize};
    use std::time::Duration;

    use common::counter::hardware_counter::HardwareCounterCell;
    use quantization::encoded_storage::TestEncodedStorageBuilder;
    use quantization::encoded_vectors::{DistanceType, EncodedVectors, VectorParameters};
    use quantization::encoded_vectors_pq::{self, EncodedVectorsPQ};
    use rand::{RngExt, SeedableRng};

    use crate::metrics::{dot_similarity, l1_similarity, l2_similarity};

    const VECTORS_COUNT: usize = 513;
    const VECTOR_DIM: usize = 65;
    const ERROR: f32 = VECTOR_DIM as f32 * 0.05;

    #[test]
    fn test_pq_dot() {
        let mut rng = rand::rngs::StdRng::seed_from_u64(42);
        let mut vector_data: Vec<Vec<_>> = vec![];
        for _ in 0..VECTORS_COUNT {
            vector_data.push((0..VECTOR_DIM).map(|_| rng.random()).collect());
        }
        let query: Vec<_> = (0..VECTOR_DIM).map(|_| rng.random()).collect();

        let vector_parameters = VectorParameters {
            dim: VECTOR_DIM,
            deprecated_count: None,
            distance_type: DistanceType::Dot,
            invert: false,
        };
        let quantized_vector_size =
            encoded_vectors_pq::get_quantized_vector_size(&vector_parameters, 1);
        let encoded = EncodedVectorsPQ::encode(
            vector_data.iter(),
            TestEncodedStorageBuilder::new(None, quantized_vector_size),
            &vector_parameters,
            VECTORS_COUNT,
            1,
            1,
            None,
            &AtomicBool::new(false),
        )
        .unwrap();
        let query_u8 = encoded.encode_query(&query);

        let counter = HardwareCounterCell::new();
        for (index, vector) in vector_data.iter().enumerate() {
            let score = encoded.score_point(&query_u8, index as u32, &counter);
            let orginal_score = dot_similarity(&query, vector);
            assert!((score - orginal_score).abs() < ERROR);
        }
    }

    #[test]
    fn test_pq_l2() {
        let mut rng = rand::rngs::StdRng::seed_from_u64(42);
        let mut vector_data: Vec<Vec<_>> = vec![];
        for _ in 0..VECTORS_COUNT {
            vector_data.push((0..VECTOR_DIM).map(|_| rng.random()).collect());
        }
        let query: Vec<_> = (0..VECTOR_DIM).map(|_| rng.random()).collect();

        let vector_parameters = VectorParameters {
            dim: VECTOR_DIM,
            deprecated_count: None,
            distance_type: DistanceType::L2,
            invert: false,
        };
        let quantized_vector_size =
            encoded_vectors_pq::get_quantized_vector_size(&vector_parameters, 1);
        let encoded = EncodedVectorsPQ::encode(
            vector_data.iter(),
            TestEncodedStorageBuilder::new(None, quantized_vector_size),
            &vector_parameters,
            VECTORS_COUNT,
            1,
            1,
            None,
            &AtomicBool::new(false),
        )
        .unwrap();
        let query_u8 = encoded.encode_query(&query);

        let counter = HardwareCounterCell::new();
        for (index, vector) in vector_data.iter().enumerate() {
            let score = encoded.score_point(&query_u8, index as u32, &counter);
            let orginal_score = l2_similarity(&query, vector);
            assert!((score - orginal_score).abs() < ERROR);
        }
    }

    #[test]
    fn test_pq_l1() {
        let mut rng = rand::rngs::StdRng::seed_from_u64(42);
        let mut vector_data: Vec<Vec<_>> = vec![];
        for _ in 0..VECTORS_COUNT {
            vector_data.push((0..VECTOR_DIM).map(|_| rng.random()).collect());
        }
        let query: Vec<_> = (0..VECTOR_DIM).map(|_| rng.random()).collect();

        let vector_parameters = VectorParameters {
            dim: VECTOR_DIM,
            deprecated_count: None,
            distance_type: DistanceType::L1,
            invert: false,
        };
        let quantized_vector_size =
            encoded_vectors_pq::get_quantized_vector_size(&vector_parameters, 1);
        let encoded = EncodedVectorsPQ::encode(
            vector_data.iter(),
            TestEncodedStorageBuilder::new(None, quantized_vector_size),
            &vector_parameters,
            VECTORS_COUNT,
            1,
            1,
            None,
            &AtomicBool::new(false),
        )
        .unwrap();
        let query_u8 = encoded.encode_query(&query);

        let counter = HardwareCounterCell::new();
        for (index, vector) in vector_data.iter().enumerate() {
            let score = encoded.score_point(&query_u8, index as u32, &counter);
            let orginal_score = l1_similarity(&query, vector);
            assert!((score - orginal_score).abs() < ERROR);
        }
    }

    #[test]
    fn test_pq_dot_inverted() {
        let mut rng = rand::rngs::StdRng::seed_from_u64(42);
        let mut vector_data: Vec<Vec<_>> = vec![];
        for _ in 0..VECTORS_COUNT {
            vector_data.push((0..VECTOR_DIM).map(|_| rng.random()).collect());
        }
        let query: Vec<_> = (0..VECTOR_DIM).map(|_| rng.random()).collect();

        let vector_parameters = VectorParameters {
            dim: VECTOR_DIM,
            deprecated_count: None,
            distance_type: DistanceType::Dot,
            invert: true,
        };
        let quantized_vector_size =
            encoded_vectors_pq::get_quantized_vector_size(&vector_parameters, 1);
        let encoded = EncodedVectorsPQ::encode(
            vector_data.iter(),
            TestEncodedStorageBuilder::new(None, quantized_vector_size),
            &vector_parameters,
            VECTORS_COUNT,
            1,
            1,
            None,
            &AtomicBool::new(false),
        )
        .unwrap();
        let query_u8 = encoded.encode_query(&query);

        let counter = HardwareCounterCell::new();
        for (index, vector) in vector_data.iter().enumerate() {
            let score = encoded.score_point(&query_u8, index as u32, &counter);
            let orginal_score = -dot_similarity(&query, vector);
            assert!((score - orginal_score).abs() < ERROR);
        }
    }

    #[test]
    fn test_pq_l2_inverted() {
        let mut rng = rand::rngs::StdRng::seed_from_u64(42);
        let mut vector_data: Vec<Vec<_>> = vec![];
        for _ in 0..VECTORS_COUNT {
            vector_data.push((0..VECTOR_DIM).map(|_| rng.random()).collect());
        }
        let query: Vec<_> = (0..VECTOR_DIM).map(|_| rng.random()).collect();

        let vector_parameters = VectorParameters {
            dim: VECTOR_DIM,
            deprecated_count: None,
            distance_type: DistanceType::L2,
            invert: true,
        };
        let quantized_vector_size =
            encoded_vectors_pq::get_quantized_vector_size(&vector_parameters, 1);
        let encoded = EncodedVectorsPQ::encode(
            vector_data.iter(),
            TestEncodedStorageBuilder::new(None, quantized_vector_size),
            &vector_parameters,
            VECTORS_COUNT,
            1,
            1,
            None,
            &AtomicBool::new(false),
        )
        .unwrap();
        let query_u8 = encoded.encode_query(&query);

        let counter = HardwareCounterCell::new();
        for (index, vector) in vector_data.iter().enumerate() {
            let score = encoded.score_point(&query_u8, index as u32, &counter);
            let orginal_score = -l2_similarity(&query, vector);
            assert!((score - orginal_score).abs() < ERROR);
        }
    }

    #[test]
    fn test_pq_l1_inverted() {
        let mut rng = rand::rngs::StdRng::seed_from_u64(42);
        let mut vector_data: Vec<Vec<_>> = vec![];
        for _ in 0..VECTORS_COUNT {
            vector_data.push((0..VECTOR_DIM).map(|_| rng.random()).collect());
        }
        let query: Vec<_> = (0..VECTOR_DIM).map(|_| rng.random()).collect();

        let vector_parameters = VectorParameters {
            dim: VECTOR_DIM,
            deprecated_count: None,
            distance_type: DistanceType::L1,
            invert: true,
        };
        let quantized_vector_size =
            encoded_vectors_pq::get_quantized_vector_size(&vector_parameters, 1);
        let encoded = EncodedVectorsPQ::encode(
            vector_data.iter(),
            TestEncodedStorageBuilder::new(None, quantized_vector_size),
            &vector_parameters,
            VECTORS_COUNT,
            1,
            1,
            None,
            &AtomicBool::new(false),
        )
        .unwrap();
        let query_u8 = encoded.encode_query(&query);

        let counter = HardwareCounterCell::new();
        for (index, vector) in vector_data.iter().enumerate() {
            let score = encoded.score_point(&query_u8, index as u32, &counter);
            let orginal_score = -l1_similarity(&query, vector);
            assert!((score - orginal_score).abs() < ERROR);
        }
    }

    #[test]
    fn test_pq_dot_internal() {
        let mut rng = rand::rngs::StdRng::seed_from_u64(42);
        let mut vector_data: Vec<Vec<_>> = vec![];
        for _ in 0..VECTORS_COUNT {
            vector_data.push((0..VECTOR_DIM).map(|_| rng.random()).collect());
        }

        let vector_parameters = VectorParameters {
            dim: VECTOR_DIM,
            deprecated_count: None,
            distance_type: DistanceType::Dot,
            invert: false,
        };
        let quantized_vector_size =
            encoded_vectors_pq::get_quantized_vector_size(&vector_parameters, 1);
        let encoded = EncodedVectorsPQ::encode(
            vector_data.iter(),
            TestEncodedStorageBuilder::new(None, quantized_vector_size),
            &vector_parameters,
            VECTORS_COUNT,
            1,
            1,
            None,
            &AtomicBool::new(false),
        )
        .unwrap();

        let counter = HardwareCounterCell::new();
        for i in 1..VECTORS_COUNT {
            let score = encoded.score_internal(0, i as u32, &counter);
            let orginal_score = dot_similarity(&vector_data[0], &vector_data[i]);
            assert!((score - orginal_score).abs() < ERROR);
        }
    }

    #[test]
    fn test_pq_dot_inverted_internal() {
        let mut rng = rand::rngs::StdRng::seed_from_u64(42);
        let mut vector_data: Vec<Vec<_>> = vec![];
        for _ in 0..VECTORS_COUNT {
            vector_data.push((0..VECTOR_DIM).map(|_| rng.random()).collect());
        }

        let vector_parameters = VectorParameters {
            dim: VECTOR_DIM,
            deprecated_count: None,
            distance_type: DistanceType::Dot,
            invert: true,
        };
        let quantized_vector_size =
            encoded_vectors_pq::get_quantized_vector_size(&vector_parameters, 1);
        let encoded = EncodedVectorsPQ::encode(
            vector_data.iter(),
            TestEncodedStorageBuilder::new(None, quantized_vector_size),
            &vector_parameters,
            VECTORS_COUNT,
            1,
            1,
            None,
            &AtomicBool::new(false),
        )
        .unwrap();

        let counter = HardwareCounterCell::new();
        for i in 1..VECTORS_COUNT {
            let score = encoded.score_internal(0, i as u32, &counter);
            let orginal_score = -dot_similarity(&vector_data[0], &vector_data[i]);
            assert!((score - orginal_score).abs() < ERROR);
        }
    }

    // ignore this test because it requires long time
    // this test should be started separately of with `--test-threads=1` flag
    // because `num_threads::num_threads()` is used to check that all encode threads finished
    #[ignore]
    #[test]
    fn test_encode_panic() {
        let mut rng = rand::rngs::StdRng::seed_from_u64(42);
        let mut vector_data: Vec<Vec<_>> = vec![];
        for _ in 0..VECTORS_COUNT {
            vector_data.push((0..VECTOR_DIM).map(|_| rng.random()).collect());
        }

        for i in 0.. {
            let counter = AtomicUsize::new(0);
            let panic_index = i * VECTORS_COUNT / 3;

            let start_num_threads = num_threads::num_threads();
            let vector_data = vector_data.clone();

            let vector_parameters = VectorParameters {
                dim: VECTOR_DIM,
                deprecated_count: None,
                distance_type: DistanceType::Dot,
                invert: false,
            };
            let quantized_vector_size =
                encoded_vectors_pq::get_quantized_vector_size(&vector_parameters, 1);

            let result = std::thread::spawn(move || {
                EncodedVectorsPQ::encode(
                    vector_data.iter().inspect(|_| {
                        let cnt = counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                        if cnt == panic_index {
                            panic!("test panic")
                        }
                        if cnt > panic_index {
                            // after panic add start sleeping to simulate large amount of data
                            std::thread::sleep(Duration::from_micros(100));
                        }
                    }),
                    TestEncodedStorageBuilder::new(None, quantized_vector_size),
                    &vector_parameters,
                    VECTORS_COUNT,
                    1,
                    5,
                    None,
                    &AtomicBool::new(false),
                )
                .unwrap()
            })
            .join();

            if result.is_ok() {
                // no panic, panic_index is too big, all panic cases are handled
                return;
            }

            // some time required to finish encoding threads
            std::thread::sleep(Duration::from_millis(50));

            // check that all threads are finished
            assert_eq!(num_threads::num_threads(), start_num_threads);

            println!("Finished iteration {i}");
        }
    }
}