Can now encode and decode MarkLogic internal vectors

Author: BillFarber

marklogic/vector_util.py

@@ -0,0 +1,36 @@
+import base64
+import struct
+from typing import List
+
+
+class VectorUtil:
+    """
+    Supports encoding and decoding vectors using the same approach as the vec:base64-encode and vec:base64-decode
+    functions supported by the MarkLogic server.
+    """
+
+    @staticmethod
+    def base64_encode(vector: List[float]) -> str:
+        """
+        Encodes a list of floats as a base64 string compatible with MarkLogic's vec:base64-encode.
+        """
+        dimensions = len(vector)
+        # version (int32, 0) + dimensions (int32) + floats (little-endian)
+        buffer = struct.pack("<ii", 0, dimensions) + struct.pack(
+            "<" + "f" * dimensions, *vector
+        )
+        return base64.b64encode(buffer).decode("ascii")
+
+    @staticmethod
+    def base64_decode(encoded_vector: str) -> List[float]:
+        """
+        Decodes a base64 string to a list of floats compatible with MarkLogic's vec:base64-decode.
+        """
+        buffer = base64.b64decode(encoded_vector)
+        version, dimensions = struct.unpack("<ii", buffer[:8])
+        if version != 0:
+            raise ValueError(f"Unsupported vector version: {version}")
+        floats = struct.unpack(
+            "<" + "f" * dimensions, buffer[8 : 8 + 4 * dimensions]
+        )
+        return list(floats)

test-app/docker-compose.yml

@@ -1,9 +1,11 @@
-name: marklogic_python
+# name: docker-tests-marklogic_python
+name: docker-tests-marklogic_python-12
 
 services:
 
   marklogic:
-    image: "progressofficial/marklogic-db:11.3.0-ubi"
+    # image: "progressofficial/marklogic-db:11.3.0-ubi"
+    image: "ml-docker-db-dev-tierpoint.bed-artifactory.bedford.progress.com/marklogic/marklogic-server-ubi-rootless:latest-12"
     platform: linux/amd64
     environment:
       - INSTALL_CONVERTERS=true

tests/test_vector_util.py

@@ -0,0 +1,55 @@
+import math
+import ast
+from marklogic.vector_util import VectorUtil
+from marklogic import Client
+
+VECTOR = [3.14, 1.59, 2.65]
+EXPECTED_BASE64 = "AAAAAAMAAADD9UhAH4XLP5qZKUA="
+ACCEPTABLE_DELTA = 0.0001
+
+
+def test_encode_and_decode_with_python():
+    encoded = VectorUtil.base64_encode(VECTOR)
+    assert encoded == EXPECTED_BASE64
+
+    decoded = VectorUtil.base64_decode(encoded)
+    assert len(decoded) == len(VECTOR)
+    for a, b in zip(decoded, VECTOR):
+        assert abs(a - b) < ACCEPTABLE_DELTA
+
+
+def test_decode_known_base64():
+    decoded = VectorUtil.base64_decode(EXPECTED_BASE64)
+    assert len(decoded) == len(VECTOR)
+    for a, b in zip(decoded, VECTOR):
+        assert abs(a - b) < ACCEPTABLE_DELTA
+
+
+def test_encode_and_decode_with_server(client: Client):
+    """
+    Encode a vector in Python, decode it on the MarkLogic server, and check the result.
+    """
+    encoded = VectorUtil.base64_encode(VECTOR)
+    assert encoded == EXPECTED_BASE64
+
+    # Use MarkLogic's eval endpoint to decode the vector on the server
+    xquery = f"vec:base64-decode('{encoded}')"
+    binary_result = client.eval(xquery=xquery)
+    float_list = ast.literal_eval(binary_result[0].decode("utf-8"))
+    assert len(float_list) == len(VECTOR)
+    for a, b in zip(float_list, VECTOR):
+        assert math.isclose(a, b, abs_tol=ACCEPTABLE_DELTA)
+
+
+def test_encode_with_server_and_decode_with_python(client: Client):
+    """
+    Encode a vector on the MarkLogic server, decode it in Python, and check the result.
+    """
+    xquery = "vec:base64-encode(vec:vector((3.14, 1.59, 2.65)))"
+    encoded = client.eval(xquery=xquery)[0]
+    assert encoded == EXPECTED_BASE64
+
+    decoded = VectorUtil.base64_decode(encoded)
+    assert len(decoded) == len(VECTOR)
+    for a, b in zip(decoded, VECTOR):
+        assert math.isclose(a, b, abs_tol=ACCEPTABLE_DELTA)