qdrant-edge 0.8.0

use std::collections::{HashMap, HashSet};
use std::fmt;
use std::num::NonZeroUsize;
use std::sync::Arc;

use crate::segment::common::BYTES_IN_KB;
use crate::segment::data_types::modifier::Modifier;
use crate::segment::index::sparse_index::sparse_index_config::{SparseIndexConfig, SparseIndexType};
use crate::segment::types::{
    Distance, HnswConfig, Indexes, Memory, MultiVectorConfig, PayloadStorageType,
    QuantizationConfig, SegmentConfig, SparseVectorDataConfig, SparseVectorStorageType,
    VectorDataConfig, VectorNameBuf, VectorStorageDatatype, VectorStorageType,
};

pub const TEMP_SEGMENTS_PATH: &str = "temp_segments";
pub const DEFAULT_MAX_SEGMENT_PER_CPU_KB: usize = 256_000;
pub const DEFAULT_INDEXING_THRESHOLD_KB: usize = 10_000;
pub const DEFAULT_DELETED_THRESHOLD: f64 = 0.2;
pub const DEFAULT_VACUUM_MIN_VECTOR_NUMBER: usize = 1000;

/// Extra configuration for dense vectors, applied on top of the plain config during optimization.
#[derive(Debug, Clone, PartialEq)]
pub struct DenseVectorOptimizerConfig {
    pub on_disk: Option<bool>,
    pub memory: Option<Memory>,
    pub hnsw_config: HnswConfig,
    pub quantization_config: Option<QuantizationConfig>,
}

impl DenseVectorOptimizerConfig {
    /// Requested memory placement of the original vector storage, resolving the new `memory`
    /// parameter against the deprecated `on_disk` flag. `None` if neither is configured.
    pub fn memory_placement(&self) -> Option<Memory> {
        Memory::resolve(self.memory, self.on_disk.map(Memory::from_on_disk))
    }
}

/// Extra configuration for sparse vectors, applied on top of the plain config during optimization.
#[derive(Debug, Clone, PartialEq)]
pub struct SparseVectorOptimizerConfig {
    pub on_disk: Option<bool>,
    pub memory: Option<Memory>,
}

impl SparseVectorOptimizerConfig {
    /// Requested memory placement of the sparse index, resolving the new `memory` parameter
    /// against the deprecated `on_disk` flag. `None` if neither is configured.
    pub fn memory_placement(&self) -> Option<Memory> {
        Memory::resolve(self.memory, self.on_disk.map(Memory::from_on_disk_heap))
    }
}

/// Live read of the vector names currently present in the collection schema.
///
/// Unlike the rest of [`SegmentOptimizerConfig`], which is a frozen snapshot taken when the
/// optimizer was built, this reads the *current* schema each time it is called. Optimization needs
/// the live view to tell a vector name that was deleted from the collection (and should be pruned
/// when rebuilding old segments) from one that was just created but is not yet in this optimizer's
/// frozen target (the CreateVectorName race, which must cancel instead). Wrapped in a newtype so
/// `SegmentOptimizerConfig` can keep deriving `Debug`.
#[derive(Clone)]
pub struct LiveVectorNamesProvider(Arc<dyn Fn() -> HashSet<VectorNameBuf> + Send + Sync>);

impl LiveVectorNamesProvider {
    pub fn new(read: impl Fn() -> HashSet<VectorNameBuf> + Send + Sync + 'static) -> Self {
        Self(Arc::new(read))
    }

    pub fn get(&self) -> HashSet<VectorNameBuf> {
        self.0()
    }
}

impl fmt::Debug for LiveVectorNamesProvider {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("LiveVectorNamesProvider")
            .finish_non_exhaustive()
    }
}

