Repository: qdrant/qdrant
File: lib/quantization/src/quantile.rs (Rust)

Code before:
    count: usize,
    quantile: f32,
    num_threads: usize,
    stopped: &AtomicBool,
) -> Result<Vec<(f32, f32)>, EncodingError> {
    debug_assert!(quantile > 0.5 && quantile <= 1.0);

    // In case of max quantile, return min-max per dimension
    if quantile >= 1.0 {
        return find_min_max_interval_per_coordinate(vector_data, dim, count, stopped);
    }

    // Otherwise, use P Square algorithm to estimate quantile intervals
    find_interval_per_coordinate_p2(vector_data, dim, count, quantile, num_threads, stopped)
}

fn find_min_max_interval_per_coordinate<'a>(
    vector_data: impl Iterator<Item = impl AsRef<[f32]> + 'a> + Clone,
    dim: usize,
    count: usize,
    stopped: &AtomicBool,
) -> Result<Vec<(f32, f32)>, EncodingError> {
    let mut result = Vec::with_capacity(dim);
    for _ in 0..dim {
        result.push((f32::MAX, f32::MIN));
    }

    let selected_vectors = take_random_vectors(vector_data, dim, count, P2_SAMPLE_SIZE, stopped)?;

    for vector in selected_vectors.chunks_exact(dim) {
        for ((min, max), &value) in result.iter_mut().zip(vector.iter()) {
            *min = min.min(value);
            *max = max.max(value);
        }
    }

    for min_max in result.iter_mut() {
        if min_max.0 == f32::MAX || min_max.1 == f32::MIN {
            *min_max = (0.0, 0.0);
        }
    }

    Ok(result)
}

fn find_interval_per_coordinate_p2<'a>(
    vector_data: impl Iterator<Item = impl AsRef<[f32]> + 'a> + Clone,
    dim: usize,
    count: usize,
    quantile: f32,
    num_threads: usize,

Review comment (suggestion):
```suggestion
    let result = vec![(f32::MAX, f32::MIN); dim];
```

Code after:
    )))
}

pub fn find_interval_per_coordinate<'a>(
    vector_data: impl Iterator<Item = impl AsRef<[f32]> + Send + Sync + 'a> + Clone,
    dim: usize,
    count: usize,
    quantile: f32,
    num_threads: usize,
    stopped: &AtomicBool,
) -> Result<Vec<(f32, f32)>, EncodingError> {
    debug_assert!(quantile > 0.5 && quantile <= 1.0);

    // In case of max quantile, return min-max per dimension
    if quantile >= 1.0 {
        return find_min_max_interval_per_coordinate(vector_data, dim, count, stopped);
    }

    // Otherwise, use P Square algorithm to estimate quantile intervals
    find_interval_per_coordinate_p2(vector_data, dim, count, quantile, num_threads, stopped)
}

fn find_min_max_interval_per_coordinate<'a>(
    vector_data: impl Iterator<Item = impl AsRef<[f32]> + 'a> + Clone,
    dim: usize,
    count: usize,
    stopped: &AtomicBool,
) -> Result<Vec<(f32, f32)>, EncodingError> {
    let mut result = vec![(f32::MAX, f32::MIN); dim];

    let selected_vectors = take_random_vectors(vector_data, count, SAMPLE_SIZE, stopped)?;

    for vector in selected_vectors {
        if stopped.load(Ordering::Relaxed) {
            return Err(EncodingError::Stopped);
        }

        for ((min, max), &value) in result.iter_mut().zip(vector.as_ref().iter()) {
            *min = min.min(value);
            *max = max.max(value);
        }
    }

    for min_max in result.iter_mut() {
        if min_max.0 == f32::MAX || min_max.1 == f32::MIN {
            *min_max = (0.0, 0.0);
        }
    }

    Ok(result)
}