ronantakizawa/codereview-bench huggingface dataset:
Repository: qdrant/qdrant
File: lib/shard/src/query/query_context.rs (Rust)

Code before:
    const DEFAULT_CAPACITY: usize = 3;
    let mut idf_vectors: SmallVec<[&VectorName; DEFAULT_CAPACITY]> = Default::default();

    // check vector names existing
    for req in batch_request {
        let vector_name = req.query.get_vector_name();
        let is_idf_required = check_idf_required(vector_name);
        if is_idf_required && !idf_vectors.contains(&vector_name) {
            idf_vectors.push(vector_name);
        }
    }

    let mut query_context = QueryContext::new(search_optimized_threshold_kb, hw_measurement_acc)
        .with_is_stopped(is_stopped_guard.get_is_stopped());

    for search_request in batch_request {
        search_request
            .query
            .iterate_sparse(|vector_name, sparse_vector| {
                if idf_vectors.contains(&vector_name) {
                    query_context.init_idf(vector_name, &sparse_vector.indices);
                }
            })
    }

    query_context
}

pub fn fill_query_context(
    mut query_context: QueryContext,
    segments: LockedSegmentHolder,
) -> Option<QueryContext> {
    let segments = segments.read();

    if segments.is_empty() {
        return None;
    }

    let segments = segments.non_appendable_then_appendable_segments();
    for locked_segment in segments {
        let segment = locked_segment.get();
        let segment_guard = segment.read();
        segment_guard.fill_query_context(&mut query_context);
    }
    Some(query_context)
}


Review comment (performance):
Though, why can't it be simplified to something like this?


```suggestion
    let mut query_context = QueryContext::new(search_optimized_threshold_kb, hw_measurement_acc)
        .with_is_stopped(is_stopped_guard.get_is_stopped());

    for search_request in batch_request {
        search_request
            .query
            .iterate_sparse(|vector_name, sparse_vector| {
                if check_idf_required(&vector_name) {
                    query_context.init_idf(vector_name, &sparse_vector.indices);
                }
            })
    }

    query_context
```

Code after:
        .with_is_stopped(is_stopped_guard.get_is_stopped());

    for search_request in batch_request {
        search_request
            .query
            .iterate_sparse(|vector_name, sparse_vector| {
                if check_idf_required(vector_name) {
                    query_context.init_idf(vector_name, &sparse_vector.indices);
                }
            })
    }

    query_context
}

pub fn fill_query_context(
    mut query_context: QueryContext,
    segments: LockedSegmentHolder,
) -> Option<QueryContext> {
    let segments = segments.read();

    if segments.is_empty() {
        return None;
    }

    let segments = segments.non_appendable_then_appendable_segments();
    for locked_segment in segments {
        let segment = locked_segment.get();
        let segment_guard = segment.read();
        segment_guard.fill_query_context(&mut query_context);
    }
    Some(query_context)
}
Repository: qdrant/qdrant
File: lib/collection/src/shards/dummy_shard.rs (Rust)

Code before:

    pub fn get_size_stats(&self) -> SizeStats {
        SizeStats::default()
    }

    pub fn estimate_cardinality(
        &self,
        _: Option<&Filter>,
    ) -> CollectionResult<CardinalityEstimation> {
        self.dummy()
    }

    fn dummy<T>(&self) -> CollectionResult<T> {
        Err(CollectionError::service_error(self.message.clone()))
    }
}

#[async_trait]
impl ShardOperation for DummyShard {
    async fn update(
        &self,
        op: OperationWithClockTag,
        _: bool,
        _: HwMeasurementAcc,
    ) -> CollectionResult<UpdateResult> {
        match &op.operation {
            CollectionUpdateOperations::FieldIndexOperation(_) => Ok(UpdateResult {
                operation_id: None,
                status: UpdateStatus::Acknowledged,
                clock_tag: None,
            }),
            _ => self.dummy(),
        }
    }

