From b07bd81af4e75a2e2e85b646092bbd284b64ec7e Mon Sep 17 00:00:00 2001 From: Roman Titov Date: Tue, 24 Jan 2023 20:24:55 +0100 Subject: [PATCH] Implemented response resolver (#1370, #1381) - Implement `Resolve` trait that "merges" multiple `retrieve`/`scroll_by`/`search` responses into a single response, to ensure it's consistent across multiple nodes in the cluster - Implement tests for the `Resolve` trait implementation - Add a few additional derives required for the `Resolve` trait implementation Co-authored-by: Andrey Vasnetsov --- Cargo.lock | 1 + lib/collection/Cargo.toml | 1 + lib/collection/src/operations/types.rs | 2 +- lib/collection/src/shards/mod.rs | 1 + lib/collection/src/shards/replica_set.rs | 2 +- lib/collection/src/shards/resolve.rs | 526 +++++++++++++++++++++++ lib/segment/src/data_types/vectors.rs | 5 +- lib/segment/src/types.rs | 2 +- 8 files changed, 534 insertions(+), 6 deletions(-) create mode 100644 lib/collection/src/shards/resolve.rs diff --git a/Cargo.lock b/Cargo.lock index c8d0571f3e..80281eb999 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -829,6 +829,7 @@ dependencies = [ "tar", "tempfile", "thiserror", + "tinyvec", "tokio", "tonic", "tower", diff --git a/lib/collection/Cargo.toml b/lib/collection/Cargo.toml index 200c307f90..b163a5a851 100644 --- a/lib/collection/Cargo.toml +++ b/lib/collection/Cargo.toml @@ -24,6 +24,7 @@ rmp-serde = "~1.1" wal = { git = "https://github.com/qdrant/wal.git", rev = "0519d97c1f806e2787abcfecf96d13d87556e8c0"} ordered-float = "3.4" hashring = "0.3.0" +tinyvec = { version = "1.6.0", features = ["alloc"] } tokio = {version = "~1.24", features = ["full"]} futures = "0.3.25" diff --git a/lib/collection/src/operations/types.rs b/lib/collection/src/operations/types.rs index 5b139ac990..44c47d3c8d 100644 --- a/lib/collection/src/operations/types.rs +++ b/lib/collection/src/operations/types.rs @@ -61,7 +61,7 @@ pub enum OptimizersStatus { } /// Point data -#[derive(Debug, Deserialize, Serialize, JsonSchema)] +#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, JsonSchema)] #[serde(rename_all = "snake_case")] pub struct Record { /// Id of the point diff --git a/lib/collection/src/shards/mod.rs b/lib/collection/src/shards/mod.rs index effae3e9a8..11f8975fb9 100644 --- a/lib/collection/src/shards/mod.rs +++ b/lib/collection/src/shards/mod.rs @@ -8,6 +8,7 @@ pub mod proxy_shard; pub mod remote_shard; #[allow(dead_code)] pub mod replica_set; +pub mod resolve; pub mod shard; pub mod shard_config; pub mod shard_holder; diff --git a/lib/collection/src/shards/replica_set.rs b/lib/collection/src/shards/replica_set.rs index 64425225ad..8d77b981f0 100644 --- a/lib/collection/src/shards/replica_set.rs +++ b/lib/collection/src/shards/replica_set.rs @@ -565,7 +565,7 @@ impl ShardReplicaSet { /// 3 - Fallbacks to all remaining shards if the optimisations fails. /// It does not report failing peer_ids to the consensus. pub async fn execute_read_operation<'a, F, Fut, Res>( - &'_ self, + &self, read_operation: F, local: &'a Option, remotes: &'a [RemoteShard], diff --git a/lib/collection/src/shards/resolve.rs b/lib/collection/src/shards/resolve.rs new file mode 100644 index 0000000000..e5fb796a10 --- /dev/null +++ b/lib/collection/src/shards/resolve.rs @@ -0,0 +1,526 @@ +use std::collections::{HashMap, HashSet}; +use std::hash; + +use segment::types::ScoredPoint; +use tinyvec::TinyVec; + +use crate::operations::types::Record; + +#[derive(Copy, Clone, Debug, Eq, PartialEq)] +pub enum ResolveCondition { + All, + Majority, +} + +pub trait Resolve: Sized { + fn resolve(responses: Vec, condition: ResolveCondition) -> Self; +} + +impl Resolve for Vec { + fn resolve(records: Vec, condition: ResolveCondition) -> Self { + let mut resolved = Resolver::resolve(records, |record| record.id, PartialEq::eq, condition); + resolved.sort_unstable_by_key(|record| record.id); + resolved + } +} + +impl Resolve for Vec> { + fn resolve(batches: Vec, condition: ResolveCondition) -> Self { + // batches: > + // transpose to > + + let batches = transpose(batches); + + batches + .into_iter() + .map(|points| { + let mut resolved = + Resolver::resolve(points, |point| point.id, scored_points_eq, condition); + + resolved.sort_unstable(); + resolved + }) + .collect() + } +} + +fn transpose(vec: Vec>) -> Vec> { + if vec.is_empty() { + return Vec::new(); + } + + let len = vec[0].len(); + + let mut iters: Vec<_> = vec.into_iter().map(IntoIterator::into_iter).collect(); + + (0..len) + .map(|_| iters.iter_mut().filter_map(Iterator::next).collect()) + .collect() +} + +fn scored_points_eq(this: &ScoredPoint, other: &ScoredPoint) -> bool { + this.id == other.id + && this.score == other.score + && this.vector == other.vector + && this.payload == other.payload +} + +struct Resolver<'a, Item, Id, Ident, Cmp> { + items: HashMap>, + identify: Ident, + compare: Cmp, +} + +type ResolverRecords<'a, Item> = TinyVec<[ResolverRecord<'a, Item>; RESOLVER_RECORDS_CAPACITY]>; + +const RESOLVER_RECORDS_CAPACITY: usize = 5; // Expected number of replicas + +impl<'a, Item, Id, Ident, Cmp> Resolver<'a, Item, Id, Ident, Cmp> +where + Id: Eq + hash::Hash, + Ident: Fn(&Item) -> Id, + Cmp: Fn(&Item, &Item) -> bool, +{ + pub fn resolve( + items: Vec>, + identify: Ident, + compare: Cmp, + condition: ResolveCondition, + ) -> Vec { + let resolution_count = match condition { + ResolveCondition::All => items.len(), + ResolveCondition::Majority => items.len() / 2 + 1, + }; + + let mut resolver = Resolver::new(items.first().map_or(0, Vec::len), identify, compare); + resolver.add_all(&items); + + // Select coordinates of accepted items, avoiding copying + let resolved_items: HashSet<_> = resolver + .items + .into_iter() + .filter_map(|(_, points)| { + points + .into_iter() + .find(|point| point.count >= resolution_count) + .map(|point| (point.row, point.index)) + }) + .collect(); + + // Shortcut if everything is consistent: return first items, avoiding filtering + let is_consistent = resolved_items.len() == items.first().map_or(0, Vec::len) + && resolved_items.iter().all(|&(row, _)| row == 0); + + if is_consistent { + items.into_iter().next().unwrap_or_default() + } else { + items + .into_iter() + .enumerate() + .flat_map(|(row, items)| { + items + .into_iter() + .enumerate() + .map(move |(index, item)| (row, index, item)) + }) + .filter_map(|(row, index, item)| { + if resolved_items.contains(&(row, index)) { + Some(item) + } else { + None + } + }) + .collect() + } + } + + fn new(capacity: usize, identify: Ident, compare: Cmp) -> Self { + Self { + items: HashMap::with_capacity(capacity), + identify, + compare, + } + } + + fn add_all(&mut self, items: I) + where + I: IntoIterator, + I::Item: IntoIterator, + { + for (row, items) in items.into_iter().enumerate() { + for (index, item) in items.into_iter().enumerate() { + self.add((self.identify)(item), item, row, index); + } + } + } + + fn add(&mut self, id: Id, item: &'a Item, row: usize, index: usize) { + let points = self.items.entry(id).or_default(); + + for point in points.iter_mut() { + if (self.compare)(item, point.item.unwrap()) { + point.count += 1; + return; + } + } + + points.push(ResolverRecord::new(item, row, index)); + } +} + +struct ResolverRecord<'a, T> { + item: Option<&'a T>, + row: usize, + index: usize, + count: usize, +} + +impl<'a, T> Default for ResolverRecord<'a, T> { + fn default() -> Self { + Self { + item: None, + row: 0, + index: 0, + count: 0, + } + } +} + +impl<'a, T> ResolverRecord<'a, T> { + fn new(item: &'a T, row: usize, index: usize) -> Self { + Self { + item: Some(item), + row, + index, + count: 1, + } + } +} + +#[cfg(test)] +mod test { + use std::fmt; + + use segment::types::ScoreType; + + use super::*; + + #[rustfmt::skip] + fn resolve_scored_points_batch_4_data() -> [Vec; 3] { + [ + vec![ + point(14, 0.0), point(17, 0.1), point(15, 0.1), + point(13, 0.2), point(11, 0.2), point(12, 0.3), + point(18, 0.3), point(16, 0.4), point(10, 0.5), + ], + vec![ + point(23, 0.0), point(21, 0.1), point(25, 0.2), + point(22, 0.2), point(20, 0.3), point(24, 0.3), + ], + vec![ + point(30, 0.1), point(31, 0.1), point(32, 0.1), + point(33, 0.2), point(34, 0.2), point(35, 0.3), + ], + ] + } + + fn point(id: u64, score: ScoreType) -> ScoredPoint { + ScoredPoint { + id: id.into(), + version: 1, + score, + payload: None, + vector: None, + } + } + + #[rustfmt::skip] + fn resolve_scored_points_batch_4_input() -> Vec>> { + let [batch1, batch2, batch3] = resolve_scored_points_batch_4_data(); + + vec![ + vec![ + batch(&batch1, [remove(2), remove(3)]), + batch(&batch2, [remove(0), remove(3)]), + batch(&batch3, [remove(4), remove(5)]), + ], + + vec![ + batch(&batch1, [remove(1), modify(3)]), + batch(&batch2, [modify(0), remove(2)]), + batch(&batch3, [remove(3), modify(5)]), + ], + + vec![ + batch(&batch1, [remove(1), modify(4)]), + batch(&batch2, [modify(3), remove(5)]), + batch(&batch3, [remove(2), modify(5)]), + ], + + vec![ + batch1, + batch2, + batch3, + ], + ] + } + + fn batch(batch: &[ScoredPoint], mut actions: [Action; N]) -> Vec { + let mut batch = batch.to_owned(); + + actions.sort_unstable_by_key(|action| action.index()); + + let mut removed = Vec::new(); + + for action in actions.into_iter() { + let offset = removed + .iter() + .filter(|&&removed| removed <= action.index()) + .count(); + + match action { + Action::Remove(index) => { + batch.remove(index - offset); + removed.push(index); + } + + Action::Modify(index) => { + batch[index - offset].score += 1.0; + } + } + } + + batch + } + + #[derive(Copy, Clone, Debug)] + enum Action { + Remove(usize), + Modify(usize), + } + + impl Action { + pub fn index(self) -> usize { + match self { + Self::Remove(index) => index, + Self::Modify(index) => index, + } + } + } + + fn remove(index: usize) -> Action { + Action::Remove(index) + } + + fn modify(index: usize) -> Action { + Action::Modify(index) + } + + #[test] + fn resolve_scored_points_batch_4_all() { + let [mut batch1, mut batch2, mut batch3] = resolve_scored_points_batch_4_data(); + + batch1.remove(4); + batch1.remove(3); + batch1.remove(2); + batch1.remove(1); + + batch2.remove(5); + batch2.remove(3); + batch2.remove(2); + batch2.remove(0); + + batch3.remove(5); + batch3.remove(4); + batch3.remove(3); + batch3.remove(2); + + test_resolve( + resolve_scored_points_batch_4_input(), + [batch1, batch2, batch3], + ResolveCondition::All, + ); + } + + #[test] + fn resolve_scored_points_batch_4_majority() { + let [mut batch1, mut batch2, mut batch3] = resolve_scored_points_batch_4_data(); + + batch1.remove(3); + batch1.remove(1); + + batch2.remove(3); + batch2.remove(0); + + batch3.remove(5); + + test_resolve( + resolve_scored_points_batch_4_input(), + [batch1, batch2, batch3], + ResolveCondition::Majority, + ); + } + + fn data_simple() -> [i32; 9] { + [1, 2, 3, 4, 5, 6, 7, 8, 9] + } + + #[rustfmt::skip] + fn input_2() -> [Vec; 2] { + [ + vec![1, 2, 3, 6, 7, 9, 11, 12, 13], + vec![ 3, 4, 5, 6, 8, 10, 11, ], + ] + } + + fn expected_2() -> [i32; 3] { + [3, 6, 11] + } + + #[rustfmt::skip] + fn input_3() -> [Vec; 3] { + [ + vec![1, 2, 6, 7, 8, 11, 13, 14, 15, ], + vec![ 2, 3, 4, 7, 9, 10, 13, 14, 16, ], + vec![ 4, 5, 6, 7, 9, 11, 12, 14, 17], + ] + } + + fn expected_3_all() -> [i32; 2] { + [7, 14] + } + + fn expected_3_majority() -> [i32; 8] { + [2, 4, 6, 7, 9, 11, 13, 14] + } + + #[rustfmt::skip] + fn input_4() -> [Vec; 4] { + [ + vec![1, 2, 3, 9, 11, 12, 13, 14, 16, 19, 21, 22, 24, 27, 29], + vec![ 2, 3, 4, 5, 6, 12, 13, 15, 17, 19, 22, 24, 26, 27, 28, ], + vec![ 3, 5, 6, 7, 8, 9, 13, 15, 16, 18, 20, 22, 26, 27, 28, ], + vec![ 6, 8, 9, 10, 11, 12, 13, 16, 18, 19, 21, 23, 26, 27, 29], + ] + } + + fn expected_4_all() -> [i32; 2] { + [13, 27] + } + + fn expected_4_majority() -> [i32; 10] { + [3, 6, 9, 12, 13, 16, 19, 22, 26, 27] + } + + #[test] + fn resolve_0_all() { + resolve_0(ResolveCondition::All); + } + + #[test] + fn resolve_0_majority() { + resolve_0(ResolveCondition::Majority); + } + + fn resolve_0(condition: ResolveCondition) { + test_resolve_simple(Vec::>::new(), Vec::new(), condition); + } + + #[test] + fn resolve_simple_all() { + for replicas in 1..=5 { + resolve_simple(replicas, ResolveCondition::All); + } + } + + #[test] + fn resolve_simple_majority() { + for replicas in 1..=5 { + resolve_simple(replicas, ResolveCondition::All); + } + } + + fn resolve_simple(replicas: usize, condition: ResolveCondition) { + let input: Vec<_> = (0..replicas).map(|_| data_simple()).collect(); + let expected = data_simple(); + + test_resolve_simple(input, expected, condition) + } + + #[test] + fn resolve_2_all() { + test_resolve_simple(input_2(), expected_2(), ResolveCondition::All); + } + + #[test] + fn resolve_2_majority() { + test_resolve_simple(input_2(), expected_2(), ResolveCondition::Majority); + } + + #[test] + fn resolve_3_all() { + test_resolve_simple(input_3(), expected_3_all(), ResolveCondition::All); + } + + #[test] + fn resolve_3_majority() { + test_resolve_simple(input_3(), expected_3_majority(), ResolveCondition::Majority); + } + + #[test] + fn resolve_4_all() { + test_resolve_simple(input_4(), expected_4_all(), ResolveCondition::All); + } + + #[test] + fn resolve_4_majority() { + test_resolve_simple(input_4(), expected_4_majority(), ResolveCondition::Majority); + } + + fn test_resolve(input: Vec, expected: E, condition: ResolveCondition) + where + T: Resolve + Clone + PartialEq + fmt::Debug, + E: fmt::Debug, + { + assert_eq!(T::resolve(input, condition), expected); + } + + fn test_resolve_simple(input: I, expected: E, condition: ResolveCondition) + where + I: IntoIterator, + I::Item: IntoIterator, + E: IntoIterator, + { + test_resolve(simple_input(input), simple_expected(expected), condition); + } + + fn simple_input(input: I) -> Vec> + where + I: IntoIterator, + I::Item: IntoIterator, + { + input + .into_iter() + .map(|items| items.into_iter().map(Val).collect()) + .collect() + } + + fn simple_expected(expected: E) -> Vec + where + E: IntoIterator, + { + expected.into_iter().map(Val).collect() + } + + #[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] + struct Val(i32); + + impl Resolve for Vec { + fn resolve(values: Vec, condition: ResolveCondition) -> Self { + let mut resolved = Resolver::resolve(values, |val| val.0, PartialEq::eq, condition); + + resolved.sort_unstable(); + resolved + } + } +} diff --git a/lib/segment/src/data_types/vectors.rs b/lib/segment/src/data_types/vectors.rs index 48b7a1f1a8..6eff8f3351 100644 --- a/lib/segment/src/data_types/vectors.rs +++ b/lib/segment/src/data_types/vectors.rs @@ -23,9 +23,8 @@ pub fn only_default_vector(vec: &[VectorElementType]) -> NamedVectors { } /// Full vector data per point separator with single and multiple vector modes -#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)] -#[serde(rename_all = "snake_case")] -#[serde(untagged)] +#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, JsonSchema)] +#[serde(untagged, rename_all = "snake_case")] pub enum VectorStruct { Single(VectorType), Multi(HashMap), diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index eada6afbcd..d37a92cfae 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -455,7 +455,7 @@ impl TryFrom for GeoPoint { } } -#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq)] +#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize, JsonSchema)] pub struct Payload(pub Map); impl Payload {