From 5dd3c1935ece9810c5fb8d347022100aabfb2005 Mon Sep 17 00:00:00 2001 From: Andrey Vasnetsov Date: Sun, 14 Mar 2021 01:01:53 +0100 Subject: [PATCH] WIP: struct payload index texts --- lib/segment/src/index/plain_payload_index.rs | 30 ++----- lib/segment/src/index/struct_payload_index.rs | 68 +++------------ .../segment_constructor.rs | 16 +++- lib/segment/src/types.rs | 4 +- lib/segment/tests/payload_index_test.rs | 85 +++++++++++++++++++ 5 files changed, 121 insertions(+), 82 deletions(-) create mode 100644 lib/segment/tests/payload_index_test.rs diff --git a/lib/segment/src/index/plain_payload_index.rs b/lib/segment/src/index/plain_payload_index.rs index ddcbb7229f..9553dc826e 100644 --- a/lib/segment/src/index/plain_payload_index.rs +++ b/lib/segment/src/index/plain_payload_index.rs @@ -38,7 +38,12 @@ impl PlainPayloadIndex { path: &Path, ) -> OperationResult { let config_path = PayloadConfig::get_config_path(path); - let config = PayloadConfig::load(&config_path)?; + let config = if config_path.exists() { + PayloadConfig::load(&config_path)? + } else { + PayloadConfig::default() + }; + let index = PlainPayloadIndex { condition_checker, @@ -47,26 +52,9 @@ impl PlainPayloadIndex { path: path.to_owned() }; - Ok(index) - } - - pub fn new( - condition_checker: Arc>, - vector_storage: Arc>, - path: &Path, - config: Option - ) -> OperationResult { - create_dir_all(path)?; - let payload_config = config.unwrap_or_default(); - - let index = PlainPayloadIndex { - condition_checker, - vector_storage, - config: payload_config, - path: path.to_owned() - }; - - index.save_config()?; + if !index.config_path().exists() { + index.save_config()? + } Ok(index) } diff --git a/lib/segment/src/index/struct_payload_index.rs b/lib/segment/src/index/struct_payload_index.rs index 52eed4702b..ac2e607f32 100644 --- a/lib/segment/src/index/struct_payload_index.rs +++ b/lib/segment/src/index/struct_payload_index.rs @@ -32,7 +32,6 @@ pub struct StructPayloadIndex { field_indexes: IndexesMap, config: PayloadConfig, path: PathBuf, - total_points: usize, } impl StructPayloadIndex { @@ -122,10 +121,13 @@ impl StructPayloadIndex { payload: Arc>, id_mapper: Arc>, path: &Path, - total_points: usize, ) -> OperationResult { let config_path = PayloadConfig::get_config_path(path); - let config = PayloadConfig::load(&config_path)?; + let config = if config_path.exists() { + PayloadConfig::load(&config_path)? + } else { + PayloadConfig::default() + }; let mut index = StructPayloadIndex { condition_checker, @@ -134,10 +136,14 @@ impl StructPayloadIndex { id_mapper, field_indexes: Default::default(), config, - path: path.to_owned(), - total_points, + path: path.to_owned() }; + if !index.config_path().exists() { + // Save default config + index.save_config()? + } + index.load_all_fields()?; Ok(index) @@ -206,33 +212,6 @@ impl StructPayloadIndex { Ok(()) } - pub fn new( - condition_checker: Arc>, - vector_storage: Arc>, - payload: Arc>, - id_mapper: Arc>, - path: &Path, - config: Option, - total_points: usize, - ) -> OperationResult { - create_dir_all(path)?; - 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(), - total_points, - }; - - payload_index.build_all_fields()?; - - Ok(payload_index) - } - fn save(&self) -> OperationResult<()> { let file = File::create(self.path.as_path())?; serde_cbor::to_writer(file, &self.field_indexes) @@ -241,7 +220,7 @@ impl StructPayloadIndex { } pub fn total_points(&self) -> usize { - self.total_points + self.vector_storage.borrow().vector_count() } } @@ -332,31 +311,10 @@ impl PayloadIndex for StructPayloadIndex { }) .flat_map(|x| x) .collect(); - let matched_points = preselected.into_iter() + let matched_points = preselected.into_iter() .filter(|i| condition_checker.check(*i, query)) .collect_vec(); Box::new(matched_points.into_iter()) }; } } - -#[cfg(test)] -mod tests { - use tempdir::TempDir; - - use crate::payload_storage::simple_payload_storage::SimplePayloadStorage; - - use super::*; - - #[test] - fn test_index_save_and_load() { - 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/segment_constructor/segment_constructor.rs b/lib/segment/src/segment_constructor/segment_constructor.rs index 486aa7ac8b..6691562338 100644 --- a/lib/segment/src/segment_constructor/segment_constructor.rs +++ b/lib/segment/src/segment_constructor/segment_constructor.rs @@ -4,7 +4,7 @@ use crate::vector_storage::simple_vector_storage::SimpleVectorStorage; use crate::payload_storage::simple_payload_storage::SimplePayloadStorage; use crate::index::plain_payload_index::{PlainPayloadIndex, PlainIndex}; use crate::query_planner::simple_query_planner::SimpleQueryPlanner; -use crate::types::{SegmentType, SegmentConfig, Indexes, SegmentState, SeqNumberType, StorageType}; +use crate::types::{SegmentType, SegmentConfig, Indexes, SegmentState, SeqNumberType, StorageType, PayloadIndexType}; use std::sync::{Arc, Mutex}; use atomic_refcell::AtomicRefCell; use crate::payload_storage::query_checker::SimpleConditionChecker; @@ -15,6 +15,8 @@ use crate::entry::entry_point::{OperationResult, OperationError}; use std::io::Read; use crate::vector_storage::memmap_vector_storage::MemmapVectorStorage; use crate::vector_storage::vector_storage::VectorStorage; +use crate::index::struct_payload_index::StructPayloadIndex; +use crate::index::index::PayloadIndex; fn sp(t: T) -> Arc> { Arc::new(AtomicRefCell::new(t)) } @@ -42,9 +44,15 @@ fn create_segment(version: SeqNumberType, segment_path: &Path, config: &SegmentC id_mapper.clone(), )); - let payload_index = sp(PlainPayloadIndex::new( - condition_checker, vector_storage.clone(), &payload_index_path, None - )?); + let payload_index: Arc> = match config.payload_index.unwrap_or_default() { + PayloadIndexType::Plain => sp(PlainPayloadIndex::open(condition_checker, vector_storage.clone(), &payload_index_path)?), + PayloadIndexType::Struct => sp(StructPayloadIndex::open( + condition_checker, + vector_storage.clone(), + payload_storage.clone(), + id_mapper.clone(), + &payload_index_path)?), + }; let index = sp(match config.index { Indexes::Plain { .. } => PlainIndex::new(vector_storage.clone(), payload_index, config.distance), diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index 925949c814..a9f20d4b37 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -133,9 +133,9 @@ impl Default for Indexes { #[serde(tag = "type", content = "options")] /// Type of payload index pub enum PayloadIndexType { - /// Store vectors in memory and use persistence storage only if vectors are changed + /// Do not index anything, just keep of what should be indexed later Plain, - /// Use memmap to store vectors, a little slower than `InMemory`, but requires little RAM + /// Build payload index. Index is saved on disc, but index itself is in RAM Struct, } diff --git a/lib/segment/tests/payload_index_test.rs b/lib/segment/tests/payload_index_test.rs new file mode 100644 index 0000000000..51168d357a --- /dev/null +++ b/lib/segment/tests/payload_index_test.rs @@ -0,0 +1,85 @@ +#[cfg(test)] +mod tests { + use super::*; + use rand::prelude::ThreadRng; + use rand::seq::SliceRandom; + use segment::types::{PayloadType, VectorElementType, SegmentConfig, Indexes, PayloadIndexType, Distance, StorageType}; + use rand::Rng; + use tempdir::TempDir; + use segment::segment_constructor::segment_constructor::build_segment; + + const ADJECTIVE: Vec = vec![ + "jobless".to_string(), + "rightful".to_string(), + "breakable".to_string(), + "impartial".to_string(), + "shocking".to_string(), + "faded".to_string(), + "phobic".to_string(), + "overt".to_string(), + "like".to_string(), + "wide-eyed".to_string(), + "broad".to_string(), + ]; + + const NOUN: Vec = vec![ + "territory".to_string(), + "jam".to_string(), + "neck".to_string(), + "chicken".to_string(), + "cap".to_string(), + "kiss".to_string(), + "veil".to_string(), + "trail".to_string(), + "size".to_string(), + "digestion".to_string(), + "rod".to_string(), + "seed".to_string(), + ]; + + fn random_keyword(rnd_gen: &mut ThreadRng) -> String { + let random_adj = ADJECTIVE.choose(rnd_gen).unwrap(); + let random_noun = NOUN.choose(rnd_gen).unwrap(); + format!("{} {}", random_adj, random_noun) + } + + fn random_keyword_payload(rnd_gen: &mut ThreadRng) -> PayloadType { + PayloadType::Keyword(vec![random_keyword(rnd_gen)]) + } + + fn random_int_payload(rnd_gen: &mut ThreadRng) -> PayloadType { + let val1: i64 = rnd_gen.gen_range(0..500); + let val2: i64 = rnd_gen.gen_range(0..500); + let val3: i64 = rnd_gen.gen_range(0..500); + + PayloadType::Integer(vec![val1, val2, val3]) + } + + fn random_vector(rnd_gen: &mut ThreadRng, size: usize) -> Vec { + (0..size).map(|_| rnd_gen.gen()).collect() + } + + #[test] + fn test_struct_payload_index() { + // Compare search with plain and struct indexes + + let dir = TempDir::new("segment_dir").unwrap(); + + let dim = 5; + + let mut config = SegmentConfig { + vector_size: dim, + index: Indexes::Plain {}, + payload_index: Some(PayloadIndexType::Plain), + storage_type: StorageType::InMemory, + distance: Distance::Dot, + }; + + let mut plain_segment = build_segment(path, &config).unwrap(); + config.payload_index = Some(PayloadIndexType::Struct); + let mut struct_segment = build_segment(path, &config).unwrap(); + + // ToDo: Init both segments with same data + // ToDo: Compare indexed and un-indexed search results + } +} \ No newline at end of file