tests hnsw multi vector indices only on supported search module versions

Author: justin-cechmanek

tests/conftest.py

@@ -232,89 +232,6 @@ def sample_datetimes():
     }
 
 
-@pytest.fixture
-def OG(sample_datetimes):
-    return [
-        {
-            "user": "john",
-            "age": 18,
-            "job": "engineer",
-            "description": "engineers conduct trains that ride on train tracks",
-            "last_updated": sample_datetimes["low"].timestamp(),
-            "credit_score": "high",
-            "location": "-122.4194,37.7749",
-            "user_embedding": [0.1, 0.1, 0.5],
-            "image_embedding": [0.1, 0.1, 0.1, 0.1, 0.1],
-        },
-        {
-            "user": "mary",
-            "age": 14,
-            "job": "doctor",
-            "description": "a medical professional who treats diseases and helps people stay healthy",
-            "last_updated": sample_datetimes["low"].timestamp(),
-            "credit_score": "low",
-            "location": "-122.4194,37.7749",
-            "user_embedding": [0.1, 0.1, 0.5],
-            "image_embedding": [0.1, 0.2, 0.3, 0.4, 0.5],
-        },
-        {
-            "user": "nancy",
-            "age": 94,
-            "job": "doctor",
-            "description": "a research scientist specializing in cancers and diseases of the lungs",
-            "last_updated": sample_datetimes["mid"].timestamp(),
-            "credit_score": "high",
-            "location": "-122.4194,37.7749",
-            "user_embedding": [0.7, 0.1, 0.5],
-            "image_embedding": [0.1, 0.1, 0.3, 0.3, 0.5],
-        },
-        {
-            "user": "tyler",
-            "age": 100,
-            "job": "engineer",
-            "description": "a software developer with expertise in mathematics and computer science",
-            "last_updated": sample_datetimes["mid"].timestamp(),
-            "credit_score": "high",
-            "location": "-110.0839,37.3861",
-            "user_embedding": [0.1, 0.4, 0.5],
-            "image_embedding": [-0.1, -0.2, -0.3, -0.4, -0.5],
-        },
-        {
-            "user": "tim",
-            "age": 12,
-            "job": "dermatologist",
-            "description": "a medical professional specializing in diseases of the skin",
-            "last_updated": sample_datetimes["mid"].timestamp(),
-            "credit_score": "high",
-            "location": "-110.0839,37.3861",
-            "user_embedding": [0.4, 0.4, 0.5],
-            "image_embedding": [-0.1, 0.0, 0.6, 0.0, -0.9],
-        },
-        {
-            "user": "taimur",
-            "age": 15,
-            "job": "CEO",
-            "description": "high stress, but financially rewarding position at the head of a company",
-            "last_updated": sample_datetimes["high"].timestamp(),
-            "credit_score": "low",
-            "location": "-110.0839,37.3861",
-            "user_embedding": [0.6, 0.1, 0.5],
-            "image_embedding": [1.1, 1.2, -0.3, -4.1, 5.0],
-        },
-        {
-            "user": "joe",
-            "age": 35,
-            "job": "dentist",
-            "description": "like the tooth fairy because they'll take your teeth, but you have to pay them!",
-            "last_updated": sample_datetimes["high"].timestamp(),
-            "credit_score": "medium",
-            "location": "-110.0839,37.3861",
-            "user_embedding": [-0.1, -0.1, -0.5],
-            "image_embedding": [-0.8, 2.0, 3.1, 1.5, -1.6],
-        },
-    ]
-
-
 @pytest.fixture
 def sample_data(sample_datetimes):
     return [

tests/integration/test_aggregation.py

@@ -50,7 +50,7 @@ def index(multi_vector_data, redis_url, worker_id):
                     "attrs": {
                         "dims": 6,
                         "distance_metric": "cosine",
-                        "algorithm": "hnsw",
+                        "algorithm": "flat",
                         "datatype": "bfloat16",
                     },
                 },
@@ -581,6 +581,64 @@ def test_multivector_query_datatypes(index):
         )
 
 
-def test_multivector_query_broadcasting(index):
+def test_multivector_query_mixed_index(index):
+    # test that we can do multi vector queries on indices with both a 'flat' and 'hnsw' index
     skip_if_redis_version_below(index.client, "7.2.0")
-    pass
+    try:
+        index.schema.remove_field("audio_embedding")
+        index.schema.add_field(
+            {
+                "name": "audio_embedding",
+                "type": "vector",
+                "attrs": {
+                    "dims": 6,
+                    "distance_metric": "cosine",
+                    "algorithm": "hnsw",
+                    "datatype": "bfloat16",
+                },
+            },
+        )
+
+    except:
+        pytest.skip("Required Redis modules not available or version too low")
+
+    vectors = [[0.1, 0.2, 0.5], [1.2, 0.3, -0.4, 0.7, 0.2, -0.3]]
+    vector_fields = ["user_embedding", "audio_embedding"]
+    return_fields = [
+        "distance_0",
+        "distance_1",
+        "score_0",
+        "score_1",
+        "user_embedding",
+        "audio_embedding",
+    ]
+
+    # changing the weights does indeed change the result order
+    multi_query = MultiVectorQuery(
+        vectors=vectors,
+        vector_field_names=vector_fields,
+        return_fields=return_fields,
+        dtypes=["float32", "bfloat16"],
+    )
+    results = index.query(multi_query)
+
+    for i in range(1, len(results)):
+        assert results[i]["combined_score"] <= results[i - 1]["combined_score"]
+
+    # verify we're doing the combined score math correctly
+    weights = [-1.322, 0.851]
+    multi_query = MultiVectorQuery(
+        vectors=vectors,
+        vector_field_names=vector_fields,
+        return_fields=return_fields,
+        dtypes=["float32", "bfloat16"],
+        weights=weights,
+    )
+
+    results = index.query(multi_query)
+    assert results
+    for r in results:
+        score = float(r["score_0"]) * weights[0] + float(r["score_1"]) * weights[1]
+        assert (
+            float(r["combined_score"]) - score <= 0.0001
+        )  # allow for small floating point error