    /// Forward read-only `scroll_by` to `wrapped_shard`
    async fn scroll_by(
        &self,
        _: Arc<ScrollRequestInternal>,
        _: &Handle,
        _: Option<Duration>,
        _: HwMeasurementAcc,
    ) -> CollectionResult<Vec<RecordInternal>> {
        self.dummy()
    }

    async fn local_scroll_by_id(
        &self,
        _: Option<ExtendedPointId>,
        _: usize,
        _: &WithPayloadInterface,

Review comment (suggestion):
I'd like a comment here


```suggestion
        match &op.operation {
            // Allow (and ignore) field operations here as they'll be applied again when the shard is recovered
```

Code after:

    pub fn get_size_stats(&self) -> SizeStats {
        SizeStats::default()
    }

    pub fn estimate_cardinality(
        &self,
        _: Option<&Filter>,
    ) -> CollectionResult<CardinalityEstimation> {
        self.dummy()
    }

    fn dummy<T>(&self) -> CollectionResult<T> {
        Err(CollectionError::service_error(self.message.clone()))
    }
}

#[async_trait]
impl ShardOperation for DummyShard {
    async fn update(
        &self,
        op: OperationWithClockTag,
        _: bool,
        _: HwMeasurementAcc,
    ) -> CollectionResult<UpdateResult> {
        match &op.operation {
            CollectionUpdateOperations::PointOperation(_) => self.dummy(),
            CollectionUpdateOperations::VectorOperation(_) => self.dummy(),
            CollectionUpdateOperations::PayloadOperation(_) => self.dummy(),

            // Allow (and ignore) field index operations. Field index schema is stored in collection
            // config, and indices will be created (if needed) when dummy shard is recovered.
            CollectionUpdateOperations::FieldIndexOperation(_) => Ok(UpdateResult {
                operation_id: None,
                status: UpdateStatus::Acknowledged,
                clock_tag: None,
            }),
        }
    }

    /// Forward read-only `scroll_by` to `wrapped_shard`
    async fn scroll_by(
        &self,
        _: Arc<ScrollRequestInternal>,
        _: &Handle,
        _: Option<Duration>,
        _: HwMeasurementAcc,
    ) -> CollectionResult<Vec<RecordInternal>> {
        self.dummy()
    }
Repository: qdrant/qdrant
File: lib/common/cancel/src/blocking.rs (Rust)

Code before:
use super::*;

/// # Cancel safety
///
/// This function is cancel safe.
///
/// If cancelled, the cancellation token provided to the `task` will be triggered automatically.
pub async fn spawn_cancel_on_drop<Out, Task>(task: Task) -> Result<Out, Error>
where
    Task: FnOnce(CancellationToken) -> Out + Send + 'static,
    Out: Send + 'static,
{
    let cancel = CancellationToken::new();

    let task = {
        let cancel = cancel.child_token();
        move || task(cancel)
    };

    let guard = cancel.drop_guard();
    let output = tokio::task::spawn_blocking(task).await?;
    guard.disarm();

    Ok(output)
}

/// # Cancel safety
///
/// This function is cancel safe.
///
/// If cancelled without triggering the cancellation token, the `task` will still run to completion.
///
/// This function *will* return early, and the `task` *may* return early by triggering the
/// cancellation token.
pub async fn spawn_cancel_on_token<Out, Task>(
    cancel: CancellationToken,
    task: Task,
) -> Result<Out, Error>
where
    Task: FnOnce(CancellationToken) -> Out + Send + 'static,
    Out: Send + 'static,
{
    let task = {
        let cancel = cancel.child_token();
        move || task(cancel)
    };

    let output = future::cancel_on_token(cancel, tokio::task::spawn_blocking(task)).await??;

    Ok(output)
}


Review comment (refactor):
I think we can simplify this. If token is triggered, `cancel_on_token` will simply exit early, which will drop `AbortOnDropHandle`, which will abort the task.

```suggestion
    let output = future::cancel_on_token(cancel, AbortOnDropHandle::new(tokio::task::spawn_blocking(task))).await??;
```

Code after:
    Ok(output)
}

/// # Cancel safety
///
/// This function is cancel safe.
///
/// If cancelled without triggering the cancellation token, the `task` will still run to completion.
///
/// This function *will* return early, and the `task` *may* never run or return early by triggering
/// the cancellation token.
pub async fn spawn_cancel_on_token<Out, Task>(
    cancel: CancellationToken,
    task: Task,
) -> Result<Out, Error>
where
    Task: FnOnce(CancellationToken) -> Out + Send + 'static,
    Out: Send + 'static,
{
    let task = {
        let cancel = cancel.child_token();
        move || task(cancel)
    };

    let handle = tokio::task::spawn_blocking(task);
    let output = future::cancel_and_abort_on_token(cancel, handle).await?;

    Ok(output)
}
Repository: qdrant/qdrant
File: lib/quantization/src/encoded_vectors_u8.rs (Rust)

Code before:
use std::alloc::Layout;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};

