Vastly improve performance

Author: NoahStapp

.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/__init__.py

@@ -1324,7 +1324,7 @@ def decode_iter(
         elements = data[position : position + obj_size]
         position += obj_size
 
-        yield _bson_to_dict(elements, opts)  # type:ignore[misc, type-var]
+        yield _bson_to_dict(elements, opts)  # type:ignore[misc]
 
 
 @overload
@@ -1370,7 +1370,7 @@ def decode_file_iter(
             raise InvalidBSON("cut off in middle of objsize")
         obj_size = _UNPACK_INT_FROM(size_data, 0)[0] - 4
         elements = size_data + file_obj.read(max(0, obj_size))
-        yield _bson_to_dict(elements, opts)  # type:ignore[type-var, arg-type, misc]
+        yield _bson_to_dict(elements, opts)  # type:ignore[arg-type, misc]
 
 
 def is_valid(bson: bytes) -> bool:

bson/_cbsonmodule.c

@@ -207,7 +207,7 @@ static PyObject* _test_long_long_to_str(PyObject* self, PyObject* args) {
  *
  * Returns a new ref */
 static PyObject* _error(char* name) {
-    PyObject* error;
+    PyObject* error = NULL;
     PyObject* errors = PyImport_ImportModule("bson.errors");
     if (!errors) {
         return NULL;
@@ -279,7 +279,7 @@ static PyObject* datetime_from_millis(long long millis) {
      * micros = diff * 1000                       111000
      * Resulting in datetime(1, 1, 1, 1, 1, 1, 111000) -- the expected result
      */
-    PyObject* datetime;
+    PyObject* datetime = NULL;
     int diff = (int)(((millis % 1000) + 1000) % 1000);
     int microseconds = diff * 1000;
     Time64_T seconds = (millis - diff) / 1000;
@@ -294,7 +294,7 @@ static PyObject* datetime_from_millis(long long millis) {
                                           timeinfo.tm_sec,
                                           microseconds);
     if(!datetime) {
-        PyObject *etype, *evalue, *etrace;
+        PyObject *etype = NULL, *evalue = NULL, *etrace = NULL;
 
         /*
         * Calling _error clears the error state, so fetch it first.
@@ -350,8 +350,8 @@ static PyObject* datetime_ms_from_millis(PyObject* self, long long millis){
         return NULL;
     }
 
-    PyObject* dt;
-    PyObject* ll_millis;
+    PyObject* dt = NULL;
+    PyObject* ll_millis = NULL;
 
     if (!(ll_millis = PyLong_FromLongLong(millis))){
         return NULL;
@@ -1790,7 +1790,7 @@ static PyObject* _cbson_dict_to_bson(PyObject* self, PyObject* args) {
     PyObject* result;
     unsigned char check_keys;
     unsigned char top_level = 1;
-    PyObject* options_obj;
+    PyObject* options_obj = NULL;
     codec_options_t options;
     buffer_t buffer;
     PyObject* raw_bson_document_bytes_obj;
@@ -2512,8 +2512,8 @@ static PyObject* get_value(PyObject* self, PyObject* name, const char* buffer,
      * Wrap any non-InvalidBSON errors in InvalidBSON.
      */
     if (PyErr_Occurred()) {
-        PyObject *etype, *evalue, *etrace;
-        PyObject *InvalidBSON;
+        PyObject *etype = NULL, *evalue = NULL, *etrace = NULL;
+        PyObject *InvalidBSON = NULL;
 
         /*
          * Calling _error clears the error state, so fetch it first.
@@ -2585,8 +2585,8 @@ static int _element_to_dict(PyObject* self, const char* string,
     if (!*name) {
         /* If NULL is returned then wrap the UnicodeDecodeError
            in an InvalidBSON error */
-        PyObject *etype, *evalue, *etrace;
-        PyObject *InvalidBSON;
+        PyObject *etype = NULL, *evalue = NULL, *etrace = NULL;
+        PyObject *InvalidBSON = NULL;
 
         PyErr_Fetch(&etype, &evalue, &etrace);
         if (PyErr_GivenExceptionMatches(etype, PyExc_Exception)) {
@@ -2620,7 +2620,7 @@ static PyObject* _cbson_element_to_dict(PyObject* self, PyObject* args) {
     /* TODO: Support buffer protocol */
     char* string;
     PyObject* bson;
-    PyObject* options_obj;
+    PyObject* options_obj = NULL;
     codec_options_t options;
     unsigned position;
     unsigned max;
@@ -2732,7 +2732,7 @@ static PyObject* _cbson_bson_to_dict(PyObject* self, PyObject* args) {
     int32_t size;
     Py_ssize_t total_size;
     const char* string;
-    PyObject* bson;
+    PyObject* bson = NULL;
     codec_options_t options;
     PyObject* result = NULL;
     PyObject* options_obj;

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, Sequence, Tuple, Type, Union
 from uuid import UUID
 
 """Tools for representing BSON binary data.
@@ -191,21 +194,75 @@ class UuidRepresentation:
 """
 
 
+VECTOR_SUBTYPE = 9
+"""**(BETA)** BSON binary subtype for densely packed vector data.
+
+.. versionadded:: 4.10
+"""
+
+
 USER_DEFINED_SUBTYPE = 128
 """BSON binary subtype for any user defined structure.
 """
 
 
+class BinaryVectorDtype(Enum):
+    """**(BETA)** 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.10
+    """
+
+    INT8 = b"\x03"
+    FLOAT32 = b"\x27"
+    PACKED_BIT = b"\x10"
+
+
+@dataclass
+class BinaryVector:
+    """**(BETA)** Vector of numbers along with metadata for binary interoperability.
+    .. versionadded:: 4.10
+    """
+
+    __slots__ = ("data", "dtype", "padding")
+
+    def __init__(self, data: Sequence[float | int], dtype: BinaryVectorDtype, padding: int = 0):
+        """
+        :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.
+        """
+        self.data = data
+        self.dtype = dtype
+        self.padding = padding
+
+
 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`.
+    **(BETA)** 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 +275,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.10
+       **(BETA)** Addition of vector subtype.
     """
 
     _type_marker = 5
@@ -337,6 +397,89 @@ 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: int = 0,
+    ) -> Binary:
+        """**(BETA)** 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.10
+        """
+        if dtype == BinaryVectorDtype.INT8:  # pack ints in [-128, 127] as signed int8
+            format_str = "b"
+            if padding:
+                raise ValueError(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"
+            if padding:
+                raise ValueError(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) -> BinaryVector:
+        """**(BETA)** From the Binary, create a list of numbers, along with dtype and padding.
+
+        :return: BinaryVector
+
+        .. versionadded:: 4.10
+        """
+
+        if self.subtype != VECTOR_SUBTYPE:
+            raise ValueError(f"Cannot decode subtype {self.subtype} as a vector.")
+
+        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
+            if n_bytes % 4:
+                raise ValueError(
+                    "Corrupt data. N bytes for a float32 vector must be a multiple of 4."
+                )
+            dtype_format = "f"
+            format_string = f"<{n_values}{dtype_format}"
+            vector = list(struct.unpack_from(format_string, self, position))
+            return BinaryVector(vector, dtype, padding)
+
+        elif dtype == BinaryVectorDtype.PACKED_BIT:
+            # data packed as uint8
+            dtype_format = "B"
+            format_string = f"<{n_values}{dtype_format}"
+            unpacked_uint8s = list(struct.unpack_from(format_string, self, position))
+            return BinaryVector(unpacked_uint8s, dtype, padding)
+
+        else:
+            raise NotImplementedError("Binary Vector dtype %s not yet supported" % dtype.name)
+
     @property
     def subtype(self) -> int:
         """Subtype of this binary data."""

bson/decimal128.py

@@ -223,7 +223,7 @@ def __init__(self, value: _VALUE_OPTIONS) -> None:
                     "from list or tuple. Must have exactly 2 "
                     "elements."
                 )
-            self.__high, self.__low = value  # type: ignore
+            self.__high, self.__low = value
         else:
             raise TypeError(f"Cannot convert {value!r} to Decimal128")
 

bson/json_util.py

@@ -324,7 +324,7 @@ def __new__(
                 "JSONOptions.datetime_representation must be one of LEGACY, "
                 "NUMBERLONG, or ISO8601 from DatetimeRepresentation."
             )
-        self = cast(JSONOptions, super().__new__(cls, *args, **kwargs))  # type:ignore[arg-type]
+        self = cast(JSONOptions, super().__new__(cls, *args, **kwargs))
         if json_mode not in (JSONMode.LEGACY, JSONMode.RELAXED, JSONMode.CANONICAL):
             raise ValueError(
                 "JSONOptions.json_mode must be one of LEGACY, RELAXED, "

bson/son.py

@@ -68,7 +68,7 @@ def __init__(
         self.update(kwargs)
 
     def __new__(cls: Type[SON[_Key, _Value]], *args: Any, **kwargs: Any) -> SON[_Key, _Value]:
-        instance = super().__new__(cls, *args, **kwargs)  # type: ignore[type-var]
+        instance = super().__new__(cls, *args, **kwargs)
         instance.__keys = []
         return instance
 

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:

doc/async-tutorial.rst

@@ -1,6 +1,11 @@
 Async Tutorial
 ==============
 
+.. warning:: This API is currently in beta, meaning the classes, methods,
+   and behaviors described within may change before the full release.
+   If you come across any bugs during your use of this API,
+   please file a Jira ticket in the "Python Driver" project at https://jira.mongodb.org/browse/PYTHON.
+
 .. code-block:: pycon
 
   from pymongo import AsyncMongoClient