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,9 +1616,19 @@ 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 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)
+
     def clean(self, data: Any) -> Any:
+        # this method does the same as the one in the parent class, but it
+        # avoids comparisons that break when data is a numpy array
         if data is not None:
             data = self.deserialize(data)
         if (data is None or len(data) == 0) and self._required:

elasticsearch/dsl/utils.py

@@ -612,10 +612,17 @@ def to_dict(self, skip_empty: bool = True) -> Dict[str, Any]:
             if skip_empty:
                 # don't serialize empty values
                 # careful not to include numeric zeros
-                # the "is" operator is used below because it is the only comparison
-                # that works for numpy arrays
-                if v is [] or v is {} or v is None:
-                    continue
+                try:
+                    if v in ([], {}, None):
+                        continue
+                except ValueError:
+                    # the above fails when v is a numpy array
+                    # try using len() instead
+                    try:
+                        if len(v) == 0:
+                            continue
+                    except TypeError:
+                        pass
 
             out[k] = v
         return out

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,9 +422,18 @@ 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 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)
+
     def clean(self, data: Any) -> Any:
+        # this method does the same as the one in the parent class, but it
+        # avoids comparisons that break when data is a numpy array
         if data is not None:
             data = self.deserialize(data)
         if (data is None or len(data) == 0) and self._required: