Vector index

.github/workflows/docs.yaml

@@ -19,11 +19,22 @@ jobs:
       - name: Create ArangoDB Docker container
         run: >
           docker create --name arango -p 8529:8529 -e ARANGO_ROOT_PASSWORD=passwd -v "$(pwd)/tests/static/":/tests/static
-          arangodb/arangodb:latest --server.jwt-secret-keyfile=/tests/static/keyfile
+          arangodb/enterprise:latest --server.jwt-secret-keyfile=/tests/static/keyfile --vector-index=true
 
       - name: Start ArangoDB Docker container
         run: docker start arango
 
+      - name: Wait for ArangoDB
+        run: |
+          for _ in {1..30}; do
+            if curl --fail --silent --user root:passwd http://localhost:8529/_api/version > /dev/null; then
+              exit 0
+            fi
+            sleep 1
+          done
+          docker logs arango
+          exit 1
+
       - name: Set up Python
         uses: actions/setup-python@v7
         with:

arango/collection.py

@@ -1274,17 +1274,29 @@ def response_handler(resp: Response) -> Optional[Json]:
     # Index Management #
     ####################
 
-    def indexes(self) -> Result[Jsons]:
+    def indexes(
+        self, with_stats: bool = False, with_hidden: bool = False
+    ) -> Result[Jsons]:
         """Return the collection indexes.
 
+        :param with_stats: Whether to include figures and estimates in the result.
+        :type with_stats: bool
+        :param with_hidden: Whether to include hidden indexes in the result.
+        :type with_hidden: bool
         :return: Collection indexes.
         :rtype: [dict]
         :raise arango.exceptions.IndexListError: If retrieval fails.
         """
+        params: Params = {"collection": self.name}
+        if with_stats:
+            params["withStats"] = True
+        if with_hidden:
+            params["withHidden"] = True
+
         request = Request(
             method="get",
             endpoint="/_api/index",
-            params={"collection": self.name},
+            params=params,
         )
 
         def response_handler(resp: Response) -> Jsons:

arango/formatter.py

@@ -73,6 +73,8 @@ def format_index(body: Json, formatter: bool = True) -> Json:
         result["cacheEnabled"] = body["cacheEnabled"]
     if "legacyPolygons" in body:
         result["legacyPolygons"] = body["legacyPolygons"]
+    if "figures" in body:
+        result["figures"] = body["figures"]
     if "estimates" in body:
         result["estimates"] = body["estimates"]
     if "analyzer" in body:
@@ -123,6 +125,8 @@ def format_index(body: Json, formatter: bool = True) -> Json:
         result["error_message"] = body["errorMessage"]
     if "trainingState" in body:
         result["training_state"] = body["trainingState"]
+    if "shards" in body:
+        result["shards"] = body["shards"]
 
     return verify_format(body, result)
 

docs/foxx.rst

@@ -13,7 +13,7 @@ information, refer to `ArangoDB manual`_.
 
 **Example:**
 
-.. testcode::
+.. code-block:: python
 
     from arango import ArangoClient
 

docs/indexes.rst

@@ -62,4 +62,43 @@ on fields ``_from`` and ``_to``. For more information on indexes, refer to
     # Delete the last index from the collection.
     cities.delete_index(index['id'])
 