use common::counter::hardware_counter::HardwareCounterCell;
use common::typelevel::True;
use common::types::PointOffsetType;
use fs_err as fs;
use io::file_operations::atomic_save_json;
use memory::mmap_type::MmapFlusher;
use serde::{Deserialize, Serialize};

use crate::EncodingError;
use crate::encoded_storage::{EncodedStorage, EncodedStorageBuilder};
use crate::encoded_vectors::{
    DistanceType, EncodedVectors, VectorParameters, validate_vector_parameters,
};
use crate::quantile::{find_min_max_from_iter, find_quantile_interval};

pub const ALIGNMENT: usize = 16;

#[derive(Clone, Default, PartialEq, Debug)]
pub enum ScalarQuantizationMethod {
    #[default]
    Uint8,
    // Future methods can be added here
}

pub struct EncodedVectorsU8<TStorage: EncodedStorage> {
    encoded_vectors: TStorage,
    metadata: Metadata,
    metadata_path: Option<PathBuf>,
}

pub struct EncodedQueryU8 {
    offset: f32,
    encoded_query: Vec<u8>,
}

#[derive(Serialize, Deserialize)]
struct Metadata {
    actual_dim: usize,
    alpha: f32,
    offset: f32,
    multiplier: f32,
    vector_parameters: VectorParameters,
}

Review comment (suggestion):
It seems `Default` is not really used now.  But if you have a good reason to keep it, it is OK.  I think it can be dangerous.

Code after:
use std::alloc::Layout;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};

use common::counter::hardware_counter::HardwareCounterCell;
use common::typelevel::True;
use common::types::PointOffsetType;
use fs_err as fs;
use io::file_operations::atomic_save_json;
use memory::mmap_type::MmapFlusher;
use serde::{Deserialize, Serialize};

use crate::EncodingError;
use crate::encoded_storage::{EncodedStorage, EncodedStorageBuilder};
use crate::encoded_vectors::{
    DistanceType, EncodedVectors, VectorParameters, validate_vector_parameters,
};
use crate::quantile::{find_min_max_from_iter, find_quantile_interval};

pub const ALIGNMENT: usize = 16;

#[derive(Clone, PartialEq, Debug)]
pub enum ScalarQuantizationMethod {
    Int8,
    // Future methods can be added here
}

pub struct EncodedVectorsU8<TStorage: EncodedStorage> {
    encoded_vectors: TStorage,
    metadata: Metadata,
    metadata_path: Option<PathBuf>,
}

pub struct EncodedQueryU8 {
    offset: f32,
    encoded_query: Vec<u8>,
}

#[derive(Serialize, Deserialize)]
struct Metadata {
    actual_dim: usize,
    alpha: f32,
    offset: f32,
    multiplier: f32,
    vector_parameters: VectorParameters,
}
Repository: qdrant/qdrant
File: lib/segment/src/common/flags/buffered_dynamic_flags.rs (Rust)

