perf: buffer accumulation in BatchMessage.send_body() (1.6-1.8x speedup, us improvement, depends on PR #790) - #791
Conversation
5eec7ed to
b1d1cd0
Compare
b1d1cd0 to
62f91eb
Compare
Replace the per-parameter write_value(f, param) loop in _QueryMessage._write_query_params() with a buffer accumulation approach: list.append + b"".join + single f.write(). This reduces the number of f.write() calls from 2*N+1 to 1, which is significant for vector workloads with large parameters. Also removes the redundant ExecuteMessage._write_query_params() pass-through override to avoid extra MRO lookup per call. Includes 14 unit tests covering normal, NULL, UNSET, empty, large vector, and mixed parameter scenarios for both ExecuteMessage and QueryMessage. Includes a benchmark script (benchmarks/bench_execute_write_params.py).
Replace per-write_value()/write_byte()/write_short() calls in BatchMessage.send_body() with buffer accumulation (list.append + b"".join + single f.write()), reducing f.write() calls from Q*(4 + 2*P) + footer to 1 for Q queries with P params each. Benchmark results (Python 3.14, Cython .so, 50K iters, best of 3, quiet machine): Scenario Before After Speedup 10 queries x 2 params (128D vec) 8364 ns 4475 ns 1.87x 10 queries x 2 params (768D vec) 8081 ns 5516 ns 1.47x 50 queries x 2 params (128D vec) 32368 ns 16271 ns 1.99x 10 queries x 10 text params 19138 ns 9051 ns 2.11x 50 queries x 10 text params 86845 ns 40020 ns 2.17x 10 unprepared x 2 params 8666 ns 4252 ns 2.04x Also updates test_batch_message_with_keyspace to use BytesIO for byte-level verification (compatible with single-write output). Adds 7 batch-specific unit tests covering prepared, unprepared, mixed, empty, many-query, NULL/UNSET, and vector parameter scenarios. Includes benchmark script benchmarks/bench_batch_send_body.py.
Replace per-call int32_pack(-1) and int32_pack(-2) with module-level _INT32_NEG1 and _INT32_NEG2 constants. Avoids redundant struct packing on every null or unset parameter in the inner write_value loop. Benchmark: ~11% speedup on the parameter serialization loop for a typical 12-param mix of values, nulls, and unsets.
62f91eb to
cd14384
Compare
|
Warning Review limit reached
Next review available in: 55 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: QUIET Plan: Pro Plus Run ID: 📒 Files selected for processing (5)
Comment |
|
Rebased onto current `master` (this branch was stacked directly on #790's single commit, which is unchanged as commit 1 of 3 here; the 83 intervening master commits included the `protocol_features` plumbing from #770/#913 work, which touches the exact same `_write_query_params`/`BatchMessage.send_body` signatures). Conflict resolution notes:
Verified: all 46 protocol tests + 32 connection tests + full unit suite (741 passed / 88 skipped) pass, including with the Cython-compiled `.so` built locally (confirmed `BatchMessage.send_body` is a `cyfunction` at test time). No unresolved review threads found on this PR. Byte-for-byte equivalence of the buffer-accumulation path against the original `write_*` helpers double-checked by hand against their implementations. Force-pushed the rebase; still a draft pending #790. |
There was a problem hiding this comment.
Pull request overview
Optimizes protocol parameter and batch serialization by accumulating encoded parts and writing them in fewer operations.
Changes:
- Buffers query parameters and complete batch bodies before writing.
- Adds precomputed NULL/UNSET markers and removes a redundant override.
- Adds serialization tests and performance benchmarks.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
cassandra/protocol.py |
Implements buffered query and batch serialization. |
tests/unit/test_protocol.py |
Adds byte-level serialization coverage. |
benchmarks/bench_protocol_write_value.py |
Benchmarks precomputed marker constants. |
benchmarks/bench_execute_write_params.py |
Compares query parameter serialization approaches. |
benchmarks/bench_batch_send_body.py |
Benchmarks optimized batch serialization. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| _p(_u16(len(ks))) | ||
| _p(ks) | ||
|
|
||
| f.write(b"".join(parts)) |
|
|
||
| PROTO_VERSION = 4 | ||
| ITERATIONS = 50_000 | ||
| REPEATS = 3 |
| ITERATIONS = 500_000 | ||
| REPEATS = 5 |
Summary
Replace per-call
write_value()/write_byte()/write_short()inBatchMessage.send_body()with buffer accumulation (list.append+b"".join+ singlef.write()), reducingf.write()calls fromQ*(4 + 2*P) + footerto 1 for Q queries with P params each.Depends on PR #790 (
perf/buffer-accum-write-params).What changed
cassandra/protocol.pyBatchMessage.send_body()-- Full buffer accumulation for the entire message: batch header, per-query framing (prepared/unprepared), all parameters (with NULL/UNSET/str handling), and trailer.Pre-computed constants --
_INT32_NULLand_INT32_UNSETas module-level constants to avoid repeatedint32_pack()calls in the hot loop.tests/unit/test_protocol.pyAdded 7 new batch-specific test methods: prepared queries, unprepared queries, mixed, empty batch, many queries (50), NULL/UNSET params, and vector params.
Benchmark
Measured with
min()oftimeit.repeat(repeat=7, number=50_000)on a quiet machine (load <3), Cython.socompiled, before/after rebuild on same machine.Consistent 1.6-1.8x speedup across all batch scenarios. Larger batches with more params see the greatest absolute savings (35+ us saved for 50q x 10p).
Tests