Facets in internal service (qdrant#4790)

* add `facet` to shard trait

* Add internal service

* gen grpc docs

docs/grpc/docs.md

@@ -138,6 +138,8 @@
     - [DiscoverInput](#qdrant-DiscoverInput)
     - [DiscoverPoints](#qdrant-DiscoverPoints)
     - [DiscoverResponse](#qdrant-DiscoverResponse)
+    - [FacetValue](#qdrant-FacetValue)
+    - [FacetValueHit](#qdrant-FacetValueHit)
     - [FieldCondition](#qdrant-FieldCondition)
     - [Filter](#qdrant-Filter)
     - [GeoBoundingBox](#qdrant-GeoBoundingBox)
@@ -2407,6 +2409,37 @@ The JSON representation for `Value` is a JSON value.
 
 
 
+<a name="qdrant-FacetValue"></a>
+
+### FacetValue
+
+
+
+| Field | Type | Label | Description |
+| ----- | ---- | ----- | ----------- |
+| string_value | [string](#string) |  | String value from the facet |
+
+
+
+
+
+
+<a name="qdrant-FacetValueHit"></a>
+
+### FacetValueHit
+
+
+
+| Field | Type | Label | Description |
+| ----- | ---- | ----- | ----------- |
+| value | [FacetValue](#qdrant-FacetValue) |  | Value from the facet |
+| count | [uint64](#uint64) |  | Number of points with this value |
+
+
+
+
+
+
 <a name="qdrant-FieldCondition"></a>
 
 ### FieldCondition

lib/api/build.rs

@@ -304,6 +304,7 @@ fn configure_validation(builder: Builder) -> Builder {
             ("SyncPoints.collection_name", "length(min = 1, max = 255)"),
             ("QueryBatchPointsInternal.collection_name", "length(min = 1, max = 255)"),
             ("QueryBatchPointsInternal.timeout", "custom = \"crate::grpc::validate::validate_u64_range_min_1\""),
+            ("FacetCountsInternal.collection_name", "length(min = 1, max = 255)"),
         ], &[])
         // Service: raft_service.proto
         .validates(&[

lib/api/src/grpc/conversions.rs

@@ -19,10 +19,11 @@ use uuid::Uuid;
 use super::qdrant::raw_query::RawContextPair;
 use super::qdrant::{
     raw_query, start_from, BinaryQuantization, BoolIndexParams, CompressionRatio,
-    DatetimeIndexParams, DatetimeRange, Direction, FieldType, FloatIndexParams, GeoIndexParams,
-    GeoLineString, GroupId, KeywordIndexParams, LookupLocation, MultiVectorComparator,
-    MultiVectorConfig, OrderBy, OrderValue, Range, RawVector, RecommendStrategy, SearchPointGroups,
-    SearchPoints, ShardKeySelector, SparseIndices, StartFrom, UuidIndexParams, WithLookup,
+    DatetimeIndexParams, DatetimeRange, Direction, FacetValue, FacetValueHit, FieldType,
+    FloatIndexParams, GeoIndexParams, GeoLineString, GroupId, KeywordIndexParams, LookupLocation,
+    MultiVectorComparator, MultiVectorConfig, OrderBy, OrderValue, Range, RawVector,
+    RecommendStrategy, SearchPointGroups, SearchPoints, ShardKeySelector, SparseIndices, StartFrom,
+    UuidIndexParams, WithLookup,
 };
 use crate::grpc::models::{CollectionsResponse, VersionInfo};
 use crate::grpc::qdrant::condition::ConditionOneOf;
@@ -2225,3 +2226,36 @@ impl From<LookupLocation> for rest::LookupLocation {
         }
     }
 }
+
+impl TryFrom<FacetValueHit> for segment::data_types::facets::FacetValueHit {
+    type Error = Status;
+
+    fn try_from(hit: FacetValueHit) -> Result<Self, Self::Error> {
+        let value = hit
+            .value
+            .ok_or_else(|| Status::internal("expected FacetValueHit to have a value"))?;
+
+        Ok(Self {
+            value: segment::data_types::facets::FacetValue::try_from(value)?,
+            count: hit.count as usize,
+        })
+    }
+}
+
+impl TryFrom<FacetValue> for segment::data_types::facets::FacetValue {
+    type Error = Status;
+
+    fn try_from(value: FacetValue) -> Result<Self, Self::Error> {
+        use segment::data_types::facets as segment;
+
+        use super::qdrant::facet_value::Variant;
+
+        let variant = value
+            .variant
+            .ok_or_else(|| Status::internal("expected FacetValue to have a value"))?;
+
+        Ok(match variant {
+            Variant::StringValue(value) => segment::FacetValue::Keyword(value),
+        })
+    }
+}

lib/api/src/grpc/proto/points.proto

@@ -586,6 +586,17 @@ message QueryPointGroups {
   optional ShardKeySelector shard_key_selector = 17; // Specify in which shards to look for the points, if not specified - look in all shards
 }
 
+message FacetValue {
+    oneof variant {
+        string string_value = 1; // String value from the facet
+    }
+}
+
+message FacetValueHit {
+    FacetValue value = 1; // Value from the facet
+    uint64 count = 2; // Number of points with this value
+}
+
 message PointsUpdateOperation {
   message PointStructList {
     repeated PointStruct points = 1;

lib/api/src/grpc/proto/points_internal_service.proto

@@ -23,6 +23,7 @@ service PointsInternal {
   rpc Recommend (RecommendPointsInternal) returns (RecommendResponse) {}
   rpc Get (GetPointsInternal) returns (GetResponse) {}
   rpc QueryBatch (QueryBatchPointsInternal) returns (QueryBatchResponseInternal) {}
+  rpc Facet(FacetCountsInternal) returns (FacetResponseInternal) {}
 }
 
 
@@ -287,3 +288,16 @@ message QueryBatchResponseInternal {
   repeated QueryResultInternal results = 1;
   double time = 2; // Time spent to process
 }
+
+message FacetCountsInternal {
+    string collection_name = 1;
+    string key = 2;
+    optional Filter filter = 3;
+    uint64 limit = 4;
+    uint32 shard_id = 5;
+}
+
+message FacetResponseInternal {
+    repeated FacetValueHit hits = 1;
+    double time = 2; // Time spent to process
+}

lib/api/src/grpc/qdrant.rs

@@ -5017,6 +5017,35 @@ pub struct QueryPointGroups {
 #[derive(serde::Serialize)]
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
+pub struct FacetValue {
+    #[prost(oneof = "facet_value::Variant", tags = "1")]
+    pub variant: ::core::option::Option<facet_value::Variant>,
+}
+/// Nested message and enum types in `FacetValue`.
+pub mod facet_value {
+    #[derive(serde::Serialize)]
+    #[allow(clippy::derive_partial_eq_without_eq)]
+    #[derive(Clone, PartialEq, ::prost::Oneof)]
+    pub enum Variant {
+        /// String value from the facet
+        #[prost(string, tag = "1")]
+        StringValue(::prost::alloc::string::String),
+    }
+}
+#[derive(serde::Serialize)]
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct FacetValueHit {
+    /// Value from the facet
+    #[prost(message, optional, tag = "1")]
+    pub value: ::core::option::Option<FacetValue>,
+    /// Number of points with this value
+    #[prost(uint64, tag = "2")]
+    pub count: u64,
+}
+#[derive(serde::Serialize)]
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
 pub struct PointsUpdateOperation {
     #[prost(
         oneof = "points_update_operation::Operation",
@@ -8787,6 +8816,33 @@ pub struct QueryBatchResponseInternal {
     #[prost(double, tag = "2")]
     pub time: f64,
 }
+#[derive(serde::Serialize)]
+#[derive(validator::Validate)]
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct FacetCountsInternal {
+    #[prost(string, tag = "1")]
+    #[validate(length(min = 1, max = 255))]
+    pub collection_name: ::prost::alloc::string::String,
+    #[prost(string, tag = "2")]
+    pub key: ::prost::alloc::string::String,
+    #[prost(message, optional, tag = "3")]
+    pub filter: ::core::option::Option<Filter>,
+    #[prost(uint64, tag = "4")]
+    pub limit: u64,
+    #[prost(uint32, tag = "5")]
+    pub shard_id: u32,
+}
+#[derive(serde::Serialize)]
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct FacetResponseInternal {
+    #[prost(message, repeated, tag = "1")]
+    pub hits: ::prost::alloc::vec::Vec<FacetValueHit>,
+    /// Time spent to process
+    #[prost(double, tag = "2")]
+    pub time: f64,
+}
 /// Generated client implementations.
 pub mod points_internal_client {
     #![allow(unused_variables, dead_code, missing_docs, clippy::let_unit_value)]
@@ -9287,6 +9343,31 @@ pub mod points_internal_client {
                 .insert(GrpcMethod::new("qdrant.PointsInternal", "QueryBatch"));
             self.inner.unary(req, path, codec).await
         }
+        pub async fn facet(
+            &mut self,
+            request: impl tonic::IntoRequest<super::FacetCountsInternal>,
+        ) -> std::result::Result<
+            tonic::Response<super::FacetResponseInternal>,
+            tonic::Status,
+        > {
+            self.inner
+                .ready()
+                .await
+                .map_err(|e| {
+                    tonic::Status::new(
+                        tonic::Code::Unknown,
+                        format!("Service was not ready: {}", e.into()),
+                    )
+                })?;
+            let codec = tonic::codec::ProstCodec::default();
+            let path = http::uri::PathAndQuery::from_static(
+                "/qdrant.PointsInternal/Facet",
+            );
+            let mut req = request.into_request();
+            req.extensions_mut()
+                .insert(GrpcMethod::new("qdrant.PointsInternal", "Facet"));
+            self.inner.unary(req, path, codec).await
+        }
     }
 }
 /// Generated server implementations.
@@ -9406,6 +9487,13 @@ pub mod points_internal_server {
             tonic::Response<super::QueryBatchResponseInternal>,
             tonic::Status,
         >;
+        async fn facet(
+            &self,
+            request: tonic::Request<super::FacetCountsInternal>,
+        ) -> std::result::Result<
+            tonic::Response<super::FacetResponseInternal>,
+            tonic::Status,
+        >;
     }
     #[derive(Debug)]
     pub struct PointsInternalServer<T: PointsInternal> {
@@ -10278,6 +10366,52 @@ pub mod points_internal_server {
                     };
                     Box::pin(fut)
                 }
+                "/qdrant.PointsInternal/Facet" => {
+                    #[allow(non_camel_case_types)]
+                    struct FacetSvc<T: PointsInternal>(pub Arc<T>);
+                    impl<
+                        T: PointsInternal,
+                    > tonic::server::UnaryService<super::FacetCountsInternal>
+                    for FacetSvc<T> {
+                        type Response = super::FacetResponseInternal;
+                        type Future = BoxFuture<
+                            tonic::Response<Self::Response>,
+                            tonic::Status,
+                        >;
+                        fn call(
+                            &mut self,
+                            request: tonic::Request<super::FacetCountsInternal>,
+                        ) -> Self::Future {
+                            let inner = Arc::clone(&self.0);
+                            let fut = async move {
+                                <T as PointsInternal>::facet(&inner, request).await
+                            };
+                            Box::pin(fut)
+                        }
+                    }
+                    let accept_compression_encodings = self.accept_compression_encodings;
+                    let send_compression_encodings = self.send_compression_encodings;
+                    let max_decoding_message_size = self.max_decoding_message_size;
+                    let max_encoding_message_size = self.max_encoding_message_size;
+                    let inner = self.inner.clone();
+                    let fut = async move {
+                        let inner = inner.0;
+                        let method = FacetSvc(inner);
+                        let codec = tonic::codec::ProstCodec::default();
+                        let mut grpc = tonic::server::Grpc::new(codec)
+                            .apply_compression_config(
+                                accept_compression_encodings,
+                                send_compression_encodings,
+                            )
+                            .apply_max_message_size_config(
+                                max_decoding_message_size,
+                                max_encoding_message_size,
+                            );
+                        let res = grpc.unary(method, req).await;
+                        Ok(res)
+                    };
+                    Box::pin(fut)
+                }
                 _ => {
                     Box::pin(async move {
                         Ok(

lib/collection/src/shards/dummy_shard.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use async_trait::async_trait;
+use segment::data_types::facets::{FacetRequest, FacetResponse};
 use segment::data_types::order_by::OrderBy;
 use segment::types::{
     ExtendedPointId, Filter, ScoredPoint, WithPayload, WithPayloadInterface, WithVector,
@@ -110,4 +111,12 @@ impl ShardOperation for DummyShard {
     ) -> CollectionResult<Vec<ShardQueryResponse>> {
         self.dummy()
     }
+
+    async fn facet(
+        &self,
+        _: Arc<FacetRequest>,
+        _search_runtime_handle: &Handle,
+    ) -> CollectionResult<FacetResponse> {
+        self.dummy()
+    }
 }

lib/collection/src/shards/forward_proxy_shard.rs

@@ -4,6 +4,7 @@ use std::time::Duration;
 
 use async_trait::async_trait;
 use common::types::TelemetryDetail;
+use segment::data_types::facets::{FacetRequest, FacetResponse};
 use segment::data_types::order_by::OrderBy;
 use segment::types::{
     ExtendedPointId, Filter, PointIdType, ScoredPoint, WithPayload, WithPayloadInterface,
@@ -378,4 +379,13 @@ impl ShardOperation for ForwardProxyShard {
             .query_batch(requests, search_runtime_handle, timeout)
             .await
     }
+
+    async fn facet(
+        &self,
+        request: Arc<FacetRequest>,
+        search_runtime_handle: &Handle,
+    ) -> CollectionResult<FacetResponse> {
+        let local_shard = &self.wrapped_shard;
+        local_shard.facet(request, search_runtime_handle).await
+    }
 }

lib/collection/src/shards/local_shard/facet.rs

@@ -10,7 +10,7 @@ use crate::collection_manager::holders::segment_holder::LockedSegment;
 use crate::operations::types::CollectionResult;
 
 impl LocalShard {
-    pub async fn facet(
+    pub async fn do_facet(
         &self,
         request: Arc<FacetRequest>,
         search_runtime_handle: &Handle,

lib/collection/src/shards/local_shard/shard_ops.rs

@@ -2,6 +2,7 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use async_trait::async_trait;
+use segment::data_types::facets::{FacetRequest, FacetResponse};
 use segment::data_types::order_by::OrderBy;
 use segment::types::{
     ExtendedPointId, Filter, ScoredPoint, WithPayload, WithPayloadInterface, WithVector,
@@ -207,4 +208,13 @@ impl ShardOperation for LocalShard {
         self.do_planned_query(planned_query, search_runtime_handle, timeout)
             .await
     }
+
+    async fn facet(
+        &self,
+        request: Arc<FacetRequest>,
+        search_runtime_handle: &Handle,
+    ) -> CollectionResult<FacetResponse> {
+        let hits = self.do_facet(request, search_runtime_handle).await?;
+        Ok(FacetResponse { hits })
+    }
 }