Move numpy dense vector to a subclass

Author: miguelgrinberg

elasticsearch/dsl/__init__.py

@@ -73,6 +73,7 @@
     MatchOnlyText,
     Murmur3,
     Nested,
+    NumpyDenseVector,
     Object,
     Passthrough,
     Percolator,
@@ -189,6 +190,7 @@
     "Murmur3",
     "Nested",
     "NestedFacet",
+    "NumpyDenseVector",
     "Object",
     "Passthrough",
     "Percolator",

elasticsearch/dsl/field.py

@@ -1555,8 +1555,6 @@ class DenseVector(Field):
     :arg dynamic:
     :arg fields:
     :arg synthetic_source_keep:
-    :arg use_numpy: if set to ``True``, deserialize as a numpy array.
-    :arg dtype: The numpy data type to use as a string, when ``use_numpy`` is ``True``. The default is "float32".
     """
 
     name = "dense_vector"
@@ -1589,8 +1587,6 @@ def __init__(
         synthetic_source_keep: Union[
             Literal["none", "arrays", "all"], "DefaultType"
         ] = DEFAULT,
-        use_numpy: bool = False,
-        dtype: str = "float32",
         **kwargs: Any,
     ):
         if dims is not DEFAULT:
@@ -1618,20 +1614,31 @@ 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
-        self._dtype = dtype
         super().__init__(*args, **kwargs)
 
+
+class NumpyDenseVector(DenseVector):
+    """A dense vector field that uses numpy arrays.
+
+    Accepts the same arguments as class ``DenseVector`` plus:
+
+    :arg dtype: The numpy data type to use for the array. If not given, numpy will select the type based on the data.
+    """
+
+    def __init__(self, *args: Any, dtype: Optional[type] = None, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self._dtype = dtype
+
     def deserialize(self, data: Any) -> Any:
-        if self._use_numpy and isinstance(data, list):
+        if isinstance(data, list):
             import numpy as np
 
-            return np.array(data, dtype=getattr(np, self._dtype))
+            return np.array(data, dtype=self._dtype)
         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
+        # this method does the same as the one in the parent classes, but it
+        # avoids comparisons that do not work for numpy arrays
         if data is not None:
             data = self.deserialize(data)
         if (data is None or len(data) == 0) and self._required:

examples/quotes/backend/quotes.py

@@ -24,8 +24,8 @@ class Quote(AsyncBaseESModel):
     embedding: Annotated[
         np.ndarray,
         PlainSerializer(lambda v: v.tolist()),
-        dsl.DenseVector(use_numpy=True)
-    ] = Field(init=False, default_factory=lambda: np.array([]))
+        dsl.NumpyDenseVector(dtype=np.float32)
+    ] = Field(init=False, default_factory=lambda: np.array([], dtype=np.float32))
 
     class Config:
         arbitrary_types_allowed = True

test_elasticsearch/test_dsl/test_integration/_async/test_document.py

@@ -48,6 +48,7 @@
     Mapping,
     MetaField,
     Nested,
+    NumpyDenseVector,
     Object,
     Q,
     RankFeatures,
@@ -866,29 +867,33 @@ class Doc(AsyncDocument):
         float_vector: List[float] = mapped_field(DenseVector())
         byte_vector: List[int] = mapped_field(DenseVector(element_type="byte"))
         bit_vector: List[int] = mapped_field(DenseVector(element_type="bit"))
-        numpy_float_vector: np.ndarray = mapped_field(DenseVector(use_numpy=True))
+        numpy_float_vector: np.ndarray = mapped_field(NumpyDenseVector())
 
         class Index:
             name = "vectors"
 
     await Doc._index.delete(ignore_unavailable=True)
     await Doc.init()
 
+    test_float_vector = [1.0, 1.2, 2.3]
+    test_byte_vector = [12, 23, 34, 45]
+    test_bit_vector = [18, -43, -112]
+
     doc = Doc(
-        float_vector=[1.0, 1.2, 2.3],
-        byte_vector=[12, 23, 34, 45],
-        bit_vector=[18, -43, -112],
-        numpy_float_vector=np.array([3.1, 2.25, 1.0]),
+        float_vector=test_float_vector,
+        byte_vector=test_byte_vector,
+        bit_vector=test_bit_vector,
+        numpy_float_vector=np.array(test_float_vector),
     )
     await doc.save(refresh=True)
 
     docs = await Doc.search().execute()
     assert len(docs) == 1
-    assert [round(v, 1) for v in docs[0].float_vector] == doc.float_vector
-    assert docs[0].byte_vector == doc.byte_vector
-    assert docs[0].bit_vector == doc.bit_vector
+    assert [round(v, 1) for v in docs[0].float_vector] == test_float_vector
+    assert docs[0].byte_vector == test_byte_vector
+    assert docs[0].bit_vector == test_bit_vector
     assert type(docs[0].numpy_float_vector) is np.ndarray
-    assert np.array_equal(docs[0].numpy_float_vector, doc.numpy_float_vector)
+    assert [round(v, 1) for v in docs[0].numpy_float_vector] == test_float_vector
 
 
 @pytest.mark.anyio

test_elasticsearch/test_dsl/test_integration/_sync/test_document.py

@@ -47,6 +47,7 @@
     Mapping,
     MetaField,
     Nested,
+    NumpyDenseVector,
     Object,
     Q,
     RankFeatures,
@@ -854,29 +855,33 @@ class Doc(Document):
         float_vector: List[float] = mapped_field(DenseVector())
         byte_vector: List[int] = mapped_field(DenseVector(element_type="byte"))
         bit_vector: List[int] = mapped_field(DenseVector(element_type="bit"))
-        numpy_float_vector: np.ndarray = mapped_field(DenseVector(use_numpy=True))
+        numpy_float_vector: np.ndarray = mapped_field(NumpyDenseVector())
 
         class Index:
             name = "vectors"
 
     Doc._index.delete(ignore_unavailable=True)
     Doc.init()
 
+    test_float_vector = [1.0, 1.2, 2.3]
+    test_byte_vector = [12, 23, 34, 45]
+    test_bit_vector = [18, -43, -112]
+
     doc = Doc(
-        float_vector=[1.0, 1.2, 2.3],
-        byte_vector=[12, 23, 34, 45],
-        bit_vector=[18, -43, -112],
-        numpy_float_vector=np.array([3.1, 2.25, 1.0]),
+        float_vector=test_float_vector,
+        byte_vector=test_byte_vector,
+        bit_vector=test_bit_vector,
+        numpy_float_vector=np.array(test_float_vector),
     )
     doc.save(refresh=True)
 
     docs = Doc.search().execute()
     assert len(docs) == 1
-    assert [round(v, 1) for v in docs[0].float_vector] == doc.float_vector
-    assert docs[0].byte_vector == doc.byte_vector
-    assert docs[0].bit_vector == doc.bit_vector
+    assert [round(v, 1) for v in docs[0].float_vector] == test_float_vector
+    assert docs[0].byte_vector == test_byte_vector
+    assert docs[0].bit_vector == test_bit_vector
     assert type(docs[0].numpy_float_vector) is np.ndarray
-    assert np.array_equal(docs[0].numpy_float_vector, doc.numpy_float_vector)
+    assert [round(v, 1) for v in docs[0].numpy_float_vector] == test_float_vector
 
 
 @pytest.mark.sync

utils/templates/field.py.tpl

@@ -217,10 +217,6 @@ 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.
-    :arg dtype: The numpy data type to use as a string, when ``use_numpy`` is ``True``. The default is "float32".
-    {% endif %}
     """
     name = "{{ k.field }}"
     {% if k.coerced %}
