From c50482d53f80dbaebd3556bb25eae47443aeffaf Mon Sep 17 00:00:00 2001 From: Andrey Vasnetsov Date: Mon, 8 Mar 2021 01:00:27 +0100 Subject: [PATCH] WIP: query_points in struct payload index --- .../src/index/field_index/field_index.rs | 4 +- .../src/index/field_index/map_index.rs | 4 +- lib/segment/src/index/field_index/mod.rs | 4 +- .../src/index/field_index/numeric_index.rs | 2 +- lib/segment/src/index/index.rs | 12 +- lib/segment/src/index/plain_payload_index.rs | 25 +++- lib/segment/src/index/query_estimator.rs | 31 +---- lib/segment/src/index/struct_payload_index.rs | 108 ++++++++++++++++-- lib/segment/src/spaces/tools.rs | 1 - 9 files changed, 136 insertions(+), 55 deletions(-) diff --git a/lib/segment/src/index/field_index/field_index.rs b/lib/segment/src/index/field_index/field_index.rs index 46cd6c45b6..bd56ebba2b 100644 --- a/lib/segment/src/index/field_index/field_index.rs +++ b/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 + '_>; + fn filter(&self, condition: &FieldCondition) -> Option + '_>>; fn estimate_cardinality(&self, condition: &FieldCondition) -> Option; } @@ -39,7 +39,7 @@ impl FieldIndex { impl PayloadFieldIndex for FieldIndex { - fn filter(&self, condition: &FieldCondition) -> Box + '_> { + fn filter(&self, condition: &FieldCondition) -> Option + '_>> { self.get_payload_field_index().filter(condition) } diff --git a/lib/segment/src/index/field_index/map_index.rs b/lib/segment/src/index/field_index/map_index.rs index d2f9d5d8ff..af3c29ed04 100644 --- a/lib/segment/src/index/field_index/map_index.rs +++ b/lib/segment/src/index/field_index/map_index.rs @@ -57,7 +57,7 @@ impl PersistedMapIndex { } impl PayloadFieldIndex for PersistedMapIndex { - fn filter(&self, condition: &FieldCondition) -> Box> { + fn filter(&self, condition: &FieldCondition) -> Option>> { unimplemented!() } @@ -76,7 +76,7 @@ impl PayloadFieldIndex for PersistedMapIndex { } impl PayloadFieldIndex for PersistedMapIndex { - fn filter(&self, condition: &FieldCondition) -> Box> { + fn filter(&self, condition: &FieldCondition) -> Option>> { unimplemented!() } diff --git a/lib/segment/src/index/field_index/mod.rs b/lib/segment/src/index/field_index/mod.rs index af116e83e4..4cdf92b03d 100644 --- a/lib/segment/src/index/field_index/mod.rs +++ b/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) + Ids(HashSet) } #[derive(Debug)] diff --git a/lib/segment/src/index/field_index/numeric_index.rs b/lib/segment/src/index/field_index/numeric_index.rs index ebee21128e..95bebc9c50 100644 --- a/lib/segment/src/index/field_index/numeric_index.rs +++ b/lib/segment/src/index/field_index/numeric_index.rs @@ -100,7 +100,7 @@ impl PersistedNumericIndex { impl PayloadFieldIndex for PersistedNumericIndex { - fn filter(&self, condition: &FieldCondition) -> Box> { + fn filter(&self, condition: &FieldCondition) -> Option>> { unimplemented!() } diff --git a/lib/segment/src/index/index.rs b/lib/segment/src/index/index.rs index 0e6c0deb2d..bc709ef0be 100644 --- a/lib/segment/src/index/index.rs +++ b/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; + fn query_points(&self, query: &Filter) -> Box + '_>; } diff --git a/lib/segment/src/index/plain_payload_index.rs b/lib/segment/src/index/plain_payload_index.rs index bacbeae7c4..ddcbb7229f 100644 --- a/lib/segment/src/index/plain_payload_index.rs +++ b/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>, @@ -87,7 +90,23 @@ impl PayloadIndex for PlainPayloadIndex { self.save_config() } - fn query_points(&self, query: &Filter) -> Vec { + 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 + '_> { 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 { 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) diff --git a/lib/segment/src/index/query_estimator.rs b/lib/segment/src/index/query_estimator.rs index d342596ed4..b6d249024b 100644 --- a/lib/segment/src/index/query_estimator.rs +++ b/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(estimator: &F, conditions: &Vec, 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) - } -} diff --git a/lib/segment/src/index/struct_payload_index.rs b/lib/segment/src/index/struct_payload_index.rs index d0b5962dfa..52eed4702b 100644 --- a/lib/segment/src/index/struct_payload_index.rs +++ b/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>; pub struct StructPayloadIndex { condition_checker: Arc>, + vector_storage: Arc>, payload: Arc>, + id_mapper: Arc>, 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 { self.field_indexes.get(&condition.key).and_then(|indexes| { let mut result_estimation: Option = 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 + '_>> { + 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>, + vector_storage: Arc>, payload: Arc>, + id_mapper: Arc>, path: &Path, total_points: usize, ) -> OperationResult { @@ -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>, + vector_storage: Arc>, payload: Arc>, + id_mapper: Arc>, path: &Path, config: Option, 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 { - 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 = 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 + '_> { + // 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 = 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)) + // } } \ No newline at end of file diff --git a/lib/segment/src/spaces/tools.rs b/lib/segment/src/spaces/tools.rs index 09c1937d21..e5530b8f92 100644 --- a/lib/segment/src/spaces/tools.rs +++ b/lib/segment/src/spaces/tools.rs @@ -69,7 +69,6 @@ pub fn peek_top_scores_iterable(scores: I, top: usize, distan pub fn peek_top_scores(scores: &[E], top: usize, distance: &Distance) -> Vec { return peek_top_scores_iterable(scores.iter().cloned(), top, distance) - } pub fn mertic_object(distance: &Distance) -> Box {