(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. |
…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>
75b9b75 to
c9219e2
Compare
|
Warning Review limit reached
Next review available in: 39 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
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
| 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. |
| serialized_size = cls._subtype_serial_size | ||
| if serialized_size is not None: | ||
| expected_byte_size = serialized_size * cls.vector_size |
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.