(improvement) Add VectorType support to numpy_parser for 2D array parsing - #731
(improvement) Add VectorType support to numpy_parser for 2D array parsing#731mykaul wants to merge 3 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds native VectorType handling to the Cython NumPy row parser so vector columns can be parsed into 2D NumPy (masked) arrays instead of falling back to object arrays, improving performance for vector/embedding workloads.
Changes:
- Extend
cassandra/numpy_parser.pyxto allocate 2D arrays forVectorTypecolumns and correctly advance/mark masks for 2D shapes. - Add unit tests in
tests/unit/test_numpy_parser.pycovering several numeric vector subtypes and mixed scalar+vector results.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 12 comments.
| File | Description |
|---|---|
| cassandra/numpy_parser.pyx | Adds VectorType 2D array allocation + 2D mask stride handling; updates NULL-mask write logic. |
| tests/unit/test_numpy_parser.py | Introduces unit tests for vector parsing into 2D NumPy arrays. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
6474dae to
e73ef19
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:
📝 WalkthroughWalkthroughThe NumPy parser now supports Sequence Diagram(s)sequenceDiagram
participant BytesIOReader
participant NumpyParser
participant make_array
participant NumPyArray
BytesIOReader->>NumpyParser: provide row payloads
NumpyParser->>make_array: allocate VectorType result
make_array->>NumPyArray: create 2D masked or object array
NumpyParser->>NumPyArray: validate and copy vector data
NumpyParser->>NumPyArray: apply NULL mask and zero data
Suggested labels: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Comment |
|
Cleaned this up: rebased onto current master (zero conflicts), and closed out the review feedback. What was already fixed (previous commit
What this pass added on top:
This is still a draft/perf-experiment PR, so flagging rather than blocking: the same unchecked-length-before-memcpy pattern that motivated the bounds check here also seems to recur in PR #689's Cython deserializers. Might be worth a shared bounds-checking convention/helper addressed once across both PRs rather than fixing it ad hoc per PR - leaving that decision to a human rather than filing it myself. |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (3)
cassandra/numpy_parser.pyx:146
- The new 2-D vector buffer is later passed to
make_native_byteorder(), wherearr.byteswap()uses its defaultinplace=False. On little-endian hosts this allocates and copies a second full vector result, doubling peak data-buffer memory for the large embedding workloads this path targets. Because this array is newly allocated and exclusively owned, byte-swapping it in place before changing the dtype would avoid that extra full-size copy.
a = np.ma.empty((array_size, vector_size), dtype=dtype)
tests/unit/test_numpy_parser.py:32
- This broad
ImportErrorhandler also catches failures importingcassandra.numpy_parseror its Cython dependencies and then marks NumPy unavailable, causing every new test to be skipped. A broken or missing parser extension can therefore make this suite pass without exercising the PR. Gate imports using the repository'sHAVE_CYTHON/HAVE_NUMPYflags, but let parser import failures surface when both dependencies are expected.
except ImportError:
HAVE_NUMPY = False
cassandra/numpy_parser.pyx:189
- A NULL numeric/vector value only masks the row; its backing bytes remain whatever
np.ma.empty()left in the allocation. The masked array exposes those bytes through its public.dataattribute, so each NULL vector can return a full row of stale process-heap data—the same information leak the undersized-payload test is intended to prevent. Zero the row before setting its mask.
else:
memset(<char *>arr.mask_ptr, 1, arr.mask_stride)
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
cassandra/numpy_parser.pyx (1)
182-195: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGood fix — strict size equality closes the buffer-overflow/under-copy hole.
The
buf.size != arr.strideguard correctly rejects both oversized and undersized payloads beforememcpy, preventing heap overflow (oversized) and stale/uninitialized data (undersized). Themask_stride-based pointer advance also correctly generalizes the previous fixed-increment logic to support 2D vector masks.One small improvement: the error message only includes the column index, not its name, making it harder to diagnose which column failed for wide result sets.
💡 Include column name in the error message
if buf.size != arr.stride: raise ValueError( - "Column %d: received %d bytes but array stride is %d " + "Column %d (%s): received %d bytes but array stride is %d " "(payload must exactly match the expected element size)" % - (i, buf.size, arr.stride)) + (i, desc.colnames[i], buf.size, arr.stride))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cassandra/numpy_parser.pyx` around lines 182 - 195, Update the ValueError raised in the buffer-size validation within the result parsing loop to include the failing column’s name alongside its index, while preserving the existing size details and strict equality check. Use the column-name symbol already available in the surrounding parsing context; do not alter the memcpy or pointer advancement logic.tests/unit/test_numpy_parser.py (1)
411-424: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider an end-to-end fallback test through
parse_rows().This test only checks
make_array()'s allocation shape/dtype for an unsupported subtype. It doesn't exercise the actualunpack_rowobject-path (from_binarydeserialization +Py_INCREFstorage) for a VectorType column with an unsupported subtype, which is the path that matters at runtime.🧪 Suggested additional test
def test_unsupported_subtype_full_parse_round_trip(self): """Test full parse_rows() round trip for an unsupported vector subtype""" vector_type = self._create_vector_type(cqltypes.UTF8Type, 2) # build a buffer using VectorType.serialize()/protocol wire format # and assert result["vec"] is a 1D object array of lists🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/test_numpy_parser.py` around lines 411 - 424, Add an end-to-end test alongside test_unsupported_subtype_falls_back_to_object_array that builds a VectorType UTF8 payload using its serialization/protocol wire format, parses it through parse_rows(), and verifies the resulting “vec” column is a 1D object array containing the expected deserialized lists. Exercise the unpack_row object path rather than testing make_array() directly.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@cassandra/numpy_parser.pyx`:
- Around line 182-195: Update the ValueError raised in the buffer-size
validation within the result parsing loop to include the failing column’s name
alongside its index, while preserving the existing size details and strict
equality check. Use the column-name symbol already available in the surrounding
parsing context; do not alter the memcpy or pointer advancement logic.
In `@tests/unit/test_numpy_parser.py`:
- Around line 411-424: Add an end-to-end test alongside
test_unsupported_subtype_falls_back_to_object_array that builds a VectorType
UTF8 payload using its serialization/protocol wire format, parses it through
parse_rows(), and verifies the resulting “vec” column is a 1D object array
containing the expected deserialized lists. Exercise the unpack_row object path
rather than testing make_array() directly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: ee4bb2b7-c632-4736-a43e-0846f115f04c
📒 Files selected for processing (2)
cassandra/numpy_parser.pyxtests/unit/test_numpy_parser.py
Extend NumpyParser to handle VectorType columns by creating 2D NumPy arrays (rows × vector_dimension) instead of object arrays. This enables zero-copy parsing for vector embeddings in ML/AI workloads. Features: - Detects VectorType via vector_size and subtype attributes - Creates 2D masked arrays for numeric vector subtypes with a fixed per-element wire width (float, double, int32, bigint) - Falls back to object arrays for subtypes without a fixed wire width, including smallint (ShortType does not override serial_size(), so Cassandra/Scylla vint-prefixes each element inside a vector instead of using a fixed 2-byte field -- mapping it to the fast path would crash on a stride mismatch instead of safely falling back) - Handles endianness conversion for both 1D and 2D arrays, byteswapping the freshly-allocated (exclusively owned) array in place instead of taking a full extra copy - Zeroes the data bytes of masked/NULL rows so no uninitialized heap memory is exposed for a NULL vector - Pre-allocates result arrays for efficiency Supported vector types (fast path): - Vector<float> → 2D float32 array - Vector<double> → 2D float64 array - Vector<int> → 2D int32 array - Vector<bigint> → 2D int64 array Vector<smallint> falls back to a 1D object array (see above). Adds comprehensive test coverage for all supported vector types, mixed column queries, and large vector dimensions (384-element embeddings), including round-trip tests that serialize via the actual production VectorType.serialize() rather than hand-rolled wire bytes. Signed-off-by: Yaniv Kaul <yaniv.kaul@scylladb.com>
- Add buffer size guard before memcpy in unpack_row() to prevent overflow - Remove dead mask_true constant and unused uint8_t cimport - Fix copyright header (DataStax -> ScyllaDB) - Replace hard-coded little-endian dtypes with native numpy dtypes - Remove unused Mock import - Add test for NULL vector mask handling - Add test for unsupported subtype fallback to object array
e73ef19 to
83a5989
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (2)
cassandra/numpy_parser.pyx:75
- Removing
ShortTypefrom this shared map also changes ordinarysmallintresult columns:make_array()below now takes itsKeyErrorfallback and returns a 1-D object array instead of the existing maskedint16array. KeepShortTypein the scalar mapping and exclude it only when selecting theVectorTypefast path.
# Note: ShortType (smallint) is intentionally absent from this table. Unlike
# LongType/Int32Type/FloatType/DoubleType/CounterColumnType, ShortType does
# not override serial_size() in cassandra.cqltypes, so it has no fixed
cassandra/numpy_parser.pyx:78
- The PR description promises a 2-D
int16fast path for supported numeric vectors, but this exclusion makesVectorType<smallint>return a 1-D object array, and the new test explicitly locks in that fallback. Either implement the advertisedint16path (including the correct wire-size contract) or update the PR description to remove that claim.
# Note: ShortType (smallint) is intentionally absent from this table. Unlike
# LongType/Int32Type/FloatType/DoubleType/CounterColumnType, ShortType does
# not override serial_size() in cassandra.cqltypes, so it has no fixed
# per-element width when serialized *inside* a VectorType: the server
# vint-prefixes each smallint element rather than encoding it as a plain
# 2-byte big-endian field. Mapping it here would make the vector fast-path
coderabbitai
Bot
added
area/Driver_-_python-driver
enhancement
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/unit/test_numpy_parser.py (2)
316-357: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider covering mask-stride advancement with a NULL in both columns.
This mixed-column test only uses non-NULL values, so the separate
mask_ptr/mask_strideadvancement inunpack_row(cassandra/numpy_parser.pyxlines 210-212) isn't actually exercised — a scalar column hasmask_stride == 1while the vector column hasmask_stride == vector_size. A row where bothidandvecare NULL (size = -1) would catch stride desynchronization between the two columns.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/test_numpy_parser.py` around lines 316 - 357, Extend test_mixed_columns_with_vectors to include a row where both id and vec are NULL, encoding each column with size = -1, and update the expected outputs to verify both columns preserve the NULL values while retaining correct shapes and alignment. Ensure the assertions exercise mask-stride advancement between the scalar id column and vector vec column.
96-108: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: extract the parser/ParseDesc boilerplate into a helper.
This 10-line block is repeated near-verbatim in ~10 tests, varying only in colnames/coltypes/deserializers.
♻️ Suggested helper
def _parse(self, serialized, colnames, coltypes, deserializers=None): desc = ParseDesc( colnames=colnames, coltypes=coltypes, column_encryption_policy=None, coldescs=None, deserializers=deserializers or obj_array([None] * len(colnames)), protocol_version=5, ) return NumpyParser().parse_rows(BytesIOReader(serialized), desc)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/test_numpy_parser.py` around lines 96 - 108, Extract the repeated NumpyParser and ParseDesc setup into a test helper, such as _parse, in the test class or module. Have it accept serialized data, column names, column types, and optional deserializers, supply the existing default deserializers when omitted, and return parse_rows’ result; update the affected tests to use this helper.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cassandra/numpy_parser.pyx`:
- Around line 73-83: Restore ShortType handling in the vector decoding flow
without adding it to the fixed-width fast-path table. Add a dedicated
per-element decoder for VectorType<smallint, N> that follows its length-prefixed
wire representation and produces the required 2D int16 masked array, while
preserving the existing fallback behavior for other variable-width subtypes and
avoiding the unpack_row stride assumption.
---
Nitpick comments:
In `@tests/unit/test_numpy_parser.py`:
- Around line 316-357: Extend test_mixed_columns_with_vectors to include a row
where both id and vec are NULL, encoding each column with size = -1, and update
the expected outputs to verify both columns preserve the NULL values while
retaining correct shapes and alignment. Ensure the assertions exercise
mask-stride advancement between the scalar id column and vector vec column.
- Around line 96-108: Extract the repeated NumpyParser and ParseDesc setup into
a test helper, such as _parse, in the test class or module. Have it accept
serialized data, column names, column types, and optional deserializers, supply
the existing default deserializers when omitted, and return parse_rows’ result;
update the affected tests to use this helper.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: f479ba91-f46a-4bf6-b58d-446177f1a701
📒 Files selected for processing (2)
cassandra/numpy_parser.pyxtests/unit/test_numpy_parser.py
The memcpy fast-path guard in unpack_row() only rejected oversized payloads (buf.size > arr.stride), which prevents the buffer-overflow case but let undersized payloads slip through, leaving the remainder of the row filled with uninitialized/stale heap bytes instead of raising. Tighten the check to buf.size != arr.stride so both directions are rejected, and update the error message accordingly, including the column name for easier debugging. Add tests covering both the oversized and undersized payload cases, since neither had test coverage before. Also revert the arrays = [] / .append() -> [None]*n / index-assign change in make_arrays(), which was functionally identical diff noise unrelated to the VectorType/2D-array feature. Restore ShortType (smallint) to _cqltype_to_numpy: an earlier fix on this PR removed it from that table entirely to stop the VectorType fast-path from assuming a fixed 2-byte stride for vint-prefixed vector elements, but the table is shared with the plain scalar top-level column dispatch in make_array(), where smallint genuinely *is* fixed-width (no length prefix). Removing the entry regressed scalar smallint columns onto the slow object-array fallback. Fix: keep ShortType in the shared table, and instead guard the VectorType.subtype call site with `subtype.serial_size() is not None` (the same check VectorType.serialize()/deserialize() use) before consulting the table, so only the vector-specific dispatch excludes non-fixed-width subtypes. This also protects any subtype added to the table in the future. Add tests for the restored scalar smallint fast path (both make_array() directly and an end-to-end parse_rows() round-trip). Also fix tests/unit/test_numpy_parser.py's HAVE_NUMPY gating: a bare `except ImportError` around the numpy_parser/bytesio/etc. imports could not distinguish "numpy genuinely not installed" from "the numpy_parser Cython extension is broken despite numpy being present", silently skipping the whole test module in both cases. Gate instead on cassandra.cython_deps.HAVE_CYTHON and HAVE_NUMPY (plus the VERIFY_CYTHON override), matching the `numpytest` convention in tests/unit/cython/utils.py, and import unconditionally once those are satisfied so a broken extension surfaces as a real test failure.
83a5989 to
8f16aba
Compare
|
Follow-up fix: a Copilot review comment on this PR was suppressed due to low confidence (never became a formal resolvable thread), but it correctly caught a regression in the previous commit here:
Confirmed directly in code: Fix pushed:
Note on CodeRabbit's separate suggestion to "restore the required int16 vector support" via a dedicated per-element decoder for Also fixed a separate, previously-open review concern in Verified: full |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
cassandra/numpy_parser.pyx:184
- The PR description promises a native
int162D array for supported numeric vector subtypes, but this guard excludesShortTypebecause itsserial_size()isNone, and the new test explicitly expects an object-array fallback. Either removeint16support from the PR details or implement the required framed-element conversion so the advertised behavior matches the code.
if subtype.serial_size() is not None:
Summary
(num_rows, vector_dimension)instead of falling back to object arraysDetails
Changes to
cassandra/numpy_parser.pyx:make_array()detects VectorType columns and creates 2Dnp.ma.empty((array_size, vector_size), dtype=...)arrays with the correct numeric dtype (float32, float64, int32, int64, int16)ArrDescextended withmask_stridefield to handle 2D mask arrays where stride =vector_dimensionbools (not 1 bool)unpack_row()uses directmemcpyof the full vector payload (e.g., 3072 bytes for float[768]) into the pre-allocated 2D array buffermake_native_byteorder()handles bulk byte-swap on any-dimensional arrays transparentlyResult: For a query returning N rows of
Vector<float, 768>, the NumpyParser produces an(N, 768)float32 array directly from wire bytes — the fastest possible path when consuming results as numpy arrays.Tests: Comprehensive unit tests in
tests/unit/test_numpy_parser.pycovering all supported numeric subtypes, NULL handling, mask strides, and unsupported type fallback.This commit is fully independent — it only modifies
numpy_parser.pyxand adds a new test file.