qdrant-edge 0.8.0

use std::cmp::max;

use ahash::{AHashMap, AHashSet};
use crate::common::fixed_length_priority_queue::FixedLengthPriorityQueue;
use crate::common::types::ScoreType;
use crate::segment::types::{PointIdType, ScoredPoint, SeqNumberType};

/// Avoid excessive memory allocation and allocation failures on huge limits
const LARGEST_REASONABLE_ALLOCATION_SIZE: usize = 1_048_576;

pub struct SearchResultAggregator {
    queue: Option<FixedLengthPriorityQueue<ScoredPoint>>,
    seen: AHashSet<PointIdType>, // Point ids seen
}

impl SearchResultAggregator {
    pub fn new(limit: usize) -> Self {
        SearchResultAggregator {
            queue: if limit > 0 {
                Some(FixedLengthPriorityQueue::new(limit))
            } else {
                None
            },
            seen: AHashSet::with_capacity(limit.min(LARGEST_REASONABLE_ALLOCATION_SIZE)),
        }
    }

    pub fn push(&mut self, point: ScoredPoint) {
        let Some(queue) = self.queue.as_mut() else {
            return;
        };

        // Only add unseen points
        if self.seen.insert(point.id) {
            queue.push(point);
        }
    }

    pub fn into_vec(self) -> Vec<ScoredPoint> {
        self.queue
            .map(|queue| queue.into_sorted_vec())
            .unwrap_or_default()
    }

    pub fn lowest(&self) -> Option<&ScoredPoint> {
        self.queue.as_ref().and_then(|queue| queue.top())
    }
}

pub struct BatchResultAggregator {
    // result aggregators for each batched request
    batch_aggregators: Vec<SearchResultAggregator>,
    // Store max version for each point id to exclude outdated points from the result
    point_versions: AHashMap<PointIdType, SeqNumberType>,
}

impl BatchResultAggregator {
    pub fn new(tops: impl IntoIterator<Item = usize>) -> Self {
        let mut merged_results_per_batch = vec![];
        for top in tops {
            merged_results_per_batch.push(SearchResultAggregator::new(top));
        }

        BatchResultAggregator {
            batch_aggregators: merged_results_per_batch,
            point_versions: AHashMap::new(),
        }
    }

    pub fn update_point_versions<'a>(
        &mut self,
        all_searches_results: impl IntoIterator<Item = &'a ScoredPoint>,
    ) {
        for point in all_searches_results {
            let point_id = point.id;
            self.point_versions
                .entry(point_id)
                .and_modify(|version| *version = max(*version, point.version))
                .or_insert(point.version);
        }
    }

    /// Updates the specific batch result aggregator with the new points
    /// Point must be:
    /// - not seen before
    /// - not outdated (not less than the version stored in point_versions)
    ///
    /// WARN: Must be called after `update_point_versions`, so that `point_versions` is up to date
    pub fn update_batch_results(
        &mut self,
        batch_id: usize,
        search_results: impl IntoIterator<Item = ScoredPoint>,
    ) {
        let aggregator = &mut self.batch_aggregators[batch_id];
        for scored_point in search_results {
            debug_assert!(self.point_versions.contains_key(&scored_point.id));
            if let Some(point_max_version) = self.point_versions.get(&scored_point.id)
                && scored_point.version >= *point_max_version
            {
                aggregator.push(scored_point);
            }
        }
    }

    /// Return lowest acceptable score for given batch id
    pub fn batch_lowest_scores(&self, batch_id: usize) -> Option<ScoreType> {
        let batch_scores = &self.batch_aggregators[batch_id];
        batch_scores.lowest().map(|x| x.score)
    }

    pub fn into_topk(self) -> Vec<Vec<ScoredPoint>> {
        self.batch_aggregators
            .into_iter()
            .map(|x| x.into_vec())
            .collect()
    }
}