diff --git a/lib/api/src/grpc/conversions.rs b/lib/api/src/grpc/conversions.rs index 9ee5966125..a142604918 100644 --- a/lib/api/src/grpc/conversions.rs +++ b/lib/api/src/grpc/conversions.rs @@ -1345,14 +1345,14 @@ impl TryFrom for segment::types::HasIdCondition { .into_iter() .map(|p| p.try_into()) .collect::>()?; - Ok(Self { has_id: set }) + Ok(Self::from(set)) } } impl From for HasIdCondition { fn from(value: segment::types::HasIdCondition) -> Self { let segment::types::HasIdCondition { has_id } = value; - let set: Vec = has_id.into_iter().map(|p| p.into()).collect(); + let set: Vec = has_id.into_inner().into_iter().map(PointId::from).collect(); Self { has_id: set } } } diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index d75e7605ee..c4a2ecc70b 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -38,6 +38,7 @@ use crate::index::sparse_index::sparse_index_config::SparseIndexConfig; use crate::json_path::JsonPath; use crate::spaces::metric::MetricPostProcessing; use crate::spaces::simple::{CosineMetric, DotProductMetric, EuclidMetric, ManhattanMetric}; +use crate::utils::maybe_arc::MaybeArc; pub type PayloadKeyType = JsonPath; pub type PayloadKeyTypeRef<'a> = &'a JsonPath; @@ -2608,7 +2609,7 @@ impl From for IsEmptyCondition { #[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq)] pub struct HasIdCondition { #[schemars(schema_with = "HashSet::::json_schema")] - pub has_id: AHashSet, + pub has_id: MaybeArc>, } /// Filter points which have specific vector assigned @@ -2623,17 +2624,30 @@ impl From for HasVectorCondition { } } +/// Threshold determining when to use an `Arc` in `HasIdCondition` if the condition includes many points. +/// Since we're cloning filters quite a lot, using an Arc for larger conditions reduces risk of memory leaks +/// and potentially improves performance in some places. +const HAS_ID_CONDITION_ARC_THRESHOLD: usize = 1_000; + impl From> for HasIdCondition { fn from(has_id: AHashSet) -> Self { - HasIdCondition { has_id } + if has_id.len() > HAS_ID_CONDITION_ARC_THRESHOLD { + HasIdCondition { + has_id: MaybeArc::arc(has_id), + } + } else { + HasIdCondition { + has_id: MaybeArc::no_arc(has_id), + } + } } } impl FromIterator for HasIdCondition { fn from_iter>(iter: T) -> Self { - HasIdCondition { - has_id: iter.into_iter().collect(), - } + let items: AHashSet<_> = iter.into_iter().collect(); + // Arc-Threshold applies here, since we're reusing the From implementation from AHashSet. + Self::from(items) } } diff --git a/lib/segment/src/utils/maybe_arc.rs b/lib/segment/src/utils/maybe_arc.rs new file mode 100644 index 0000000000..12260facf0 --- /dev/null +++ b/lib/segment/src/utils/maybe_arc.rs @@ -0,0 +1,111 @@ +use std::ops::Deref; +use std::sync::Arc; + +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +// Structure that acts as `T` most of the time but allows to interchange being wrapped within an `Arc` or not. +// This is helpful, when a variable can become memory-intensive but must remain the ability to get cloned. +#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq)] +#[serde(untagged)] // Make this type transparent when de/serializing and always deserialize as `NoArc`, since it's the first enum kind that matches. +pub enum MaybeArc { + NoArc(T), + Arc(Arc), +} + +impl MaybeArc { + /// Create a new `MaybeArc` wrapper that uses an `Arc` internally. + #[inline] + pub fn arc(t: T) -> Self { + Self::Arc(Arc::new(t)) + } + + /// Create a new `MaybeArc` wrapper that doesn't use an `Arc` internally. + #[inline] + pub fn no_arc(t: T) -> Self { + Self::NoArc(t) + } + + /// Returns `true` if the value is wrapped around an `Arc`. + pub fn is_arc(&self) -> bool { + matches!(self, Self::Arc(..)) + } +} + +impl AsRef for MaybeArc { + #[inline] + fn as_ref(&self) -> &T { + self + } +} + +impl MaybeArc { + /// Converts the `MaybeArc` back to `T`, potentially cloning the inner value + /// in case it's an `Arc` that has existing references. + #[inline] + pub fn into_inner(self) -> T { + match self { + Self::Arc(a) => Arc::unwrap_or_clone(a), + Self::NoArc(a) => a, + } + } +} + +impl Deref for MaybeArc { + type Target = T; + + #[inline] + fn deref(&self) -> &Self::Target { + match self { + Self::Arc(a) => a, + Self::NoArc(a) => a, + } + } +} + +impl FromIterator for MaybeArc +where + T: FromIterator, +{ + fn from_iter>(iter: U) -> Self { + let inner = T::from_iter(iter); + + // Using `NoArc` as default implementation to stay as close as possible to the type `T` and + // don't accidentally introducing overhead. + // A caller can always manually create the `MaybeArc` if using an `Arc` is preferred. + MaybeArc::NoArc(inner) + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_serializing() { + let original = String::from("42"); + + let ma_original = MaybeArc::arc(original.clone()); + + let encoded = serde_json::to_string(&ma_original).unwrap(); + + let decoded: MaybeArc = serde_json::from_str(&encoded).unwrap(); + assert_eq!(decoded.as_ref(), &original); + assert!(!decoded.is_arc()); // Always using `NoArc` to deserialize. + + // `MaybeArc` can be deserialized as inner type, since information about arc is not serialized. + let decoded: String = serde_json::from_str(&encoded).unwrap(); + assert_eq!(decoded, original); + } + + #[test] + fn test_deserializing() { + let original = String::from("42"); + let encoded = serde_json::to_string(&original).unwrap(); + + // Any type can be deserialized as `MaybeArc`, defaulting to `NoArc`. + let decoded: MaybeArc = serde_json::from_str(&encoded).unwrap(); + assert_eq!(decoded.as_ref(), &original); + assert!(!decoded.is_arc()); + } +} diff --git a/lib/segment/src/utils/mod.rs b/lib/segment/src/utils/mod.rs index 43a210b359..2ab223ec54 100644 --- a/lib/segment/src/utils/mod.rs +++ b/lib/segment/src/utils/mod.rs @@ -1,5 +1,6 @@ pub mod fmt; pub mod fs; +pub mod maybe_arc; pub mod mem; pub mod path; pub mod scored_point_ties;