From 0d9c32468544dbf5e60532528a17ad8dad844edb Mon Sep 17 00:00:00 2001
From: Phil Barber <Phil.Barber@marklogic.com>
Date: Thu, 10 Jul 2025 20:08:27 -0400
Subject: [PATCH] Can now encode and decode MarkLogic internal vectors

---
 marklogic/vectors.py        | 45 ++++++++++++++++++++++++++++++
 test-app/docker-compose.yml |  4 +--
 tests/test_vectors.py       | 55 +++++++++++++++++++++++++++++++++++++
 3 files changed, 102 insertions(+), 2 deletions(-)
 create mode 100644 marklogic/vectors.py
 create mode 100644 tests/test_vectors.py

diff --git a/marklogic/vectors.py b/marklogic/vectors.py
new file mode 100644
index 0000000..f166f63
--- /dev/null
+++ b/marklogic/vectors.py
@@ -0,0 +1,45 @@
+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)
+        if len(buffer) < 8:
+            raise ValueError(
+                "Buffer is too short to contain version and dimensions."
+            )
+        version, dimensions = struct.unpack("<ii", buffer[:8])
+        if version != 0:
+            raise ValueError(f"Unsupported vector version: {version}")
+        expected_length = 8 + 4 * dimensions
+        if len(buffer) < expected_length:
+            raise ValueError(
+                f"Buffer is too short for the specified dimensions: expected {expected_length}, got {len(buffer)}"
+            )
+        floats = struct.unpack(
+            "<" + "f" * dimensions, buffer[8 : 8 + 4 * dimensions]
+        )
+        return list(floats)
diff --git a/test-app/docker-compose.yml b/test-app/docker-compose.yml
index 1dda1d0..825ae3d 100644
--- a/test-app/docker-compose.yml
+++ b/test-app/docker-compose.yml
@@ -1,9 +1,9 @@
-name: marklogic_python
+name: docker-tests-marklogic_python
 
 services:
 
   marklogic:
-    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
diff --git a/tests/test_vectors.py b/tests/test_vectors.py
new file mode 100644
index 0000000..02c4e76
--- /dev/null
+++ b/tests/test_vectors.py
@@ -0,0 +1,55 @@
+import math
+import ast
+from marklogic.vectors 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)