qdrant-edge 0.8.0

pub mod immutable_bool_index;
pub mod mutable_bool_index;
pub mod read_only_bool_index;
mod read_ops;

use crate::common::counter::hardware_accumulator::HwMeasurementAcc;
use crate::common::counter::hardware_counter::HardwareCounterCell;
use crate::common::types::PointOffsetType;
use crate::common::universal_io::MmapFile;
pub use immutable_bool_index::ImmutableBoolIndex;
pub use mutable_bool_index::MutableBoolIndex;
pub use read_only_bool_index::ReadOnlyBoolIndex;
pub use read_ops::{BoolConditionChecker, BoolIndexRead};
use serde_json::Value as JsonValue;

use super::facet_index::FacetIndex;
use super::{PayloadFieldIndex, PayloadFieldIndexRead, ValueIndexer};
use crate::segment::common::flags::roaring_flags::RoaringFlags;
use crate::segment::common::operation_error::{OperationError, OperationResult};
use crate::segment::data_types::facets::{FacetHit, FacetValue, FacetValueRef};
use crate::segment::index::condition_checker::ConditionCheckerEnum;
use crate::segment::index::payload_config::IndexMutability;
use crate::segment::index::query_optimization::rescore_formula::value_retriever::VariableRetrieverFn;
use crate::segment::types::FieldCondition;

pub enum BoolIndex {
    Mutable(MutableBoolIndex),
    Immutable(ImmutableBoolIndex),
}

impl From<MutableBoolIndex> for BoolIndex {
    #[inline]
    fn from(index: MutableBoolIndex) -> Self {
        BoolIndex::Mutable(index)
    }
}

impl From<ImmutableBoolIndex> for BoolIndex {
    #[inline]
    fn from(index: ImmutableBoolIndex) -> Self {
        BoolIndex::Immutable(index)
    }
}

impl BoolIndex {
    pub fn get_mutability_type(&self) -> IndexMutability {
        match self {
            BoolIndex::Mutable(_) => IndexMutability::Mutable,
            BoolIndex::Immutable(_) => IndexMutability::Immutable,
        }
    }

    /// Produce a closure that maps a point id to its indexed bool
    /// values as JSON `Value`s. Used by `FieldIndex::value_retriever`.
    pub fn value_retriever<'a>(
        &'a self,
        hw_counter: &'a HardwareCounterCell,
    ) -> OperationResult<VariableRetrieverFn<'a>> {
        read_ops::value_retriever(self, hw_counter)
    }
}

impl BoolIndexRead for BoolIndex {
    type Flags = RoaringFlags<MmapFile>;

    fn trues_flags(&self) -> &Self::Flags {
        match self {
            BoolIndex::Mutable(index) => index.trues_flags(),
            BoolIndex::Immutable(index) => index.trues_flags(),
        }
    }

    fn falses_flags(&self) -> &Self::Flags {
        match self {
            BoolIndex::Mutable(index) => index.falses_flags(),
            BoolIndex::Immutable(index) => index.falses_flags(),
        }
    }

    fn indexed_count(&self) -> OperationResult<usize> {
        match self {
            BoolIndex::Mutable(index) => index.indexed_count(),
            BoolIndex::Immutable(index) => index.indexed_count(),
        }
    }

    fn telemetry_index_type(&self) -> &'static str {
        match self {
            BoolIndex::Mutable(index) => index.telemetry_index_type(),
            BoolIndex::Immutable(index) => index.telemetry_index_type(),
        }
    }

    fn trues_count(&self) -> OperationResult<usize> {
        match self {
            BoolIndex::Mutable(index) => index.trues_count(),
            BoolIndex::Immutable(index) => index.trues_count(),
        }
    }

    fn falses_count(&self) -> OperationResult<usize> {
        match self {
            BoolIndex::Mutable(index) => index.falses_count(),
            BoolIndex::Immutable(index) => index.falses_count(),
        }
    }
}

impl PayloadFieldIndexRead for BoolIndex {
    fn count_indexed_points(&self) -> OperationResult<usize> {
        self.indexed_count()
    }

