qdrant-edge 0.8.0

use std::collections::{BTreeSet, HashMap};
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;

use ahash::AHashMap;
use crate::common::counter::hardware_counter::HardwareCounterCell;
use crate::common::types::{DeferredBehavior, TelemetryDetail};
use uuid::Uuid;

use crate::segment::common::Flusher;
use crate::segment::common::operation_error::{OperationError, OperationResult, SegmentFailedState};
use crate::segment::data_types::build_index_result::BuildFieldIndexResult;
use crate::segment::data_types::facets::{FacetParams, FacetValue};
use crate::segment::data_types::named_vectors::NamedVectors;
use crate::segment::data_types::order_by::{OrderBy, OrderValue};
use crate::segment::data_types::query_context::{FormulaContext, QueryContext, SegmentQueryContext};
use crate::segment::data_types::segment_record::{SegmentRecord, SegmentRecordRaw};
use crate::segment::data_types::vector_name_config::VectorNameConfig;
use crate::segment::data_types::vectors::{QueryVector, VectorInternal};
use crate::segment::entry::snapshot_entry::SnapshotEntry;
use crate::segment::index::field_index::{CardinalityEstimation, FieldIndex};
use crate::segment::json_path::JsonPath;
use crate::segment::telemetry::SegmentTelemetry;
use crate::segment::types::{
    ExtendedPointId, Filter, Payload, PayloadFieldSchema, PayloadKeyType, PayloadKeyTypeRef,
    PointIdType, ScoredPoint, SearchParams, SegmentConfig, SegmentInfo, SegmentType, SeqNumberType,
    VectorName, VectorNameBuf, WithPayload, WithVector,
};

/// Define all operations on segment that do not require mutable access.
///
/// Assume all operations are idempotent - which means that no matter how many times an operation
/// is executed - the storage state will be the same.
pub trait ReadSegmentEntry {
    fn is_proxy(&self) -> bool;

    /// Get version of specified point
    ///
    /// Returns `None` if point does not exist or is soft-deleted.
    fn point_version(&self, point_id: PointIdType) -> Option<SeqNumberType>;

    #[allow(clippy::too_many_arguments)]
    fn search_batch(
        &self,
        vector_name: &VectorName,
        query_vectors: &[&QueryVector],
        with_payload: &WithPayload,
        with_vector: &WithVector,
        filter: Option<&Filter>,
        top: usize,
        params: Option<&SearchParams>,
        query_context: &SegmentQueryContext,
    ) -> OperationResult<Vec<Vec<ScoredPoint>>>;

