Updated vector index support

arangoasync/typings.py

@@ -1060,6 +1060,10 @@ def cache_enabled(self) -> Optional[bool]:
     def legacy_polygons(self) -> Optional[bool]:
         return self._data.get("legacyPolygons")
 
+    @property
+    def figures(self) -> Optional[Json]:
+        return self._data.get("figures")
+
     @property
     def estimates(self) -> Optional[bool]:
         return self._data.get("estimates")
@@ -1140,6 +1144,10 @@ def error_message(self) -> Optional[str]:
     def training_state(self) -> Optional[str]:
         return self._data.get("trainingState")
 
+    @property
+    def shards(self) -> Optional[Json]:
+        return self._data.get("shards")
+
     @staticmethod
     def compatibility_formatter(data: Json) -> Json:
         """python-arango compatibility formatter."""
@@ -1170,6 +1178,8 @@ def compatibility_formatter(data: Json) -> Json:
             result["storedValues"] = data["storedValues"]
         if "legacyPolygons" in data:
             result["legacyPolygons"] = data["legacyPolygons"]
+        if "figures" in data:
+            result["figures"] = data["figures"]
         if "estimates" in data:
             result["estimates"] = data["estimates"]
         if "analyzer" in data:
@@ -1202,6 +1212,8 @@ def compatibility_formatter(data: Json) -> Json:
             result["error_message"] = data["errorMessage"]
         if "trainingState" in data:
             result["training_state"] = data["trainingState"]
+        if "shards" in data:
+            result["shards"] = data["shards"]
         return result
 
     def format(self, formatter: Optional[Formatter] = None) -> Json:

docs/indexes.rst

@@ -50,4 +50,43 @@ on fields ``_from`` and ``_to``. For more information on indexes, refer to
         # Delete the last index from the collection.
         await cities.delete_index(index["id"])
 
+        # Insert documents with vector embeddings.
+        await cities.insert_many([
+            {
+                "_key": f"city{i}",
+                "continent": f"continent{i}",
+                "country": f"country{i}",
+                "population": i,
+                "coordinates": [float(i % 180), float(i % 90)],
+                "embedding": [float(i), float(i % 7), float(i % 11), 1.0],
+            }
+            for i in range(100)
+        ])
+
+        # Let ArangoDB determine the number of vector-index centroids.
+        vector_index = await cities.add_index(
+            type="vector",
+            fields=["embedding"],
+            options={
+                "name": "vector_index",
+                "params": {
+                    "metric": "cosine",
+                    "dimension": 4,
+                },
+            },
+        )
+
+        # Index creation may succeed even if vector training fails.
+        if vector_index.training_state != "ready":
+            raise RuntimeError(
+                vector_index.error_message or "Vector index is not ready"
+            )
+
+Omitted or scaling-object ``nLists``, ``numberOfDocsPerCentroid``, factory
+placeholders such as ``IVF{},Flat``, and successful-but-unusable creation
+behavior require ArangoDB 3.12.10 or later. A successful creation response
+means that the index exists, but callers should check ``training_state`` before
+using it. If training fails permanently, the state is ``"unusable"`` and
+``error_message`` describes the failure.
+
 See :class:`arangoasync.collection.StandardCollection` for API specification.

tests/test_collection.py

@@ -1,6 +1,7 @@
 import asyncio
 
 import pytest
