universal-query: Expose batch query in REST (#4497)

Exposes the ability to query in batch

docs/redoc/master/openapi.json

@@ -5186,6 +5186,16 @@
             "schema": {
               "$ref": "#/components/schemas/ReadConsistency"
             }
+          },
+          {
+            "name": "timeout",
+            "in": "query",
+            "description": "If set, overrides global timeout for this request. Unit is seconds.",
+            "required": false,
+            "schema": {
+              "type": "integer",
+              "minimum": 1
+            }
           }
         ],
         "responses": {
@@ -5237,6 +5247,107 @@
           }
         }
       }
+    },
+    "/collections/{collection_name}/points/query/batch": {
+      "post": {
+        "tags": [
+          "points"
+        ],
+        "summary": "Query points in batch",
+        "description": "Universally query points in batch. This endpoint covers all capabilities of search, recommend, discover, filters. But also enables hybrid and multi-stage queries.",
+        "operationId": "query_batch_points",
+        "requestBody": {
+          "description": "Describes the queries to make to the collection",
+          "content": {
+            "application/json": {
+              "schema": {
+                "$ref": "#/components/schemas/QueryRequestBatch"
+              }
+            }
+          }
+        },
+        "parameters": [
+          {
+            "name": "collection_name",
+            "in": "path",
+            "description": "Name of the collection to query",
+            "required": true,
+            "schema": {
+              "type": "string"
+            }
+          },
+          {
+            "name": "consistency",
+            "in": "query",
+            "description": "Define read consistency guarantees for the operation",
+            "required": false,
+            "schema": {
+              "$ref": "#/components/schemas/ReadConsistency"
+            }
+          },
+          {
+            "name": "timeout",
+            "in": "query",
+            "description": "If set, overrides global timeout for this request. Unit is seconds.",
+            "required": false,
+            "schema": {
+              "type": "integer",
+              "minimum": 1
+            }
+          }
+        ],
+        "responses": {
+          "default": {
+            "description": "error",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/ErrorResponse"
+                }
+              }
+            }
+          },
+          "4XX": {
+            "description": "error",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/ErrorResponse"
+                }
+              }
+            }
+          },
+          "200": {
+            "description": "successful operation",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "type": "object",
+                  "properties": {
+                    "time": {
+                      "type": "number",
+                      "format": "float",
+                      "description": "Time spent to process this request"
+                    },
+                    "status": {
+                      "type": "string"
+                    },
+                    "result": {
+                      "type": "array",
+                      "items": {
+                        "type": "array",
+                        "items": {
+                          "$ref": "#/components/schemas/ScoredPoint"
+                        }
+                      }
+                    }
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
     }
   },
   "openapi": "3.0.1",
@@ -11827,6 +11938,20 @@
             ]
           }
         ]
+      },
+      "QueryRequestBatch": {
+        "type": "object",
+        "required": [
+          "searches"
+        ],
+        "properties": {
+          "searches": {
+            "type": "array",
+            "items": {
+              "$ref": "#/components/schemas/QueryRequest"
+            }
+          }
+        }
       }
     }
   }

lib/api/src/rest/schema.rs

@@ -203,6 +203,12 @@ pub struct QueryRequest {
     pub shard_key: Option<ShardKeySelector>,
 }
 
