feat: NumPy-accelerated vector serialization in VectorType (huge improvements in some use cases 60us -> 0.6us) - #793
Conversation
|
CC @swasik |
fba36f3 to
9e4bf4a
Compare
mykaul
changed the title
mykaul
changed the title
9e4bf4a to
5d891a8
Compare
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
|
Rebased onto current
Also rebuilt the Cython/numpy extensions from scratch and ran the full unit suite: CI: all real checks (Docs build, Integration tests across 3.11-3.14t, security/snyk) are green. "Test wheels building" shows No unresolved review threads / pending review requests at this time. Net result: no code changes were needed beyond the rebase itself (force-pushed with |
There was a problem hiding this comment.
Pull request overview
Note
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Adds optional NumPy-accelerated serialization paths for VectorType to significantly reduce per-vector overhead in bulk workloads, plus documentation and benchmarking to validate the improvements.
Changes:
- Add NumPy ndarray fast path and a bytes/bytearray passthrough path to
VectorType.serialize() - Add
VectorType.serialize_numpy_bulk()for batch conversion from 2-D NumPy arrays - Add unit tests, documentation updates, and a standalone benchmark script
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
cassandra/cqltypes.py |
Implements NumPy dtype caching, ndarray fast path, bytes passthrough, and serialize_numpy_bulk() |
tests/unit/test_types.py |
Adds NumPy-focused unit tests for correctness, error handling, and bulk serialization |
docs/performance.rst |
Documents the new fast paths and provides usage guidance |
benchmarks/bench_vector_numpy_serialize.py |
Adds a reproducible benchmark for list vs NumPy vs bulk vs passthrough serialization |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Add three fast paths to VectorType.serialize() for the common case of fixed-size numeric subtypes (float, double, int, bigint): 1. bytes/bytearray passthrough – skip all conversion when the caller already holds a correctly-sized blob (e.g. from serialize_numpy_bulk). 2. NumPy ndarray fast path – convert a 1-D numpy array to big-endian bytes via asarray(dtype=...).tobytes() instead of 768+ individual struct.pack + BytesIO.write calls. 3. serialize_numpy_bulk() classmethod – byte-swap an entire 2-D array (N rows × dim columns) once and slice the raw buffer, yielding one bytes object per row with zero per-element overhead. Benchmarks on 768-dim float32 vectors show 70-300× speedups depending on the path, directly benefiting bulk-insert workloads such as loading embeddings from Parquet files (VectorDBBench use case). NumPy remains an optional dependency; all new code is guarded by try/except ImportError. Variable-size subtypes (smallint, tinyint, text, etc.) are excluded and continue to use the original element-by-element path. Unit tests cover correctness, round-trip fidelity, error handling, and fallback behavior for all three paths.
Standalone benchmark comparing the four serialization paths for VectorType across dimensions (128, 768, 1536) and batch sizes (1, 100, 10000): - list (element-by-element) – baseline - numpy (per-row ndarray) – single-row fast path - bulk (serialize_numpy_bulk) – batch fast path - bytes passthrough (bind) – pre-serialized blob Includes auto-calibrated iteration counts and correctness verification.
Add a new section to docs/performance.rst covering the three fast serialization paths (single-row ndarray, bulk serialize_numpy_bulk, bytes passthrough) with usage examples and supported subtype list.
5d891a8 to
1d24619
Compare
There was a problem hiding this comment.
🟡 Not ready to approve
Variable-width byte input does not meet the stated safety contract, and the documented example is not runnable as written.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Review details
Comments suppressed due to low confidence (5)
cassandra/cqltypes.py:1581
- The public method's
Raisessection omits two implemented failures: aTypeErrorfor unsafe dtype conversion and aValueErrorwhen the input is not a 2-D ndarray. Documenting only column-count validation leaves callers with an incomplete API contract.
TypeError
If NumPy is not available or ``cls._numpy_dtype`` is ``None``
(subtype not supported for the fast path).
ValueError
If the second dimension of *vectors* does not match
cassandra/cqltypes.py:1526
- Variable-width vectors still accept or mishandle raw bytes because this condition simply falls through when
expected is None. For example, avector<tinyint, 4>receives a four-byte value as four elements, while other byte lengths raiseValueError; neither behavior provides the explicitTypeErrorpromised by the PR's safety contract. Reject bytes before the element-by-element path when no validated serialized length exists.
This issue also appears on line 1577 of the same file.
if expected is not None and isinstance(v, (bytes, bytearray)):
docs/performance.rst:68
- This example raises
NameErrorbecause it callslookup_casstype()below without importing it. Import the helper alongsideVectorTypeso the documented bulk workflow is runnable.
from cassandra.cqltypes import VectorType
docs/performance.rst:90
smallintandtinyintare fixed-width CQL numeric types, so describing them as variable-length is inaccurate. They fall back here because they are not in the new NumPy dtype map (and currently do not exposeserial_size()), not because their values have variable width.
**Supported subtypes** – the fast paths are available for ``float``
(``>f4``), ``double`` (``>f8``), ``int`` (``>i4``), and ``bigint``
(``>i8``). Variable-length subtypes (``smallint``, ``tinyint``, text
types, etc.) fall back to the original element-by-element serialization
automatically.
benchmarks/bench_vector_numpy_serialize.py:24
- The implementation no longer slices one whole-array
bytesbuffer; it callstobytes()separately for each row. Update this benchmark description so it measures and documents the implementation actually in this PR.
3. bulk path – serialize_numpy_bulk() on a 2-D array (one byte-swap for
the entire batch, then bytes slicing)
- Files reviewed: 4/4 changed files
- Comments generated: 0 new
- Review effort level: Balanced
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
Summary
VectorTypefor fixed-size numeric subtypes (float,double,int,bigint), delivering 70–300× speedups for high-dimensional vector insertsserialize_numpy_bulk()classmethod for batch workloads (e.g. loading embeddings from Parquet files in VectorDBBench)try/except ImportErrorMotivation
The VectorDBBench use case (https://github.com/scylladb/VectorDBBench) pipelines data as
Parquet → PyArrow/NumPy → wire. The existingVectorType.serialize()performed 768+ individualstruct.pack('>f', val)+BytesIO.write()calls per vector, making serialization the bottleneck for bulk inserts.What's new
Three fast paths in
VectorType.serialize():serialize_numpy_bulk()), return it directly with zero conversionnp.asarray(v, dtype='>f4').tobytes()instead of per-element struct.packserialize_numpy_bulk()classmethod – byte-swap an entire 2-D array(N_rows × dim)once and slice the raw buffer into alist[bytes]Benchmark results
768-dim float32, 100-row batches (Python 3.14, NumPy 2.3, best of 3):
In absolute terms: serializing 100 × 768-dim float32 vectors drops from ~6.2 ms (baseline) to ~60 µs (bulk+passthrough end-to-end) — a 103× improvement. The bottleneck was 76,800 individual
struct.pack('>f', val)+BytesIO.write()calls (768 elements × 100 rows); NumPy replaces all of that with a single array byte-swap + buffer slice.768-dim float32, single vector (latency per call):
Per-vector latency: a single 768-dim vector goes from ~65 µs to under 1 µs via the NumPy path, or 237 ns if pre-serialized bytes are reused.
Safety
smallint,tinyint, text) are excluded from the fast path and continue to use the original element-by-element serializationbytesfor a variable-size subtype now raisesTypeErrorexplicitly (rather than silently falling through)_numpy_dtypeis only set whensubtype.serial_size() is not Noneand the typename is in the 4-entry_NUMPY_DTYPE_MAPCommits
docs/performance.rst