[DRIVERS-2926] [PYTHON-4577] BSON Binary Vector Subtype Support

.evergreen/resync-specs.sh

@@ -76,6 +76,9 @@ do
     atlas-data-lake-testing|data_lake)
       cpjson atlas-data-lake-testing/tests/ data_lake
       ;;
+    bson-binary-vector|bson_binary_vector)
+      cpjson bson-binary-vector/tests/ bson_binary_vector
+      ;;
     bson-corpus|bson_corpus)
       cpjson bson-corpus/tests/ bson_corpus
       ;;

bson/binary.py

@@ -13,7 +13,10 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Tuple, Type, Union
+import struct
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Optional, Sequence, Tuple, Type, Union
 from uuid import UUID
 
 """Tools for representing BSON binary data.
@@ -191,21 +194,76 @@ class UuidRepresentation:
 """
 
 
+VECTOR_SUBTYPE = 9
+"""BSON binary subtype for densely packed vector data.
+
+.. versionadded:: 4.9
+"""
+
+
 USER_DEFINED_SUBTYPE = 128
 """BSON binary subtype for any user defined structure.
 """
 
 
+class BinaryVectorDtype(Enum):
+    """Datatypes of vector subtype.
+
+    :param FLOAT32: (0x27) Pack list of :class:`float` as float32
+    :param INT8: (0x03) Pack list of :class:`int` in [-128, 127] as signed int8
+    :param PACKED_BIT: (0x10) Pack list of :class:`int` in [0, 255] as unsigned uint8
+
+    The `PACKED_BIT` value represents a special case where vector values themselves
+    can only be of two values (0 or 1) but these are packed together into groups of 8,
+    a byte. In Python, these are displayed as ints in range [0, 255]
+
+    Each value is of type bytes with a length of one.
+
+    .. versionadded:: 4.9
+    """
+
+    INT8 = b"\x03"
+    FLOAT32 = b"\x27"
+    PACKED_BIT = b"\x10"
+
+
+# Map from bytes to enum value, for decoding.
+DTYPE_FROM_HEX = {key.value: key for key in BinaryVectorDtype}
+
+
+@dataclass
+class BinaryVector:
+    """Vector of numbers along with metadata for binary interoperability.
+
+    :param data: Sequence of numbers representing the mathematical vector.
+    :param dtype:  The data type stored in binary
+    :param padding: The number of bits in the final byte that are to be ignored
+      when a vector element's size is less than a byte
+      and the length of the vector is not a multiple of 8.
+
+    .. versionadded:: 4.9
+    """
+
+    data: Sequence[float | int]
+    dtype: BinaryVectorDtype
+    padding: Optional[int] = 0
+
+
 class Binary(bytes):
     """Representation of BSON binary data.
 
-    This is necessary because we want to represent Python strings as
-    the BSON string type. We need to wrap binary data so we can tell
+    We want to represent Python strings as the BSON string type.
+    We need to wrap binary data so that we can tell
     the difference between what should be considered binary data and
     what should be considered a string when we encode to BSON.
 
-    Raises TypeError if `data` is not an instance of :class:`bytes`
-    or `subtype` is not an instance of :class:`int`.
+    Subtype 9 provides a space-efficient representation of 1-dimensional vector data.
+    Its data is prepended with two bytes of metadata.
+    The first (dtype) describes its data type, such as float32 or int8.
+    The second (padding) prescribes the number of bits to ignore in the final byte.
+    This is relevant when the element size of the dtype is not a multiple of 8.
+
+    Raises TypeError if subtype` is not an instance of :class:`int`.
     Raises ValueError if `subtype` is not in [0, 256).
 
     .. note::
@@ -218,7 +276,10 @@ class Binary(bytes):
         to use
 
     .. versionchanged:: 3.9
-      Support any bytes-like type that implements the buffer protocol.
+       Support any bytes-like type that implements the buffer protocol.
+
+    .. versionchanged:: 4.9
+       Addition of vector subtype.
     """
 
     _type_marker = 5
@@ -337,6 +398,90 @@ def as_uuid(self, uuid_representation: int = UuidRepresentation.STANDARD) -> UUI
             f"cannot decode subtype {self.subtype} to {UUID_REPRESENTATION_NAMES[uuid_representation]}"
         )
 
