WIP: query_points in struct payload index

lib/segment/src/index/field_index/field_index.rs

@@ -6,7 +6,7 @@ use crate::types::{FieldCondition, FloatPayloadType, IntPayloadType, PayloadType
 
 pub trait PayloadFieldIndex {
     /// Get iterator over points fitting given `condition`
-    fn filter(&self, condition: &FieldCondition) -> Box<dyn Iterator<Item=PointOffsetType> + '_>;
+    fn filter(&self, condition: &FieldCondition) -> Option<Box<dyn Iterator<Item=PointOffsetType> + '_>>;
 
     fn estimate_cardinality(&self, condition: &FieldCondition) -> Option<CardinalityEstimation>;
 }
@@ -39,7 +39,7 @@ impl FieldIndex {
 
 impl PayloadFieldIndex for FieldIndex {
 
-    fn filter(&self, condition: &FieldCondition) -> Box<dyn Iterator<Item=usize> + '_> {
+    fn filter(&self, condition: &FieldCondition) -> Option<Box<dyn Iterator<Item=PointOffsetType> + '_>> {
         self.get_payload_field_index().filter(condition)
     }
 

lib/segment/src/index/field_index/map_index.rs

@@ -57,7 +57,7 @@ impl<N: Hash + Eq + Clone> PersistedMapIndex<N> {
 }
 
 impl PayloadFieldIndex for PersistedMapIndex<String> {
-    fn filter(&self, condition: &FieldCondition) -> Box<dyn Iterator<Item=usize>> {
+    fn filter(&self, condition: &FieldCondition) -> Option<Box<dyn Iterator<Item=PointOffsetType>>> {
         unimplemented!()
     }
 
@@ -76,7 +76,7 @@ impl PayloadFieldIndex for PersistedMapIndex<String> {
 }
 
 impl PayloadFieldIndex for PersistedMapIndex<IntPayloadType> {
-    fn filter(&self, condition: &FieldCondition) -> Box<dyn Iterator<Item=usize>> {
+    fn filter(&self, condition: &FieldCondition) -> Option<Box<dyn Iterator<Item=PointOffsetType>>> {
         unimplemented!()
     }
 

lib/segment/src/index/field_index/mod.rs

@@ -1,6 +1,6 @@
 
 
-use crate::types::{FieldCondition, PointIdType};
+use crate::types::{FieldCondition, PointOffsetType};
 use std::collections::HashSet;
 
 pub mod numeric_index;
@@ -12,7 +12,7 @@ pub mod index_selector;
 #[derive(Debug, Clone)]
 pub enum PrimaryCondition {
     Condition(FieldCondition),
-    Ids(HashSet<PointIdType>)
+    Ids(HashSet<PointOffsetType>)
 }
 
 #[derive(Debug)]

lib/segment/src/index/field_index/numeric_index.rs

@@ -100,7 +100,7 @@ impl<N: ToPrimitive + Clone> PersistedNumericIndex<N> {
 
 
 impl<N: ToPrimitive + Clone> PayloadFieldIndex for PersistedNumericIndex<N> {
-    fn filter(&self, condition: &FieldCondition) -> Box<dyn Iterator<Item=usize>> {
+    fn filter(&self, condition: &FieldCondition) -> Option<Box<dyn Iterator<Item=PointOffsetType>>> {
         unimplemented!()
     }
 

lib/segment/src/index/index.rs

@@ -15,12 +15,7 @@ pub trait Index {
 
 
     /// Force internal index rebuild.
-    fn build_index(&mut self) -> OperationResult<()> ;
-}
-
-pub trait QueryEstimator {
-    /// Estimate amount of points (min, max) which satisfies filtering condition.
-    fn estimate_cardinality(&self, query: &Filter) -> CardinalityEstimation;
+    fn build_index(&mut self) -> OperationResult<()>;
 }
 
 pub trait PayloadIndex {
@@ -33,6 +28,9 @@ pub trait PayloadIndex {
     /// Remove index
     fn drop_index(&mut self, field: &PayloadKeyType) -> OperationResult<()>;
 
+    /// Estimate amount of points (min, max) which satisfies filtering condition.
+    fn estimate_cardinality(&self, query: &Filter) -> CardinalityEstimation;
+
     /// Return list of all point ids, which satisfy filtering criteria
-    fn query_points(&self, query: &Filter) -> Vec<PointOffsetType>;
+    fn query_points(&self, query: &Filter) -> Box<dyn Iterator<Item=PointOffsetType> + '_>;
 }

lib/segment/src/index/plain_payload_index.rs

@@ -9,6 +9,9 @@ use crate::entry::entry_point::OperationResult;
 use crate::index::payload_config::PayloadConfig;
 use std::path::{Path, PathBuf};
 use std::fs::create_dir_all;
+use crate::index::field_index::CardinalityEstimation;
+use itertools::Itertools;
+
 
 pub struct PlainPayloadIndex {
     condition_checker: Arc<AtomicRefCell<dyn ConditionChecker>>,
@@ -87,7 +90,23 @@ impl PayloadIndex for PlainPayloadIndex {
         self.save_config()
     }
 
-    fn query_points(&self, query: &Filter) -> Vec<PointOffsetType> {
+    fn estimate_cardinality(&self, query: &Filter) -> CardinalityEstimation {
+        let mut matched_points = 0;
+        let condition_checker = self.condition_checker.borrow();
+        for i in self.vector_storage.borrow().iter_ids() {
+            if condition_checker.check(i, query) {
+                matched_points += 1;
+            }
+        }
+        CardinalityEstimation {
+            primary_clauses: vec![],
+            min: matched_points,
+            exp: matched_points,
+            max: matched_points
+        }
+    }
+
+    fn query_points(&self, query: &Filter) -> Box<dyn Iterator<Item=PointOffsetType> + '_> {
         let mut matched_points = vec![];
         let condition_checker = self.condition_checker.borrow();
         for i in self.vector_storage.borrow().iter_ids() {
@@ -95,7 +114,7 @@ impl PayloadIndex for PlainPayloadIndex {
                 matched_points.push(i);
             }
         }
-        return matched_points;
+        return Box::new(matched_points.into_iter());
     }
 }
 
@@ -131,7 +150,7 @@ impl Index for PlainIndex {
     ) -> Vec<ScoredPointOffset> {
         match filter {
             Some(filter) => {
-                let filtered_ids = self.payload_index.borrow().query_points(filter);
+                let filtered_ids = self.payload_index.borrow().query_points(filter).collect_vec();
                 self.vector_storage.borrow().score_points(vector, &filtered_ids, top, &self.distance)
             }
             None => self.vector_storage.borrow().score_all(vector, top, &self.distance)

lib/segment/src/index/query_estimator.rs

@@ -1,4 +1,4 @@
-use crate::index::index::{QueryEstimator, PayloadIndex};
+use crate::index::index::PayloadIndex;
 use crate::types::{Filter, Condition};
 use crate::index::field_index::{CardinalityEstimation, PrimaryCondition};
 use crate::index::struct_payload_index::StructPayloadIndex;
@@ -120,32 +120,3 @@ fn estimate_must_not<F>(estimator: &F, conditions: &Vec<Condition>, total: usize
     let must_not_estimations = conditions.iter().map(estimate).collect_vec();
     combine_must_estimations(&must_not_estimations, total)
 }
-
-
-impl QueryEstimator for StructPayloadIndex {
-    fn estimate_cardinality(&self, query: &Filter) -> CardinalityEstimation {
-        let total = self.total_points();
-
-        let estimator = |condition: &Condition| {
-            match condition {
-                Condition::Filter(_) => panic!("Unexpected branching"),
-                Condition::HasId(ids) => CardinalityEstimation {
-                    primary_clauses: vec![PrimaryCondition::Ids(ids.clone())],
-                    min: 0,
-                    exp: ids.len(),
-                    max: ids.len(),
-                },
-                Condition::Field(field_condition) => self
-                    .estimate_field_condition(field_condition)
-                    .unwrap_or(CardinalityEstimation {
-                        primary_clauses: vec![],
-                        min: 0,
-                        exp: self.total_points() / 2,
-                        max: self.total_points(),
-                    }),
-            }
-        };
-
-        estimate_filter(&estimator, query, total)
-    }
-}

lib/segment/src/index/struct_payload_index.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::fs::{create_dir_all, File, remove_file};
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
@@ -13,8 +13,12 @@ use crate::index::field_index::index_selector::index_selector;
 use crate::index::index::PayloadIndex;
 use crate::index::payload_config::PayloadConfig;
 use crate::payload_storage::payload_storage::{ConditionChecker, PayloadStorage};
-use crate::types::{Filter, PayloadKeyType, FieldCondition};
-use crate::index::field_index::CardinalityEstimation;
+use crate::types::{Filter, PayloadKeyType, FieldCondition, Condition, PointOffsetType};
+use crate::index::field_index::{CardinalityEstimation, PrimaryCondition};
+use crate::index::query_estimator::estimate_filter;
+use crate::vector_storage::vector_storage::VectorStorage;
+use std::iter::FromIterator;
+use crate::id_mapper::id_mapper::IdMapper;
 
 pub const PAYLOAD_FIELD_INDEX_PATH: &str = "fields";
 
@@ -22,7 +26,9 @@ type IndexesMap = HashMap<PayloadKeyType, Vec<FieldIndex>>;
 
 pub struct StructPayloadIndex {
     condition_checker: Arc<AtomicRefCell<dyn ConditionChecker>>,
+    vector_storage: Arc<AtomicRefCell<dyn VectorStorage>>,
     payload: Arc<AtomicRefCell<dyn PayloadStorage>>,
+    id_mapper: Arc<AtomicRefCell<dyn IdMapper>>,
     field_indexes: IndexesMap,
     config: PayloadConfig,
     path: PathBuf,
@@ -30,20 +36,33 @@ pub struct StructPayloadIndex {
 }
 
 impl StructPayloadIndex {
-
     pub fn estimate_field_condition(&self, condition: &FieldCondition) -> Option<CardinalityEstimation> {
         self.field_indexes.get(&condition.key).and_then(|indexes| {
             let mut result_estimation: Option<CardinalityEstimation> = None;
             for index in indexes {
                 result_estimation = index.estimate_cardinality(condition);
                 if result_estimation.is_some() {
-                    break
+                    break;
                 }
             }
             result_estimation
         })
     }
 
+    fn query_field(&self, field_condition: &FieldCondition) -> Option<Box<dyn Iterator<Item=PointOffsetType> + '_>> {
+        let indexes = self.field_indexes
+            .get(&field_condition.key)
+            .and_then(|indexes|
+                indexes
+                    .iter()
+                    .map(|field_index| field_index.filter(field_condition))
+                    .skip_while(|filter_iter| filter_iter.is_none())
+                    .next()
+                    .map(|filter_iter| filter_iter.unwrap())
+            );
+        indexes
+    }
+
     fn config_path(&self) -> PathBuf {
         PayloadConfig::get_config_path(&self.path)
     }
@@ -99,7 +118,9 @@ impl StructPayloadIndex {
 
 
     pub fn open(condition_checker: Arc<AtomicRefCell<dyn ConditionChecker>>,
+                vector_storage: Arc<AtomicRefCell<dyn VectorStorage>>,
                 payload: Arc<AtomicRefCell<dyn PayloadStorage>>,
+                id_mapper: Arc<AtomicRefCell<dyn IdMapper>>,
                 path: &Path,
                 total_points: usize,
     ) -> OperationResult<Self> {
@@ -108,7 +129,9 @@ impl StructPayloadIndex {
 
         let mut index = StructPayloadIndex {
             condition_checker,
+            vector_storage,
             payload,
+            id_mapper,
             field_indexes: Default::default(),
             config,
             path: path.to_owned(),
@@ -185,7 +208,9 @@ impl StructPayloadIndex {
 
     pub fn new(
         condition_checker: Arc<AtomicRefCell<dyn ConditionChecker>>,
+        vector_storage: Arc<AtomicRefCell<dyn VectorStorage>>,
         payload: Arc<AtomicRefCell<dyn PayloadStorage>>,
+        id_mapper: Arc<AtomicRefCell<dyn IdMapper>>,
         path: &Path,
         config: Option<PayloadConfig>,
         total_points: usize,
@@ -194,7 +219,9 @@ impl StructPayloadIndex {
         let payload_config = config.unwrap_or_default();
         let mut payload_index = Self {
             condition_checker,
+            vector_storage,
             payload,
+            id_mapper,
             field_indexes: Default::default(),
             config: payload_config,
             path: path.to_owned(),
@@ -247,8 +274,69 @@ impl PayloadIndex for StructPayloadIndex {
         Ok(())
     }
 
-    fn query_points(&self, query: &Filter) -> Vec<usize> {
-        unimplemented!()
+    fn estimate_cardinality(&self, query: &Filter) -> CardinalityEstimation {
+        let total = self.total_points();
+
+        let estimator = |condition: &Condition| {
+            match condition {
+                Condition::Filter(_) => panic!("Unexpected branching"),
+                Condition::HasId(ids) => {
+                    let id_mapper_ref = self.id_mapper.borrow();
+                    let mapped_ids: HashSet<PointOffsetType> = ids.iter()
+                        .filter_map(|external_id| id_mapper_ref.internal_id(*external_id))
+                        .collect();
+                    let num_ids = mapped_ids.len();
+                    CardinalityEstimation {
+                        primary_clauses: vec![PrimaryCondition::Ids(mapped_ids)],
+                        min: 0,
+                        exp: num_ids,
+                        max: num_ids,
+                    }
+                }
+                Condition::Field(field_condition) => self
+                    .estimate_field_condition(field_condition)
+                    .unwrap_or(CardinalityEstimation {
+                        primary_clauses: vec![],
+                        min: 0,
+                        exp: self.total_points() / 2,
+                        max: self.total_points(),
+                    }),
+            }
+        };
+
+        estimate_filter(&estimator, query, total)
+    }
+
+    fn query_points(&self, query: &Filter) -> Box<dyn Iterator<Item=PointOffsetType> + '_> {
+        // Assume query is already estimated to be small enough so we can iterate over all matched ids
+        let query_cardinality = self.estimate_cardinality(query);
+        let condition_checker = self.condition_checker.borrow();
+        let vector_storage_ref = self.vector_storage.borrow();
+        let full_scan_iterator = vector_storage_ref.iter_ids(); // Should not be used if filter restricted by indexed fields
+        return if query_cardinality.primary_clauses.is_empty() {
+            // Worst case: query expected to return few matches, but index can't be used
+            let matched_points = full_scan_iterator
+                .filter(|i| condition_checker.check(*i, query))
+                .collect_vec();
+
+            Box::new(matched_points.into_iter())
+        } else {
+            // CPU-optimized strategy here: points are made unique before applying other filters.
+            let preselected: HashSet<PointOffsetType> = query_cardinality.primary_clauses.iter()
+                .map(|clause| {
+                    match clause {
+                        PrimaryCondition::Condition(field_condition) => self.query_field(field_condition)
+                            .unwrap_or(vector_storage_ref.iter_ids() /* index is not built */),
+                        PrimaryCondition::Ids(ids) => Box::new(ids.iter().cloned())
+                    }
+                })
+                .flat_map(|x| x)
+                .collect();
+            let  matched_points = preselected.into_iter()
+                .filter(|i| condition_checker.check(*i, query))
+                .collect_vec();
+            Box::new(matched_points.into_iter())
+        };
     }
 }
 
@@ -265,4 +353,10 @@ mod tests {
         let dir = TempDir::new("storage_dir").unwrap();
         let mut storage = SimplePayloadStorage::open(dir.path()).unwrap();
     }
+
+    // #[test]
+    // fn test_flat_map() {
+    //     let a = vec![vec![1,2,3], vec![4,5,6], vec![7,7,7]];
+    //     a.iter().flat_map(|x| x.iter()).for_each(|x| println!("{}", x))
+    // }
 }

lib/segment/src/spaces/tools.rs

@@ -69,7 +69,6 @@ pub fn peek_top_scores_iterable<I, E: Ord + Clone>(scores: I, top: usize, distan
 
 pub fn peek_top_scores<E: Ord + Clone>(scores: &[E], top: usize, distance: &Distance) -> Vec<E> {
     return peek_top_scores_iterable(scores.iter().cloned(), top, distance)
-
 }
 
 pub fn mertic_object(distance: &Distance) -> Box<dyn Metric> {