Code before:

            // Keep the guard till the end of the flush to prevent concurrent drop/flushes
            let _flushing_lock_guard = flushing_lock.lock();

            // lock for the entire flushing process
            let mut flags_guard = flags_arc.lock();

            // resize if needed
            if required_len > flags_guard.len() {
                flags_guard.set_len(required_len)?;
            }

            for (index, value) in updates {
                flags_guard.set(index as usize, value);
            }

            flags_guard.flusher()()?;
            Ok(())
        })
    }
}

impl Drop for BufferedDynamicFlags {
    fn drop(&mut self) {
        // Wait for all background flush operations to finish, and cancel future flushes
        _ = self.flushing_lock.lock();
    }
}

#[cfg(test)]
mod tests {

    use common::types::PointOffsetType;
    use rand::rngs::StdRng;
    use rand::{Rng, SeedableRng};

    use crate::common::flags::buffered_dynamic_flags::BufferedDynamicFlags;
    use crate::common::flags::dynamic_mmap_flags::DynamicMmapFlags;

    #[test]
    fn test_buffered_flags_growth_persistence() {
        let dir = tempfile::Builder::new()
            .prefix("buffered_flags_growth")
            .tempdir()
            .unwrap();

        // Start with smaller flags
        {
            let mut mmap_flags = DynamicMmapFlags::open(dir.path(), false).unwrap();
            mmap_flags.set_len(3).unwrap();
            mmap_flags.set(0, true);

Review comment (suggestion):
I don't think this works.

When we obtain this lock, the `Arc`s are still alive. And so a flusher thread is still able to upgrade, lock and flush after this drop function if the timing is right.

Code after:
            }

            let Some(flags_arc) = flags_arc.upgrade() else {
                log::debug!("skipping flushing on deleted storage");
                return Ok(());
            };

            // lock for the entire flushing process
            let mut flags_guard = flags_arc.lock();

            // resize if needed
            if required_len > flags_guard.len() {
                flags_guard.set_len(required_len)?;
            }

            for (index, value) in updates {
                flags_guard.set(index as usize, value);
            }

            flags_guard.flusher()()?;

            Ok(())
        })
    }
}

impl Drop for BufferedDynamicFlags {
    fn drop(&mut self) {
        // Wait for all background flush operations to finish, and cancel future flushes
        *self.is_alive_flush_lock.lock() = false;
    }
}

#[cfg(test)]
mod tests {

    use common::types::PointOffsetType;
    use rand::rngs::StdRng;
    use rand::{Rng, SeedableRng};

    use crate::common::flags::buffered_dynamic_flags::BufferedDynamicFlags;
    use crate::common::flags::dynamic_mmap_flags::DynamicMmapFlags;

