fix(local): load pydantic v1 pickled points under pydantic v2 (#481)

.gitignore

@@ -15,3 +15,4 @@ local_cache/*/*
 docs/source/examples/local_cache/*
 docs/source/examples/path/to/db/*
 .venv
+.venv/

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:

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]