pydantic v1 => v2

Makefile

@@ -1,89 +1,8 @@
-PYTHON ?= /home/rowan/.local/share/uv/tools/pip/bin/python
-BENCH = $(PYTHON) -m vectordb_bench.cli.vectordbbench
-COMMON = --drop-old --skip-search-concurrent
-ANTFLY_NUM_PER_BATCH ?= 5000
-
-help:
-	@echo "VectorDBBench — Antfly Integration"
-	@echo ""
-	@echo "Benchmarks (50K vectors, 1536 dim):"
-	@echo "  make bench-antfly-50k    Antfly SPANN"
-	@echo "  make bench-qdrant-50k    Qdrant HNSW"
-	@echo "  make bench-milvus-50k    Milvus AutoIndex"
-	@echo "  make bench-all-50k       All three (50K)"
-	@echo ""
-	@echo "Benchmarks (1M vectors, 768 dim):"
-	@echo "  make bench-antfly-1m     Antfly SPANN"
-	@echo "  make bench-qdrant-1m     Qdrant HNSW"
-	@echo "  make bench-milvus-1m     Milvus AutoIndex"
-	@echo "  make bench-all-1m        All three (1M)"
-	@echo ""
-	@echo "Infrastructure:"
-	@echo "  make start-qdrant        Start Qdrant (Docker)"
-	@echo "  make start-milvus        Start Milvus (Docker)"
-	@echo "  make stop-all            Stop Docker containers"
-	@echo ""
-	@echo "Dev:"
-	@echo "  make unittest            Run unit tests"
-	@echo "  make format / lint       Code formatting"
-
-# --- 50K benchmarks ---
-bench-antfly-50k:
-	NUM_PER_BATCH=$(ANTFLY_NUM_PER_BATCH) $(BENCH) antflyaknn --host localhost --port 8080 --num-shards 1 \
-		--case-type Performance1536D50K $(COMMON) --db-label antfly-local
-
-bench-qdrant-50k:
-	$(BENCH) qdrantlocal --url http://localhost:6333 \
-		--case-type Performance1536D50K $(COMMON) --db-label qdrant-local \
-		--m 16 --ef-construct 128 --hnsw-ef 64
-
-bench-milvus-50k:
-	$(BENCH) milvusautoindex --uri http://localhost:19530 \
-		--case-type Performance1536D50K $(COMMON) --db-label milvus-local
-
-bench-all-50k: bench-antfly-50k bench-qdrant-50k bench-milvus-50k
-
-# --- 1M benchmarks ---
-bench-antfly-1m:
-	NUM_PER_BATCH=$(ANTFLY_NUM_PER_BATCH) $(BENCH) antflyaknn --host localhost --port 8080 --num-shards 1 \
-		--case-type Performance768D1M $(COMMON) --db-label antfly-local
-
-bench-qdrant-1m:
-	$(BENCH) qdrantlocal --url http://localhost:6333 \
-		--case-type Performance768D1M $(COMMON) --db-label qdrant-local \
-		--m 16 --ef-construct 128 --hnsw-ef 64
-
-bench-milvus-1m:
-	$(BENCH) milvusautoindex --uri http://localhost:19530 \
-		--case-type Performance768D1M $(COMMON) --db-label milvus-local
-
-bench-all-1m: bench-antfly-1m bench-qdrant-1m bench-milvus-1m
-
-# --- Infrastructure ---
-start-qdrant:
-	docker run -d --name qdrant_bench -p 6333:6333 -p 6334:6334 qdrant/qdrant:latest
-
-start-milvus:
-	docker run -d --name milvus_bench -p 19530:19530 -p 9091:9091 \
-		-e ETCD_USE_EMBED=true -e COMMON_STORAGETYPE=local \
-		milvusdb/milvus:latest milvus run standalone
-
-stop-all:
-	-docker stop qdrant_bench milvus_bench 2>/dev/null
-	-docker rm qdrant_bench milvus_bench 2>/dev/null
-
-# --- Dev ---
-unittest:
-	PYTHONPATH=`pwd` python3 -m pytest tests/test_dataset.py::TestDataSet::test_download_small -svv
-
-format:
-	PYTHONPATH=`pwd` python3 -m black vectordb_bench
-	PYTHONPATH=`pwd` python3 -m ruff check vectordb_bench --fix
-
 lint:
 	PYTHONPATH=`pwd` python3 -m black vectordb_bench --check
 	PYTHONPATH=`pwd` python3 -m ruff check vectordb_bench
 
