qdrant-edge 0.8.0

use std::collections::HashSet;
use std::sync::atomic::AtomicBool;

use crate::common::counter::hardware_accumulator::HwMeasurementAcc;
use crate::common::types::DeferredBehavior;
use itertools::Itertools as _;
use rand::RngExt;
use rand::distr::weighted::WeightedIndex;
use rand::rngs::StdRng;
use crate::segment::common::operation_error::{OperationError, OperationResult};
use crate::segment::data_types::order_by::{Direction, OrderBy};
use crate::segment::entry::ReadSegmentEntry;
use crate::segment::types::*;
use crate::shard::query::scroll::{QueryScrollRequestInternal, ScrollOrder};
use crate::shard::retrieve::record_internal::RecordInternal;
use crate::shard::retrieve::retrieve_blocking::retrieve_over;
use crate::shard::scroll::ScrollRequestInternal;

use crate::edge::read_view::{EdgeReadView, ReadSegmentHandle};

impl<H: ReadSegmentHandle> EdgeReadView<H> {
    pub(crate) fn scroll(
        &self,
        request: ScrollRequestInternal,
    ) -> OperationResult<(Vec<RecordInternal>, Option<PointIdType>)> {
        let ScrollRequestInternal {
            offset,
            limit,
            filter,
            with_payload,
            with_vector,
            order_by,
        } = request;

        let limit = limit.unwrap_or(ScrollRequestInternal::default_limit());
        let with_payload = with_payload.unwrap_or(ScrollRequestInternal::default_with_payload());

        match order_by.map(OrderBy::from) {
            None => {
                let limit_plus_one = limit.saturating_add(1);
                let mut records = self.scroll_by_id(
                    offset,
                    limit_plus_one,
                    &with_payload,
                    &with_vector,
                    filter.as_ref(),
                    HwMeasurementAcc::disposable_edge(),
                )?;
                let next_offset = if records.len() > limit {
                    let last_record = records.pop().unwrap();
                    Some(last_record.id)
                } else {
                    None
                };
                Ok((records, next_offset))
            }

            Some(order_by) => {
                if offset.is_some() {
                    return Err(OperationError::validation_error(
                        "Offset is not supported when ordering by field",
                    ));
                }
                let records = self.scroll_by_field(
                    limit,
                    &with_payload,
                    &with_vector,
                    filter.as_ref(),
                    &order_by,
                    HwMeasurementAcc::disposable_edge(),
                )?;
                Ok((records, None))
            }
        }
    }

    pub(crate) fn query_scroll(
        &self,
        request: &QueryScrollRequestInternal,
    ) -> OperationResult<Vec<ScoredPoint>> {
        let QueryScrollRequestInternal {
            limit,
            with_vector,
            filter,
            scroll_order,
            with_payload,
        } = request;

        let records = match scroll_order {
            ScrollOrder::ById => self.scroll_by_id(
                None,
                *limit,
                with_payload,
                with_vector,
                filter.as_ref(),
                HwMeasurementAcc::disposable_edge(),
            )?,
            ScrollOrder::ByField(order_by) => self.scroll_by_field(
                *limit,
                with_payload,
                with_vector,
                filter.as_ref(),
                order_by,
                HwMeasurementAcc::disposable_edge(),
            )?,
            ScrollOrder::Random => self.scroll_randomly(
                *limit,
                with_payload,
                with_vector,
                filter.as_ref(),
                HwMeasurementAcc::disposable_edge(),
            )?,
        };

        let point_results = records
            .into_iter()
            .map(|record| ScoredPoint {
                id: record.id,
                version: 0,
                score: 1.0,
                payload: record.payload,
                vector: record.vector,
                shard_key: record.shard_key,
                order_value: record.order_value,
            })
            .collect();

        Ok(point_results)
    }

    fn scroll_by_id(
        &self,
        offset: Option<ExtendedPointId>,
        limit: usize,
        with_payload_interface: &WithPayloadInterface,
        with_vector: &WithVector,
        filter: Option<&Filter>,
        hw_measurement_acc: HwMeasurementAcc,
    ) -> OperationResult<Vec<RecordInternal>> {
        let per_segment = self.par_map_segments(|segment| {
            segment.read_segment().read_filtered(
                offset,
                Some(limit),
                filter,
                &AtomicBool::new(false),
                &hw_measurement_acc.get_counter_cell(),
                DeferredBehavior::VisibleOnly,
            )
        })?;

        let point_ids: Vec<_> = per_segment
            .into_iter()
            .flatten()
            .sorted()
            .dedup()
            .take(limit)
            .collect();

        let mut points = retrieve_over(
            self.segment_arcs(),
            &point_ids,
            &WithPayload::from(with_payload_interface),
            with_vector,
            &AtomicBool::new(false),
            hw_measurement_acc,
            DeferredBehavior::VisibleOnly,
        )?;

        let ordered_points = point_ids
            .iter()
            .filter_map(|point_id| points.remove(point_id))
            .collect();

        Ok(ordered_points)
    }

