Optimize process search results (#8163)

lib/segment/src/entry/entry_point.rs

@@ -3,6 +3,7 @@ use std::path::PathBuf;
 use std::sync::Arc;
 use std::sync::atomic::AtomicBool;
 
+use ahash::AHashMap;
 use common::counter::hardware_counter::HardwareCounterCell;
 use common::types::TelemetryDetail;
 use uuid::Uuid;
@@ -21,9 +22,9 @@ use crate::index::field_index::{CardinalityEstimation, FieldIndex};
 use crate::json_path::JsonPath;
 use crate::telemetry::SegmentTelemetry;
 use crate::types::{
-    Filter, Payload, PayloadFieldSchema, PayloadKeyType, PayloadKeyTypeRef, PointIdType,
-    ScoredPoint, SearchParams, SegmentConfig, SegmentInfo, SegmentType, SeqNumberType, VectorName,
-    VectorNameBuf, WithPayload, WithVector,
+    ExtendedPointId, Filter, Payload, PayloadFieldSchema, PayloadKeyType, PayloadKeyTypeRef,
+    PointIdType, ScoredPoint, SearchParams, SegmentConfig, SegmentInfo, SegmentType, SeqNumberType,
+    VectorName, VectorNameBuf, WithPayload, WithVector,
 };
 
 /// Define all operations which can be performed with non-appendable Segment or Segment-like entity.
@@ -91,7 +92,7 @@ pub trait NonAppendableSegmentEntry: SnapshotEntry {
         with_vector: &WithVector,
         hw_counter: &HardwareCounterCell,
         is_stopped: &AtomicBool,
-    ) -> OperationResult<Vec<SegmentRecord>>;
+    ) -> OperationResult<AHashMap<ExtendedPointId, SegmentRecord>>;
 
     /// Retrieve payload for the point
     /// If not found, return empty payload

lib/segment/src/segment/entry.rs

@@ -30,9 +30,9 @@ use crate::json_path::JsonPath;
 use crate::payload_storage::PayloadStorage;
 use crate::telemetry::SegmentTelemetry;
 use crate::types::{
-    Filter, Payload, PayloadFieldSchema, PayloadIndexInfo, PayloadKeyType, PayloadKeyTypeRef,
-    PointIdType, ScoredPoint, SearchParams, SegmentConfig, SegmentInfo, SegmentType, SeqNumberType,
-    VectorDataInfo, VectorName, VectorNameBuf, WithPayload, WithVector,
+    ExtendedPointId, Filter, Payload, PayloadFieldSchema, PayloadIndexInfo, PayloadKeyType,
+    PayloadKeyTypeRef, PointIdType, ScoredPoint, SearchParams, SegmentConfig, SegmentInfo,
+    SegmentType, SeqNumberType, VectorDataInfo, VectorName, VectorNameBuf, WithPayload, WithVector,
 };
 use crate::vector_storage::VectorStorage;
 
@@ -173,7 +173,7 @@ impl NonAppendableSegmentEntry for Segment {
         with_vector: &WithVector,
         hw_counter: &HardwareCounterCell,
         is_stopped: &AtomicBool,
-    ) -> OperationResult<Vec<SegmentRecord>> {
+    ) -> OperationResult<AHashMap<ExtendedPointId, SegmentRecord>> {
         let mut records = AHashMap::with_capacity(point_ids.len());
 
         let mut update_record_vector =
@@ -244,7 +244,7 @@ impl NonAppendableSegmentEntry for Segment {
             point_record.payload = payload;
         }
 
-        Ok(records.into_values().collect())
+        Ok(records)
     }
 
     fn iter_points(&self) -> Box<dyn Iterator<Item = PointIdType>> {

lib/segment/src/segment/search.rs

@@ -1,6 +1,5 @@
 use std::sync::atomic::AtomicBool;
 
-use ahash::AHashMap;
 use common::counter::hardware_counter::HardwareCounterCell;
 use common::types::ScoredPointOffset;
 
@@ -45,17 +44,13 @@ impl Segment {
             })
             .unzip();
 
-        let mut segment_records: AHashMap<_, _> = self
-            .retrieve(
-                &point_ids,
-                with_payload,
-                with_vector,
-                hw_counter,
-                is_stopped,
-            )?
-            .into_iter()
-            .map(|record| (record.id, record))
-            .collect();
+        let mut segment_records = self.retrieve(
+            &point_ids,
+            with_payload,
+            with_vector,
+            hw_counter,
+            is_stopped,
+        )?;
 
         let mut results = Vec::with_capacity(point_ids.len());
 

lib/shard/src/proxy_segment/segment_entry.rs

@@ -4,6 +4,7 @@ use std::path::PathBuf;
 use std::sync::Arc;
 use std::sync::atomic::AtomicBool;
 
+use ahash::AHashMap;
 use common::counter::hardware_counter::HardwareCounterCell;
 use common::types::TelemetryDetail;
 use segment::common::Flusher;
@@ -215,7 +216,7 @@ impl NonAppendableSegmentEntry for ProxySegment {
         with_vector: &WithVector,
         hw_counter: &HardwareCounterCell,
         is_stopped: &AtomicBool,
-    ) -> OperationResult<Vec<SegmentRecord>> {
+    ) -> OperationResult<AHashMap<ExtendedPointId, SegmentRecord>> {
         let filtered_point_ids: Vec<PointIdType> = point_ids
             .iter()
             .copied()

lib/shard/src/retrieve/retrieve_blocking.rs

@@ -45,7 +45,7 @@ pub fn retrieve_blocking(
             *version_entry.or_default() = version;
         }
 
-        for record in segment.retrieve(
+        for (id, record) in segment.retrieve(
             &newer_version_points,
             with_payload,
             with_vector,
@@ -53,7 +53,7 @@ pub fn retrieve_blocking(
             is_stopped,
         )? {
             // We expect all points to be found since we already checked their versions
-            point_records.insert(record.id, RecordInternal::from(record));
+            point_records.insert(id, RecordInternal::from(record));
             applied += 1;
         }
 

lib/shard/src/update.rs

@@ -464,8 +464,8 @@ pub fn sync_points(
                 segment.retrieve(ids, &with_payload, &with_vector, hw_counter, &is_stopped)?;
             let mut updated = 0;
 
-            for stored_record in stored_records {
-                let point = id_to_point.get(&stored_record.id).unwrap();
+            for (id, stored_record) in stored_records {
+                let point = id_to_point.get(&id).unwrap();
                 if !point.is_equal_to(&stored_record) {
                     points_to_update.push(*point);
                     updated += 1;