numpy deserialization

Author: miguelgrinberg

elasticsearch/dsl/field.py

@@ -1555,6 +1555,7 @@ class DenseVector(Field):
     :arg dynamic:
     :arg fields:
     :arg synthetic_source_keep:
+    :arg use_numpy: if set to ``True``, deserialize as a numpy array.
     """
 
     name = "dense_vector"
@@ -1587,6 +1588,7 @@ def __init__(
         synthetic_source_keep: Union[
             Literal["none", "arrays", "all"], "DefaultType"
         ] = DEFAULT,
+        use_numpy: bool = False,
         **kwargs: Any,
     ):
         if dims is not DEFAULT:
@@ -1614,6 +1616,7 @@ def __init__(
         self._element_type = kwargs.get("element_type", "float")
         if self._element_type in ["float", "byte"]:
             kwargs["multi"] = True
+        self._use_numpy = use_numpy
         super().__init__(*args, **kwargs)
 
     def clean(self, data: Any) -> Any:
@@ -1623,6 +1626,13 @@ def clean(self, data: Any) -> Any:
             raise ValidationException("Value required for this field.")
         return data
 
+    def deserialize(self, data: Any) -> Any:
+        if self._use_numpy and isinstance(data, list):
+            import numpy as np
+
+            return np.array(data)
+        return super().deserialize(data)
+
 
 class Double(Float):
     """

examples/quotes/backend/quotes.py

@@ -5,22 +5,30 @@
 from typing import Annotated
 
 from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel, Field, ValidationError
+import numpy as np
+from pydantic import BaseModel, Field, PlainSerializer
 from sentence_transformers import SentenceTransformer
 
-from elasticsearch import NotFoundError
+from elasticsearch import NotFoundError, OrjsonSerializer
 from elasticsearch.dsl.pydantic import AsyncBaseESModel
 from elasticsearch import dsl
 
 model = SentenceTransformer("all-MiniLM-L6-v2")
-dsl.async_connections.create_connection(hosts=[os.environ['ELASTICSEARCH_URL']])
+dsl.async_connections.create_connection(hosts=[os.environ['ELASTICSEARCH_URL']], serializer=OrjsonSerializer())
 
 
 class Quote(AsyncBaseESModel):
     quote: str
     author: Annotated[str, dsl.Keyword()]
     tags: Annotated[list[str], dsl.Keyword()]
-    embedding: Annotated[list[float], dsl.DenseVector()] = Field(init=False, default=[])
+    embedding: Annotated[
+        np.ndarray,
+        PlainSerializer(lambda v: v.tolist()),
+        dsl.DenseVector(use_numpy=True)
+    ] = Field(init=False, default_factory=lambda: np.array([]))
+
+    class Config:
+        arbitrary_types_allowed = True
 
     class Index:
         name = 'quotes'

utils/templates/field.py.tpl

@@ -217,6 +217,9 @@ class {{ k.name }}({{ k.parent }}):
             {% endfor %}
         {% endfor %}
     {% endif %}
+    {% if k.field == "dense_vector" %}
+    :arg use_numpy: if set to ``True``, deserialize as a numpy array.
+    {% endif %}
     """
     name = "{{ k.field }}"
     {% if k.coerced %}
@@ -246,6 +249,9 @@ class {{ k.name }}({{ k.parent }}):
         {{ arg.name }}: {{ arg.type }} = DEFAULT,
             {% endif %}
         {% endfor %}
+        {% if k.field == "dense_vector" %}
+        use_numpy: bool = False,
+        {% endif %}
         **kwargs: Any
     ):
         {% for arg in k.args %}
@@ -416,6 +422,7 @@ class {{ k.name }}({{ k.parent }}):
         self._element_type = kwargs.get("element_type", "float")
         if self._element_type in ["float", "byte"]:
             kwargs["multi"] = True
+        self._use_numpy = use_numpy
         super().__init__(*args, **kwargs)
 
     def clean(self, data: Any) -> Any:
@@ -424,6 +431,12 @@ class {{ k.name }}({{ k.parent }}):
         if (data is None or len(data) == 0) and self._required:
             raise ValidationException("Value required for this field.")
         return data
+
+    def deserialize(self, data: Any) -> Any: 
+        if self._use_numpy and isinstance(data, list):
+            import numpy as np
+            return np.array(data)
+        return super().deserialize(data)
         {% elif k.field == "scaled_float" %}
         if 'scaling_factor' not in kwargs:
             if len(args) > 0: