Fix ANN flaky tests (#246)

Fix ANN LIMIT tests as the USearch library does not guarantee to return
all the items, but only those nearly matching the expected vector.

Refactor the `get_query_results` as it's a better practice not to
allocate a whole vector of result rows, but work with iterator.
It also helps to serialize the results into expected types.

Fix ANN expected results test to check for a valid PKs and if the rows
returned are the ones that was inserted.

Refs: VECTOR-221
Fixes: VECTOR-239

Author: knowack1

crates/validator/src/common.rs

@@ -10,7 +10,7 @@ use crate::vector_store_cluster::VectorStoreClusterExt;
 use httpclient::HttpClient;
 use scylla::client::session::Session;
 use scylla::client::session_builder::SessionBuilder;
-use scylla::value::Row;
+use scylla::response::query_result::QueryRowsResult;
 use std::net::Ipv4Addr;
 use std::sync::Arc;
 use std::time::Duration;
@@ -97,17 +97,13 @@ where
     .unwrap_or_else(|_| panic!("Timeout on: {msg}"))
 }
 
-pub(crate) async fn get_query_results(query: String, session: &Session) -> Vec<Row> {
+pub(crate) async fn get_query_results(query: String, session: &Session) -> QueryRowsResult {
     session
         .query_unpaged(query, ())
         .await
         .expect("failed to run query")
         .into_rows_result()
         .expect("failed to get rows")
-        .rows()
-        .unwrap()
-        .collect::<Result<Vec<Row>, _>>()
-        .expect("failed to decode rows")
 }
 
 pub(crate) async fn create_keyspace(session: &Session) -> String {

crates/validator/src/tests/ann.rs

@@ -5,8 +5,8 @@
 
 use crate::common::*;
 use crate::tests::*;
+use std::collections::HashMap;
 use std::time::Duration;
-use tracing::debug;
 use tracing::info;
 
 pub(crate) async fn new() -> TestCase {
@@ -39,17 +39,23 @@ async fn ann_query_returns_expected_results(actors: TestActors) {
     let keyspace = create_keyspace(&session).await;
     let table = create_table(&session, "pk INT PRIMARY KEY, v VECTOR<FLOAT, 3>", None).await;
 
-    // Insert 1000 vectors
+    // Create a map of pk -> embedding
+    let mut embeddings: HashMap<i32, Vec<f32>> = HashMap::new();
     for i in 0..1000 {
-        let embedding: Vec<f32> = vec![
+        let embedding = vec![
             if i < 100 { 0.0 } else { (i % 3) as f32 },
             if i < 100 { 0.0 } else { (i % 5) as f32 },
             if i < 100 { 0.0 } else { (i % 7) as f32 },
         ];
+        embeddings.insert(i, embedding);
+    }
+
+    // Insert 1000 vectors from the map
+    for (pk, embedding) in &embeddings {
         session
             .query_unpaged(
                 format!("INSERT INTO {table} (pk, v) VALUES (?, ?)"),
-                (i, embedding),
+                (pk, embedding),
             )
             .await
             .expect("failed to insert data");
@@ -65,24 +71,25 @@ async fn ann_query_returns_expected_results(actors: TestActors) {
     .await;
 
     // Check if the query returns the expected results (recall at least 85%)
-    let rows = get_query_results(
-        format!("SELECT pk FROM {table} ORDER BY v ANN OF [0.0, 0.0, 0.0] LIMIT 100"),
+    let results = get_query_results(
+        format!("SELECT pk, v FROM {table} ORDER BY v ANN OF [0.0, 0.0, 0.0] LIMIT 100"),
         &session,
     )
     .await;
-    assert_eq!(rows.len(), 100);
-    let correct = rows
-        .iter()
-        .filter(|row| {
-            let pk: i32 = row.columns[0].as_ref().unwrap().as_int().unwrap();
-            pk < 100
-        })
-        .count();
-    debug!("Number of matching results: {}", correct);
-    assert!(
-        correct >= 85,
-        "Expected more than 85 matching results, got {correct}"
-    );
+    let rows = results
+        .rows::<(i32, Vec<f32>)>()
+        .expect("failed to get rows");
+    assert!(rows.rows_remaining() <= 100);
+    for row in rows {
+        let row = row.expect("failed to get row");
+        let (pk, v) = row;
+        assert!(
+            embeddings.contains_key(&pk),
+            "pk {pk} not found in embeddings"
+        );
+        let expected = embeddings.get(&pk).unwrap();
+        assert_eq!(&v, expected, "Returned vector does not match for pk={pk}");
+    }
 
     // Drop keyspace
     session
@@ -124,19 +131,25 @@ async fn ann_query_respects_limit(actors: TestActors) {
     .await;
 
     // Check if queries return the expected number of results
-    let rows = get_query_results(
+    let results = get_query_results(
         format!("SELECT * FROM {table} ORDER BY v ANN OF [0.0, 0.0, 0.0] LIMIT 10"),
         &session,
     )
     .await;
-    assert_eq!(rows.len(), 10);
+    let rows = results
+        .rows::<(i32, Vec<f32>)>()
+        .expect("failed to get rows");
+    assert!(rows.rows_remaining() <= 10);
 
-    let rows = get_query_results(
+    let results = get_query_results(
         format!("SELECT * FROM {table} ORDER BY v ANN OF [0.0, 0.0, 0.0] LIMIT 1000"),
         &session,
     )
     .await;
-    assert_eq!(rows.len(), 10); // Should return only 10, as there are only 10 vectors
+    let rows = results
+        .rows::<(i32, Vec<f32>)>()
+        .expect("failed to get rows");
+    assert!(rows.rows_remaining() <= 10); // Should return only 10, as there are only 10 vectors
 
     // Check if LIMIT over 1000 fails
     session
@@ -186,20 +199,25 @@ async fn ann_query_respects_limit_over_1000_vectors(actors: TestActors) {
     .await;
 
     // Check if queries return the expected number of results
-    let rows = get_query_results(
+    let results = get_query_results(
         format!("SELECT * FROM {table} ORDER BY v ANN OF [0.0, 0.0, 0.0] LIMIT 10"),
         &session,
     )
     .await;
-    assert_eq!(rows.len(), 10);
-
-    // Due to VECTOR-221 the test fails. Uncomment after fixing.
-    // let rows = get_query_results(
-    //     format!("SELECT * FROM {table} ORDER BY v ANN OF [0.0, 0.0, 0.0] LIMIT 1000"),
-    //     &session,
-    // )
-    // .await;
-    // assert_eq!(rows.len(), 1000);
+    let rows = results
+        .rows::<(i32, Vec<f32>)>()
+        .expect("failed to get rows");
+    assert!(rows.rows_remaining() <= 10);
+
+    let results = get_query_results(
+        format!("SELECT * FROM {table} ORDER BY v ANN OF [0.0, 0.0, 0.0] LIMIT 1000"),
+        &session,
+    )
+    .await;
+    let rows = results
+        .rows::<(i32, Vec<f32>)>()
+        .expect("failed to get rows");
+    assert!(rows.rows_remaining() <= 1000);
 
     // Check if LIMIT over 1000 fails
     session