+#[derive(Debug, Serialize, Deserialize, JsonSchema, Validate)]
+pub struct QueryRequestBatch {
+    #[validate]
+    pub searches: Vec<QueryRequest>,
+}
+
 #[derive(Debug, Serialize, Deserialize, JsonSchema)]
 #[serde(untagged)]
 pub enum QueryInterface {

openapi/openapi-main.ytt.yaml

@@ -628,9 +628,53 @@ paths:
           required: false
           schema:
             $ref: "#/components/schemas/ReadConsistency"
-        #! TODO(universal-query): add timeout
+        - name: timeout
+          in: query
+          description: If set, overrides global timeout for this request. Unit is seconds.
+          required: false
+          schema:
+            type: integer
+            minimum: 1
 
       responses: #@ response(array(reference("ScoredPoint")))
+  
+  /collections/{collection_name}/points/query/batch:
+    post:
+      tags:
+        - points
+      summary: Query points in batch
+      description: Universally query points in batch. This endpoint covers all capabilities of search, recommend, discover, filters. But also enables hybrid and multi-stage queries.
+      operationId: query_batch_points
+      requestBody: 
+        description: Describes the queries to make to the collection
+        content:
+          application/json:
+            schema:
+              $ref: "#/components/schemas/QueryRequestBatch"
+              
+
+      parameters:
+        - name: collection_name
+          in: path
+          description: Name of the collection to query
+          required: true
+          schema:
+            type: string
+        - name: consistency
+          in: query
+          description: Define read consistency guarantees for the operation
+          required: false
+          schema:
+            $ref: "#/components/schemas/ReadConsistency"
+        - name: timeout
+          in: query
+          description: If set, overrides global timeout for this request. Unit is seconds.
+          required: false
+          schema:
+            type: integer
+            minimum: 1
+
+      responses: #@ response(array(array(reference("ScoredPoint"))))
 
 components:
   securitySchemes:

src/actix/api/query_api.rs

@@ -1,7 +1,8 @@
 use actix_web::{post, web, Responder};
 use actix_web_validator::{Json, Path, Query};
-use api::rest::QueryRequest;
+use api::rest::{QueryRequest, QueryRequestBatch};
 use collection::operations::shard_selector_internal::ShardSelectorInternal;
+use collection::operations::universal_query::collection_query::CollectionQueryRequest;
 use itertools::Itertools;
 use storage::content_manager::errors::StorageError;
 use storage::dispatcher::Dispatcher;
@@ -53,6 +54,59 @@ async fn query_points(
     .await
 }
 
+#[post("/collections/{name}/points/query/batch")]
+async fn query_points_batch(
+    dispatcher: web::Data<Dispatcher>,
+    collection: Path<CollectionPath>,
+    request: Json<QueryRequestBatch>,
+    params: Query<ReadParams>,
+    ActixAccess(access): ActixAccess,
+) -> impl Responder {
+    helpers::time(async move {
+        let QueryRequestBatch { searches } = request.into_inner();
+
+        let batch = searches
+            .into_iter()
+            .map(|request| {
+                let QueryRequest {
+                    internal,
+                    shard_key,
+                } = request;
+
+                let request = CollectionQueryRequest::from(internal);
+                let shard_selection = match shard_key {
+                    None => ShardSelectorInternal::All,
+                    Some(shard_keys) => shard_keys.into(),
+                };
+
+                (request, shard_selection)
+            })
+            .collect();
+
+        let res = dispatcher
+            .toc(&access)
+            .query_batch(
+                &collection.name,
+                batch,
+                params.consistency,
+                access,
+                params.timeout(),
+            )
+            .await?
+            .into_iter()
+            .map(|response| {
+                response
+                    .into_iter()
+                    .map(api::rest::ScoredPoint::from)
+                    .collect_vec()
+            })
+            .collect_vec();
+
+        Ok(res)
+    })
+    .await
+}
+
 pub fn config_query_api(cfg: &mut web::ServiceConfig) {
     cfg.service(query_points);
 }

src/common/metrics.rs

@@ -24,6 +24,7 @@ const REST_ENDPOINT_WHITELIST: &[&str] = &[
     "/collections/{name}/points/discover/batch",
     "/collections/{name}/points/payload",
     "/collections/{name}/points/query",
+    "/collections/{name}/points/query/batch",
     "/collections/{name}/points/recommend",
     "/collections/{name}/points/recommend/batch",
     "/collections/{name}/points/search",
@@ -40,6 +41,7 @@ const GRPC_ENDPOINT_WHITELIST: &[&str] = &[
     "/qdrant.Points/DiscoverBatch",
     "/qdrant.Points/OverwritePayload",
     "/qdrant.Points/Query",
+    "/qdrant.Points/QueryBatch",
     "/qdrant.Points/Recommend",
     "/qdrant.Points/RecommendBatch",
     "/qdrant.Points/Search",

src/schema_generator.rs

@@ -1,5 +1,5 @@
 use api::grpc::models::{CollectionsResponse, VersionInfo};
-use api::rest::{QueryRequest, Record, ScoredPoint};
+use api::rest::{QueryRequest, QueryRequestBatch, Record, ScoredPoint};
 use collection::operations::cluster_ops::ClusterOperations;
 use collection::operations::consistency_params::ReadConsistency;
 use collection::operations::payload_ops::{DeletePayload, SetPayload};
@@ -82,6 +82,7 @@ struct AllDefinitions {
     bc: VersionInfo,
     bd: CollectionExistence,
     be: QueryRequest,
+    bf: QueryRequestBatch,
 }
 
 fn save_schema<T: JsonSchema>() {

tests/consensus_tests/auth_tests/test_jwt_access.py

@@ -527,6 +527,9 @@ ACTION_ACCESS = {
     "query_points": EndpointAccess(
         True, True, True, "POST /collections/{collection_name}/points/query", "qdrant.Points/Query"
     ),
+    "query_batch_points": EndpointAccess(
+        True, True, True, "POST /collections/{collection_name}/points/query/batch", # "qdrant.Points/QueryBatch"
+    ),
     ### Service ###
     "root": EndpointAccess(True, True, True, "GET /", "qdrant.Qdrant/HealthCheck"),
     "readyz": EndpointAccess(True, True, True, "GET /readyz", "grpc.health.v1.Health/Check"),
@@ -1725,6 +1728,27 @@ def test_query_points():
             },
         },
     )
+
+def test_query_batch_points():
+    check_access(
+        "query_batch_points",
+        rest_request={ "searches": [{ "query": [0.1, 0.2, 0.3, 0.4] }] },
+        path_params={"collection_name": COLL_NAME},
+        grpc_request={
+            "collection_name": COLL_NAME, 
+            "query_points": [
+                { 
+                    "query": {
+                        "nearest": {
+                            "dense": {
+                                "data": [0.1, 0.2, 0.3, 0.4]
+                            }
+                        }
+                    }
+                }
+            ]
+        }
+    )
     
 
 def test_root():

tests/openapi_consistency_check.sh

@@ -36,7 +36,7 @@ fi
 rm -f ./docs/redoc/master/.diff.openapi.json
 
 NUMBER_OF_APIS=$(cat ./docs/redoc/master/openapi.json | jq '[.paths[] | length] | add')
-EXPECTED_NUMBER_OF_APIS=66
+EXPECTED_NUMBER_OF_APIS=67
 
 if [ "$NUMBER_OF_APIS" -ne "$EXPECTED_NUMBER_OF_APIS" ]; then
     echo "ERROR: It looks like the total number of APIs has changed."