diff --git a/cassandra/cqltypes.py b/cassandra/cqltypes.py index 99018eef03..8cf013a4aa 100644 --- a/cassandra/cqltypes.py +++ b/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'} diff --git a/tests/unit/test_types.py b/tests/unit/test_types.py index 11aab2748d..1f869c99dd 100644 --- a/tests/unit/test_types.py +++ b/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]], \