feat: Upgrade pydantic to v2

Signed-off-by: yangxuan <xuan.yang@zilliz.com>

Author: XuanYang-cn

.gitignore

@@ -10,7 +10,6 @@ build/
 venv/
 .venv/
 .idea/
-results/
 logs/
 
 # Worktrees

install/requirements_py3.11.txt

@@ -20,7 +20,7 @@ psutil
 polars
 plotly
 environs
-pydantic<v2
+pydantic>=2.0,<3
 scikit-learn
 pymilvus
 clickhouse_connect

pyproject.toml

@@ -37,7 +37,7 @@ dependencies = [
     "polars",
     "plotly",
     "environs",
-    "pydantic<v2",
+    "pydantic>=2.0,<3",
     "scikit-learn",
     "pymilvus", # with pandas, numpy
     "hdrhistogram>=0.10.1",

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
 

vectordb_bench/backend/clients/api.py

@@ -2,7 +2,7 @@
 from contextlib import contextmanager
 from enum import StrEnum
 
-from pydantic import BaseModel, SecretStr, validator
+from pydantic import BaseModel, model_validator
 
 from vectordb_bench.backend.filter import Filter, FilterOp
 
@@ -90,13 +90,18 @@ def common_long_configs() -> list[str]:
     def to_dict(self) -> dict:
         raise NotImplementedError
 
-    @validator("*")
-    def not_empty_field(cls, v: any, field: any):
-        if field.name in cls.common_short_configs() or field.name in cls.common_long_configs():
-            return v
-        if not v and isinstance(v, str | SecretStr):
-            raise ValueError("Empty string!")
-        return v
+    @model_validator(mode="before")
+    @classmethod
+    def not_empty_field(cls, data: any) -> any:
+        if not isinstance(data, dict):
+            return data
+        skip = set(cls.common_short_configs()) | set(cls.common_long_configs())
+        for field_name, v in data.items():
+            if field_name in skip:
+                continue
+            if isinstance(v, str) and not v:
+                raise ValueError("Empty string!")
+        return data
 
 
 class DBCaseConfig(ABC):

vectordb_bench/backend/clients/aws_opensearch/config.py

@@ -1,7 +1,7 @@
 import logging
 from enum import Enum
 
-from pydantic import BaseModel, SecretStr, validator
+from pydantic import BaseModel, SecretStr, model_validator
 
 from ..api import DBCaseConfig, DBConfig, MetricType
 
@@ -32,17 +32,18 @@ def to_dict(self) -> dict:
             "timeout": 600,
         }
 
-    @validator("*")
-    def not_empty_field(cls, v: any, field: any):
-        if (
-            field.name in cls.common_short_configs()
-            or field.name in cls.common_long_configs()
-            or field.name in ["user", "password", "host"]
-        ):
-            return v
-        if isinstance(v, str | SecretStr) and len(v) == 0:
-            raise ValueError("Empty string!")
-        return v
+    @model_validator(mode="before")
+    @classmethod
+    def not_empty_field(cls, data: any) -> any:
+        if not isinstance(data, dict):
+            return data
+        skip = set(cls.common_short_configs()) | set(cls.common_long_configs()) | {"user", "password", "host"}
+        for field_name, v in data.items():
+            if field_name in skip:
+                continue
+            if isinstance(v, str) and not v:
+                raise ValueError("Empty string!")
+        return data
 
 
 class AWSOS_Engine(Enum):

vectordb_bench/backend/clients/chroma/config.py

@@ -6,7 +6,7 @@
 
 class ChromaConfig(DBConfig):
     user: str | None = None
-    password: SecretStr | None
+    password: SecretStr | None = None
     host: SecretStr = "localhost"
     port: int = 8000
 

vectordb_bench/backend/clients/cockroachdb/config.py

@@ -75,16 +75,16 @@ class CockroachDBIndexParam(TypedDict):
 
     metric: str
     index_creation_with_options: Sequence[dict[str, Any]]
-    min_partition_size: int | None
-    max_partition_size: int | None
-    build_beam_size: int | None
+    min_partition_size: int | None = None
+    max_partition_size: int | None = None
+    build_beam_size: int | None = None
 
 
 class CockroachDBSearchParam(TypedDict):
     """Search parameters for CockroachDB vector queries."""
 
     metric_fun_op: LiteralString
-    vector_search_beam_size: int | None
+    vector_search_beam_size: int | None = None
 
 
 class CockroachDBSessionCommands(TypedDict):

vectordb_bench/backend/clients/doris/config.py

@@ -1,6 +1,6 @@
 import logging
 
-from pydantic import BaseModel, SecretStr, validator
+from pydantic import BaseModel, SecretStr, model_validator
 
 from ..api import DBCaseConfig, DBConfig, MetricType
 
@@ -17,9 +17,10 @@ class DorisConfig(DBConfig):
     db_name: str = "test"
     ssl: bool = False
 
-    @validator("*")
-    def not_empty_field(cls, v: any, field: any):
-        return v
+    @model_validator(mode="before")
+    @classmethod
+    def not_empty_field(cls, data: any) -> any:
+        return data
 
     def to_dict(self) -> dict:
         pwd_str = self.password.get_secret_value()