Sparse idf dot (#4126)

* introduce QueryContext, which accumulates runtime info needed for executing search

* fmt

* propagate query context into segment internals

* [WIP] prepare idf stats for search query context

* Split SparseVector and RemmapedSparseVector to guarantee we will not mix them up on the type level

* implement filling of the query context with IDF statistics

* implement re-weighting of the sparse query with idf

* fmt

* update idf param only if explicitly specified (more consistent with diff param update

* replace idf bool with modifier enum, improve further extensibility

* test and fixes

* Update lib/collection/src/operations/types.rs

Co-authored-by: Arnaud Gourlay <arnaud.gourlay@gmail.com>

* review fixes

* fmt

---------

Co-authored-by: Arnaud Gourlay <arnaud.gourlay@gmail.com>

docs/grpc/docs.md

@@ -79,6 +79,7 @@
     - [CompressionRatio](#qdrant-CompressionRatio)
     - [Datatype](#qdrant-Datatype)
     - [Distance](#qdrant-Distance)
+    - [Modifier](#qdrant-Modifier)
     - [PayloadSchemaType](#qdrant-PayloadSchemaType)
     - [QuantizationType](#qdrant-QuantizationType)
     - [ReplicaState](#qdrant-ReplicaState)
@@ -1213,6 +1214,7 @@ Note: 1kB = 1 vector of size 256. |
 | Field | Type | Label | Description |
 | ----- | ---- | ----- | ----------- |
 | index | [SparseIndexConfig](#qdrant-SparseIndexConfig) | optional | Configuration of sparse index |
+| modifier | [Modifier](#qdrant-Modifier) | optional | If set - apply modifier to the vector values |
 
 
 
@@ -1504,6 +1506,18 @@ Note: 1kB = 1 vector of size 256. |
 
 
 
+<a name="qdrant-Modifier"></a>
+
+### Modifier
+
+
+| Name | Number | Description |
+| ---- | ------ | ----------- |
+| None | 0 |  |
+| Idf | 1 | Apply Inverse Document Frequency |
+
+
+
 <a name="qdrant-PayloadSchemaType"></a>
 
 ### PayloadSchemaType

docs/redoc/master/openapi.json

@@ -5711,6 +5711,17 @@
                 "nullable": true
               }
             ]
+          },
+          "modifier": {
+            "description": "Configures addition value modifications for sparse vectors. Default: none",
+            "anyOf": [
+              {
+                "$ref": "#/components/schemas/Modifier"
+              },
+              {
+                "nullable": true
+              }
+            ]
           }
         }
       },
@@ -5732,6 +5743,14 @@
           }
         }
       },
+      "Modifier": {
+        "description": "If used, include weight modification, which will be applied to sparse vectors at query time: None - no modification (default) Idf - inverse document frequency, based on statistics of the collection",
+        "type": "string",
+        "enum": [
+          "none",
+          "idf"
+        ]
+      },
       "HnswConfig": {
         "description": "Config of HNSW index",
         "type": "object",

lib/api/src/grpc/proto/collections.proto

@@ -45,8 +45,14 @@ message VectorsConfigDiff {
   }
 }
 
+enum Modifier {
+    None = 0;
+    Idf = 1; // Apply Inverse Document Frequency
+}
+
 message SparseVectorParams {
   optional SparseIndexConfig index = 1; // Configuration of sparse index
+  optional Modifier modifier = 2; // If set - apply modifier to the vector values
 }
 
 message SparseVectorConfig {

lib/api/src/grpc/qdrant.rs

@@ -112,6 +112,9 @@ pub struct SparseVectorParams {
     /// Configuration of sparse index
     #[prost(message, optional, tag = "1")]
     pub index: ::core::option::Option<SparseIndexConfig>,
+    /// If set - apply modifier to the vector values
+    #[prost(enumeration = "Modifier", optional, tag = "2")]
+    pub modifier: ::core::option::Option<i32>,
 }
 #[derive(serde::Serialize)]
 #[allow(clippy::derive_partial_eq_without_eq)]
@@ -1111,6 +1114,34 @@ impl Datatype {
 #[derive(serde::Serialize)]
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
 #[repr(i32)]
+pub enum Modifier {
+    None = 0,
+    /// Apply Inverse Document Frequency
+    Idf = 1,
+}
+impl Modifier {
+    /// String value of the enum field names used in the ProtoBuf definition.
+    ///
+    /// The values are not transformed in any way and thus are considered stable
+    /// (if the ProtoBuf definition does not change) and safe for programmatic use.
+    pub fn as_str_name(&self) -> &'static str {
+        match self {
+            Modifier::None => "None",
+            Modifier::Idf => "Idf",
+        }
+    }
+    /// Creates an enum from field names used in the ProtoBuf definition.
+    pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
+        match value {
+            "None" => Some(Self::None),
+            "Idf" => Some(Self::Idf),
+            _ => None,
+        }
+    }
+}
+#[derive(serde::Serialize)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
+#[repr(i32)]
 pub enum Distance {
     UnknownDistance = 0,
     Cosine = 1,

lib/collection/src/collection/search.rs

@@ -9,6 +9,7 @@ use segment::types::{ExtendedPointId, Order, ScoredPoint, WithPayloadInterface,
 
 use super::Collection;
 use crate::operations::consistency_params::ReadConsistency;
+use crate::operations::query_enum::QueryEnum;
 use crate::operations::shard_selector_internal::ShardSelectorInternal;
 use crate::operations::types::*;
 

lib/collection/src/collection_manager/holders/proxy_segment.rs

@@ -9,6 +9,7 @@ use parking_lot::{RwLock, RwLockUpgradableReadGuard};
 use segment::common::operation_error::{OperationResult, SegmentFailedState};
 use segment::data_types::named_vectors::NamedVectors;
 use segment::data_types::order_by::OrderingValue;
+use segment::data_types::query_context::QueryContext;
 use segment::data_types::vectors::{QueryVector, Vector};
 use segment::entry::entry_point::SegmentEntry;
 use segment::index::field_index::CardinalityEstimation;
@@ -238,7 +239,7 @@ impl ProxySegment {
             top,
             params,
             &false.into(),
-            usize::MAX,
+            &Default::default(),
         )?;
 
         Ok(result.into_iter().next().unwrap())
@@ -272,7 +273,7 @@ impl SegmentEntry for ProxySegment {
         top: usize,
         params: Option<&SearchParams>,
         is_stopped: &AtomicBool,
-        search_optimized_threshold_kb: usize,
+        query_context: &QueryContext,
     ) -> OperationResult<Vec<Vec<ScoredPoint>>> {
         let deleted_points = self.deleted_points.read();
 
@@ -296,7 +297,7 @@ impl SegmentEntry for ProxySegment {
                 top,
                 params,
                 is_stopped,
-                search_optimized_threshold_kb,
+                query_context,
             )?
         } else {
             self.wrapped_segment.get().read().search_batch(
@@ -308,7 +309,7 @@ impl SegmentEntry for ProxySegment {
                 top,
                 params,
                 is_stopped,
-                search_optimized_threshold_kb,
+                query_context,
             )?
         };
         let mut write_results = self.write_segment.get().read().search_batch(
@@ -320,7 +321,7 @@ impl SegmentEntry for ProxySegment {
             top,
             params,
             is_stopped,
-            search_optimized_threshold_kb,
+            query_context,
         )?;
         for (index, write_result) in write_results.iter_mut().enumerate() {
             wrapped_results[index].append(write_result)
@@ -854,6 +855,14 @@ impl SegmentEntry for ProxySegment {
     fn get_telemetry_data(&self, detail: TelemetryDetail) -> SegmentTelemetry {
         self.wrapped_segment.get().read().get_telemetry_data(detail)
     }
+
+    fn fill_query_context(&self, query_context: &mut QueryContext) {
+        // Information from temporary segment is not too important for query context
+        self.wrapped_segment
+            .get()
+            .read()
+            .fill_query_context(query_context)
+    }
 }
 
 #[cfg(test)]
@@ -992,7 +1001,7 @@ mod tests {
                 10,
                 None,
                 &false.into(),
-                10_000,
+                &Default::default(),
             )
             .unwrap();
 
@@ -1047,7 +1056,7 @@ mod tests {
                 10,
                 None,
                 &false.into(),
-                10_000,
+                &Default::default(),
             )
             .unwrap();
 
@@ -1112,7 +1121,7 @@ mod tests {
                 10,
                 None,
                 &false.into(),
-                10_000,
+                &Default::default(),
             )
             .unwrap();