    fn filter<'a>(
        &'a self,
        condition: &'a FieldCondition,
        hw_counter: &'a HardwareCounterCell,
    ) -> OperationResult<Option<Box<dyn Iterator<Item = PointOffsetType> + 'a>>> {
        read_ops::filter(self, condition, hw_counter)
    }

    fn estimate_cardinality(
        &self,
        condition: &FieldCondition,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<Option<super::CardinalityEstimation>> {
        read_ops::estimate_cardinality(self, condition, hw_counter)
    }

    fn for_each_payload_block(
        &self,
        threshold: usize,
        key: crate::segment::types::PayloadKeyType,
        f: &mut dyn FnMut(super::PayloadBlockCondition) -> OperationResult<()>,
    ) -> OperationResult<()> {
        read_ops::for_each_payload_block(self, threshold, key, f)
    }

    fn condition_checker<'a>(
        &'a self,
        condition: &FieldCondition,
        hw_acc: HwMeasurementAcc,
    ) -> OperationResult<Option<ConditionCheckerEnum<'a>>> {
        match self {
            BoolIndex::Mutable(index) => index.condition_checker(condition, hw_acc),
            BoolIndex::Immutable(index) => index.condition_checker(condition, hw_acc),
        }
    }
}

impl PayloadFieldIndex for BoolIndex {
    fn wipe(self) -> OperationResult<()> {
        match self {
            BoolIndex::Mutable(index) => index.wipe(),
            BoolIndex::Immutable(index) => index.wipe(),
        }
    }

    fn flusher(&self) -> crate::segment::common::Flusher {
        match self {
            BoolIndex::Mutable(index) => index.flusher(),
            BoolIndex::Immutable(index) => index.flusher(),
        }
    }

    fn files(&self) -> Vec<std::path::PathBuf> {
        BoolIndexRead::files(self)
    }

    fn immutable_files(&self) -> Vec<std::path::PathBuf> {
        match self {
            BoolIndex::Mutable(index) => index.immutable_files(),
            BoolIndex::Immutable(index) => index.immutable_files(),
        }
    }
}

impl FacetIndex for BoolIndex {
    fn unique_values_count(&self) -> usize {
        // A boolean field has at most two distinct values (true/false).
        // We could compute the actual number cheaply, but this upper bound is
        // enough for the facet-strategy decision (which only checks whether
        // the cardinality vastly exceeds the user limit).
        2
    }

