(improvement) Optimize VectorType deserialization with struct.unpack and numpy (us level improvements - 2-13x speedup - Python path only!)

cassandra/cqltypes.py

@@ -51,6 +51,25 @@
                                vints_pack, vints_unpack, uvint_unpack, uvint_pack)
 from cassandra import util
 
+# NOTE: deliberately not importing this from cassandra.cython_deps. That
+# module imports cassandra.row_parser (the Cython row parser), which in turn
+# imports cassandra.deserializers, which imports this module (cqltypes) back.
+# If cqltypes is what pulls in cassandra.cython_deps, and cassandra.cython_deps
+# (or anything upstream of it, like cassandra.row_parser) happens to be the
+# first "cassandra.*" submodule imported in the process, that cycle closes
+# on a partially-initialized cassandra.cython_deps module: its own
+# `from cassandra.row_parser import ...` is still executing, so it hasn't
+# even set HAVE_CYTHON/HAVE_NUMPY yet, and this module's plain `import`
+# statement would raise an uncaught ImportError -- which cython_deps'
+# try/except then swallows, permanently (mis)recording HAVE_CYTHON as False
+# for the rest of the process, even though Cython is actually available.
+# Doing our own independent numpy probe here avoids creating that cycle.
+try:
+    import numpy as np
+    HAVE_NUMPY = True
+except ImportError:
+    HAVE_NUMPY = False
+
 _little_endian_flag = 1  # we always serialize LE
 import ipaddress
 
@@ -1432,47 +1451,106 @@ class VectorType(_CassandraType):
     typename = 'org.apache.cassandra.db.marshal.VectorType'
     vector_size = 0
     subtype = None
+    _vector_struct = None  # Cached struct.Struct for bulk deserialization
+    _struct_format_map = {}  # Populated after FloatType etc. are defined
+    _numpy_dtype = None  # Cached numpy dtype string for large vector deserialization
+    _subtype_serial_size = None  # Cached subtype.serial_size() (computed once in apply_parameters)
+    _serial_size = None  # Cached serial_size() for the full vector (subtype_serial_size * vector_size)
 
     @classmethod
     def serial_size(cls):
-        serialized_size = cls.subtype.serial_size()
-        return cls.vector_size * serialized_size if serialized_size is not None else None
+        return cls._serial_size
+
 
     @classmethod
     def apply_parameters(cls, params, names):
         assert len(params) == 2
         subtype = lookup_casstype(params[0])
         vsize = params[1]
-        return type('%s(%s)' % (cls.cass_parameterized_type_with([]), vsize), (cls,), {'vector_size': vsize, 'subtype': subtype})
+        # Cache subtype serial_size and full vector serial_size to avoid
+        # repeated method dispatch in serialize/deserialize hot paths.
+        subtype_ss = subtype.serial_size()
+        vec_ss = vsize * subtype_ss if subtype_ss is not None else None
+        # Cache a struct.Struct for bulk deserialization of known numeric types.
+        # Only ever consider this for subtypes that genuinely have a fixed
+        # serialized size (subtype_ss is not None): the wire format for any
+        # type without a fixed size is vint-length-prefixed per element (e.g.
+        # smallint/tinyint on Cassandra 5.0, see ShortType/ByteType above), and
+        # treating it as fixed-width would misparse real vector data.
+        vector_struct = None
+        numpy_dtype = None
+        if subtype_ss is not None:
+            for base_type, fmt_char in cls._struct_format_map.items():
+                if subtype is base_type or (isinstance(subtype, type) and issubclass(subtype, base_type)):
+                    vector_struct = struct.Struct(f'>{vsize}{fmt_char}')
+                    numpy_dtype = cls._numpy_dtype_map.get(fmt_char)
+                    break
+        return type('%s(%s)' % (cls.cass_parameterized_type_with([]), vsize), (cls,),
+                     {'vector_size': vsize, 'subtype': subtype, '_vector_struct': vector_struct,
+                      '_numpy_dtype': numpy_dtype, '_subtype_serial_size': subtype_ss, '_serial_size': vec_ss})
 
     @classmethod
     def deserialize(cls, byts, protocol_version):