    #[test]
    fn test_buffered_flags_growth_persistence() {
        let dir = tempfile::Builder::new()
            .prefix("buffered_flags_growth")
            .tempdir()
            .unwrap();

        // Start with smaller flags
Repository: qdrant/qdrant
File: src/actix/api/service_api.rs (Rust)

Code before:
    pub fn timeout(&self) -> Option<Duration> {
        self.timeout.map(Duration::from_secs)
    }
}

#[get("/telemetry")]
fn telemetry(
    telemetry_collector: Data<Mutex<TelemetryCollector>>,
    params: Query<TelemetryParam>,
    ActixAccess(access): ActixAccess,
) -> impl Future<Output = HttpResponse> {
    helpers::time(async move {
        let anonymize = params.anonymize.unwrap_or(false);
        let details_level = params
            .details_level
            .map_or(DetailsLevel::Level0, Into::into);

        let detail = TelemetryDetail {
            level: details_level,
            histograms: false,
        };
        let telemetry_collector = telemetry_collector.lock().await;
        let telemetry_data = telemetry_collector
            .prepare_data(&access, detail, params.timeout())
            .await?;
        drop(telemetry_collector);
        let telemetry_data = if anonymize {
            telemetry_data.anonymize()
        } else {
            telemetry_data
        };
        Ok(telemetry_data)
    })
}

#[derive(Deserialize, Serialize, JsonSchema, Validate)]
pub struct MetricsParam {
    pub anonymize: Option<bool>,
    #[validate(range(min = 1))]
    pub timeout: Option<u64>,
}

impl MetricsParam {
    pub fn timeout(&self) -> Option<Duration> {
        self.timeout.map(Duration::from_secs)
    }
}

#[get("/metrics")]
async fn metrics(
    telemetry_collector: Data<Mutex<TelemetryCollector>>,

Review comment (bug):
Maybe just?

```suggestion
let telemetry_data = telemetry_collector
    .lock()
    .await
    .prepare_data(&access, detail, params.timeout())
    .await?;

```

(note that this suggestion diff looks broken 🤷‍♂️ )

Code after:
    pub fn timeout(&self) -> Option<Duration> {
        self.timeout.map(Duration::from_secs)
    }
}

#[get("/telemetry")]
fn telemetry(
    telemetry_collector: Data<Mutex<TelemetryCollector>>,
    params: Query<TelemetryParam>,
    ActixAccess(access): ActixAccess,
) -> impl Future<Output = HttpResponse> {
    helpers::time(async move {
        let anonymize = params.anonymize.unwrap_or(false);
        let details_level = params
            .details_level
            .map_or(DetailsLevel::Level0, Into::into);

        let detail = TelemetryDetail {
            level: details_level,
            histograms: false,
        };
        let telemetry_data = telemetry_collector
            .lock()
            .await
            .prepare_data(&access, detail, params.timeout())
            .await?;
        let telemetry_data = if anonymize {
            telemetry_data.anonymize()
        } else {
            telemetry_data
        };
        Ok(telemetry_data)
    })
}

#[derive(Deserialize, Serialize, JsonSchema, Validate)]
pub struct MetricsParam {
    pub anonymize: Option<bool>,
    #[validate(range(min = 1))]
    pub timeout: Option<u64>,
}

impl MetricsParam {
    pub fn timeout(&self) -> Option<Duration> {
        self.timeout.map(Duration::from_secs)
    }
}

#[get("/metrics")]
async fn metrics(
    telemetry_collector: Data<Mutex<TelemetryCollector>>,
Repository: qdrant/qdrant
File: lib/shard/src/operations/staging.rs (Rust)

Code before:
//! Staging-only operations for testing and debugging purposes.
//!
//! This module contains operations that are only available when the `staging` feature is enabled.

use std::time::Duration;

use ordered_float::OrderedFloat;
use serde::{Deserialize, Serialize};

/// Test operation that introduces an artificial delay.
#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, Hash)]
pub struct TestDelayUpsertPointsOperation {
    /// Duration of the delay in seconds.
    pub duration: OrderedFloat<f64>,
}

impl TestDelayUpsertPointsOperation {
    /// Execute the delay operation (blocking).
    pub fn execute(&self) {
        let duration_secs = self.duration.into_inner();
        log::debug!("TestDelayUpsertPoints: sleeping for {duration_secs}s");
        std::thread::sleep(Duration::from_secs_f64(duration_secs));
        log::debug!("TestDelayUpsertPoints: finished sleeping");
    }
}


Review comment (bug):
The `Duration` type itself it quite good at formatting time, I suggest to use that.

Demo: https://play.rust-lang.org/?version=stable&mode=debug&edition=2024&gist=029f33c3a44399848df807365dab9d3c

```suggestion
        let duration = Duration::from_secs_f64(self.duration.into_inner().max(0.0));
        log::debug!("TestDelayUpsertPoints: sleeping for {duration:?}");
        std::thread::sleep(duration);
        log::debug!("TestDelayUpsertPoints: finished sleeping");
```

This also applies Rabbit's fix.

Code after:
//! Staging-only operations for testing and debugging purposes.
//!
//! This module contains operations that are only available when the `staging` feature is enabled.

use std::time::Duration;

use ordered_float::OrderedFloat;
use serde::{Deserialize, Serialize};

/// Test operation that introduces an artificial delay.
#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, Hash)]
pub struct TestDelayOperation {
    /// Duration of the delay in seconds.
    pub duration: OrderedFloat<f64>,
}

impl TestDelayOperation {
    /// Create a new TestDelayOperation with the given duration in seconds.
    pub fn new(duration_secs: f64) -> Self {
        Self {
            duration: OrderedFloat(duration_secs),
        }
    }

    /// Execute the delay operation (blocking).
    pub fn execute(&self) {
        let duration_secs = self.duration.into_inner();
        if duration_secs < 0.0 {
            log::warn!("TestDelay: negative duration {duration_secs}s, skipping");
            return;
        }
        let delay = Duration::from_secs_f64(duration_secs);
        log::debug!("TestDelay: sleeping for {delay:?}");
        std::thread::sleep(delay);
        log::debug!("TestDelay: finished sleeping");
    }
}
Repository: qdrant/qdrant
File: lib/quantization/src/quantile.rs (Rust)

Code before:
    let mut result = Vec::with_capacity(dim);
    for (p_square_min, p_square_max) in p_squares {
        result.push((
            p_square_min.estimate() as f32,
            p_square_max.estimate() as f32,
        ));
    }
    Ok(result)
}

fn take_random_vectors<'a>(
    vector_data: impl Iterator<Item = impl AsRef<[f32]> + 'a>,
    count: usize,
    sample_size: usize,
) -> (usize, impl Iterator<Item = impl AsRef<[f32]> + 'a>) {
    let slice_size = std::cmp::min(count, sample_size);
    let permutor = Permutor::new(count as u64);
    let mut selected_vectors: Vec<usize> = permutor.map(|i| i as usize).take(slice_size).collect();
    selected_vectors.sort_unstable();

    (
        slice_size,
        vector_data
            .enumerate()
            .filter_map(move |(vector_index, vector)| {
                if selected_vectors.binary_search(&vector_index).is_ok() {
                    Some(vector)
                } else {
                    None
                }
            }),
    )
}