    /// Rescore results with a formula that can reference payload values.
    ///
    /// A deleted bitslice is passed to exclude points from a wrapped segment.
    fn rescore_with_formula(
        &self,
        formula_ctx: Arc<FormulaContext>,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<Vec<ScoredPoint>>;

    fn vector(
        &self,
        vector_name: &VectorName,
        point_id: PointIdType,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<Option<VectorInternal>>;

    /// Like [`ReadSegmentEntry::vector`], but with explicit deferred semantics.
    ///
    /// With [`DeferredBehavior::WithDeferred`] this resolves the latest head of
    /// the point, including a deferred head that is invisible to ordinary reads.
    fn vector_with_behavior(
        &self,
        vector_name: &VectorName,
        point_id: PointIdType,
        deferred_behavior: DeferredBehavior,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<Option<VectorInternal>>;

    fn all_vectors(
        &self,
        point_id: PointIdType,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<NamedVectors<'_>>;

    /// Reads Records from the segment, according to specified selectors and a list of point ids.
    ///
    /// WARNING:
    /// This function may return fewer records than requested, if some points are not found.
    /// Order of returned records is not guaranteed to match order of requested point ids.
    fn retrieve(
        &self,
        point_ids: &[PointIdType],
        with_payload: &WithPayload,
        with_vector: &WithVector,
        hw_counter: &HardwareCounterCell,
        is_stopped: &AtomicBool,
        deferred_behavior: DeferredBehavior,
    ) -> OperationResult<AHashMap<ExtendedPointId, SegmentRecord>>;

    /// Byte-blob analogue of [`ReadSegmentEntry::retrieve`]: returns vectors as
    /// storage-native bytes ([`SegmentRecordRaw`]) to avoid a lossy round-trip
    /// when relocating points (copy-on-write moves, shard transfer).
    ///
    /// Like `retrieve`, may return fewer records than requested and in any order.
    fn retrieve_raw(
        &self,
        point_ids: &[PointIdType],
        with_payload: &WithPayload,
        with_vector: &WithVector,
        hw_counter: &HardwareCounterCell,
        is_stopped: &AtomicBool,
        deferred_behavior: DeferredBehavior,
    ) -> OperationResult<AHashMap<ExtendedPointId, SegmentRecordRaw>>;

    /// Retrieve payload for the point
    /// If not found, return empty payload
    fn payload(
        &self,
        point_id: PointIdType,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<Payload>;

    /// Paginate over points which satisfies filtering condition starting with `offset` id including.
    ///
    /// Cancelled by `is_stopped` flag.
    fn read_filtered(
        &self,
        offset: Option<PointIdType>,
        limit: Option<usize>,
        filter: Option<&Filter>,
        is_stopped: &AtomicBool,
        hw_counter: &HardwareCounterCell,
        deferred_behavior: DeferredBehavior,
    ) -> OperationResult<Vec<PointIdType>>;

    /// Return points which satisfies filtering condition ordered by the `order_by.key` field,
    /// starting with `order_by.start_from` value including.
    ///
    /// Will fail if there is no index for the order_by key.
    /// Cancelled by `is_stopped` flag.
    fn read_ordered_filtered<'a>(
        &'a self,
        limit: Option<usize>,
        filter: Option<&'a Filter>,
        order_by: &'a OrderBy,
        is_stopped: &AtomicBool,
        hw_counter: &HardwareCounterCell,
        deferred_behavior: DeferredBehavior,
    ) -> OperationResult<Vec<(OrderValue, PointIdType)>>;

    /// Return random points which satisfies filtering condition.
    ///
    /// Cancelled by `is_stopped` flag.
    fn read_random_filtered(
        &self,
        limit: usize,
        filter: Option<&Filter>,
        is_stopped: &AtomicBool,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<Vec<PointIdType>>;

    /// Read points in [from; to) range
    fn read_range(&self, from: Option<PointIdType>, to: Option<PointIdType>) -> Vec<PointIdType>;

    /// Return all unique values for the given key.
    fn unique_values(
        &self,
        key: &JsonPath,
        filter: Option<&Filter>,
        is_stopped: &AtomicBool,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<BTreeSet<FacetValue>>;

    /// Return the largest counts for the given facet request.
    fn facet(
        &self,
        request: &FacetParams,
        is_stopped: &AtomicBool,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<HashMap<FacetValue, usize>>;

    /// Check if there is point with `point_id` in this segment.
    ///
    /// Soft deleted points are excluded. `deferred_behavior` selects whether a
    /// deferred-only point counts as present.
    fn has_point(&self, point_id: PointIdType, deferred_behavior: DeferredBehavior) -> bool;

    /// Estimate available point count in this segment for given filter.
    fn estimate_point_count<'a>(
        &'a self,
        filter: Option<&'a Filter>,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<CardinalityEstimation>;

    /// Names of all vectors in this segment, sorted.
    fn vector_names(&self) -> Vec<VectorNameBuf>;

    /// Whether this segment is completely empty in terms of points
    ///
    /// The segment is considered to not be empty if it contains any points, even if deleted.
    /// Deleted points still have a version which may be important at time of recovery. Deciding
    /// this by just the reported point count is not reliable in case a proxy segment is used.
    ///
    /// Payload indices or type of storage are not considered here.
    fn is_empty(&self) -> bool;

    /// Number of available points
    ///
    /// - excludes soft deleted points
    /// - includes deferred points.
    fn available_point_count(&self) -> usize;

    /// Number of deleted points
    fn deleted_point_count(&self) -> usize;

    /// Similar to `available_point_count()` but excludes all deferred points.
    fn available_point_count_without_deferred(&self) -> usize;

    /// Size of all available vectors in storage
    fn available_vectors_size_in_bytes(&self, vector_name: &VectorName) -> OperationResult<usize>;

    /// Max value from all `available_vectors_size_in_bytes`
    fn max_available_vectors_size_in_bytes(&self) -> OperationResult<usize> {
        let mut max_size = 0;
        for vector_name in self.vector_names() {
            let inner_size = self.available_vectors_size_in_bytes(&vector_name)?;
            max_size = std::cmp::max(max_size, inner_size);
        }
        Ok(max_size)
    }

    /// Get segment uuid
    fn segment_uuid(&self) -> Uuid;

    /// Get segment type
    fn segment_type(&self) -> SegmentType;

    /// Get current stats of the segment
    fn info(&self) -> OperationResult<SegmentInfo>;

    /// Get size related stats of the segment.
    /// This returns `SegmentInfo` with some non size-related data (like `schema`) unset to improve performance.
    fn size_info(&self) -> SegmentInfo;

    /// Get segment configuration
    fn config(&self) -> &SegmentConfig;

    /// Whether this segment is appendable
    ///
    /// Returns appendable state of outer most segment. If this is a proxy segment, this shadows
    /// the appendable state of the wrapped segment.
    fn is_appendable(&self) -> bool;

    /// Get indexed fields
    fn get_indexed_fields(&self) -> HashMap<PayloadKeyType, PayloadFieldSchema>;

    // Get collected telemetry data of segment
    fn get_telemetry_data(&self, detail: TelemetryDetail) -> OperationResult<SegmentTelemetry>;

    fn fill_query_context(&self, query_context: &mut QueryContext) -> OperationResult<()>;

    /// Check whether the point is marked as deferred in the segment
    fn point_is_deferred(&self, point_id: PointIdType) -> bool;

    /// Returns external IDs of all deferred points in the segment
    fn deferred_point_ids(&self) -> Vec<PointIdType>;

    /// Returns the amount of non-deleted deferred points.
    ///
    /// Note: This value can return `0` with `has_deferred_points()` returning `true`.
    /// This is because this function returns the *non-deleted* deferred points.
    fn deferred_point_count(&self) -> usize;

    /// Returns `true` if there is at least one point that is hidden (deferred).
    /// Non-appendable segments always return `false` as they can't have deferred points.
    ///
    /// Note: the deferred point can be deleted and this function would still return `true`.
    fn has_deferred_points(&self) -> bool;
}

/// Segment with storage.
pub trait StorageSegmentEntry: ReadSegmentEntry + SnapshotEntry {
    /// Get current update version of the segment
    fn version(&self) -> SeqNumberType;

    /// Checks if segment errored during last operations
    fn check_error(&self) -> Option<SegmentFailedState>;

    /// Get current persistent version of the segment
    fn persistent_version(&self) -> SeqNumberType;

    /// Returns a function, which when called, will flush all pending changes to disk.
    /// If there are currently no changes to flush, returns None.
    /// If `force` is true, will return a flusher even if there are no changes to flush.
    fn flusher(&self, force: bool) -> Option<Flusher>;

    /// Immediately flush all changes to disk and return persisted version.
    /// Blocks the current thread.
    fn flush(&self, force: bool) -> OperationResult<SeqNumberType> {
        if let Some(flusher) = self.flusher(force) {
            flusher()?;
        }
        Ok(self.persistent_version())
    }

    /// Removes all persisted data and forces to destroy segment
    fn drop_data(self) -> OperationResult<()>;

    /// Path to data, owned by segment
    fn data_path(&self) -> PathBuf;
}

/// Define all operations which can be performed with non-appendable Segment or Segment-like entity.
///
/// Assume all operations are idempotent - which means that no matter how many times an operation
/// is executed - the storage state will be the same.
pub trait NonAppendableSegmentEntry: StorageSegmentEntry {
    fn delete_point(
        &mut self,
        op_num: SeqNumberType,
        point_id: PointIdType,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<bool>;

    /// Delete field index, if exists
    fn delete_field_index(
        &mut self,
        op_num: SeqNumberType,
        key: PayloadKeyTypeRef,
    ) -> OperationResult<bool>;

    /// Delete field index, if exists and doesn't match the schema
    fn delete_field_index_if_incompatible(
        &mut self,
        op_num: SeqNumberType,
        key: PayloadKeyTypeRef,
        field_schema: &PayloadFieldSchema,
    ) -> OperationResult<bool>;

    /// Build the field index for the key and schema, if not built before.
    fn build_field_index(
        &self,
        op_num: SeqNumberType,
        key: PayloadKeyTypeRef,
        field_type: &PayloadFieldSchema,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<BuildFieldIndexResult>;

    /// Apply a built index. Returns whether it was actually applied or not.
    fn apply_field_index(
        &mut self,
        op_num: SeqNumberType,
        key: PayloadKeyType,
        field_schema: PayloadFieldSchema,
        field_index: Vec<FieldIndex>,
    ) -> OperationResult<bool>;

    /// Create index for a payload field, if not exists
    fn create_field_index(
        &mut self,
        op_num: SeqNumberType,
        key: PayloadKeyTypeRef,
        field_schema: Option<&PayloadFieldSchema>,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<bool> {
        let Some(field_schema) = field_schema else {
            // Legacy case, where we tried to automatically detect the schema for the field.
            // We don't do this anymore, as it is not reliable.
            return Err(OperationError::TypeInferenceError {
                field_name: key.clone(),
            });
        };

        self.delete_field_index_if_incompatible(op_num, key, field_schema)?;

        let (schema, indexes) =
            match self.build_field_index(op_num, key, field_schema, hw_counter)? {
                BuildFieldIndexResult::SkippedByVersion => {
                    return Ok(false);
                }
                BuildFieldIndexResult::AlreadyExists => {
                    return Ok(false);
                }
                BuildFieldIndexResult::IncompatibleSchema => {
                    // This is a service error, as we should have just removed the old index
                    // So it should not be possible to get this error
                    return Err(OperationError::service_error(format!(
                        "Incompatible schema for field index on field {key}",
                    )));
                }
                BuildFieldIndexResult::Built { schema, indexes } => (schema, indexes),
            };

        self.apply_field_index(op_num, key.to_owned(), schema, indexes)
    }

    /// Create a new named vector in the segment.
    /// For appendable segments: creates a real, writable vector storage + plain index.
    /// For immutable segments: creates a placeholder (empty) vector storage.
    /// Returns Ok(false) if the vector already exists (idempotent).
    fn create_vector_name(
        &mut self,
        op_num: SeqNumberType,
        vector_name: &VectorName,
        vector_config: &VectorNameConfig,
    ) -> OperationResult<bool>;

    /// Delete a named vector from the segment.
    /// Removes vector storage, index, and quantization data.
    /// Removes the vector from segment config.
    /// Returns Ok(false) if the vector does not exist (idempotent).
    fn delete_vector_name(
        &mut self,
        op_num: SeqNumberType,
        vector_name: &VectorName,
    ) -> OperationResult<bool>;
}

/// Define mutable operations which can be performed with Segment or Segment-like entity.
///
/// Assume all operations are idempotent - which means that no matter how many times an operation
/// is executed - the storage state will be the same.
pub trait SegmentEntry: NonAppendableSegmentEntry {
    fn upsert_point(
        &mut self,
        op_num: SeqNumberType,
        point_id: PointIdType,
        vectors: NamedVectors,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<bool>;

    /// Byte-blob analogue of [`SegmentEntry::upsert_point`]: vector values are
    /// storage-native bytes in the exact form returned by
    /// [`ReadSegmentEntry::retrieve_raw`], so requantized (e.g. TurboQuant)
    /// vectors relocate without a lossy decode/re-encode round-trip.
    ///
    /// The bytes carry no encoding/version tag: the target segment must have
    /// the same vector configuration (kind, datatype, dim) as the source. The
    /// bytes are inserted as-is, without preprocessing — they were already
    /// preprocessed (e.g. cosine-normalized) when first ingested.
    fn upsert_point_raw(
        &mut self,
        op_num: SeqNumberType,
        point_id: PointIdType,
        vectors: &[(VectorNameBuf, Vec<u8>)],
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<bool>;

    /// Upsert a complete point in a single operation: storage-native raw
    /// vectors (the [`ReadSegmentEntry::retrieve_raw`] form, same contract as
    /// [`SegmentEntry::upsert_point_raw`]), decoded vectors overriding them
    /// name-by-name, and the full payload. Named vectors present in neither
    /// list are deleted.
    ///
    /// This is the copy-on-write move primitive: it is equivalent to
    /// `upsert_point_raw` + `update_vectors` + `set_full_payload`, but writes
    /// the point once. On append-only segments each of those steps clones the
    /// whole point to a fresh internal id, so issuing them separately turns
    /// one moved point into a chain of immediately-dead slots.
    fn upsert_moved_point(
        &mut self,
        op_num: SeqNumberType,
        point_id: PointIdType,
        raw_vectors: &[(VectorNameBuf, Vec<u8>)],
        updated_vectors: NamedVectors,
        payload: &Payload,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<bool>;

    fn update_vectors(
        &mut self,
        op_num: SeqNumberType,
        point_id: PointIdType,
        vectors: NamedVectors,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<bool>;

    fn delete_vector(
        &mut self,
        op_num: SeqNumberType,
        point_id: PointIdType,
        vector_name: &VectorName,
    ) -> OperationResult<bool>;

    fn set_payload(
        &mut self,
        op_num: SeqNumberType,
        point_id: PointIdType,
        payload: &Payload,
        key: &Option<JsonPath>,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<bool>;

    fn set_full_payload(
        &mut self,
        op_num: SeqNumberType,
        point_id: PointIdType,
        full_payload: &Payload,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<bool>;

    fn delete_payload(
        &mut self,
        op_num: SeqNumberType,
        point_id: PointIdType,
        key: PayloadKeyTypeRef,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<bool>;

    fn clear_payload(
        &mut self,
        op_num: SeqNumberType,
        point_id: PointIdType,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<bool>;
}