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

@@ -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"

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"

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"] }

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;

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<Record> 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<PseudoId>,
+    collection_by_name: F,
+    read_consistency: Option<ReadConsistency>,
+    shard_selection: Option<ShardId>,
+) -> CollectionResult<HashMap<PseudoId, Lookup>>
+where
+    F: FnOnce(String) -> Fut,
+    Fut: Future<Output = Option<RwLockReadGuard<'a, Collection>>>,
+{
+    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)
+}

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<GroupId> 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<PseudoId> 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<PseudoId> for PointIdType {
+    type Error = ConversionError;
+
+    fn try_from(value: PseudoId) -> Result<Self, Self::Error> {
+        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<PointIdType> 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<u64> for PseudoId {
+    fn from(id: u64) -> Self {
+        PseudoId::NumberU64(id)
+    }
+}
+
+impl From<i64> for PseudoId {
+    fn from(id: i64) -> Self {
+        PseudoId::NumberI64(id)
+    }
+}
+
+impl From<String> 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())
+    }
+}

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<Collection>,
+    read_consistency: Option<ReadConsistency>,
+    shard_selection: Option<ShardId>,
+}
+
+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<PseudoId>) {
+    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<PseudoId>) {
+    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<PseudoId>) {
+    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()
+    );
+}

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<GroupId> 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<Self, Self::Error> {
         match value {
             serde_json::Value::String(s) => Ok(Self::String(s.to_string())),

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<ScoredPoint>,
-    /// 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<bool> 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")]