/// This configuration contains all necessary information to build an optimized segment.
#[derive(Debug, Clone)]
pub struct SegmentOptimizerConfig {
    pub payload_storage_type: PayloadStorageType,
    /// Configuration of dense vectors, as it should be for a plain segment (without any optimization).
    pub plain_dense_vector_config: HashMap<VectorNameBuf, VectorDataConfig>,
    /// Configuration of sparse vectors, as it should be for a plain segment (without any optimization).
    pub plain_sparse_vector_config: HashMap<VectorNameBuf, SparseVectorDataConfig>,
    /// Extra configuration for dense vectors, which _might_ be applied during optimization,
    /// depending on the segment state.
    pub dense_vector: HashMap<VectorNameBuf, DenseVectorOptimizerConfig>,
    /// Extra configuration for sparse vectors, which _might_ be applied during optimization,
    /// depending on the segment state.
    pub sparse_vector: HashMap<VectorNameBuf, SparseVectorOptimizerConfig>,
    /// Live read of the collection's vector names, when wired in via
    /// [`SegmentOptimizerConfig::with_live_vector_names`]. `None` if no live source is available.
    pub live_vector_names: Option<LiveVectorNamesProvider>,
}

impl SegmentOptimizerConfig {
    pub fn plain_segment_config(&self) -> SegmentConfig {
        SegmentConfig {
            vector_data: self.plain_dense_vector_config.clone(),
            sparse_vector_data: self.plain_sparse_vector_config.clone(),
            payload_storage_type: self.payload_storage_type,
        }
    }

    pub fn new(
        payload_storage_type: PayloadStorageType,
        dense_vectors: HashMap<VectorNameBuf, DenseVectorOptimizerInput>,
        sparse_vectors: HashMap<VectorNameBuf, SparseVectorOptimizerInput>,
    ) -> SegmentOptimizerConfig {
        let (mut plain_dense_vector_config, mut dense_vector) = (HashMap::new(), HashMap::new());
        for (name, input) in dense_vectors {
            let DenseVectorOptimizerInput {
                size,
                distance,
                on_disk,
                memory,
                hnsw_config,
                quantization_config,
                multivector_config,
                datatype,
            } = input;
            let plain_memory = Memory::resolve(
                memory,
                Some(Memory::from_on_disk(on_disk.unwrap_or_default())),
            )
            .unwrap_or(Memory::Cached);
            plain_dense_vector_config.insert(
                name.clone(),
                VectorDataConfig {
                    size,
                    distance,
                    index: Indexes::Plain {},
                    storage_type: VectorStorageType::appendable_from_memory(plain_memory),
                    quantization_config: QuantizationConfig::for_appendable_segment(
                        quantization_config.as_ref(),
                    ),
                    multivector_config,
                    datatype,
                },
            );
            dense_vector.insert(
                name,
                DenseVectorOptimizerConfig {
                    on_disk,
                    memory,
                    hnsw_config,
                    quantization_config,
                },
            );
        }

        let (mut plain_sparse_vector_config, mut sparse_vector) = (HashMap::new(), HashMap::new());
        for (name, input) in sparse_vectors {
            let SparseVectorOptimizerInput {
                on_disk,
                memory,
                full_scan_threshold,
                index_datatype,
                storage_type,
                modifier,
            } = input;
            plain_sparse_vector_config.insert(
                name.clone(),
                SparseVectorDataConfig {
                    index: SparseIndexConfig {
                        full_scan_threshold,
                        index_type: SparseIndexType::MutableRam,
                        datatype: index_datatype,
                        memory,
                    },
                    storage_type,
                    modifier,
                },
            );
            sparse_vector.insert(name, SparseVectorOptimizerConfig { on_disk, memory });
        }

        SegmentOptimizerConfig {
            payload_storage_type,
            plain_dense_vector_config,
            plain_sparse_vector_config,
            dense_vector,
            sparse_vector,
            live_vector_names: None,
        }
    }

    /// Attach a live read of the collection's vector names (see [`LiveVectorNamesProvider`]).
    #[must_use]
    pub fn with_live_vector_names(mut self, provider: LiveVectorNamesProvider) -> Self {
        self.live_vector_names = Some(provider);
        self
    }