-.PHONY: help bench-antfly-50k bench-qdrant-50k bench-milvus-50k bench-all-50k \
-	bench-antfly-1m bench-qdrant-1m bench-milvus-1m bench-all-1m \
-	start-qdrant start-milvus stop-all unittest format lint
+unittest:
+	PYTHONPATH=`pwd` python3 -m pytest tests/test_dataset.py::TestDataSet::test_download_small -svv
+
+.PHONY: lint unittest

pyproject.toml

@@ -37,7 +37,7 @@ dependencies = [
     "polars",
     "plotly",
     "environs",
-    "pydantic<v2",
+    "pydantic>=2.0",
     "scikit-learn",
     "pymilvus", # with pandas, numpy, ujson
     "ujson",
@@ -79,7 +79,7 @@ all = [
     "mysql-connector-python",
     "turbopuffer[fast]",
     'zvec',
-    "endee==0.1.10", # compatible with pydantic<2
+    "endee==0.1.10", # TODO: check pydantic v2 compat
 ]
 
 qdrant          = [ "qdrant-client" ]

tests/test_bench_runner.py

@@ -52,7 +52,7 @@ def test_performance_case_no_error(self):
         )
 
         t = task_config.copy()
-        d = t.json(exclude={'db_config': {'password', 'api_key'}})
+        d = t.model_dump_json(exclude={'db_config': {'password', 'api_key'}})
         log.info(f"{d}")
 
         import ujson

vectordb_bench/backend/clients/alisql/config.py

@@ -49,8 +49,8 @@ def parse_metric(self) -> str:
 
 
 class AliSQLHNSWConfig(AliSQLIndexConfig, DBCaseConfig):
-    M: int | None
-    ef_search: int | None
+    M: int | None = None
+    ef_search: int | None = None
     index: IndexType = IndexType.HNSW
 
     def index_param(self) -> dict:

vectordb_bench/backend/clients/alloydb/config.py

@@ -43,8 +43,8 @@ class AlloyDBIndexParam(TypedDict):
     metric: str
     index_type: str
     index_creation_with_options: Sequence[dict[str, Any]]
-    maintenance_work_mem: str | None
-    max_parallel_workers: int | None
+    maintenance_work_mem: str | None = None
+    max_parallel_workers: int | None = None
 
 
 class AlloyDBSearchParam(TypedDict):
@@ -120,15 +120,15 @@ def _optionally_build_set_options(
 
 class AlloyDBScaNNConfig(AlloyDBIndexConfig):
     index: IndexType = IndexType.SCANN
-    num_leaves: int | None
-    quantizer: str | None
-    enable_pca: str | None
-    max_num_levels: int | None
-    num_leaves_to_search: int | None
-    max_top_neighbors_buffer_size: int | None
-    pre_reordering_num_neighbors: int | None
-    num_search_threads: int | None
-    max_num_prefetch_datasets: int | None
+    num_leaves: int | None = None
+    quantizer: str | None = None
+    enable_pca: str | None = None
+    max_num_levels: int | None = None
+    num_leaves_to_search: int | None = None
+    max_top_neighbors_buffer_size: int | None = None
+    pre_reordering_num_neighbors: int | None = None
+    num_search_threads: int | None = None
+    max_num_prefetch_datasets: int | None = None
     maintenance_work_mem: str | None = None
     max_parallel_workers: int | None = None