More tests, fix comparison to use enum

abrookins · abrookins · commit 977b859ab55b · 2025-04-16T09:30:23.000-07:00
diff --git a/redisvl/index/index.py b/redisvl/index/index.py
@@ -48,16 +48,18 @@
     BaseVectorQuery,
     CountQuery,
     FilterQuery,
-    HybridQuery,
 )
 from redisvl.query.filter import FilterExpression
 from redisvl.redis.connection import (
     RedisConnectionFactory,
     convert_index_info_to_schema,
 )
-from redisvl.redis.utils import convert_bytes
 from redisvl.schema import IndexSchema, StorageType
-from redisvl.schema.fields import VECTOR_NORM_MAP, VectorDistanceMetric
+from redisvl.schema.fields import (
+    VECTOR_NORM_MAP,
+    VectorDistanceMetric,
+    VectorIndexAlgorithm,
+)
 from redisvl.utils.log import get_logger
 
 logger = get_logger(__name__)
@@ -196,6 +198,15 @@ def _storage(self) -> BaseStorage:
             index_schema=self.schema
         )
 
+    def _validate_query(self, query: BaseQuery) -> None:
+        """Validate a query."""
+        if isinstance(query, VectorQuery):
+            field = self.schema.fields[query._vector_field_name]
+            if query.ef_runtime and field.attrs.algorithm != VectorIndexAlgorithm.HNSW:  # type: ignore
+                raise QueryValidationError(
+                    "Vector field using 'flat' algorithm does not support EF_RUNTIME query parameter."
+                )
+
     @property
     def name(self) -> str:
         """The name of the Redis search index."""