    /// The collection's current vector names, if a live source was wired in.
    pub fn live_vector_names(&self) -> Option<HashSet<VectorNameBuf>> {
        self.live_vector_names
            .as_ref()
            .map(LiveVectorNamesProvider::get)
    }
}

/// Per-dense-vector input for the optimizer builder.
#[derive(Debug, Clone)]
pub struct DenseVectorOptimizerInput {
    pub size: usize,
    pub distance: Distance,
    pub on_disk: Option<bool>,
    pub memory: Option<Memory>,
    pub hnsw_config: HnswConfig,
    pub quantization_config: Option<QuantizationConfig>,
    pub multivector_config: Option<MultiVectorConfig>,
    pub datatype: Option<VectorStorageDatatype>,
}

/// Per-sparse-vector input for the optimizer builder.
#[derive(Debug, Clone)]
pub struct SparseVectorOptimizerInput {
    pub on_disk: Option<bool>,
    pub memory: Option<Memory>,
    pub full_scan_threshold: Option<usize>,
    pub index_datatype: Option<VectorStorageDatatype>,
    pub storage_type: SparseVectorStorageType,
    pub modifier: Option<Modifier>,
}

/// Target segment count for the merge optimizer.
pub fn default_segment_number() -> usize {
    // Configure 1 segment per 2 CPUs, as a middle ground between
    // latency and RPS.
    let expected_segments = crate::common::cpu::get_num_cpus() / 2;
    // Do not configure less than 2 and more than 8 segments
    // until it is not explicitly requested
    expected_segments.clamp(2, 8)
}

// --- Shared optimizer threshold helpers (used by collection and edge) ---

/// Resolve number of segments: if `default_segment_number` is 0, use CPU-based default.
pub fn get_number_segments(requested_segment_number: usize) -> usize {
    if requested_segment_number == 0 {
        default_segment_number()
    } else {
        requested_segment_number
    }
}

/// Resolve indexing threshold in KB: `None` => default, `Some(0)` => disable (usize::MAX).
pub fn get_indexing_threshold_kb(indexing_threshold: Option<usize>) -> usize {
    match indexing_threshold {
        None => DEFAULT_INDEXING_THRESHOLD_KB,
        Some(0) => usize::MAX,
        Some(custom) => custom,
    }
}

/// Resolve max segment size in KB: custom value or per-thread default.
pub fn get_max_segment_size_kb(
    max_segment_size: Option<usize>,
    num_indexing_threads: usize,
) -> usize {
    if let Some(max) = max_segment_size {
        max
    } else {
        num_indexing_threads.saturating_mul(DEFAULT_MAX_SEGMENT_PER_CPU_KB)
    }
}

/// Build deferred points threshold in bytes when `prevent_unoptimized` is true.
pub fn get_deferred_points_threshold_bytes(
    prevent_unoptimized: Option<bool>,
    indexing_threshold_kb: usize,
) -> Option<NonZeroUsize> {
    (prevent_unoptimized == Some(true))
        .then(|| indexing_threshold_kb.saturating_mul(BYTES_IN_KB))
        .and_then(NonZeroUsize::new)
}

#[cfg(test)]
mod tests {
    use std::sync::Mutex;

    use super::*;

    #[test]
    fn live_vector_names_provider_reads_current_state() {
        // The provider must re-read the live source on every call rather than snapshot it once,
        // otherwise a vector deleted (or created) after optimizer construction would be missed and
        // the merge would make the wrong cancel/prune decision.
        let source = Arc::new(Mutex::new(HashSet::from(["a".to_owned(), "b".to_owned()])));
        let provider = {
            let source = source.clone();
            LiveVectorNamesProvider::new(move || source.lock().unwrap().clone())
        };

        assert_eq!(
            provider.get(),
            HashSet::from(["a".to_owned(), "b".to_owned()])
        );

        // Delete "b" from the live source: the provider must reflect it on the next read.
        source.lock().unwrap().remove("b");
        assert_eq!(provider.get(), HashSet::from(["a".to_owned()]));
    }
}