From b92d3f45cd42e045b08b6ab7bdef368bb7c93f65 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20Coss=C3=ADo?= Date: Wed, 31 May 2023 14:23:46 -0400 Subject: [PATCH] Lookup internals (#1981) * feat: add lookup_ids function - rename GroupId -> PseudoId - impl try_from PseudoId to PointIdType - create lookup_ids function - add tests * fix: keep `GroupId` name in the openapi schema * refactor: make new type for PseudoId * refactor: make new `Lookup` output for lookup_ids() * changes from review, thanks @agourlay --- Cargo.lock | 33 ++++ docs/redoc/master/openapi.json | 1 + lib/collection/Cargo.toml | 1 + lib/collection/src/lib.rs | 1 + lib/collection/src/lookup/mod.rs | 80 +++++++++ lib/collection/src/lookup/types.rs | 98 +++++++++++ lib/collection/tests/lookup_test.rs | 232 +++++++++++++++++++++++++++ lib/segment/src/data_types/groups.rs | 3 +- lib/segment/src/types.rs | 8 +- 9 files changed, 455 insertions(+), 2 deletions(-) create mode 100644 lib/collection/src/lookup/mod.rs create mode 100644 lib/collection/src/lookup/types.rs create mode 100644 lib/collection/tests/lookup_test.rs diff --git a/Cargo.lock b/Cargo.lock index 5f6cbac838..a7325d6142 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -976,6 +976,7 @@ dependencies = [ "pprof", "rand 0.8.5", "rmp-serde", + "rstest", "schemars", "segment", "semver", @@ -1635,6 +1636,12 @@ version = "0.3.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "76d3d132be6c0e6aa1534069c705a74a5997a356c0dc2f86a47765e5617c5b65" +[[package]] +name = "futures-timer" +version = "3.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e64b03909df88034c26dc1547e8970b91f98bdb65165d6a4e9110d94263dbb2c" + [[package]] name = "futures-util" version = "0.3.28" @@ -3428,6 +3435,32 @@ dependencies = [ "smallvec", ] +[[package]] +name = "rstest" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de1bb486a691878cd320c2f0d319ba91eeaa2e894066d8b5f8f117c000e9d962" +dependencies = [ + "futures", + "futures-timer", + "rstest_macros", + "rustc_version", +] + +[[package]] +name = "rstest_macros" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "290ca1a1c8ca7edb7c3283bd44dc35dd54fdec6253a3912e201ba1072018fca8" +dependencies = [ + "cfg-if", + "proc-macro2", + "quote", + "rustc_version", + "syn 1.0.107", + "unicode-ident", +] + [[package]] name = "rust-ini" version = "0.18.0" diff --git a/docs/redoc/master/openapi.json b/docs/redoc/master/openapi.json index b7c3348079..79231ac128 100644 --- a/docs/redoc/master/openapi.json +++ b/docs/redoc/master/openapi.json @@ -7827,6 +7827,7 @@ } }, "GroupId": { + "description": "Value of the group_by key, shared across all the hits in the group", "anyOf": [ { "type": "string" diff --git a/lib/collection/Cargo.toml b/lib/collection/Cargo.toml index 803724b336..07f488c768 100644 --- a/lib/collection/Cargo.toml +++ b/lib/collection/Cargo.toml @@ -9,6 +9,7 @@ edition = "2021" [dev-dependencies] tempfile = "3.5.0" criterion = "0.5" +rstest = "0.17.0" [target.'cfg(not(target_os = "windows"))'.dev-dependencies] pprof = { version = "0.11", features = ["flamegraph", "prost-codec"] } diff --git a/lib/collection/src/lib.rs b/lib/collection/src/lib.rs index 92a9dba7b9..4393a7e3f4 100644 --- a/lib/collection/src/lib.rs +++ b/lib/collection/src/lib.rs @@ -5,6 +5,7 @@ pub mod common; pub mod config; pub mod grouping; pub mod hash_ring; +pub mod lookup; pub mod operations; pub mod optimizers_builder; pub mod recommendations; diff --git a/lib/collection/src/lookup/mod.rs b/lib/collection/src/lookup/mod.rs new file mode 100644 index 0000000000..c2fa25138b --- /dev/null +++ b/lib/collection/src/lookup/mod.rs @@ -0,0 +1,80 @@ +pub mod types; + +use std::collections::HashMap; + +use futures::Future; +use itertools::Itertools; +use schemars::JsonSchema; +use segment::types::{PointIdType, WithPayloadInterface, WithVector}; +use serde::{Deserialize, Serialize}; +use tokio::sync::RwLockReadGuard; +use types::PseudoId; + +use crate::collection::Collection; +use crate::operations::consistency_params::ReadConsistency; +use crate::operations::types::{CollectionError, CollectionResult, PointRequest, Record}; +use crate::shards::shard::ShardId; + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +#[serde(untagged)] +pub enum Lookup { + None, + Single(Record), + // We may want to implement multi-record lookup in the future +} + +impl From for Lookup { + fn from(record: Record) -> Self { + Lookup::Single(record) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +pub struct LookupRequest { + #[serde(rename = "collection")] + pub collection_name: String, + pub with_payload: WithPayloadInterface, + pub with_vectors: WithVector, +} + +pub async fn lookup_ids<'a, F, Fut>( + request: LookupRequest, + values: Vec, + collection_by_name: F, + read_consistency: Option, + shard_selection: Option, +) -> CollectionResult> +where + F: FnOnce(String) -> Fut, + Fut: Future>>, +{ + let collection = collection_by_name(request.collection_name.clone()) + .await + .ok_or(CollectionError::NotFound { + what: format!("Collection {}", request.collection_name), + })?; + + let ids = values + .into_iter() + .filter_map(|v| PointIdType::try_from(v).ok()) + .collect_vec(); + + if ids.is_empty() { + return Ok(HashMap::new()); + } + + let point_request = PointRequest { + ids, + with_payload: Some(request.with_payload), + with_vector: request.with_vectors, + }; + + let result = collection + .retrieve(point_request, read_consistency, shard_selection) + .await? + .into_iter() + .map(|point| (PseudoId::from(point.id), Lookup::from(point))) + .collect(); + + Ok(result) +} diff --git a/lib/collection/src/lookup/types.rs b/lib/collection/src/lookup/types.rs new file mode 100644 index 0000000000..b13222a2ae --- /dev/null +++ b/lib/collection/src/lookup/types.rs @@ -0,0 +1,98 @@ +use std::fmt::Display; + +use segment::data_types::groups::GroupId; +use segment::types::PointIdType; +use uuid::Uuid; + +/// A value that can be used as a temporary ID +#[derive(Debug, Eq, PartialEq, Clone, Hash)] +pub enum PseudoId { + String(String), + NumberU64(u64), + NumberI64(i64), +} + +impl Display for PseudoId { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + PseudoId::String(s) => write!(f, "{}", s), + PseudoId::NumberU64(n) => write!(f, "{}", n), + PseudoId::NumberI64(n) => write!(f, "{}", n), + } + } +} + +impl From for PseudoId { + fn from(id: GroupId) -> Self { + match id { + GroupId::String(s) => Self::String(s), + GroupId::NumberU64(n) => Self::NumberU64(n), + GroupId::NumberI64(n) => Self::NumberI64(n), + } + } +} + +impl From for GroupId { + fn from(id: PseudoId) -> Self { + match id { + PseudoId::String(s) => Self::String(s), + PseudoId::NumberU64(n) => Self::NumberU64(n), + PseudoId::NumberI64(n) => Self::NumberI64(n), + } + } +} + +#[derive(Debug)] +pub enum ConversionError { + IntError(core::num::TryFromIntError), + ParseError(uuid::Error), +} + +impl TryFrom for PointIdType { + type Error = ConversionError; + + fn try_from(value: PseudoId) -> Result { + match value { + PseudoId::String(s) => Ok(PointIdType::Uuid( + Uuid::try_parse(&s).map_err(ConversionError::ParseError)?, + )), + PseudoId::NumberU64(n) => Ok(PointIdType::NumId(n)), + PseudoId::NumberI64(n) => Ok(PointIdType::NumId( + u64::try_from(n).map_err(ConversionError::IntError)?, + )), + } + } +} + +impl From for PseudoId { + fn from(id: PointIdType) -> Self { + match id { + PointIdType::NumId(n) => PseudoId::NumberU64(n), + PointIdType::Uuid(u) => PseudoId::String(u.to_string()), + } + } +} + +impl From for PseudoId { + fn from(id: u64) -> Self { + PseudoId::NumberU64(id) + } +} + +impl From for PseudoId { + fn from(id: i64) -> Self { + PseudoId::NumberI64(id) + } +} + +impl From for PseudoId { + fn from(id: String) -> Self { + PseudoId::String(id) + } +} + +impl From<&str> for PseudoId { + fn from(id: &str) -> Self { + PseudoId::String(id.to_string()) + } +} diff --git a/lib/collection/tests/lookup_test.rs b/lib/collection/tests/lookup_test.rs new file mode 100644 index 0000000000..c19a52590b --- /dev/null +++ b/lib/collection/tests/lookup_test.rs @@ -0,0 +1,232 @@ +use collection::collection::Collection; +use collection::lookup::types::PseudoId; +use collection::lookup::{lookup_ids, Lookup, LookupRequest}; +use collection::operations::consistency_params::ReadConsistency; +use collection::operations::point_ops::{Batch, WriteOrdering}; +use collection::shards::shard::ShardId; +use common::simple_collection_fixture; +use itertools::Itertools; +use rand::rngs::SmallRng; +use rand::{self, Rng, SeedableRng}; +use rstest::*; +use segment::data_types::vectors::VectorStruct; +use segment::types::{Payload, PointIdType}; +use serde_json::json; +use tempfile::Builder; +use tokio::sync::RwLock; +use uuid::Uuid; + +mod common; + +const SEED: u64 = 42; + +struct Resources { + request: LookupRequest, + collection: RwLock, + read_consistency: Option, + shard_selection: Option, +} + +async fn setup() -> Resources { + let request = LookupRequest { + collection_name: "test".to_string(), + with_payload: false.into(), + with_vectors: false.into(), + }; + + let collection_dir = Builder::new().prefix("storage").tempdir().unwrap(); + + let collection = simple_collection_fixture(collection_dir.path(), 1).await; + + let int_ids = (0..1000).map(PointIdType::from); + + let mut rng = SmallRng::seed_from_u64(SEED); + let uuids = (0..1000).map(|_| PointIdType::Uuid(Uuid::from_u128(rng.gen()))); + + let ids = int_ids.chain(uuids).collect_vec(); + + let mut rng = SmallRng::seed_from_u64(SEED); + let vectors = (0..2000) + .map(|_| rng.gen::<[f32; 4]>().to_vec()) + .collect_vec(); + + let payloads = ids + .iter() + .map(|i| Some(Payload::from(json!({ "foo": format!("bar {}", i) })))) + .collect_vec(); + + let upsert_points = collection::operations::CollectionUpdateOperations::PointOperation( + Batch { + ids, + vectors: vectors.into(), + payloads: Some(payloads), + } + .into(), + ); + + collection + .update_from_client(upsert_points, true, WriteOrdering::default()) + .await + .unwrap(); + + let read_consistency = None; + + let shard_selection = None; + + Resources { + request, + collection: RwLock::new(collection), + read_consistency, + shard_selection, + } +} + +#[tokio::test] +async fn happy_lookup_ids() { + let Resources { + mut request, + collection, + read_consistency, + shard_selection, + } = setup().await; + + let collection = collection.read().await; + + let collection_by_name = |_: String| async { Some(collection) }; + + let n = 100u64; + let ints = (0..n).map_into(); + + let mut rng = SmallRng::seed_from_u64(SEED); + let uuids = (0..n) + .map(|_| Uuid::from_u128(rng.gen()).to_string()) + .map_into(); + + let values = ints.chain(uuids).collect_vec(); + request.with_payload = true.into(); + request.with_vectors = true.into(); + + let result = lookup_ids( + request.clone(), + values.clone(), + collection_by_name, + read_consistency, + shard_selection, + ) + .await; + + assert!(result.is_ok()); + + let result = result.unwrap(); + + assert_eq!(result.len(), (n * 2) as usize); + + let mut rng = SmallRng::seed_from_u64(SEED); + + // use points 0..n and 1000..1000+n as expected vectors + let expected_vectors = (0..1000 + n) + .map(|i| (i, rng.gen::<[f32; 4]>().to_vec())) + .filter(|(i, _)| !(&n..&1000).contains(&i)) + .map(|(_, v)| v) + .map(VectorStruct::from); + + for (id_value, vector) in values.into_iter().zip(expected_vectors) { + let Lookup::Single(record) = result.get(&id_value).unwrap() else { + panic!("Expected to find record for id {}", id_value); + }; + assert_eq!(record.id, PointIdType::try_from(id_value.clone()).unwrap()); + assert_eq!( + record.payload, + Some(Payload::from(json!({ "foo": format!("bar {}", id_value) }))) + ); + assert_eq!(record.vector, Some(vector)); + } +} + +fn first_uuid() -> String { + let mut rng = SmallRng::seed_from_u64(SEED); + Uuid::from_u128(rng.gen()).to_string() +} + +#[rstest] +#[case::existing_uuid(first_uuid())] +#[case::zero_int(0i64)] +#[case::positive_int(1i64)] +#[case::existing_uint(999u64)] +fn parsable_pseudo_id_to_point_id(#[case] value: impl Into) { + let value = value.into(); + assert!(PointIdType::try_from(value).is_ok()); +} + +#[rstest] +#[case::negative_int(-1i64)] +#[case::non_uuid_string("not a uuid")] +fn non_parsable_pseudo_id_to_point_id(#[case] value: impl Into) { + let value = value.into(); + assert!(PointIdType::try_from(value).is_err()); +} + +#[rstest] +#[case::uuid(Uuid::new_v4().to_string())] +#[case::int(1001u64)] +#[tokio::test] +async fn inexisting_lookup_ids_are_ignored(#[case] value: impl Into) { + let value = value.into(); + + let Resources { + mut request, + collection, + read_consistency, + shard_selection, + } = setup().await; + + let collection = collection.read().await; + + let collection_by_name = |_: String| async { Some(collection) }; + + let values = vec![value]; + request.with_payload = true.into(); + request.with_vectors = true.into(); + + let result = lookup_ids( + request, + values, + collection_by_name, + read_consistency, + shard_selection, + ) + .await; + + assert!(result.is_ok()); + + let result = result.unwrap(); + + assert_eq!(result.len(), 0); +} + +#[tokio::test] +async fn err_when_collection_by_name_returns_none() { + let Resources { + request, + read_consistency, + shard_selection, + .. + } = setup().await; + + let collection_by_name = |_: String| async { None }; + + let result = lookup_ids( + request, + vec![], + collection_by_name, + read_consistency, + shard_selection, + ) + .await; + + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Collection test not found".to_string() + ); +} diff --git a/lib/segment/src/data_types/groups.rs b/lib/segment/src/data_types/groups.rs index 8061338f10..7de241ba0f 100644 --- a/lib/segment/src/data_types/groups.rs +++ b/lib/segment/src/data_types/groups.rs @@ -2,6 +2,7 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use serde_json::json; +/// Value of the group_by key, shared across all the hits in the group #[derive(Debug, Serialize, Deserialize, JsonSchema, Eq, PartialEq, Clone, Hash)] #[serde(untagged)] pub enum GroupId { @@ -47,7 +48,7 @@ impl From for serde_json::Value { impl TryFrom<&serde_json::Value> for GroupId { type Error = (); - /// Only allows Strings and Numbers to be converted into GroupId + /// Only allows Strings and Numbers to be converted into Scalar fn try_from(value: &serde_json::Value) -> Result { match value { serde_json::Value::String(s) => Ok(Self::String(s.to_string())), diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index 3e1f89054f..25ff5e532e 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -203,7 +203,7 @@ impl PartialEq for ScoredPoint { pub struct PointGroup { /// Scored points that have the same value of the group_by key pub hits: Vec, - /// Value of the group_by key shared by all the hits + /// Value of the group_by key, shared across all the hits in the group pub id: GroupId, } @@ -1397,6 +1397,12 @@ pub enum WithPayloadInterface { Selector(PayloadSelector), } +impl From for WithPayloadInterface { + fn from(b: bool) -> Self { + WithPayloadInterface::Bool(b) + } +} + /// Options for specifying which vector to include #[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq)] #[serde(rename_all = "snake_case")]