perf: buffer accumulation in BatchMessage.send_body() (1.6-1.8x speedup, us improvement, depends on PR #790) - #791
Conversation
5eec7ed to
b1d1cd0
Compare
mykaul
changed the title
b1d1cd0 to
62f91eb
Compare
mykaul
changed the title
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
|
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 `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