Review comment (suggestion):
It could be implemented as linear operation because both `selected_vectors`'s values and `vector_index` is increasing.

Upd: actually it was before.

Code after:
    })
}

// Take random vectors from the input iterator using `Permutor`.
fn take_random_vectors<'a>(
    vector_data: impl Iterator<Item = impl AsRef<[f32]> + 'a>,
    count: usize,
    sample_size: usize,
    stopped: &AtomicBool,
) -> Result<Vec<impl AsRef<[f32]> + 'a>, EncodingError> {
    let slice_size = std::cmp::min(count, sample_size);
    let permutor = Permutor::new(count as u64);
    let mut selected_vectors: Vec<usize> = permutor.map(|i| i as usize).take(slice_size).collect();
    selected_vectors.sort_unstable();

    let mut data_slice = Vec::with_capacity(slice_size);
    let mut selected_index: usize = 0;
    for (vector_index, vector_data) in vector_data.into_iter().enumerate() {
        if stopped.load(Ordering::Relaxed) {
            return Err(EncodingError::Stopped);
        }

        if vector_index == selected_vectors[selected_index] {
            data_slice.push(vector_data);
            selected_index += 1;
            if selected_index == slice_size {
                break;
            }
        }
    }

    Ok(data_slice)
}

#[cfg(test)]
mod tests {
    use rand::rngs::StdRng;
    use rand::{Rng, SeedableRng};
    use rstest::rstest;

    use super::*;

    #[rstest]
    #[case(1.0)]
    #[case(0.95)]
    fn test_vectors_quantile_interval(#[case] quantile: f32) {
        const COUNT: usize = 5_000;
        const DIM: usize = 4;

        let mut rng = StdRng::seed_from_u64(42);
        let mut data = Vec::with_capacity(COUNT);