Add helper function to pack dense vectors for efficient uploading

Author: miguelgrinberg

docs/sphinx/api_helpers.rst

@@ -17,6 +17,10 @@ Bulk
 ----
 .. autofunction:: bulk
 
+Dense Vector packing
+--------------------
+.. autofunction:: pack_dense_vector
+
 Scan
 ----
 .. autofunction:: scan

elasticsearch/helpers/__init__.py

@@ -23,6 +23,7 @@
     BULK_FLUSH,
     bulk,
     expand_action,
+    pack_dense_vector,
     parallel_bulk,
     reindex,
     scan,
@@ -37,6 +38,7 @@
     "expand_action",
     "streaming_bulk",
     "bulk",
+    "pack_dense_vector",
     "parallel_bulk",
     "scan",
     "reindex",

elasticsearch/helpers/actions.py

@@ -15,6 +15,7 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
+import base64
 import logging
 import queue
 import time
@@ -31,6 +32,7 @@
     Mapping,
     MutableMapping,
     Optional,
+    Sequence,
     Tuple,
     Union,
 )
@@ -708,6 +710,22 @@ def _setup_queues(self) -> None:
             pool.join()
 
 
+
+def pack_dense_vector(vector: Union["np.ndarray", Sequence[float]]) -> str:
+    """Helper function that packs a dense vector for efficient uploading.
+
+    :arg v: the list or numpy array to pack.
+    """
+    import numpy as np
+
+    if type(vector) is not np.ndarray:
+        vector = np.array(vector, dtype=np.float32)
+    elif vector.dtype != np.float32:
+        raise ValueError("Only arrays of type float32 can be packed")
+    byte_array = vector.byteswap().tobytes()
+    return base64.b64encode(byte_array).decode()
+
+
 def scan(
     client: Elasticsearch,
     query: Optional[Any] = None,

examples/quotes/backend/quotes.py

@@ -12,6 +12,7 @@
 from elasticsearch import NotFoundError, OrjsonSerializer
 from elasticsearch.dsl.pydantic import AsyncBaseESModel
 from elasticsearch import dsl
+from elasticsearch.helpers import pack_dense_vector
 
 model = SentenceTransformer("all-MiniLM-L6-v2")
 dsl.async_connections.create_connection(hosts=[os.environ['ELASTICSEARCH_URL']], serializer=OrjsonSerializer())
@@ -33,6 +34,9 @@ class Config:
     class Index:
         name = 'quotes'
 
+    def clean(self):
+        # pack the embedding for efficient uploading
+        self.embedding = pack_dense_vector(self.embedding)
 
 class Tag(BaseModel):
     tag: str

test_elasticsearch/test_dsl/test_integration/_async/test_document.py

@@ -60,6 +60,7 @@
 from elasticsearch.dsl.types import MatchQuery
 from elasticsearch.dsl.utils import AttrList
 from elasticsearch.helpers.errors import BulkIndexError
+from elasticsearch.helpers import pack_dense_vector
 
 snowball = analyzer("my_snow", tokenizer="standard", filter=["lowercase", "snowball"])
 
@@ -868,10 +869,19 @@ class Doc(AsyncDocument):
         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(NumpyDenseVector())
+        packed_float_vector: List[float] = mapped_field(DenseVector())
+        packed_numpy_float_vector: np.ndarray = mapped_field(NumpyDenseVector())
 
         class Index:
             name = "vectors"
 
+        def clean(self):
+            # pack the dense vectors before they are sent to Elasticsearch
+            self.packed_float_vector = pack_dense_vector(self.packed_float_vector)
+            self.packed_numpy_float_vector = pack_dense_vector(
+                self.packed_numpy_float_vector
+            )
+
     await Doc._index.delete(ignore_unavailable=True)
     await Doc.init()
 
@@ -884,6 +894,8 @@ class Index:
         byte_vector=test_byte_vector,
         bit_vector=test_bit_vector,
         numpy_float_vector=np.array(test_float_vector),
+        packed_float_vector=test_float_vector,
+        packed_numpy_float_vector=np.array(test_float_vector, dtype=np.float32),
     )
     await doc.save(refresh=True)
 
@@ -894,6 +906,9 @@ class Index:
     assert docs[0].bit_vector == test_bit_vector
     assert type(docs[0].numpy_float_vector) is np.ndarray
     assert [round(v, 1) for v in docs[0].numpy_float_vector] == test_float_vector
+    assert [round(v, 1) for v in docs[0].packed_float_vector] == test_float_vector
+    assert type(docs[0].packed_numpy_float_vector) is np.ndarray
+    assert [round(v, 1) for v in docs[0].packed_numpy_float_vector] == test_float_vector
 
 
 @pytest.mark.anyio

test_elasticsearch/test_dsl/test_integration/_sync/test_document.py

@@ -60,6 +60,7 @@
 from elasticsearch.dsl.types import MatchQuery
 from elasticsearch.dsl.utils import AttrList
 from elasticsearch.helpers.errors import BulkIndexError
+from elasticsearch.helpers import pack_dense_vector
 
 snowball = analyzer("my_snow", tokenizer="standard", filter=["lowercase", "snowball"])
 
@@ -856,10 +857,19 @@ class Doc(Document):
         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(NumpyDenseVector())
+        packed_float_vector: List[float] = mapped_field(DenseVector())
+        packed_numpy_float_vector: np.ndarray = mapped_field(NumpyDenseVector())
 
         class Index:
             name = "vectors"
 
+        def clean(self):
+            # pack the dense vectors before they are sent to Elasticsearch
+            self.packed_float_vector = pack_dense_vector(self.packed_float_vector)
+            self.packed_numpy_float_vector = pack_dense_vector(
+                self.packed_numpy_float_vector
+            )
+
     Doc._index.delete(ignore_unavailable=True)
     Doc.init()
 
@@ -872,6 +882,8 @@ class Index:
         byte_vector=test_byte_vector,
         bit_vector=test_bit_vector,
         numpy_float_vector=np.array(test_float_vector),
+        packed_float_vector=test_float_vector,
+        packed_numpy_float_vector=np.array(test_float_vector, dtype=np.float32),
     )
     doc.save(refresh=True)
 
@@ -882,6 +894,9 @@ class Index:
     assert docs[0].bit_vector == test_bit_vector
     assert type(docs[0].numpy_float_vector) is np.ndarray
     assert [round(v, 1) for v in docs[0].numpy_float_vector] == test_float_vector
+    assert [round(v, 1) for v in docs[0].packed_float_vector] == test_float_vector
+    assert type(docs[0].packed_numpy_float_vector) is np.ndarray
+    assert [round(v, 1) for v in docs[0].packed_numpy_float_vector] == test_float_vector
 
 
 @pytest.mark.sync