    fn scroll_by_field(
        &self,
        limit: usize,
        with_payload_interface: &WithPayloadInterface,
        with_vector: &WithVector,
        filter: Option<&Filter>,
        order_by: &OrderBy,
        hw_measurement_acc: HwMeasurementAcc,
    ) -> OperationResult<Vec<RecordInternal>> {
        let read_results = self.par_map_segments(|segment| {
            segment.read_segment().read_ordered_filtered(
                Some(limit),
                filter,
                order_by,
                &AtomicBool::new(false),
                &hw_measurement_acc.get_counter_cell(),
                DeferredBehavior::VisibleOnly,
            )
        })?;

        let (order_values, point_ids): (Vec<_>, Vec<_>) = read_results
            .into_iter()
            .kmerge_by(|a, b| match order_by.direction() {
                Direction::Asc => a <= b,
                Direction::Desc => a >= b,
            })
            .dedup()
            .take(limit)
            .unzip();

        let points = retrieve_over(
            self.segment_arcs(),
            &point_ids,
            &WithPayload::from(with_payload_interface),
            with_vector,
            &AtomicBool::new(false),
            hw_measurement_acc,
            DeferredBehavior::VisibleOnly,
        )?;

        let ordered_points = point_ids
            .iter()
            .zip(order_values)
            .filter_map(|(point_id, value)| {
                let mut record = points.get(point_id).cloned()?;
                record.order_value = Some(value);
                Some(record)
            })
            .collect();

        Ok(ordered_points)
    }

    fn scroll_randomly(
        &self,
        limit: usize,
        with_payload_interface: &WithPayloadInterface,
        with_vector: &WithVector,
        filter: Option<&Filter>,
        hw_measurement_acc: HwMeasurementAcc,
    ) -> OperationResult<Vec<RecordInternal>> {
        let per_segment = self.par_map_segments(|segment| {
            let segment = segment.read_segment();

            let point_count = segment.available_point_count_without_deferred();
            let point_ids = segment.read_random_filtered(
                limit,
                filter,
                &AtomicBool::new(false),
                &hw_measurement_acc.get_counter_cell(),
            )?;

            OperationResult::Ok((point_count, point_ids))
        })?;

        let (point_count, mut point_ids): (Vec<_>, Vec<_>) = per_segment.into_iter().unzip();

        // Shortcut if all segments are empty
        if point_count.iter().all(|&count| count == 0) {
            return Ok(Vec::new());
        }

        // Select points in a weighted fashion from each segment, depending on how many points each segment has.
        let distribution = WeightedIndex::new(point_count).map_err(|err| {
            OperationError::service_error(format!(
                "failed to create weighted index for random scroll: {err:?}"
            ))
        })?;

        let mut rng = rand::make_rng::<StdRng>();
        // Grow on demand rather than pre-allocating `limit` buckets: an
        // unbounded/huge `limit` would otherwise request an enormous allocation
        // up front (a DoS / allocator-abort vector). The loop below is bounded
        // by the actual point count, so capacity grows only to what's inserted.
        let mut random_point_ids = HashSet::new();

        // Randomly sample points in two stages
        //
        // 1. This loop iterates <= LIMIT times, and either breaks early if we
        // have enough points, or if some of the segments are exhausted.
        //
        // 2. If the segments are exhausted, we will fill up the rest of the
        // points from other segments. In total, the complexity is guaranteed to
        // be O(limit).
        while random_point_ids.len() < limit {
            let segment_idx = rng.sample(&distribution);
            let segment_point_ids = &mut point_ids[segment_idx];

            if let Some(point) = segment_point_ids.pop() {
                random_point_ids.insert(point);
            } else {
                // It seems that some segments are empty early,
                // so distribution does not make sense anymore.
                // This is only possible if segments size < limit.
                break;
            }
        }

        // If we still need more points, we will get them from the rest of the segments.
        // This is a rare case, as it seems we don't have enough points in individual segments.
        // Therefore, we can ignore "proper" distribution, as it won't be accurate anyway.
        if random_point_ids.len() < limit {
            for point_id in point_ids.into_iter().flatten() {
                random_point_ids.insert(point_id);
                if random_point_ids.len() >= limit {
                    break;
                }
            }
        }

        let random_point_ids: Vec<_> = random_point_ids.into_iter().collect();

        let random_points = retrieve_over(
            self.segment_arcs(),
            &random_point_ids,
            &WithPayload::from(with_payload_interface),
            with_vector,
            &AtomicBool::new(false),
            hw_measurement_acc,
            DeferredBehavior::VisibleOnly,
        )?
        .into_values()
        .collect();

        Ok(random_points)
    }
}