+    # Insert documents with vector embeddings.
+    cities.insert_many([
+        {
+            '_key': f'city{i}',
+            'continent': f'continent{i}',
+            'country': f'country{i}',
+            'population': i,
+            'coordinates': [float(i % 180), float(i % 90)],
+            'x': float(i),
+            'y': float(i),
+            '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 = cities.add_index({
+        'type': 'vector',
+        'fields': ['embedding'],
+        'name': 'vector_index',
+        'params': {
+            'metric': 'cosine',
+            'dimension': 4,
+        },
+    })
+
+    # Index creation may succeed even if vector training fails.
+    if vector_index.get('trainingState') != 'ready':
+        raise RuntimeError(
+            vector_index.get('errorMessage', '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 ``trainingState`` before
+using it. If training fails permanently, the state is ``"unusable"`` and
+``errorMessage`` describes the failure.
+
 See :ref:`StandardCollection` for API specification.

starter.sh

@@ -8,8 +8,8 @@
 #   ./starter.sh [single|cluster] [image[:tag]]
 # Example:
 #   ./starter.sh cluster enterprise 3.12.4
-#   ./starter.sh single enterprise-preview 4.0-nightly
-#   ./starter.sh single arangodb/enterprise-preview:4.0-nightly
+#   ./starter.sh single core-preview 4.0-nightly
+#   ./starter.sh single arangodb/enterprise-preview:3.12-nightly
 
 setup="${1:-single}"
 image="${2:-community}"

tests/test_index.py

@@ -22,11 +22,46 @@ def test_list_indexes(icol, bad_col):
     assert "sparse" in indexes[0]
     assert "unique" in indexes[0]
 
+    indexes = icol.indexes(with_stats=True)
+    assert "figures" in indexes[0]
+
     with assert_raises(IndexListError) as err:
         bad_col.indexes()
     assert err.value.error_code in {11, 1228}
 
 
+def test_list_indexes_options(icol, monkeypatch):
+    requests = []
+
+    def execute(request, response_handler):
+        requests.append(request)
+        return []
+
+    monkeypatch.setattr(icol, "_execute", execute)
+
+    assert icol.indexes() == []
+    assert requests[-1].params == {"collection": icol.name}
+
+    assert icol.indexes(with_hidden=True) == []
+    assert requests[-1].params == {
+        "collection": icol.name,
+        "withHidden": "1",
+    }
+
+    assert icol.indexes(with_stats=True) == []
+    assert requests[-1].params == {
+        "collection": icol.name,
+        "withStats": "1",
+    }
+
+    assert icol.indexes(with_stats=True, with_hidden=True) == []
+    assert requests[-1].params == {
+        "collection": icol.name,
+        "withStats": "1",
+        "withHidden": "1",
+    }
+
+
 def test_get_index(icol, bad_col):
     indexes = icol.indexes()
     for index in indexes:
@@ -253,25 +288,137 @@ def test_add_mdi_index(icol, db_version):
     icol.delete_index(result["id"])
 
 
-def test_add_vector_index(col):
+def test_add_vector_index(col, db_version):
+    # Insert vector data.
     docs = []
     for i in range(100):
-        docs.append({"_key": generate_doc_key(), "x": [1] * 128})
+        docs.append(
+            {
+                "_key": generate_doc_key(),
+                "x": [1] * 128,
+                "y": [1] * 128,
+            }
+        )
     col.insert_many(docs)
-    result = col.add_index(
+
+    # Test basic compatibility.
+    index_meta = [
+        {
+            "type": "vector",
+            "fields": ["x"],
+            "name": "vector_index_1",
+            "params": {
+                "metric": "cosine",
+                "dimension": 128,
+                "nLists": 2,
+            },
+        },
         {
+            "type": "vector",
+            "fields": ["y"],
+            "name": "vector_index_2",
+            "params": {
+                "metric": "cosine",
+                "dimension": 128,
+                "nLists": 3,
+            },
+        },
+    ]
+
+    results = [col.add_index(index_meta[0]), col.add_index(index_meta[1])]
+    assert results[0]["name"] == "vector_index_1"
+    assert results[1]["name"] == "vector_index_2"
+
+    # Test hidden shard details.
+    indexes = col.indexes(with_hidden=True)
+
+    if db_version >= version.parse("3.12.10"):
+        details = {item["id"]: item for item in indexes}
+        for result in results:
+            shards = details[result["id"]]["shards"]
+            assert shards is not None
+            for status in shards.values():
+                assert {"trainingState", "error", "resolvedNLists"} <= status.keys()
+
+    col.delete_index(results[0]["id"])
+    col.delete_index(results[1]["id"])
+
+    if db_version >= version.parse("3.12.10"):
+        # Test server-managed nLists.
+        default_index = {
+            "type": "vector",
+            "fields": ["x"],
+            "name": "vector_index_default",
+            "params": {
+                "metric": "cosine",
+                "dimension": 128,
+            },
+        }
+        scaling_n_lists = {
+            "strategy": "autoSqrt",
+            "multiplier": 1,
+            "minNLists": 2,
+            "tiers": [],
+        }
+        scaling_index = {
+            "type": "vector",
+            "fields": ["y"],
+            "name": "vector_index_scaling",
+            "params": {
+                "metric": "cosine",
+                "dimension": 128,
+                "nLists": scaling_n_lists,
+                "numberOfDocsPerCentroid": 10,
+                "factory": "IVF{},Flat",
+            },
+        }
+
+        default_result = col.add_index(default_index)
+        scaling_result = col.add_index(scaling_index)
+
+        default_n_lists = default_result["params"]["nLists"]
+        assert default_n_lists["strategy"] == "autoSqrt"
+        assert default_n_lists["multiplier"] == 4
+        assert default_n_lists["minNLists"] == 2
+        assert scaling_result["params"]["nLists"] == scaling_n_lists
+        assert scaling_result["params"]["numberOfDocsPerCentroid"] == 10
+        assert scaling_result["params"]["factory"] == "IVF{},Flat"
+
+        col.delete_index(default_result["id"])
+        col.delete_index(scaling_result["id"])
+
+        # Test unusable creation.
+        unusable_index = {
             "type": "vector",
             "fields": ["x"],
-            "name": "vector_index",
+            "name": "vector_index_unusable",
             "params": {
                 "metric": "cosine",
                 "dimension": 128,
                 "nLists": 2,
+                "factory": "IVF3,Flat",
             },
         }
-    )
-    assert result["name"] == "vector_index"
-    col.delete_index(result["id"])
+        unusable_result = col.add_index(unusable_index)
+        assert unusable_result["trainingState"] == "unusable"
+        assert unusable_result["errorMessage"]
+        col.delete_index(unusable_result["id"])
+
+        # Test invalid request failure.
+        with assert_raises(IndexCreateError) as err:
+            col.add_index(
+                {
+                    "type": "vector",
+                    "fields": ["x"],
+                    "name": "vector_index_invalid",
+                    "params": {
+                        "metric": "cosine",
+                        "dimension": 128,
+                        "nLists": 0,  # must be greater than 0
+                    },
+                }
+            )
+        assert err.value.http_code == 400
 
 
 def test_delete_index(icol, bad_col):