@@ -837,15 +848,6 @@ def batch_query(
             all_parsed.append(parsed)
         return all_parsed
 
-    def _validate_query(self, query: BaseQuery) -> None:
-        """Validate a query."""
-        if isinstance(query, VectorQuery):
-            field = self.schema.fields[query._vector_field_name]
-            if query.ef_runtime and field.attrs.algorithm != "hnsw":  # type: ignore
-                raise QueryValidationError(
-                    "Flat index does not support vector queries."
-                )
-
     def _query(self, query: BaseQuery) -> List[Dict[str, Any]]:
         """Execute a query and process results."""
         try:
@@ -1416,7 +1418,8 @@ async def _aggregate(
     ) -> List[Dict[str, Any]]:
         """Execute an aggregation query and processes the results."""
         results = await self.aggregate(
-            aggregation_query, query_params=aggregation_query.params  # type: ignore[attr-defined]
+            aggregation_query,
+            query_params=aggregation_query.params,  # type: ignore[attr-defined]
         )
         return process_aggregate_results(
             results,
@@ -1542,15 +1545,6 @@ async def batch_query(
 
         return all_parsed
 
-    def _validate_query(self, query: BaseQuery) -> None:
-        """Validate a query."""
-        if isinstance(query, VectorQuery):
-            field = self.schema.fields[query._vector_field_name]
-            if query.ef_runtime and field.attrs.algorithm != "hnsw":  # type: ignore
-                raise QueryValidationError(
-                    "Flat index does not support vector queries."
-                )
-
     async def _query(self, query: BaseQuery) -> List[Dict[str, Any]]:
         """Asynchronously execute a query and process results."""
         try:
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -301,3 +301,103 @@ def hash_preprocess(item: dict) -> dict:
 
     # clean up
     await index.delete(drop=True)
+
+
+@pytest.fixture
+async def async_hnsw_index(sample_data, redis_url):
+    """
+    A fixture that uses the "hnsw" algorithm for its vector field.
+    """
+    index = AsyncSearchIndex.from_dict(
+        {
+            "index": {
+                "name": "user_index",
+                "prefix": "v1",
+                "storage_type": "hash",
+            },
+            "fields": [
+                {"name": "description", "type": "text"},
+                {"name": "credit_score", "type": "tag"},
+                {"name": "job", "type": "text"},
+                {"name": "age", "type": "numeric"},
+                {"name": "last_updated", "type": "numeric"},
+                {"name": "location", "type": "geo"},
+                {
+                    "name": "user_embedding",
+                    "type": "vector",
+                    "attrs": {
+                        "dims": 3,
+                        "distance_metric": "cosine",
+                        "algorithm": "hnsw",
+                        "datatype": "float32",
+                    },
+                },
+            ],
+        },
+        redis_url=redis_url,
+    )
+
+    # create the index (no data yet)
+    await index.create(overwrite=True)
+
+    # Prepare and load the data
+    def hash_preprocess(item: dict) -> dict:
+        return {
+            **item,
+            "user_embedding": array_to_buffer(item["user_embedding"], "float32"),
+        }
+
+    await index.load(sample_data, preprocess=hash_preprocess)
+
+    # run the test
+    yield index
+
+
+@pytest.fixture
+def hnsw_index(sample_data, redis_url):
+    """
+    A fixture that uses the "hnsw" algorithm for its vector field.
+    """
+    index = SearchIndex.from_dict(
+        {
+            "index": {
+                "name": "user_index",
+                "prefix": "v1",
+                "storage_type": "hash",
+            },
+            "fields": [
+                {"name": "description", "type": "text"},
+                {"name": "credit_score", "type": "tag"},
+                {"name": "job", "type": "text"},
+                {"name": "age", "type": "numeric"},
+                {"name": "last_updated", "type": "numeric"},
+                {"name": "location", "type": "geo"},
+                {
+                    "name": "user_embedding",
+                    "type": "vector",
+                    "attrs": {
+                        "dims": 3,
+                        "distance_metric": "cosine",
+                        "algorithm": "hnsw",
+                        "datatype": "float32",
+                    },
+                },
+            ],
+        },
+        redis_url=redis_url,
+    )
+
+    # create the index (no data yet)
+    index.create(overwrite=True)
+
+    # Prepare and load the data
+    def hash_preprocess(item: dict) -> dict:
+        return {
+            **item,
+            "user_embedding": array_to_buffer(item["user_embedding"], "float32"),
+        }
+
+    index.load(sample_data, preprocess=hash_preprocess)
+
+    # run the test
+    yield index
diff --git a/tests/integration/test_async_search_index.py b/tests/integration/test_async_search_index.py
@@ -16,6 +16,7 @@
 from redisvl.query.query import FilterQuery
 from redisvl.redis.utils import convert_bytes
 from redisvl.schema import IndexSchema, StorageType
+from redisvl.schema.fields import VectorIndexAlgorithm
 
 fields = [{"name": "test", "type": "tag"}]
 
@@ -622,7 +623,13 @@ async def test_async_search_index_expire_keys(async_index):
 
 
 @pytest.mark.asyncio
-async def test_search_index_validates_query(async_flat_index, sample_data):
+async def test_search_index_validates_query_with_flat_algorithm(
+    async_flat_index, sample_data
+):
+    assert (
+        async_flat_index.schema.fields["user_embedding"].attrs.algorithm
+        == VectorIndexAlgorithm.FLAT
+    )
     query = VectorQuery(
         [0.1, 0.1, 0.5],
         "user_embedding",
@@ -632,3 +639,22 @@ async def test_search_index_validates_query(async_flat_index, sample_data):
     )
     with pytest.raises(QueryValidationError):
         await async_flat_index.query(query)
+
+
+@pytest.mark.asyncio
+async def test_search_index_validates_query_with_hnsw_algorithm(
+    async_hnsw_index, sample_data
+):
+    assert (
+        async_hnsw_index.schema.fields["user_embedding"].attrs.algorithm
+        == VectorIndexAlgorithm.HNSW
+    )
+    query = VectorQuery(
+        [0.1, 0.1, 0.5],
+        "user_embedding",
+        return_fields=["user", "credit_score", "age", "job", "location"],
+        num_results=7,
+        ef_runtime=100,
+    )
+    # Should not raise
+    await async_hnsw_index.query(query)
diff --git a/tests/integration/test_query.py b/tests/integration/test_query.py
@@ -21,7 +21,6 @@
     Text,
     Timestamp,
 )
-from redisvl.query.query import VectorRangeQuery
 from redisvl.redis.utils import array_to_buffer
 
 # TODO expand to multiple schema types and sync + async
diff --git a/tests/integration/test_search_index.py b/tests/integration/test_search_index.py
@@ -15,6 +15,7 @@
 from redisvl.query.query import FilterQuery
 from redisvl.redis.utils import convert_bytes
 from redisvl.schema import IndexSchema, StorageType
+from redisvl.schema.fields import VectorIndexAlgorithm
 
 fields = [
     {"name": "test", "type": "tag"},
@@ -563,7 +564,11 @@ def test_search_index_expire_keys(index):
         assert ttl <= 30
 
 
-def test_search_index_validates_query(flat_index, sample_data):
+def test_search_index_validates_query_with_flat_algorithm(flat_index, sample_data):
+    assert (
+        flat_index.schema.fields["user_embedding"].attrs.algorithm
+        == VectorIndexAlgorithm.FLAT
+    )
     query = VectorQuery(
         [0.1, 0.1, 0.5],
         "user_embedding",
@@ -573,3 +578,19 @@ def test_search_index_validates_query(flat_index, sample_data):
     )
     with pytest.raises(QueryValidationError):
         flat_index.query(query)
+
+
+def test_search_index_validates_query_with_hnsw_algorithm(hnsw_index, sample_data):
+    assert (
+        hnsw_index.schema.fields["user_embedding"].attrs.algorithm
+        == VectorIndexAlgorithm.HNSW
+    )
+    query = VectorQuery(
+        [0.1, 0.1, 0.5],
+        "user_embedding",
+        return_fields=["user", "credit_score", "age", "job", "location"],
+        num_results=7,
+        ef_runtime=100,
+    )
+    # Should not raise
+    hnsw_index.query(query)

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,6 @@`
`21`	`21`	`Text,`
`22`	`22`	`Timestamp,`
`23`	`23`	`)`
`24`		`-from redisvl.query.query import VectorRangeQuery`
`25`	`24`	`from redisvl.redis.utils import array_to_buffer`
`26`	`25`
`27`	`26`	`# TODO expand to multiple schema types and sync + async`