(improvement) Optimize VectorType deserialization with struct.unpack and numpy (us level improvements - 2-13x speedup - Python path only!) - #730
Conversation
There was a problem hiding this comment.
Pull request overview
This PR optimizes VectorType (de)serialization in cassandra/cqltypes.py by introducing bulk numeric (de)serialization via a cached struct.Struct, and an optional numpy-based deserialization fast path for larger vectors.
Changes:
- Cache a per-parameterized-vector
struct.Structto bulkunpack/packcommon numeric vector subtypes. - Add an optional numpy
frombuffer(...).tolist()deserialization fast-path for vectors withvector_size >= 32. - Refactor variable-size vector deserialization to a fixed-iteration loop with stricter bounds checks.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
c417e73 to
0535ecd
Compare
mykaul
changed the title
mykaul
changed the title
|
@mykaul This is not a draft, but review was not requested. Please either change to draft, or request review. |
It's an improvement, not a fix. I believe it's ready, but I don't want to disrupt the team. I'm not sure what to do (and I do it for fun anyway). If there's anything that I see as important - I'm not shy. |
75b9b75 to
c9219e2
Compare
|
Warning Review limit reachedNext included review available in 49 minutes. View limit detailsLimit details: You’ve used the included review currently available. You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Review configuration: ⚙️ Run configurationConfiguration used: Organization UI Review profile: QUIET Plan: Advanced Run ID: 📒 Files selected for processing (2)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Rebased onto current ShortType/ByteType fixed-width bug: present, now fixed
In practice this PR's own gating (
Second bug found: real The new module-level Other checks
All fixes were amended into the original 3 commits (not new commits) and force-pushed. |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (3)
cassandra/cqltypes.py:1456
apply_parameters()usescls._numpy_dtype_map.get(...), but the class definition doesn’t define_numpy_dtype_map(it’s only assigned later at module scope). To avoid potentialAttributeErrorduring unusual import ordering / partial initialization (and to make the class’ expected attributes explicit), initialize_numpy_dtype_map = {}alongside_struct_format_mapin the class body.
_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
tests/unit/test_types.py:419
- This test appears to live in a
unittest.TestCase-style suite, but it uses bareassertstatements. Bare asserts can be skipped with Python optimizations (-O) and typically produce less helpful failure output compared toself.assert*methods. Recommend switching toself.assertIsNone(...)/self.assertEqual(...)for consistency and clearer diagnostics.
assert ctype.subtype.serial_size() is None
assert ctype.serial_size() is None
tests/unit/test_types.py:427
- This test appears to live in a
unittest.TestCase-style suite, but it uses bareassertstatements. Bare asserts can be skipped with Python optimizations (-O) and typically produce less helpful failure output compared toself.assert*methods. Recommend switching toself.assertIsNone(...)/self.assertEqual(...)for consistency and clearer diagnostics.
assert len(data_bytes) == len(data) * (1 + packed_value_size)
assert ctype.deserialize(data_bytes, 0) == data
…ct.unpack
Add bulk deserialization using struct.unpack for common numeric vector types
instead of element-by-element deserialization. This provides significant
performance improvements, especially for small vectors and integer types.
Optimized types:
- FloatType ('>Nf' format)
- DoubleType ('>Nd' format)
- Int32Type ('>Ni' format)
- LongType ('>Nq' format)
ShortType (smallint) and ByteType (tinyint) are intentionally NOT included,
even though they have a fixed in-memory representation: real 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 like any other variable-size type. Treating
them as fixed-width here would produce a wire format a real server can't
parse.
Performance improvements (measured with CASS_DRIVER_NO_CYTHON=1):
Small vectors (3-4 elements):
Vector<float, 3> : 0.88 μs → 0.25 μs (3.58x faster)
Vector<float, 4> : 0.78 μs → 0.28 μs (2.79x faster)
Medium vectors (128 elements):
Vector<float, 128> : 4.72 μs → 4.06 μs (1.16x faster)
Vector<double, 128> : 4.83 μs → 4.01 μs (1.20x faster)
Vector<int, 128> : 2.27 μs → 1.25 μs (1.82x faster)
Large vectors (384-1536 elements):
Vector<float, 384> : 15.38 μs → 14.67 μs (1.05x faster)
Vector<float, 768> : 32.43 μs → 30.72 μs (1.06x faster)
Vector<float, 1536> : 63.74 μs → 63.24 μs (1.01x faster)
The optimization is most effective for:
- Small vectors (3-4 elements): 2.8-3.6x speedup
- Integer vectors: 1.8x speedup
- Medium-sized float/double vectors: 1.2-1.3x speedup
For very large vectors (384+ elements), the benefit is minimal as the
deserialization time is dominated by data copying rather than function
call overhead.
Variable-size subtypes and other numeric types continue to use the
element-by-element fallback path.
Signed-off-by: Yaniv Kaul <yaniv.kaul@scylladb.com>
For vectors with 32 or more elements, use numpy.frombuffer() which provides
1.3-1.5x speedup for large vectors (128+ elements) compared to struct.unpack.
The hybrid approach:
- Small vectors (< 32 elements): struct.unpack (2.8-3.6x faster than baseline)
- Large vectors (>= 32 elements): numpy.frombuffer().tolist() (1.3-1.5x faster than struct.unpack)
Threshold of 32 elements balances code complexity with performance gains.
_numpy_dtype_map has no entry for ShortType ('h'), matching
_struct_format_map: smallint vectors are variable-length on the wire on
real Cassandra 5.0, so they never take this fast path.
Probe for numpy directly (try/except import) instead of importing HAVE_NUMPY
from cassandra.cython_deps. cassandra.cython_deps imports cassandra.row_parser
(Cython row parser), which imports cassandra.deserializers, which imports
this module (cqltypes) back. If cqltypes is what first pulls in
cassandra.cython_deps, and cassandra.cython_deps (or 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
that hasn't set HAVE_CYTHON/HAVE_NUMPY yet, causing an uncaught ImportError
that cython_deps' own try/except then swallows -- permanently (and
incorrectly) recording HAVE_CYTHON as False for the rest of the process
even when Cython is available. Verified end-to-end: e.g. `import
cassandra.cython_deps` (or `tests.unit.cython.utils`, which does the same
thing) as the first cassandra import in a process previously left
HAVE_CYTHON False; with this change it correctly reports True.
Benchmark results:
- float[128]: 2.15 μs → 1.87 μs (1.15x faster)
- float[384]: 6.17 μs → 4.44 μs (1.39x faster)
- float[768]: 12.25 μs → 8.45 μs (1.45x faster)
- float[1536]: 24.44 μs → 15.77 μs (1.55x faster)
Signed-off-by: Yaniv Kaul <yaniv.kaul@scylladb.com>
…ated method dispatch Cache subtype.serial_size() and the full vector serial_size() as class attributes (_subtype_serial_size, _serial_size) during apply_parameters(). This eliminates per-call method dispatch overhead in serialize(), deserialize(), and serial_size() hot paths. serial_size() call: 99ns -> 46ns (2.2x faster) Attribute access: 54ns -> 17ns (3.2x faster) While here, compute subtype_ss before building the struct/numpy fast-path cache and gate that cache on `subtype_ss is not None` as a second line of defense: only subtypes with a genuine fixed serial_size() (FloatType, DoubleType, Int32Type, LongType) may populate _vector_struct/_numpy_dtype. This guards against ShortType/ByteType (or any future variable-length type mistakenly added to _struct_format_map) ever taking the fixed-width fast path -- real Cassandra 5.0 vint-length-prefixes smallint/tinyint vector elements, so treating them as fixed-width would misparse real vector data. Add a regression test (test_short_and_byte_vectors_use_variable_length_wire_format) that asserts ShortType/ByteType vectors serialize using vint-length-prefixed elements, not a flat fixed-width encoding. Signed-off-by: Yaniv Kaul <yaniv.kaul@scylladb.com>
c9219e2 to
134c00e
Compare
|
Re-verified the two open Copilot threads about This is confirmed by No code change needed; resolving these threads as already addressed by the existing design + regression test. |
Summary
struct.unpackfor known numeric types (float, double, int32, int64, short), caching astruct.Structobject at type-creation timenp.frombuffer().tolist()) for vectors with >= 32 elementsserial_size()results to eliminate per-call method dispatch overheadKeyErrorcatch, wrapsubtype.deserializefailures with element context and proper exception chainingPerformance (pure Python, best of 5)
Deserialization:
Vector<float, 4>Vector<float, 16>Vector<float, 128>Vector<float, 768>Vector<float, 1536>Serialization:
Vector<float, 4>Vector<float, 16>Vector<float, 128>Vector<float, 768>Vector<float, 1536>serial_size() overhead:
serial_size()call (768-dim)Details
Commit 1 -- struct.unpack optimization + variable-size path fixes:
apply_parameters()time, cache astruct.Struct('>Nf')for the vector's subtype+dimensiondeserialize()callslist(struct.unpack(byts))-- single C-level bulk unpackstruct.pack(*v)KeyErrorfrom except clause (uvint_unpackonly raisesIndexError), wrapsubtype.deserializefailures inValueErrorwith element index and proper exception chaining (from e)Commit 2 -- numpy for large vectors:
np.frombuffer(byts, dtype='>f4', count=N).tolist().tolist()batch-converts with better cache locality_numpy_dtypecached on the class at type-creation time (no per-call dict construction)Commit 3 -- serial_size caching:
subtype.serial_size()result as_subtype_serial_sizeand the full vector serial size as_serial_sizeduringapply_parameters()serial_size()returns cached value directly (no method dispatch chain)serialize()anddeserialize()usecls._subtype_serial_sizeinstead of callingcls.subtype.serial_size()each timeAll three commits modify only
cassandra/cqltypes.py. No Cython dependency.