[PERF] Validate KNN Projection Input Parameters Before Applying knn_filter (#3205)

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
- Ensure input parameters for projections are validated before applying
`knn_filter` to avoid unnecessary processing
   - Provide precise error message when scan operator inputs are invalid
   - Add unit tests to verify KNN queries and error scenarios
 - New functionality
   - ...

## Test plan
*How are these changes tested?*
 Added new unit tests
- [x] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Documentation Changes
*Are all docstrings for user-facing APIs updated if required? Do we need
to make documentation changes in the [docs
repository](https://github.com/chroma-core/docs)?*

Author: redouan-rhazouani

rust/worker/src/server.rs

@@ -143,14 +143,19 @@ impl WorkerServer {
         let collection = scan
             .collection
             .ok_or(Status::invalid_argument("Invalid Collection"))?;
+
         let collection_uuid = CollectionUuid::from_str(&collection.id)
-            .map_err(|err| Status::invalid_argument(err.to_string()))?;
+            .map_err(|_| Status::invalid_argument("Invalid Collection UUID"))?;
+
         let vector_uuid = SegmentUuid::from_str(&scan.knn_id)
-            .map_err(|err| Status::invalid_argument(err.to_string()))?;
+            .map_err(|_| Status::invalid_argument("Invalid UUID for Vector segment"))?;
+
         let metadata_uuid = SegmentUuid::from_str(&scan.metadata_id)
-            .map_err(|err| Status::invalid_argument(err.to_string()))?;
+            .map_err(|_| Status::invalid_argument("Invalid UUID for Metadata segment"))?;
+
         let record_uuid = SegmentUuid::from_str(&scan.record_id)
-            .map_err(|err| Status::invalid_argument(err.to_string()))?;
+            .map_err(|_| Status::invalid_argument("Invalid UUID for Record segment"))?;
+
         Ok((
             FetchLogOperator {
                 log_client: self.log.clone(),
@@ -213,13 +218,17 @@ impl WorkerServer {
         let scan = get_inner
             .scan
             .ok_or(Status::invalid_argument("Invalid Scan Operator"))?;
+
         let (fetch_log_operator, fetch_segment_operator) = self.decompose_proto_scan(scan)?;
+
         let filter = get_inner
             .filter
             .ok_or(Status::invalid_argument("Invalid Filter Operator"))?;
+
         let limit = get_inner
             .limit
             .ok_or(Status::invalid_argument("Invalid Scan Operator"))?;
+
         let projection = get_inner
             .projection
             .ok_or(Status::invalid_argument("Invalid Projection Operator"))?;
@@ -248,17 +257,28 @@ impl WorkerServer {
     ) -> Result<Response<KnnBatchResult>, Status> {
         let dispatcher = self.clone_dispatcher()?;
         let system = self.clone_system()?;
+
         let knn_inner = knn.into_inner();
+
         let scan = knn_inner
             .scan
             .ok_or(Status::invalid_argument("Invalid Scan Operator"))?;
+
         let (fetch_log_operator, fetch_segment_operator) = self.decompose_proto_scan(scan)?;
+
         let filter = knn_inner
             .filter
             .ok_or(Status::invalid_argument("Invalid Filter Operator"))?;
+
         let knn = knn_inner
             .knn
-            .ok_or(Status::invalid_argument("Invalid Scan Operator"))?;
+            .ok_or(Status::invalid_argument("Invalid Knn Operator"))?;
+
+        let projection = knn_inner
+            .projection
+            .ok_or(Status::invalid_argument("Invalid Projection Operator"))?;
+        let knn_projection = KnnProjectionOperator::try_from(projection)
+            .map_err(|e| Status::invalid_argument(format!("Invalid Projection Operator: {}", e)))?;
 
         if knn.embeddings.is_empty() {
             return Ok(Response::new(to_proto_knn_batch_result(Vec::new())?));
@@ -290,11 +310,6 @@ impl WorkerServer {
             }
         };
 
-        let projection = knn_inner
-            .projection
-            .ok_or(Status::invalid_argument("Invalid Projection Operator"))?;
-        let knn_projection = KnnProjectionOperator::try_from(projection)?;
-
         let knn_orchestrator_futures = from_proto_knn(knn)?
             .into_iter()
             .map(|knn| {
@@ -589,4 +604,180 @@ mod tests {
         assert!(response.is_err());
         assert_eq!(response.unwrap_err().code(), tonic::Code::InvalidArgument);
     }
+
+    fn gen_knn_request(mut scan_operator: Option<chroma_proto::ScanOperator>) -> chroma_proto::KnnPlan {
+        if scan_operator.is_none() {
+            scan_operator = Some(scan());
+        }
+        chroma_proto::KnnPlan {
+            scan: scan_operator,
+            filter: Some(chroma_proto::FilterOperator {
+                ids: None,
+                r#where: None,
+                where_document: None,
+            }),
+            knn: Some(chroma_proto::KnnOperator {
+                embeddings: vec![],
+                fetch: 0,
+            }),
+            projection: Some(chroma_proto::KnnProjectionOperator {
+                projection: Some(chroma_proto::ProjectionOperator {
+                    document: false,
+                    embedding: false,
+                    metadata: false,
+                }),
+                distance: false,
+            }),
+        }
+    }
+
+    #[tokio::test]
+    async fn validate_knn_plan_empty_embeddings() {
+        let mut executor = QueryExecutorClient::connect(run_server()).await.unwrap();
+        let response = executor.knn(gen_knn_request(None)).await;
+        assert!(response.is_ok());
+        assert_eq!(response.unwrap().into_inner().results.len(), 0);
+    }
+
+    #[tokio::test]
+    async fn validate_knn_plan_filter() {
+        let mut executor = QueryExecutorClient::connect(run_server()).await.unwrap();
+        let mut request = gen_knn_request(None);
+        request.filter = None;
+        let response = executor.knn(request).await;
+        let err = response.unwrap_err();
+        assert_eq!(err.code(), tonic::Code::InvalidArgument);
+        assert!(
+            err.message().to_lowercase().contains("filter operator"),
+            "{}",
+            err.message()
+        );
+    }
+
+    #[tokio::test]
+    async fn validate_knn_plan_knn() {
+        let mut executor = QueryExecutorClient::connect(run_server()).await.unwrap();
+        let mut request = gen_knn_request(None);
+        request.knn = None;
+        let response = executor.knn(request).await;
+        assert!(response.is_err());
+        let err = response.unwrap_err();
+        assert_eq!(err.code(), tonic::Code::InvalidArgument);
+        assert!(
+            err.message().to_lowercase().contains("knn operator"),
+            "{}",
+            err.message()
+        );
+    }
+
+    #[tokio::test]
+    async fn validate_knn_plan_projection() {
+        let mut executor = QueryExecutorClient::connect(run_server()).await.unwrap();
+        let mut request = gen_knn_request(None);
+        request.projection = None;
+        let response = executor.knn(request).await;
+        let err = response.unwrap_err();
+        assert_eq!(err.code(), tonic::Code::InvalidArgument);
+        assert!(
+            err.message().to_lowercase().contains("projection operator"),
+            "{}",
+            err.message()
+        );
+
+        let mut request = gen_knn_request(None);
+        request.projection = Some(chroma_proto::KnnProjectionOperator {
+            projection: None,
+            distance: false,
+        });
+        let response = executor.knn(request).await;
+        let err = response.unwrap_err();
+        assert_eq!(err.code(), tonic::Code::InvalidArgument);
+        assert!(
+            err.message()
+                .to_lowercase()
+                .contains("projection operator: "),
+            "{}",
+            err.message()
+        );
+    }
+
+    #[tokio::test]
+    async fn validate_knn_plan_scan() {
+        let mut executor = QueryExecutorClient::connect(run_server()).await.unwrap();
+        let mut request = gen_knn_request(None);
+        request.scan = None;
+        let response = executor.knn(request).await;
+        let err = response.unwrap_err();
+        assert_eq!(err.code(), tonic::Code::InvalidArgument);
+        assert!(
+            err.message().to_lowercase().contains("scan operator"),
+            "{}",
+            err.message()
+        );
+    }
+
+    #[tokio::test]
+    async fn validate_knn_plan_scan_collection() {
+        let mut executor = QueryExecutorClient::connect(run_server()).await.unwrap();
+        let mut scan = scan();
+        scan.collection.as_mut().unwrap().id = "Invalid-Collection-ID".to_string();
+        let response = executor.knn(gen_knn_request(Some(scan))).await;
+        assert!(response.is_err());
+        let err = response.unwrap_err();
+        assert_eq!(err.code(), tonic::Code::InvalidArgument);
+        assert!(
+            err.message().to_lowercase().contains("collection uuid"),
+            "{}",
+            err.message()
+        );
+    }
+
+    #[tokio::test]
+    async fn validate_knn_plan_scan_vector() {
+        let mut executor = QueryExecutorClient::connect(run_server()).await.unwrap();
+        // invalid vector uuid
+        let mut scan_operator = scan();
+        scan_operator.knn_id = "invalid_segment_id".to_string();
+        let response = executor.knn(gen_knn_request(Some(scan_operator))).await;
+        assert!(response.is_err());
+        let err = response.unwrap_err();
+        assert_eq!(err.code(), tonic::Code::InvalidArgument);
+        assert!(
+            err.message().to_lowercase().contains("vector"),
+            "{}",
+            err.message()
+        );
+    }
+
+    #[tokio::test]
+    async fn validate_knn_plan_scan_record() {
+        let mut executor = QueryExecutorClient::connect(run_server()).await.unwrap();
+        let mut scan_operator = scan();
+        scan_operator.record_id = "invalid_record_id".to_string();
+        let response = executor.knn(gen_knn_request(Some(scan_operator))).await;
+        assert!(response.is_err());
+        let err = response.unwrap_err();
+        assert_eq!(err.code(), tonic::Code::InvalidArgument);
+        assert!(
+            err.message().to_lowercase().contains("record"),
+            "{}",
+            err.message()
+        );
+    }
+
+    #[tokio::test]
+    async fn validate_knn_plan_scan_metadata() {
+        let mut executor = QueryExecutorClient::connect(run_server()).await.unwrap();
+        let mut scan_operator = scan();
+        scan_operator.metadata_id = "invalid_metadata_id".to_string();
+        let response = executor.knn(gen_knn_request(Some(scan_operator))).await;
+        assert!(response.is_err());
+        let err = response.unwrap_err();
+        assert_eq!(err.code(), tonic::Code::InvalidArgument);
+        assert!(
+            err.message().to_lowercase().contains("metadata"),
+            "{}",
+            err.message()
+        );
+    }
 }