-        serialized_size = cls.subtype.serial_size()
+        serialized_size = cls._subtype_serial_size
         if serialized_size is not None:
             expected_byte_size = serialized_size * cls.vector_size
             if len(byts) != expected_byte_size:
                 raise ValueError(
                     "Expected vector of type {0} and dimension {1} to have serialized size {2}; observed serialized size of {3} instead"\
                     .format(cls.subtype.typename, cls.vector_size, expected_byte_size, len(byts)))
-            indexes = (serialized_size * x for x in range(0, cls.vector_size))
-            return [cls.subtype.deserialize(byts[idx:idx + serialized_size], protocol_version) for idx in indexes]
 
+            # Optimization: bulk deserialization for common numeric types
+            # For small vectors: use cached struct.Struct (avoids per-call format string allocation)
+            # For large vectors with numpy: use numpy.frombuffer (1.3-1.5x faster for 128+ elements)
+            # Threshold at 32 elements balances simplicity with performance
+            if cls._vector_struct is not None:
+                if HAVE_NUMPY and cls.vector_size >= 32 and cls._numpy_dtype is not None:
+                    return np.frombuffer(byts, dtype=cls._numpy_dtype, count=cls.vector_size).tolist()
+                return list(cls._vector_struct.unpack(byts))
+            # Fallback: element-by-element deserialization for other fixed-size types
+            result = [None] * cls.vector_size
+            subtype_deserialize = cls.subtype.deserialize
+            offset = 0
+            for i in range(cls.vector_size):
+                result[i] = subtype_deserialize(byts[offset:offset + serialized_size], protocol_version)
+                offset += serialized_size
+            return result
+
+        # Variable-size subtype path
+        result = [None] * cls.vector_size
         idx = 0
-        rv = []
-        while (len(rv) < cls.vector_size):
+        byts_len = len(byts)
+        subtype_deserialize = cls.subtype.deserialize
+
+        for i in range(cls.vector_size):
+            if idx >= byts_len:
+                raise ValueError("Error reading additional data during vector deserialization after successfully adding {} elements"
+                    .format(i))
+
             try:
                 size, bytes_read = uvint_unpack(byts[idx:])
-                idx += bytes_read
-                rv.append(cls.subtype.deserialize(byts[idx:idx + size], protocol_version))
-                idx += size
-            except:
-                raise ValueError("Error reading additional data during vector deserialization after successfully adding {} elements"\
-                .format(len(rv)))
-
-        # If we have any additional data in the serialized vector treat that as an error as well
-        if idx < len(byts):
+            except IndexError:
+                raise ValueError("Error reading additional data during vector deserialization after successfully adding {} elements"
+                    .format(i))
+
+            idx += bytes_read
+
+            if idx + size > byts_len:
+                raise ValueError("Error reading additional data during vector deserialization after successfully adding {} elements"
+                    .format(i))
+
+            try:
+                result[i] = subtype_deserialize(byts[idx:idx + size], protocol_version)
+            except Exception as e:
+                raise ValueError("Error deserializing element {} during vector deserialization after successfully adding {} elements"
+                    .format(i, i)) from e
+            idx += size
+
+        # Check for additional data
+        if idx < byts_len:
             raise ValueError("Additional bytes remaining after vector deserialization completed")
-        return rv
+
+        return result
 
     @classmethod
     def serialize(cls, v, protocol_version):
@@ -1482,7 +1560,10 @@ def serialize(cls, v, protocol_version):
                 "Expected sequence of size {0} for vector of type {1} and dimension {0}, observed sequence of length {2}"\
                 .format(cls.vector_size, cls.subtype.typename, v_length))
 
-        serialized_size = cls.subtype.serial_size()
+        serialized_size = cls._subtype_serial_size
+        # Bulk serialization for known numeric types (symmetric with struct.unpack in deserialize)
+        if cls._vector_struct is not None and serialized_size is not None:
+            return cls._vector_struct.pack(*v)
         buf = io.BytesIO()
         for item in v:
             item_bytes = cls.subtype.serialize(item, protocol_version)
@@ -1494,3 +1575,25 @@ def serialize(cls, v, protocol_version):
     @classmethod
     def cql_parameterized_type(cls):
         return "%s<%s, %s>" % (cls.typename, cls.subtype.cql_parameterized_type(), cls.vector_size)
+
+
+# Populate VectorType._struct_format_map now that all types are defined.
+# NOTE: ShortType (smallint) and ByteType (tinyint) are intentionally excluded.
+# Although they have a fixed in-memory representation, Cassandra 5.0 does not
+# treat them as fixed-width for vector serialization (AbstractType.valueLengthIfFixed()
+# defaults to variable-length, and neither ShortType.java nor ByteType.java override
+# it) so their vector elements are vint-length-prefixed on the wire, same as any
+# other variable-size type. Only include types here that have a real fixed
+# serial_size(); apply_parameters() also gates on subtype.serial_size() is not
+# None as a second line of defense.
+VectorType._struct_format_map = {
+    FloatType: 'f',
+    DoubleType: 'd',
+    Int32Type: 'i',
+    LongType: 'q',
+}
+
+# Map struct format chars to numpy dtype strings for large vector deserialization.
+# Kept in sync with _struct_format_map above -- no entry for 'h' (ShortType),
+# since smallint vectors are variable-length on the wire (see the NOTE above).
+VectorType._numpy_dtype_map = {'f': '>f4', 'd': '>f8', 'i': '>i4', 'q': '>i8'}

tests/unit/test_types.py

@@ -22,8 +22,8 @@
 import cassandra
 from cassandra import util
 from cassandra.cqltypes import (
-    CassandraType, DateRangeType, DateType, DecimalType,
-    EmptyValue, LongType, SetType, UTF8Type,
+    ByteType, CassandraType, DateRangeType, DateType, DecimalType,
+    EmptyValue, LongType, SetType, ShortType, UTF8Type,
     cql_typename, int8_pack, int64_pack, int64_unpack, lookup_casstype,
     lookup_casstype_simple, parse_casstype_args,
     int32_pack, Int32Type, ListType, MapType, VectorType,
@@ -400,6 +400,32 @@ def test_round_trip_basic_types_without_fixed_serialized_size(self):
         self._round_trip_test([util.Duration(1,1,1), util.Duration(2,2,2), util.Duration(3,3,3)], \
             "org.apache.cassandra.db.marshal.VectorType(org.apache.cassandra.db.marshal.DurationType, 3)")
 
+    def test_short_and_byte_vectors_use_variable_length_wire_format(self):
+        # Regression test: smallint (ShortType) and tinyint (ByteType) have a
+        # fixed in-memory representation, but real Cassandra 5.0 does NOT treat
+        # them as fixed-width for vector (de)serialization -- AbstractType's
+        # valueLengthIfFixed() defaults to variable-length, and neither
+        # ShortType.java nor ByteType.java override it. So, like TimeType
+        # above, their vector elements must be vint-length-prefixed on the
+        # wire, not packed as a fixed number of bytes per element. This guards
+        # against ShortType/ByteType ever being (re-)added to a fixed-width
+        # struct/numpy fast-path map for VectorType, which would silently
+        # produce a wire format a real server can't parse.
+        for subtype, packed_value_size in ((ShortType, 2), (ByteType, 1)):
+            ctype = parse_casstype_args(
+                "org.apache.cassandra.db.marshal.VectorType(org.apache.cassandra.db.marshal.%s, 4)"
+                % subtype.__name__)
+            assert ctype.subtype.serial_size() is None
+            assert ctype.serial_size() is None
+
+            data = [3, -2, 100, -100]
+            data_bytes = ctype.serialize(data, 0)
+            # 1 vint length byte + the packed value per element -- NOT a flat
+            # packed_value_size * len(data), which is what a fixed-width fast
+            # path would (incorrectly) produce.
+            assert len(data_bytes) == len(data) * (1 + packed_value_size)
+            assert ctype.deserialize(data_bytes, 0) == data
+
     def test_round_trip_collection_types(self):
         # List (subtype of fixed size)
         self._round_trip_test([[1, 2, 3, 4], [5, 6], [7, 8, 9, 10], [11, 12]], \