+    @classmethod
+    def from_vector(
+        cls: Type[Binary],
+        vector: list[int, float],
+        dtype: BinaryVectorDtype,
+        padding: Optional[int] = 0,
+    ) -> Binary:
+        """Create a BSON :class:`~bson.binary.Binary` of Vector subtype from a list of Numbers.
+
+        To interpret the representation of the numbers, a data type must be included.
+        See :class:`~bson.binary.BinaryVectorDtype` for available types and descriptions.
+
+        The dtype and padding are prepended to the binary data's value.
+
+        :param vector: List of values
+        :param dtype: Data type of the values
+        :param padding: For fractional bytes, number of bits to ignore at end of vector.
+        :return: Binary packed data identified by dtype and padding.
+
+        .. versionadded:: 4.9
+        """
+        if dtype == BinaryVectorDtype.INT8:  # pack ints in [-128, 127] as signed int8
+            format_str = "b"
+            assert not padding, f"padding does not apply to {dtype=}"
+        elif dtype == BinaryVectorDtype.PACKED_BIT:  # pack ints in [0, 255] as unsigned uint8
+            format_str = "B"
+        elif dtype == BinaryVectorDtype.FLOAT32:  # pack floats as float32
+            format_str = "f"
+            assert not padding, f"padding does not apply to {dtype=}"
+        else:
+            raise NotImplementedError("%s not yet supported" % dtype)
+
+        metadata = struct.pack("<sB", dtype.value, padding)
+        data = struct.pack(f"{len(vector)}{format_str}", *vector)
+        return cls(metadata + data, subtype=VECTOR_SUBTYPE)
+
+    def as_vector(self, uncompressed: Optional[bool] = False) -> BinaryVector:
+        """From the Binary, create a list of numbers, along with dtype and padding.
+
+
+        :param uncompressed: If true, return the true mathematical vector.
+            This is only necessary for datatypes where padding is applicable.
+            For example, setting this to True for a PACKED_BIT vector will result
+            in a List[int] of zeros and ones.
+        :return: List of numbers, along with dtype and padding.
+
+        .. versionadded:: 4.9
+        """
+
+        position = 0
+        dtype, padding = struct.unpack_from("<sB", self, position)
+        position += 2
+        dtype = BinaryVectorDtype(dtype)
+        n_values = len(self) - position
+
+        if dtype == BinaryVectorDtype.INT8:
+            dtype_format = "b"
+            format_string = f"{n_values}{dtype_format}"
+            vector = list(struct.unpack_from(format_string, self, position))
+            return BinaryVector(vector, dtype, padding)
+
+        elif dtype == BinaryVectorDtype.FLOAT32:
+            n_bytes = len(self) - position
+            n_values = n_bytes // 4
+            assert n_bytes % 4 == 0
+            vector = list(struct.unpack_from(f"{n_values}f", self, position))
+            return BinaryVector(vector, dtype, padding)
+
+        elif dtype == BinaryVectorDtype.PACKED_BIT:
+            # data packed as uint8
+            dtype_format = "B"
+            unpacked_uint8s = list(struct.unpack_from(f"{n_values}{dtype_format}", self, position))
+            if not uncompressed:
+                return BinaryVector(unpacked_uint8s, dtype, padding)
+            else:
+                bits = []
+                for uint8 in unpacked_uint8s:
+                    bits.extend([int(bit) for bit in f"{uint8:08b}"])
+                vector = bits[:-padding] if padding else bits
+                return BinaryVector(vector, dtype, padding)
+
+        else:
+            raise NotImplementedError("Binary Vector dtype %s not yet supported" % dtype.name)
+
     @property
     def subtype(self) -> int:
         """Subtype of this binary data."""

doc/api/bson/binary.rst

@@ -21,6 +21,14 @@
    .. autoclass:: UuidRepresentation
       :members:
 
+   .. autoclass:: BinaryVectorDtype
+      :members:
+      :show-inheritance:
+
+   .. autoclass:: BinaryVector
+      :members:
+
+
    .. autoclass:: Binary(data, subtype=BINARY_SUBTYPE)
       :members:
       :show-inheritance:

test/bson_binary_vector/float32.json

@@ -0,0 +1,45 @@
+{
+  "description": "Tests of Binary subtype 9, Vectors, with dtype FLOAT32",
+  "test_key": "vector",
+  "tests": [
+    {
+      "description": "Simple Vector FLOAT32",
+      "valid": true,
+      "vector": [127.0, 7.0],
+      "dtype_hex": "0x27",
+      "dtype_alias": "FLOAT32",
+      "padding": 0,
+      "canonical_bson": "1C00000005766563746F72000A0000000927000000FE420000E04000", 
+      "canonical_extjson": "{\"vector\": {\"$binary\": {\"base64\": \"JwAAAP5CAADgQA==\", \"subType\": \"09\"}}}"
+    },
+    {
+      "description": "Empty Vector FLOAT32",
+      "valid": true,
+      "vector": [],
+      "dtype_hex": "0x27",
+      "dtype_alias": "FLOAT32",
+      "padding": 0,
+      "canonical_bson": "1400000005766563746F72000200000009270000",
+      "canonical_extjson": "{\"vector\": {\"$binary\": {\"base64\": \"JwA=\", \"subType\": \"09\"}}}"
+    },
+    {
+      "description": "Infinity Vector FLOAT32",
+      "valid": true,
+      "vector": ["-inf", 0.0, "inf"],
+      "dtype_hex": "0x27",
+      "dtype_alias": "FLOAT32",
+      "padding": 0,
+      "canonical_bson": "2000000005766563746F72000E000000092700000080FF000000000000807F00",
+      "canonical_extjson": "{\"vector\": {\"$binary\": {\"base64\": \"JwAAAID/AAAAAAAAgH8=\", \"subType\": \"09\"}}}"
+    },
+    {
+      "description": "FLOAT32 with padding",
+      "valid": false,
+      "vector": [127.0, 7.0],
+      "dtype_hex": "0x27",
+      "dtype_alias": "FLOAT32",
+      "padding": 3
+    }
+  ]
+}
+