@@ -250,10 +246,6 @@ class {{ k.name }}({{ k.parent }}):
         {{ arg.name }}: {{ arg.type }} = DEFAULT,
             {% endif %}
         {% endfor %}
-        {% if k.field == "dense_vector" %}
-        use_numpy: bool = False,
-        dtype: str = "float32",
-        {% endif %}
         **kwargs: Any
     ):
         {% for arg in k.args %}
@@ -424,19 +416,28 @@ 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
-        self._dtype = dtype
         super().__init__(*args, **kwargs)
 
+class NumpyDenseVector(DenseVector):
+    """A dense vector field that uses numpy arrays.
+    
+    Accepts the same arguments as class ``DenseVector`` plus:
+
+    :arg dtype: The numpy data type to use for the array. If not given, numpy will select the type based on the data.
+    """
+    def __init__(self, *args: Any, dtype: Optional[type] = None, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self._dtype = dtype
+
     def deserialize(self, data: Any) -> Any: 
-        if self._use_numpy and isinstance(data, list):
+        if isinstance(data, list):
             import numpy as np
-            return np.array(data, dtype=getattr(np, self._dtype))
+            return np.array(data, dtype=self._dtype)
         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
+        # this method does the same as the one in the parent classes, but it
+        # avoids comparisons that do not work for numpy arrays
         if data is not None:
             data = self.deserialize(data)
         if (data is None or len(data) == 0) and self._required: