Treat empty Antfly external indexes as ready

vectordb_bench/backend/clients/antfly/antfly.py

@@ -41,6 +41,36 @@ def _make_client(base_url: str, timeout: float) -> httpx.Client:
     )
 
 
+def _looks_like_json_response(response: httpx.Response) -> bool:
+    content_type = response.headers.get("content-type", "")
+    body = response.content.strip()
+    return (
+        response.is_success
+        and ("json" in content_type or body.startswith((b"{", b"[")))
+    )
+
+
+def _detect_metadata_base_url(host: str, port: int) -> str:
+    root = f"http://{_httpx_host(host)}:{port}"
+    candidates = (f"{root}/db/v1", f"{root}/api/v1")
+    for base_url in candidates:
+        try:
+            with httpx.Client(base_url=base_url, timeout=5) as client:
+                r = client.get("/status")
+                if _looks_like_json_response(r):
+                    return base_url
+        except Exception:
+            pass
+    try:
+        with httpx.Client(base_url=root, timeout=5) as client:
+            r = client.get("/api/v1/tables")
+            if _looks_like_json_response(r):
+                return f"{root}/api/v1"
+    except Exception:
+        pass
+    return candidates[0]
+
+
 class Antfly(VectorDB):
     def __init__(
         self,
@@ -56,13 +86,14 @@ def __init__(
         self.collection_name = collection_name
         self.dim = dim
 
-        self._metadata_base_url = (
-            f"http://{_httpx_host(db_config['host'])}:{db_config['port']}/api/v1"
+        self._metadata_base_url = _detect_metadata_base_url(
+            db_config["host"], db_config["port"]
         )
         self._store_host = _httpx_host(db_config.get("store_host") or db_config["host"])
         self._store_port = db_config.get("store_port")
         self._use_direct_store_search = bool(db_config.get("use_direct_store_search"))
         self._pack_query_vectors = bool(db_config.get("pack_query_vectors"))
+        self._legacy_wire = not self._pack_query_vectors
         self._direct_shard_id: str | None = None
         num_shards = db_config.get("num_shards", 1)
 
@@ -77,16 +108,24 @@ def __init__(
                 r = client.delete(f"/tables/{self.collection_name}")
                 log.info(f"Drop table response: {r.status_code}")
 
-            table = self._get_table_status_or_none(client)
-            if table is None:
+            def create_table_if_needed() -> None:
+                table = self._get_table_status_or_none(client)
+                if table is not None:
+                    log.info("Reusing existing table: %s", self.collection_name)
+                    return
                 r = client.post(
                     f"/tables/{self.collection_name}", json={"num_shards": num_shards}
                 )
                 log.info(f"Create table response: {r.status_code}")
                 r.raise_for_status()
-            else:
-                log.info("Reusing existing table: %s", self.collection_name)
 
+            def reset_table() -> None:
+                r = client.delete(f"/tables/{self.collection_name}")
+                log.info(f"Reset table response: {r.status_code}")
+                create_table_if_needed()
+                self._wait_for_shard_ready(client)
+
+            create_table_if_needed()
             self._wait_for_shard_ready(client)
 
             if self._get_index_status(client) is None:
@@ -97,28 +136,48 @@ def __init__(
                     **self.case_config.index_param(),
                 }
                 index_error = None
-                # Try each index type, with and without field, to handle
-                # both old binaries (require field) and new source (reject field with external).
-                for index_type in INDEX_TYPES:
-                    for extra in ({}, {"field": SOURCE_FIELD}):
-                        r = client.post(
-                            f"/tables/{self.collection_name}/indexes/{INDEX_NAME}",
-                            json={"type": index_type, **index_def, **extra},
-                        )
-                        log.info(
-                            f"Add embeddings index response ({index_type}, field={'field' in extra}): {r.status_code}"
-                        )
-                        if r.is_success:
-                            index_error = None
-                            break
+                index_created = False
+                extras = ({"field": SOURCE_FIELD}, {}) if self._legacy_wire else ({}, {"field": SOURCE_FIELD})
+                candidates = [
+                    (index_type, extra)
+                    for index_type in INDEX_TYPES
+                    for extra in extras
+                ]
+                # Treat a create-index HTTP success as provisional. Older stable
+                # binaries can accept a metadata change that the shard later
+                # rejects, so require a ready status before loading vectors.
+                for idx, (index_type, extra) in enumerate(candidates):
+                    if idx > 0:
+                        reset_table()
+                    r = client.post(
+                        f"/tables/{self.collection_name}/indexes/{INDEX_NAME}",
+                        json={"type": index_type, **index_def, **extra},
+                    )
+                    log.info(
+                        f"Add embeddings index response ({index_type}, field={'field' in extra}): {r.status_code}"
+                    )
+                    if not r.is_success:
                         index_error = r
-                    if index_error is None:
+                        continue
+                    if self._legacy_wire:
+                        index_created = True
                         break
-                if index_error is not None:
+                    if self._wait_for_index_ready(
+                        client,
+                        expected_total=0,
+                        timeout=TABLE_READY_TIMEOUT,
+                    ):
+                        index_created = True
+                        break
+                    index_error = r
+                if not index_created and index_error is not None:
                     index_error.raise_for_status()
+                if not index_created:
+                    raise RuntimeError("Antfly index was created but never became ready")
             else:
                 log.info("Reusing existing embeddings index: %s", INDEX_NAME)
-            self._wait_for_index_ready(client, expected_total=0)
+            if not self._legacy_wire:
+                self._wait_for_index_ready(client, expected_total=0)
             self._refresh_direct_search_routing(client)
         finally:
             client.close()
@@ -151,21 +210,59 @@ def _wait_for_shard_ready(self, client: httpx.Client):
     def _get_index_status(self, client: httpx.Client) -> dict | None:
         r = client.get(f"/tables/{self.collection_name}/indexes/{INDEX_NAME}")
         if r.status_code == 404:
-            return None
+            return self._get_legacy_index_status(client)
         r.raise_for_status()
+        body = r.content.strip()
+        if not body or not body.startswith(b"{"):
+            return self._get_legacy_index_status(client)
         return r.json()
 
     def _get_table_status(self, client: httpx.Client) -> dict:
+        table = self._get_table_status_or_none(client)
+        if table is None:
+            raise RuntimeError(f"Antfly table not found: {self.collection_name}")
+        return table
+
+    def _get_table_status_or_none(self, client: httpx.Client) -> dict | None:
         r = client.get(f"/tables/{self.collection_name}")
+        if r.status_code == 404:
+            return self._get_legacy_table_status(client)
         r.raise_for_status()
+        body = r.content.strip()
+        if not body or not body.startswith(b"{"):
+            return self._get_legacy_table_status(client)
         return r.json()
 
-    def _get_table_status_or_none(self, client: httpx.Client) -> dict | None:
-        r = client.get(f"/tables/{self.collection_name}")
+    def _get_legacy_table_status(self, client: httpx.Client) -> dict | None:
+        r = client.get("/tables")
         if r.status_code == 404:
             return None
         r.raise_for_status()
-        return r.json()
+        body = r.content.strip()
+        if not body or not body.startswith(b"["):
+            return None
+        for table in r.json():
+            if isinstance(table, dict) and table.get("name") == self.collection_name:
+                return table
+        return None
+
+    def _get_legacy_index_status(self, client: httpx.Client) -> dict | None:
+        r = client.get("/status")
+        if r.status_code == 404:
+            return None
+        r.raise_for_status()
+        body = r.content.strip()
+        if not body or not body.startswith(b"{"):
+            return None
+        statuses = (((r.json().get("shards") or {}).get("statuses")) or {})
+        for shard in statuses.values():
+            if not isinstance(shard, dict) or shard.get("table") != self.collection_name:
+                continue
+            indexes = ((((shard.get("info") or {}).get("shard_stats") or {}).get("indexes")) or {})
+            status = indexes.get(INDEX_NAME)
+            if isinstance(status, dict):
+                return {"status": status}
+        return None
 
     def _refresh_direct_search_routing(self, client: httpx.Client):
         if not self._use_direct_store_search:
@@ -186,21 +283,28 @@ def _index_status_is_ready(
         if payload is None:
             return False
         if status is None:
-            return expected_total == 0
+            return False
 
         rebuilding = bool(status.get("rebuilding"))
         wal_backlog = int(status.get("wal_backlog", 0) or 0)
         total_indexed = int(status.get("total_indexed", 0) or 0)
+        doc_count = int(status.get("doc_count", 0) or 0)
         has_error = bool(status.get("error"))
 
+        if expected_total == 0 and total_indexed == 0 and doc_count == 0:
+            return not has_error and wal_backlog == 0
+
         if has_error or rebuilding or wal_backlog > 0:
             return False
         return expected_total is None or total_indexed >= expected_total
 
     def _wait_for_index_ready(
-        self, client: httpx.Client, expected_total: int | None = None
-    ):
-        deadline = time.monotonic() + INDEX_READY_TIMEOUT
+        self,
+        client: httpx.Client,
+        expected_total: int | None = None,
+        timeout: int = INDEX_READY_TIMEOUT,
+    ) -> bool:
+        deadline = time.monotonic() + timeout
         last_status = None
 
         while time.monotonic() < deadline:
@@ -210,17 +314,18 @@ def _wait_for_index_ready(
                 last_status = status
                 if self._index_status_is_ready(payload, status, expected_total):
                     log.info(f"Embeddings index is ready: {status}")
-                    return
+                    return True
             except Exception as e:
                 last_status = {"error": str(e)}
             time.sleep(INDEX_READY_POLL_INTERVAL)
 
         log.warning(
             "Embeddings index readiness timeout after %ss, expected_total=%s, last_status=%s",
-            INDEX_READY_TIMEOUT,
+            timeout,
             expected_total,
             last_status,
         )
+        return False
 
     @contextmanager
     def init(self):
@@ -271,8 +376,10 @@ def _serialize_query_vector(self, vector: list[float]) -> list[float] | str:
             return self._pack_vector(vector)
         return vector
 
-    def _serialize_insert_vector(self, vector: list[float]) -> str:
-        return self._pack_vector(vector)
+    def _serialize_insert_vector(self, vector: list[float]) -> list[float] | str:
+        if self._pack_query_vectors or os.environ.get("ANTFLY_PACK_VECTORS") == "1":
+            return self._pack_vector(vector)
+        return vector
 
     def _metadata_query_body(self, query: list[float], k: int) -> dict[str, Any]:
         return {