Fix enormous memory usage in Distance Matrix API on high sample size (#6640)

* Fix memory leak in distance API

* Add tests

lib/api/src/grpc/conversions.rs

@@ -1345,14 +1345,14 @@ impl TryFrom<HasIdCondition> for segment::types::HasIdCondition {
             .into_iter()
             .map(|p| p.try_into())
             .collect::<Result<_, _>>()?;
-        Ok(Self { has_id: set })
+        Ok(Self::from(set))
     }
 }
 
 impl From<segment::types::HasIdCondition> for HasIdCondition {
     fn from(value: segment::types::HasIdCondition) -> Self {
         let segment::types::HasIdCondition { has_id } = value;
-        let set: Vec<PointId> = has_id.into_iter().map(|p| p.into()).collect();
+        let set: Vec<PointId> = has_id.into_inner().into_iter().map(PointId::from).collect();
         Self { has_id: set }
     }
 }

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<JsonPath> for IsEmptyCondition {
 #[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq)]
 pub struct HasIdCondition {
     #[schemars(schema_with = "HashSet::<PointIdType>::json_schema")]
-    pub has_id: AHashSet<PointIdType>,
+    pub has_id: MaybeArc<AHashSet<PointIdType>>,
 }
 
 /// Filter points which have specific vector assigned
@@ -2623,17 +2624,30 @@ impl From<VectorNameBuf> 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<AHashSet<PointIdType>> for HasIdCondition {
     fn from(has_id: AHashSet<PointIdType>) -> 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<PointIdType> for HasIdCondition {
     fn from_iter<T: IntoIterator<Item = PointIdType>>(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)
     }
 }
 

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<T> {
+    NoArc(T),
+    Arc(Arc<T>),
+}
+
+impl<T> MaybeArc<T> {
+    /// 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<T> AsRef<T> for MaybeArc<T> {
+    #[inline]
+    fn as_ref(&self) -> &T {
+        self
+    }
+}
+
+impl<T: Clone> MaybeArc<T> {
+    /// 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<T> Deref for MaybeArc<T> {
+    type Target = T;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        match self {
+            Self::Arc(a) => a,
+            Self::NoArc(a) => a,
+        }
+    }
+}
+
+impl<T, I> FromIterator<I> for MaybeArc<T>
+where
+    T: FromIterator<I>,
+{
+    fn from_iter<U: IntoIterator<Item = I>>(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<String> = 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<String> = serde_json::from_str(&encoded).unwrap();
+        assert_eq!(decoded.as_ref(), &original);
+        assert!(!decoded.is_arc());
+    }
+}

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;