+from packaging import version
 
 from arangoasync.errno import DATA_SOURCE_NOT_FOUND, INDEX_NOT_FOUND
 from arangoasync.exceptions import (
@@ -113,11 +114,11 @@ async def test_collection_rename(cluster, db, bad_col, docs):
         doc = await col.insert(docs[0])
         assert col.get_col_name(doc) == new_name
     finally:
-        db.delete_collection(new_name, ignore_missing=True)
+        await db.delete_collection(new_name, ignore_missing=True)
 
 
 @pytest.mark.asyncio
-async def test_collection_index(doc_col, bad_col, cluster):
+async def test_collection_index(doc_col, bad_col, cluster, db_version):
     # Create indexes
     idx1 = await doc_col.add_index(
         type="persistent",
@@ -228,36 +229,147 @@ async def test_collection_index(doc_col, bad_col, cluster):
         await bad_col.load_indexes()
     assert err.value.error_code == DATA_SOURCE_NOT_FOUND
 
-    # Create a vector index
+    # Create vector indexes using the fixed nLists format supported by older servers.
     docs = []
     for key in range(100):
-        docs.append({"_key": f"key_{key}", "embedding": [1] * 128})
+        docs.append(
+            {
+                "_key": f"key_{key}",
+                "embedding1": [1] * 128,
+                "embedding2": [1] * 128,
+            }
+        )
     await doc_col.insert_many(docs)
     idx4 = await doc_col.add_index(
         "vector",
-        ["embedding"],
+        ["embedding1"],
         {
-            "name": "vector_index",
+            "name": "vector_index_1",
             "params": {
                 "metric": "cosine",
                 "dimension": 128,
                 "nLists": 2,
             },
         },
     )
-    assert idx4.name == "vector_index"
+    idx5 = await doc_col.add_index(
+        "vector",
+        ["embedding2"],
+        {
+            "name": "vector_index_2",
+            "params": {
+                "metric": "cosine",
+                "dimension": 128,
+                "nLists": 3,
+            },
+        },
+    )
+    assert idx4.name == "vector_index_1"
+    assert idx5.name == "vector_index_2"
+
+    if db_version >= version.parse("3.12.10"):
+        # Hidden listing details expose resolved vector-index settings per shard.
+        indexes = {idx.id: idx for idx in await doc_col.indexes(with_hidden=True)}
+        for index in (idx4, idx5):
+            shards = indexes[index.id].shards
+            assert shards is not None
+            for status in shards.values():
+                assert {
+                    "trainingState",
+                    "error",
+                    "resolvedNLists",
+                } <= status.keys()
+                assert isinstance(status["resolvedNLists"], int)
 
     # Delete indexes
-    del1, del2, del3, del4 = await asyncio.gather(
+    del1, del2, del3, del4, del5 = await asyncio.gather(
         doc_col.delete_index(idx1.id),
         doc_col.delete_index(idx2.numeric_id),
         doc_col.delete_index(str(idx3.numeric_id)),
         doc_col.delete_index(idx4.id),
+        doc_col.delete_index(idx5.id),
     )
     assert del1 is True
     assert del2 is True
     assert del3 is True
     assert del4 is True
+    assert del5 is True
+
+    if db_version >= version.parse("3.12.10"):
+        # Let the server choose nLists, then supply an explicit scaling object.
+        scaling_n_lists = {
+            "strategy": "autoSqrt",
+            "multiplier": 1,
+            "minNLists": 2,
+            "tiers": [],
+        }
+        default_index = await doc_col.add_index(
+            "vector",
+            ["embedding1"],
+            {
+                "name": "vector_index_default",
+                "params": {"metric": "cosine", "dimension": 128},
+            },
+        )
+        scaling_index = await doc_col.add_index(
+            "vector",
+            ["embedding2"],
+            {
+                "name": "vector_index_scaling",
+                "params": {
+                    "metric": "cosine",
+                    "dimension": 128,
+                    "nLists": scaling_n_lists,
+                    "numberOfDocsPerCentroid": 10,
+                    "factory": "IVF{},Flat",
+                },
+            },
+        )
+
+        default_n_lists = default_index["params"]["nLists"]
+        assert default_n_lists["strategy"] == "autoSqrt"
+        assert default_n_lists["multiplier"] == 4
+        assert default_n_lists["minNLists"] == 2
+        assert scaling_index["params"]["nLists"] == scaling_n_lists
+        assert scaling_index["params"]["numberOfDocsPerCentroid"] == 10
+        assert scaling_index["params"]["factory"] == "IVF{},Flat"
+
+        await doc_col.delete_index(default_index.id)
+        await doc_col.delete_index(scaling_index.id)
+
+        # A permanent training failure still creates an unusable index.
+        unusable_index = await doc_col.add_index(
+            "vector",
+            ["embedding1"],
+            {
+                "name": "vector_index_unusable",
+                "params": {
+                    "metric": "cosine",
+                    "dimension": 128,
+                    "nLists": 2,
+                    "factory": "IVF3,Flat",
+                },
+            },
+        )
+        assert unusable_index.training_state == "unusable"
+        assert unusable_index.error_message
+        await doc_col.delete_index(unusable_index.id)
+
+        # Invalid requests continue to fail at the HTTP layer.
+        with pytest.raises(IndexCreateError) as err:
+            await doc_col.add_index(
+                "vector",
+                ["embedding1"],
+                {
+                    "name": "vector_index_invalid",
+                    "params": {
+                        "metric": "cosine",
+                        "dimension": 128,
+                        "nLists": 0,
+                    },
+                },
+            )
+        assert err.value.http_code == 400
 
     # Now, the indexes should be gone
     with pytest.raises(IndexDeleteError) as err:

tests/test_typings.py

@@ -9,6 +9,7 @@
     EdgeDefinitionOptions,
     GraphOptions,
     GraphProperties,
+    IndexProperties,
     JsonWrapper,
     KeyOptions,
     QueryCacheProperties,
@@ -449,6 +450,31 @@ def test_CollectionStatistics():
     assert stats.object_id == "69124"
 
 
+def test_IndexProperties_hidden_details():
+    shards = {
+        "s1001": {
+            "trainingState": "ready",
+            "error": False,
+            "resolvedNLists": 4,
+        }
+    }
+    figures = {"memory": 4096}
+    properties = IndexProperties(
+        {
+            "id": "products/123",
+            "fields": ["embedding"],
+            "type": "vector",
+            "figures": figures,
+            "shards": shards,
+        }
+    )
+
+    assert properties.figures == figures
+    assert properties.shards == shards
+    assert properties.format()["figures"] == figures
+    assert properties.format()["shards"] == shards
+
+
 def test_AccessToken():
     data = {
         "active": True,