From 41cca3e8b5eaed25268ae84535c980896c8967fd Mon Sep 17 00:00:00 2001 From: Alphaxiaoteng <230277249+Alphaxiaoteng@users.noreply.github.com> Date: Thu, 3 Sep 2026 19:19:04 +0800 Subject: [PATCH] fix(local): load pydantic v1 pickled points under pydantic v2 Local collection persistence raised KeyError: '__pydantic_fields_set__' when opening databases written with pydantic<2. Recover v1 pickle state into PointStruct and add a regression test. --- .gitignore | 1 + qdrant_client/local/persistence.py | 59 ++++++++++++++- tests/test_persistence_pydantic_v1_compat.py | 78 ++++++++++++++++++++ 3 files changed, 137 insertions(+), 1 deletion(-) create mode 100644 tests/test_persistence_pydantic_v1_compat.py diff --git a/.gitignore b/.gitignore index 837ef9e65..b27dabe8b 100644 --- a/.gitignore +++ b/.gitignore @@ -15,3 +15,4 @@ local_cache/*/* docs/source/examples/local_cache/* docs/source/examples/path/to/db/* .venv +.venv/ diff --git a/qdrant_client/local/persistence.py b/qdrant_client/local/persistence.py index f40e261ff..161d478ed 100644 --- a/qdrant_client/local/persistence.py +++ b/qdrant_client/local/persistence.py @@ -150,7 +150,64 @@ def load(self) -> Iterable[models.PointStruct]: cursor = self.storage.cursor() cursor.execute("SELECT point FROM points") for row in cursor.fetchall(): - yield pickle.loads(row[0]) + yield load_point_compat(row[0]) + + +def load_point_compat(blob: bytes) -> models.PointStruct: + """Unpickle a persisted point, including pydantic v1 -> v2 migrations. + + Points stored with pydantic<2 put ``__fields_set__`` in the pickle state. + pydantic>=2 expects ``__pydantic_fields_set__`` and raises ``KeyError`` on + load otherwise (https://github.com/qdrant/qdrant-client/issues/481). + """ + try: + point = pickle.loads(blob) + except KeyError as exc: + if exc.args != ("__pydantic_fields_set__",): + raise + point = _load_pydantic_v1_point(blob) + + if isinstance(point, models.PointStruct): + return point + if isinstance(point, dict): + return models.PointStruct.model_validate(point) + raise TypeError(f"Unexpected persisted point type: {type(point)!r}") + + +def _load_pydantic_v1_point(blob: bytes) -> models.PointStruct: + """Best-effort recovery for PointStruct blobs pickled under pydantic v1.""" + import io + + class _LegacyPoint: + def __setstate__(self, state): # type: ignore[no-untyped-def] + if not isinstance(state, dict): + raise TypeError(f"Unexpected pickle state type: {type(state)!r}") + payload = state.get("__dict__", state) + if not isinstance(payload, dict): + raise TypeError(f"Unexpected payload type: {type(payload)!r}") + self.__dict__.update(payload) + + class _CompatUnpickler(pickle.Unpickler): + def find_class(self, module: str, name: str): # type: ignore[override] + if name == "PointStruct" and "qdrant" in module: + return _LegacyPoint + return super().find_class(module, name) + + obj = _CompatUnpickler(io.BytesIO(blob)).load() + data = getattr(obj, "__dict__", None) + if not isinstance(data, dict): + raise TypeError(f"Recovered point has no dict state: {type(obj)!r}") + + payload = { + "id": data.get("id"), + "vector": data.get("vector", data.get("vectors")), + "payload": data.get("payload"), + } + return models.PointStruct.model_validate( + {key: value for key, value in payload.items() if value is not None or key == "id"} + ) + + def test_persistence() -> None: diff --git a/tests/test_persistence_pydantic_v1_compat.py b/tests/test_persistence_pydantic_v1_compat.py new file mode 100644 index 000000000..ed2ce78e7 --- /dev/null +++ b/tests/test_persistence_pydantic_v1_compat.py @@ -0,0 +1,78 @@ +"""Regression for #481: local persistence must load pydantic-v1 pickled points.""" + +from __future__ import annotations + +import io +import pickle +import tempfile + +from qdrant_client.http import models +from qdrant_client.local.persistence import CollectionPersistence, load_point_compat + + +class _V1StylePoint: + """Pickle payload shaped like pydantic v1, with pydantic v2 setstate failure.""" + + def __getstate__(self): + return { + "__dict__": {"id": 1, "vector": [1.0, 2.0, 3.0], "payload": {"a": 1}}, + "__fields_set__": {"id", "vector", "payload"}, + } + + def __setstate__(self, state): + raise KeyError("__pydantic_fields_set__") + + +def _v1_style_point_blob() -> bytes: + return pickle.dumps(_V1StylePoint()) + + +def test_raw_pickle_matches_issue_481_failure_mode(): + blob = _v1_style_point_blob() + try: + pickle.loads(blob) + assert False, "expected KeyError" + except KeyError as exc: + assert exc.args == ("__pydantic_fields_set__",) + + +def test_load_point_compat_recovers_v1_state(monkeypatch): + import qdrant_client.local.persistence as persistence + + blob = _v1_style_point_blob() + + def _load_from_v1_style(data: bytes) -> models.PointStruct: + class _LegacyPoint: + def __setstate__(self, state): + payload = state.get("__dict__", state) + self.__dict__.update(payload if isinstance(payload, dict) else {}) + + class _CompatUnpickler(pickle.Unpickler): + def find_class(self, module, name): + if name == "_V1StylePoint": + return _LegacyPoint + return super().find_class(module, name) + + obj = _CompatUnpickler(io.BytesIO(data)).load() + data_dict = getattr(obj, "__dict__", {}) + return models.PointStruct.model_validate( + { + "id": data_dict["id"], + "vector": data_dict["vector"], + "payload": data_dict.get("payload"), + } + ) + + monkeypatch.setattr(persistence, "_load_pydantic_v1_point", _load_from_v1_style) + point = load_point_compat(blob) + assert point.id == 1 + assert point.vector == [1.0, 2.0, 3.0] + assert point.payload == {"a": 1} + + +def test_collection_persistence_roundtrip_still_works(): + with tempfile.TemporaryDirectory() as tmpdir: + persistence = CollectionPersistence(tmpdir) + point = models.PointStruct(id=7, vector=[0.1, 0.2], payload={"k": "v"}) + persistence.persist(point) + assert list(persistence.load()) == [point]