test/bson_binary_vector/int8.json

@@ -0,0 +1,59 @@
+{
+  "description": "Tests of Binary subtype 9, Vectors, with dtype INT8",
+  "test_key": "vector",
+  "tests": [
+    {
+      "description": "Simple Vector INT8",
+      "valid": true,
+      "vector": [127, 7],
+      "dtype_hex": "0x03",
+      "dtype_alias": "INT8",
+      "padding": 0,
+      "canonical_bson": "1600000005766563746F7200040000000903007F0700",
+      "canonical_extjson": "{\"vector\": {\"$binary\": {\"base64\": \"AwB/Bw==\", \"subType\": \"09\"}}}"
+    },
+    {
+      "description": "Empty Vector INT8",
+      "valid": true,
+      "vector": [],
+      "dtype_hex": "0x03",
+      "dtype_alias": "INT8",
+      "padding": 0,
+      "canonical_bson": "1400000005766563746F72000200000009030000",
+      "canonical_extjson": "{\"vector\": {\"$binary\": {\"base64\": \"AwA=\", \"subType\": \"09\"}}}"
+    },
+    {
+      "description": "Overflow Vector INT8",
+      "valid": false,
+      "vector": [128],
+      "dtype_hex": "0x03",
+      "dtype_alias": "INT8",
+      "padding": 0
+    },
+    {
+      "description": "Underflow Vector INT8",
+      "valid": false,
+      "vector": [-129],
+      "dtype_hex": "0x03",
+      "dtype_alias": "INT8",
+      "padding": 0
+    },
+    {
+      "description": "INT8 with padding",
+      "valid": false,
+      "vector": [127, 7],
+      "dtype_hex": "0x03",
+      "dtype_alias": "INT8",
+      "padding": 3
+    },
+    {
+      "description": "INT8 with float inputs",
+      "valid": false,
+      "vector": [127.77, 7.77],
+      "dtype_hex": "0x03",
+      "dtype_alias": "INT8",
+      "padding": 0
+    }
+  ]
+}
+

test/bson_binary_vector/packed_bit.json

@@ -0,0 +1,53 @@
+{
+  "description": "Tests of Binary subtype 9, Vectors, with dtype PACKED_BIT",
+  "test_key": "vector",
+  "tests": [
+    {
+      "description": "Simple Vector PACKED_BIT",
+      "valid": true,
+      "vector": [127, 7],
+      "dtype_hex": "0x10",
+      "dtype_alias": "PACKED_BIT",
+      "padding": 0,
+      "canonical_bson": "1600000005766563746F7200040000000910007F0700",
+      "canonical_extjson": "{\"vector\": {\"$binary\": {\"base64\": \"EAB/Bw==\", \"subType\": \"09\"}}}"
+    },
+    {
+      "description": "Empty Vector PACKED_BIT",
+      "valid": true,
+      "vector": [],
+      "dtype_hex": "0x10",
+      "dtype_alias": "PACKED_BIT",
+      "padding": 0,
+      "canonical_bson": "1400000005766563746F72000200000009100000",
+      "canonical_extjson": "{\"vector\": {\"$binary\": {\"base64\": \"EAA=\", \"subType\": \"09\"}}}"
+    },
+    {
+      "description": "PACKED_BIT with padding",
+      "valid": true,
+      "vector": [127, 7],
+      "dtype_hex": "0x10",
+      "dtype_alias": "PACKED_BIT",
+      "padding": 3,
+      "canonical_bson": "1600000005766563746F7200040000000910037F0700",
+      "canonical_extjson": "{\"vector\": {\"$binary\": {\"base64\": \"EAN/Bw==\", \"subType\": \"09\"}}}"
+    },
+    {
+      "description": "Overflow Vector PACKED_BIT",
+      "valid": false,
+      "vector": [256],
+      "dtype_hex": "0x10",
+      "dtype_alias": "PACKED_BIT",
+      "padding": 0
+    },
+    {
+      "description": "Underflow Vector PACKED_BIT",
+      "valid": false,
+      "vector": [-1],
+      "dtype_hex": "0x10",
+      "dtype_alias": "PACKED_BIT",
+      "padding": 0
+    }
+  ]
+}
+