    fn for_points_values(
        &self,
        points: impl Iterator<Item = PointOffsetType>,
        _hw_counter: &HardwareCounterCell,
        mut f: impl FnMut(PointOffsetType, &mut dyn Iterator<Item = FacetValueRef<'_>>),
    ) -> OperationResult<()> {
        for point_id in points {
            let values = self.get_point_values(point_id)?;
            f(point_id, &mut values.into_iter().map(FacetValueRef::Bool));
        }
        Ok(())
    }

    fn for_each_value(
        &self,
        mut f: impl FnMut(FacetValueRef<'_>) -> OperationResult<()>,
    ) -> OperationResult<()> {
        BoolIndexRead::iter_values(self)?.try_for_each(|v| f(FacetValueRef::Bool(v)))
    }

    fn for_each_value_map(
        &self,
        hw_counter: &HardwareCounterCell,
        mut f: impl FnMut(
            FacetValueRef<'_>,
            &mut dyn Iterator<Item = PointOffsetType>,
        ) -> OperationResult<()>,
    ) -> OperationResult<()> {
        BoolIndexRead::for_each_value_map(self, hw_counter, |value, iter| {
            f(FacetValueRef::Bool(value), iter)
        })
    }

    fn for_values_map(
        &self,
        values: impl Iterator<Item = FacetValue>,
        hw_counter: &HardwareCounterCell,
        mut f: impl FnMut(FacetValue, &mut dyn Iterator<Item = PointOffsetType>) -> OperationResult<()>,
    ) -> OperationResult<()> {
        let bools = values.filter_map(|value| match value {
            FacetValue::Bool(b) => Some(b),
            FacetValue::Keyword(_) | FacetValue::Int(_) | FacetValue::Uuid(_) => None,
        });
        BoolIndexRead::for_values_map(self, bools, hw_counter, |b, iter| {
            f(FacetValue::Bool(b), iter)
        })
    }

    fn for_each_count_per_value(
        &self,
        deferred_internal_id: Option<PointOffsetType>,
        mut f: impl FnMut(FacetHit<FacetValueRef<'_>>) -> OperationResult<()>,
    ) -> OperationResult<()> {
        BoolIndexRead::for_each_count_per_value(self, deferred_internal_id, |value, count| {
            f(FacetHit {
                value: FacetValueRef::Bool(value),
                count,
            })
        })
    }
}

impl ValueIndexer for BoolIndex {
    type ValueType = bool;

    fn add_many(
        &mut self,
        id: PointOffsetType,
        values: Vec<Self::ValueType>,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<()> {
        match self {
            BoolIndex::Mutable(index) => index.add_many(id, values, hw_counter),
            BoolIndex::Immutable(_) => Err(OperationError::service_error(
                "Can't add values to immutable bool index",
            )),
        }
    }

    fn get_value(value: &JsonValue) -> Option<Self::ValueType> {
        match value {
            JsonValue::Bool(value) => Some(*value),
            JsonValue::Null
            | JsonValue::Number(_)
            | JsonValue::String(_)
            | JsonValue::Array(_)
            | JsonValue::Object(_) => None,
        }
    }

    fn remove_point(&mut self, id: PointOffsetType) -> OperationResult<()> {
        match self {
            BoolIndex::Mutable(index) => index.remove_point(id),
            BoolIndex::Immutable(index) => index.remove_point(id),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::path::Path;

    use crate::common::counter::hardware_accumulator::HwMeasurementAcc;
    use crate::common::counter::hardware_counter::HardwareCounterCell;
    use itertools::Itertools;
    use rstest::rstest;
    use serde_json::json;
    use tempfile::Builder;

    use super::immutable_bool_index::{ImmutableBoolIndex, ImmutableBoolIndexBuilder};
    use super::mutable_bool_index::{MutableBoolIndex, MutableBoolIndexBuilder};
    use crate::segment::index::field_index::{FieldIndexBuilderTrait, PayloadFieldIndex, ValueIndexer};
    use crate::segment::json_path::JsonPath;

    const FIELD_NAME: &str = "bool_field";
    const DB_NAME: &str = "test_db";

    #[derive(Debug, Copy, Clone, PartialEq, Eq)]
    enum IndexType {
        Mutable,
        Immutable,
    }

    trait BuildableIndex: PayloadFieldIndex {
        type BuilderType: FieldIndexBuilderTrait<FieldIndexType = Self>;

        fn builder(path: &Path) -> Self::BuilderType;
        fn open_at(path: &Path) -> Self;
    }

    impl BuildableIndex for MutableBoolIndex {
        type BuilderType = MutableBoolIndexBuilder;

        fn builder(path: &Path) -> Self::BuilderType {
            MutableBoolIndex::builder(path).unwrap()
        }

        fn open_at(path: &Path) -> Self {
            MutableBoolIndex::builder(path)
                .unwrap()
                .make_empty()
                .unwrap()
        }
    }

    impl BuildableIndex for ImmutableBoolIndex {
        type BuilderType = ImmutableBoolIndexBuilder;

        fn builder(path: &Path) -> Self::BuilderType {
            ImmutableBoolIndex::builder(path).unwrap()
        }

        fn open_at(path: &Path) -> Self {
            let mutable_index = MutableBoolIndex::builder(path)
                .unwrap()
                .make_empty()
                .unwrap();
            ImmutableBoolIndex::from_mutable(mutable_index).unwrap()
        }
    }

    fn match_bool(value: bool) -> crate::segment::types::FieldCondition {
        crate::segment::types::FieldCondition::new_match(
            JsonPath::new(FIELD_NAME),
            crate::segment::types::Match::Value(crate::segment::types::MatchValue {
                value: crate::segment::types::ValueVariants::Bool(value),
            }),
        )
    }

    fn bools_fixture() -> Vec<serde_json::Value> {
        vec![
            json!(true),
            json!(false),
            json!([true, false]),
            json!([false, true]),
            json!([true, true]),
            json!([false, false]),
            json!([true, false, true]),
            serde_json::Value::Null,
            json!(1),
            json!("test"),
            json!([false]),
            json!([true]),
        ]
    }

    fn filter<I: BuildableIndex>(given: serde_json::Value, match_on: bool, expected_count: usize) {
        let tmp_dir = Builder::new().prefix(DB_NAME).tempdir().unwrap();
        let mut builder = I::builder(tmp_dir.path());

        let hw_counter = HardwareCounterCell::new();

        builder.add_point(0, &[&given], &hw_counter).unwrap();

        let hw_acc = HwMeasurementAcc::new();
        let hw_counter = hw_acc.get_counter_cell();
        let index = builder.finalize().unwrap();
        let count = index
            .filter(&match_bool(match_on), &hw_counter)
            .unwrap()
            .unwrap()
            .count();

        assert_eq!(count, expected_count);
    }

    #[rstest]
    #[case(json!(true), 1)]
    #[case(json!(false), 0)]
    #[case(json!([true]), 1)]
    #[case(json!([false]), 0)]
    #[case(json!([true, false]), 1)]
    #[case(json!([false, true]), 1)]
    #[case(json!([false, false]), 0)]
    #[case(json!([true, true]), 1)]
    fn test_filter_true(
        #[case] given: serde_json::Value,
        #[case] expected_count: usize,
        #[values(IndexType::Mutable, IndexType::Immutable)] index_type: IndexType,
    ) {
        match index_type {
            IndexType::Mutable => filter::<MutableBoolIndex>(given, true, expected_count),
            IndexType::Immutable => filter::<ImmutableBoolIndex>(given, true, expected_count),
        }
    }

    #[rstest]
    #[case(json!(true), 0)]
    #[case(json!(false), 1)]
    #[case(json!([true]), 0)]
    #[case(json!([false]), 1)]
    #[case(json!([true, false]), 1)]
    #[case(json!([false, true]), 1)]
    #[case(json!([false, false]), 1)]
    #[case(json!([true, true]), 0)]
    fn test_filter_false(
        #[case] given: serde_json::Value,
        #[case] expected_count: usize,
        #[values(IndexType::Mutable, IndexType::Immutable)] index_type: IndexType,
    ) {
        match index_type {
            IndexType::Mutable => filter::<MutableBoolIndex>(given.clone(), false, expected_count),
            IndexType::Immutable => {
                filter::<ImmutableBoolIndex>(given.clone(), false, expected_count)
            }
        }
    }

    #[rstest]
    fn test_load_from_disk(
        #[values(IndexType::Mutable, IndexType::Immutable)] index_type: IndexType,
    ) {
        match index_type {
            IndexType::Mutable => load_from_disk::<MutableBoolIndex>(),
            IndexType::Immutable => load_from_disk::<ImmutableBoolIndex>(),
        }
    }

    fn load_from_disk<I: BuildableIndex>() {
        let tmp_dir = Builder::new().prefix(DB_NAME).tempdir().unwrap();
        let mut builder = I::builder(tmp_dir.path());

        let hw_counter = HardwareCounterCell::new();

        bools_fixture()
            .into_iter()
            .enumerate()
            .for_each(|(i, value)| {
                builder.add_point(i as u32, &[&value], &hw_counter).unwrap();
            });

        let index = builder.finalize().unwrap();
        index.flusher()().unwrap();
        drop(index);

        let new_index = I::open_at(tmp_dir.path());

        let hw_acc = HwMeasurementAcc::new();
        let hw_counter = hw_acc.get_counter_cell();
        let point_offsets = new_index
            .filter(&match_bool(false), &hw_counter)
            .unwrap()
            .unwrap()
            .collect_vec();
        assert_eq!(point_offsets, vec![1, 2, 3, 5, 6, 10]);

        let point_offsets = new_index
            .filter(&match_bool(true), &hw_counter)
            .unwrap()
            .unwrap()
            .collect_vec();
        assert_eq!(point_offsets, vec![0, 2, 3, 4, 6, 11]);

        assert_eq!(new_index.count_indexed_points().unwrap(), 9);
    }

    #[rstest]
    #[case(json!(false), json!(true))]
    #[case(json!([false, true]), json!(true))]
    fn test_modify_value(#[case] before: serde_json::Value, #[case] after: serde_json::Value) {
        modify_value::<MutableBoolIndex>(before, after);
    }

    /// Try to modify from falsy to only true
    fn modify_value<I: BuildableIndex + ValueIndexer>(
        before: serde_json::Value,
        after: serde_json::Value,
    ) {
        let tmp_dir = Builder::new().prefix(DB_NAME).tempdir().unwrap();
        let mut index = I::open_at(tmp_dir.path());

        let hw_cell = HardwareCounterCell::new();

        let idx = 1000;
        index.add_point(idx, &[&before], &hw_cell).unwrap();

        let hw_acc = HwMeasurementAcc::new();
        let hw_counter = hw_acc.get_counter_cell();

        let point_offsets = index
            .filter(&match_bool(false), &hw_counter)
            .unwrap()
            .unwrap()
            .collect_vec();
        assert_eq!(point_offsets, vec![idx]);

        index.add_point(idx, &[&after], &hw_cell).unwrap();

        let point_offsets = index
            .filter(&match_bool(true), &hw_counter)
            .unwrap()
            .unwrap()
            .collect_vec();
        assert_eq!(point_offsets, vec![idx]);
        let point_offsets = index
            .filter(&match_bool(false), &hw_counter)
            .unwrap()
            .unwrap()
            .collect_vec();
        assert!(point_offsets.is_empty());
    }

    #[rstest]
    fn test_indexed_count(
        #[values(IndexType::Mutable, IndexType::Immutable)] index_type: IndexType,
    ) {
        match index_type {
            IndexType::Mutable => indexed_count::<MutableBoolIndex>(),
            IndexType::Immutable => indexed_count::<ImmutableBoolIndex>(),
        }
    }

    fn indexed_count<I: BuildableIndex + PayloadFieldIndex>() {
        let tmp_dir = Builder::new().prefix(DB_NAME).tempdir().unwrap();
        let mut builder = I::builder(tmp_dir.path());

        let hw_counter = HardwareCounterCell::new();

        bools_fixture()
            .into_iter()
            .enumerate()
            .for_each(|(i, value)| {
                builder.add_point(i as u32, &[&value], &hw_counter).unwrap();
            });

        let index = builder.finalize().unwrap();

        assert_eq!(index.count_indexed_points().unwrap(), 9);
    }

    #[test]
    fn test_payload_blocks() {
        payload_blocks::<MutableBoolIndex>();
    }

    fn payload_blocks<I: BuildableIndex + ValueIndexer>() {
        let tmp_dir = Builder::new().prefix(DB_NAME).tempdir().unwrap();
        let mut index = I::open_at(tmp_dir.path());

        let hw_counter = HardwareCounterCell::new();

        bools_fixture()
            .into_iter()
            .enumerate()
            .for_each(|(i, value)| {
                index.add_point(i as u32, &[&value], &hw_counter).unwrap();
            });

        let mut blocks = Vec::new();
        index
            .for_each_payload_block(0, JsonPath::new(FIELD_NAME), &mut |block| {
                blocks.push(block);
                Ok(())
            })
            .unwrap();
        assert_eq!(blocks.len(), 2);
        assert_eq!(blocks[0].cardinality, 6);
        assert_eq!(blocks[1].cardinality, 6);
    }

    #[rstest]
    fn test_estimate_cardinality(
        #[values(IndexType::Mutable, IndexType::Immutable)] index_type: IndexType,
    ) {
        match index_type {
            IndexType::Mutable => estimate_cardinality::<MutableBoolIndex>(),
            IndexType::Immutable => estimate_cardinality::<ImmutableBoolIndex>(),
        }
    }

    fn estimate_cardinality<I: BuildableIndex>() {
        let tmp_dir = Builder::new().prefix(DB_NAME).tempdir().unwrap();
        let mut builder = I::builder(tmp_dir.path());

        let hw_counter = HardwareCounterCell::new();

        bools_fixture()
            .into_iter()
            .enumerate()
            .for_each(|(i, value)| {
                builder.add_point(i as u32, &[&value], &hw_counter).unwrap();
            });

        let hw_counter = HardwareCounterCell::new();

        let index = builder.finalize().unwrap();
        let cardinality = index
            .estimate_cardinality(&match_bool(true), &hw_counter)
            .unwrap()
            .unwrap();
        assert_eq!(cardinality.exp, 6);

        let cardinality = index
            .estimate_cardinality(&match_bool(false), &hw_counter)
            .unwrap()
            .unwrap();
        assert_eq!(cardinality.exp, 6);
    }
}