perf: replace BytesIO with b"".join() in collection serialization (us improvement, x1.07-1.24 speedup) - #787
Conversation
d9b8b3c to
17369e1
Compare
17369e1 to
88833fc
Compare
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
Fix-up pass on this PRRebased onto current Byte-for-byte equivalence verification (against the sibling concern that this overlaps with #763's null-sentinel work): I diffed the pre-change and post-change I also verified this empirically: swapped in the old BytesIO-based
Tests: CI: all 13 checks were green on the prior commit; re-triggered by this push. Self-review nit fixed: cleaned up a handful of f-strings without placeholders in the new Re: #763 overlap: confirmed these are independent branches; #763 (still open/unmerged) adds only the Amended into the existing single commit (no new commits added) and force-pushed. Still a draft. |
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.
Improve serialization performance for CQL collection/composite types by reducing per-call overhead during serialize_safe.
Changes:
- Replace
io.BytesIO()+ repeated.write()calls with list accumulation andb"".join(...)in severalserialize_safeimplementations. - Add a module-level
_INT32_NULLconstant to avoid repeatedint32_pack(-1)work. - Add a benchmark script to measure serialization performance across common types and sizes.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| cassandra/cqltypes.py | Refactors multiple serializers to use b"".join(parts) and introduces _INT32_NULL sentinel for faster null encoding. |
| benchmarks/bench_collection_serialize.py | Adds a standalone benchmark script for collection/composite serialization performance. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
88833fc to
8a9c7da
Compare
Fixed a genuine pre-existing flaky test causing CI wheel-build failuresBoth macos-arm:
|
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (5)
cassandra/cqltypes.py:970
- Joining a retained fragment list requires all encoded fields and the full final bytes object to coexist. For tuples with large serialized fields, this raises peak memory by about another payload compared with writing and releasing each field through
BytesIO, creating an avoidable OOM risk. Please preserve bounded accumulation or explicitly account for the memory tradeoff.
return b"".join(parts)
cassandra/cqltypes.py:924
- The
partslist retains every encoded key and value whilejoinallocates the complete result, adding approximately one extra serialized map payload plus two references per fragment at peak. The previousBytesIOloop released each temporary fragment after writing it, so large maps can now consume substantially more memory and may OOM. Please use bounded accumulation or otherwise account for this regression.
return b"".join(parts)
cassandra/cqltypes.py:1047
- This retains every encoded UDT field until
joinhas allocated the complete result, so all fragments and the final bytes object coexist. Large text/blob/nested fields can therefore add roughly another UDT payload to peak memory versus the previous streaming writes and can OOM serialization. Please preserve bounded accumulation or explicitly account for this memory regression.
return b"".join(parts)
cassandra/cqltypes.py:852
- Collecting every encoded fragment in
partskeeps the entire fragmented representation alive whilejoinallocates the full result. For large lists/sets this adds roughly another serialized payload plus two list references per element to peak memory, whereas the previousBytesIOpath copied each fragment and released it as it went. This can turn large collection serialization into an OOM despite the microbenchmark speedup; preserve bounded accumulation or explicitly address this memory regression.
This issue also appears in the following locations of the same file:
- line 924
- line 970
- line 1047
return b"".join(parts)
cassandra/cqltypes.py:68
- This shared sentinel is not exercised by any serialization regression test. The nearby
test_collection_null_supportonly deserializes null sentinels, while the marshalling fixtures cover only empty or non-null collections. Add exact-byte serialization cases for null list/set elements, map keys/values, tuple fields, and UDT fields so this common path is protected.
_INT32_NULL = int32_pack(-1)
dkropachev
left a comment
There was a problem hiding this comment.
[P2] Avoid retaining every serialized fragment (cassandra/cqltypes.py:852). parts keeps each payload and separately allocated length until join allocates the final result. Reproduced with 500,000 int values: a 4,000,004-byte result raises traced peak memory from 4.1 MiB with BytesIO to 123.5 MiB. Large valid lists/maps can therefore OOM the client. Preserve streaming or use bounded accumulation.
Motivation
The
serialize_safemethods for collection types (ListType,SetType,MapType,TupleType,UserType) useio.BytesIO()as an intermediate buffer — creating a BytesIO object, writing fragments with multiple.write()calls, then extracting the result with.getvalue(). This pattern incurs overhead from:.write()call (Python method lookup + C-level buffer management)int32_pack(-1)calls for null elements, recomputing the same 4-byte value each timeSummary of Changes
io.BytesIO()with list accumulation +b"".join()in fourserialize_safemethods:_SimpleParameterizedType.serialize_safe(used byListType,SetType)MapType.serialize_safeTupleType.serialize_safeUserType.serialize_safe_INT32_NULL = int32_pack(-1)as a pre-computed module-level constant, eliminating repeated packing of the null sentinel valueThe
b"".join(parts)pattern is a well-known Python idiom that avoids intermediate buffer object overhead. CPython'sbytes.join()pre-calculates the total output size and copies all fragments in a single pass, whereas BytesIO must manage a growable internal buffer with potential reallocations.How It Was Tested
pytest tests/unit/test_types.py(62 passed),pytest tests/unit/test_query.py tests/unit/test_cluster.py(36 passed)Benchmarks
Micro-benchmarks measuring end-to-end
serialize_safeperformance (Python 3.14, timeit):(Re-measured on the same machine as the original benchmark, 5 runs each, mean ns/call, CV 0.6%-3.6%, before/after ranges non-overlapping in every case — apples-to-apples against the current
serialize_safecode onorigin/mastervia the sametimeitharness. The table above supersedes the original numbers below, which were captured ~3 months earlier; the gap is most likely explained by unrelated changes to this benchmark script itself since then (e.g. theif __name__ == "__main__":wrapping added later changes the script's own variable-access bytecode, shifting fixed per-call overhead for both "before" and "after" alike) rather than any change to the actual optimization, which remains real and reproducible.)The pattern also produces simpler, more idiomatic Python code.