use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::Arc; use atomic_refcell::AtomicRefCell; use bitvec::prelude::{BitSlice, BitVec}; use common::types::PointOffsetType; use rand::prelude::StdRng; use rand::SeedableRng; use super::payload_fixtures::BOOL_KEY; use crate::common::operation_error::OperationResult; use crate::common::Flusher; use crate::fixtures::payload_fixtures::{ generate_diverse_payload, FLT_KEY, GEO_KEY, INT_KEY, STR_KEY, TEXT_KEY, }; use crate::id_tracker::IdTracker; use crate::index::plain_payload_index::PlainPayloadIndex; use crate::index::struct_payload_index::StructPayloadIndex; use crate::index::PayloadIndex; use crate::payload_storage::in_memory_payload_storage::InMemoryPayloadStorage; use crate::payload_storage::query_checker::SimpleConditionChecker; use crate::payload_storage::PayloadStorage; use crate::types::{PayloadSchemaType, PointIdType, SeqNumberType}; /// Warn: Use for tests only /// /// This struct mimics the interface of `PointsIterator` and `IdTracker` only for basic cases #[derive(Debug)] pub struct FixtureIdTracker { ids: Vec, deleted: BitVec, deleted_count: usize, } impl FixtureIdTracker { pub fn new(num_points: usize) -> Self { Self { ids: (0..num_points).map(|x| x as PointOffsetType).collect(), deleted: BitVec::repeat(false, num_points), deleted_count: 0, } } } impl IdTracker for FixtureIdTracker { fn internal_version(&self, _internal_id: PointOffsetType) -> Option { Some(0) } fn set_internal_version( &mut self, _internal_id: PointOffsetType, _version: SeqNumberType, ) -> OperationResult<()> { Ok(()) } fn internal_id(&self, external_id: PointIdType) -> Option { Some(match external_id { PointIdType::NumId(id) => { assert!(id < self.ids.len() as u64); id as PointOffsetType } PointIdType::Uuid(_) => unreachable!(), }) } fn external_id(&self, internal_id: PointOffsetType) -> Option { assert!(internal_id < self.ids.len() as PointOffsetType); Some(PointIdType::NumId(u64::from(internal_id))) } fn set_link( &mut self, _external_id: PointIdType, _internal_id: PointOffsetType, ) -> OperationResult<()> { Ok(()) } fn drop(&mut self, external_id: PointIdType) -> OperationResult<()> { let internal_id = self.internal_id(external_id).unwrap() as usize; if !self.deleted.replace(internal_id, true) { self.deleted_count += 1; } Ok(()) } fn iter_external(&self) -> Box + '_> { Box::new( self.ids .iter() .copied() .map(|id| PointIdType::NumId(u64::from(id))), ) } fn iter_internal(&self) -> Box + '_> { Box::new(self.ids.iter().copied()) } fn iter_from( &self, external_id: Option, ) -> Box + '_> { let start = match external_id { None => 0, Some(id) => match id { PointIdType::NumId(num) => num, PointIdType::Uuid(_) => unreachable!(), }, } as PointOffsetType; Box::new( self.ids .iter() .copied() .skip_while(move |x| *x < start) .map(|x| (PointIdType::NumId(u64::from(x)), x)), ) } fn iter_random(&self) -> Box + '_> { unimplemented!("Not used for tests yet") } fn total_point_count(&self) -> usize { self.ids.len() } fn deleted_point_count(&self) -> usize { self.deleted_count } fn iter_ids(&self) -> Box + '_> { Box::new( self.ids .iter() .copied() .filter(|id| !self.is_deleted_point(*id)), ) } fn mapping_flusher(&self) -> Flusher { Box::new(|| Ok(())) } fn versions_flusher(&self) -> Flusher { Box::new(|| Ok(())) } fn is_deleted_point(&self, key: PointOffsetType) -> bool { let key = key as usize; if key >= self.deleted.len() { return true; } self.deleted[key] } fn deleted_point_bitslice(&self) -> &BitSlice { &self.deleted } fn cleanup_versions(&mut self) -> OperationResult<()> { Ok(()) } fn name(&self) -> &'static str { "fixture id tracker" } fn files(&self) -> Vec { vec![] } } /// Creates in-memory payload storage and fills it with random points /// /// # Arguments /// /// * `num_points` - how many random points to insert /// /// # Result /// /// Payload storage fixture /// pub fn create_payload_storage_fixture(num_points: usize, seed: u64) -> InMemoryPayloadStorage { let mut payload_storage = InMemoryPayloadStorage::default(); let mut rng = StdRng::seed_from_u64(seed); for id in 0..num_points { let payload = generate_diverse_payload(&mut rng); payload_storage .set(id as PointOffsetType, &payload) .unwrap(); } payload_storage } /// Function generates `PlainPayloadIndex` with random payload for testing /// /// # Arguments /// /// * `path` - temp directory path /// * `num_points` - how many payloads generate? /// /// # Result /// /// `PlainPayloadIndex` /// pub fn create_plain_payload_index(path: &Path, num_points: usize, seed: u64) -> PlainPayloadIndex { let payload_storage = create_payload_storage_fixture(num_points, seed); let id_tracker = Arc::new(AtomicRefCell::new(FixtureIdTracker::new(num_points))); let condition_checker = Arc::new(SimpleConditionChecker::new( Arc::new(AtomicRefCell::new(payload_storage.into())), id_tracker.clone(), HashMap::new(), )); PlainPayloadIndex::open(condition_checker, id_tracker, path).unwrap() } /// Function generates `StructPayloadIndex` with random payload for testing. /// It will also create indexes for payloads /// /// # Arguments /// /// * `path` - temp directory path /// * `num_points` - how many payloads generate? /// /// # Result /// /// `StructPayloadIndex` /// pub fn create_struct_payload_index( path: &Path, num_points: usize, seed: u64, ) -> StructPayloadIndex { let payload_storage = Arc::new(AtomicRefCell::new( create_payload_storage_fixture(num_points, seed).into(), )); let id_tracker = Arc::new(AtomicRefCell::new(FixtureIdTracker::new(num_points))); let mut index = StructPayloadIndex::open( payload_storage, id_tracker, std::collections::HashMap::new(), path, true, ) .unwrap(); index .set_indexed(&STR_KEY.parse().unwrap(), PayloadSchemaType::Keyword) .unwrap(); index .set_indexed(&INT_KEY.parse().unwrap(), PayloadSchemaType::Integer) .unwrap(); index .set_indexed(&FLT_KEY.parse().unwrap(), PayloadSchemaType::Float) .unwrap(); index .set_indexed(&GEO_KEY.parse().unwrap(), PayloadSchemaType::Geo) .unwrap(); index .set_indexed(&TEXT_KEY.parse().unwrap(), PayloadSchemaType::Text) .unwrap(); index .set_indexed(&BOOL_KEY.parse().unwrap(), PayloadSchemaType::Bool) .unwrap(); index }