Skip to content

perf: replace BytesIO with b"".join() in collection serialization (us improvement, x1.07-1.24 speedup) - #787

Open
mykaul wants to merge 1 commit into
scylladb:masterfrom
mykaul:perf/collection-serialize-join
Open

perf: replace BytesIO with b"".join() in collection serialization (us improvement, x1.07-1.24 speedup)#787
mykaul wants to merge 1 commit into
scylladb:masterfrom
mykaul:perf/collection-serialize-join

Conversation

@mykaul

@mykaul mykaul commented Apr 3, 2026

Copy link
Copy Markdown

Motivation

The serialize_safe methods for collection types (ListType, SetType, MapType, TupleType, UserType) use io.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:

  1. BytesIO object allocation on every serialization call
  2. Multiple method dispatch for each .write() call (Python method lookup + C-level buffer management)
  3. Repeated int32_pack(-1) calls for null elements, recomputing the same 4-byte value each time

Summary of Changes

  • Replace io.BytesIO() with list accumulation + b"".join() in four serialize_safe methods:
    • _SimpleParameterizedType.serialize_safe (used by ListType, SetType)
    • MapType.serialize_safe
    • TupleType.serialize_safe
    • UserType.serialize_safe
  • Add _INT32_NULL = int32_pack(-1) as a pre-computed module-level constant, eliminating repeated packing of the null sentinel value

The b"".join(parts) pattern is a well-known Python idiom that avoids intermediate buffer object overhead. CPython's bytes.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.

Note: PR #763 on this repo adds only the _INT32_NULL constant. This PR is a superset that also replaces BytesIO with b"".join() across all four collection/composite type serializers.

How It Was Tested

  • All unit tests pass: pytest tests/unit/test_types.py (62 passed), pytest tests/unit/test_query.py tests/unit/test_cluster.py (36 passed)
  • Correctness verified: before and after produce identical serialized bytes for all collection types

Benchmarks

Micro-benchmarks measuring end-to-end serialize_safe performance (Python 3.14, timeit):

Scenario Before (ns/call) After (ns/call) Speedup
List 10 elements 2274.7 1970.2 1.16x
List 100 elements 17051.2 15348.1 1.11x
List 1000 elements 156725.7 146230.7 1.07x
List 100 elements (33% null) 13962.1 11305.2 1.24x
Set 10 elements 2310.1 2023.2 1.14x
Set 100 elements 17604.5 15656.2 1.12x
Map 10 entries 4734.3 4209.7 1.13x
Map 100 entries 38962.5 35588.8 1.10x
Tuple 10 fields 2536.7 2246.4 1.13x
UDT 5 fields 1677.7 1510.7 1.11x

(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_safe code on origin/master via the same timeit harness. 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. the if __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.

@mykaul
mykaul force-pushed the perf/collection-serialize-join branch from d9b8b3c to 17369e1 Compare April 3, 2026 22:08
@mykaul mykaul changed the title perf: replace BytesIO with b"".join() in collection serialization perf: replace BytesIO with b"".join() in collection serialization (us improvement, x1.25-2.65 speedup) Apr 7, 2026
Copilot AI review requested due to automatic review settings July 29, 2026 20:31
@mykaul
mykaul force-pushed the perf/collection-serialize-join branch from 17369e1 to 88833fc Compare July 29, 2026 20:31
@coderabbitai

coderabbitai Bot commented Jul 29, 2026

Copy link
Copy Markdown
🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Description check ⚠️ Warning The description covers motivation, implementation, testing, and benchmarks but omits the repository's required pre-review checklist. Add the complete pre-review checklist, mark each applicable item, and include any required Fixes annotation or state why it does not apply.
✅ Passed checks (4 passed)
Check name Status Explanation
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Title check ✅ Passed The title clearly identifies the main serialization performance change and reports its benchmark result.

Comment @coderabbitai help to get the list of available commands.

@mykaul

mykaul commented Jul 29, 2026

Copy link
Copy Markdown
Author

Fix-up pass on this PR

Rebased onto current origin/master (was 84 commits behind; rebase applied cleanly, no conflicts) and re-verified the change.

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 cassandra/cqltypes.py line-by-line and confirmed the change is a pure mechanical refactor — every buf.write(X) becomes parts.append(X) in the same order, and buf.getvalue() becomes b"".join(parts). Since BytesIO.getvalue() after sequential writes is definitionally the concatenation of the written fragments in order, output equivalence holds by construction.

I also verified this empirically: swapped in the old BytesIO-based serialize_safe implementations and compared serialized output against the new join-based implementations for ListType, MapType, TupleType, and UserType across these cases — all byte-for-byte identical:

  • empty collection
  • single element
  • elements needing the null sentinel (including all-null and mixed null/non-null)
  • large collections (10,000-element list, 2,000-entry map)

Tests: tests/unit/test_types.py — 62 passed, 1 skipped (pre-existing, unrelated). Full tests/unit/ — 720 passed, 88 skipped, 0 failed.

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 benchmarks/bench_collection_serialize.py (e.g. print(f"...")print("...") where there was no interpolation) — no functional change, just style.

Re: #763 overlap: confirmed these are independent branches; #763 (still open/unmerged) adds only the _INT32_NULL constant, which this PR's master base did not yet include prior to rebase. No conflicts between the two as they stand. If #763 merges first, rebasing this PR onto master afterward should be a trivial no-op on that hunk since the added constant is identical.

Amended into the existing single commit (no new commits added) and force-pushed. Still a draft.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 and b"".join(...) in several serialize_safe implementations.
  • Add a module-level _INT32_NULL constant to avoid repeated int32_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.

Comment thread cassandra/cqltypes.py Outdated
Comment thread benchmarks/bench_collection_serialize.py
Comment thread benchmarks/bench_collection_serialize.py Outdated
Copilot AI review requested due to automatic review settings July 30, 2026 09:55
@mykaul
mykaul force-pushed the perf/collection-serialize-join branch from 88833fc to 8a9c7da Compare July 30, 2026 09:55
@mykaul

mykaul commented Jul 30, 2026

Copy link
Copy Markdown
Author

Fixed a genuine pre-existing flaky test causing CI wheel-build failures

Both macos-arm and (separately) linux wheel-build jobs were failing on this PR's CI, but neither failure was actually caused by this PR's own change (the BytesIO -> b"".join() collection-serialization rewrite).

macos-arm: tests/unit/test_types.py::TypeTests::test_datetype (fixed here)

FAILED tests/unit/test_types.py::TypeTests::test_datetype - AssertionError: assert b'\x00\x00\x01\x9f\xb0\xe0.E' == b'\x00\x00\x01\x9f\xb0\xe0.D'

Root cause: the test computed now_datetime and now_timestamp independently from the same time.time() float:

now_time_seconds = time.time()
now_datetime = datetime.datetime.fromtimestamp(now_time_seconds, tz=datetime.timezone.utc)
now_timestamp = now_time_seconds * 1e3

datetime.fromtimestamp() rounds the fractional seconds to the nearest microsecond internally, while DateType.serialize() truncates a raw numeric timestamp towards zero via int(timestamp). When the fractional part of now_time_seconds sits very close to a millisecond boundary, these two independent rounding paths can disagree by 1ms, producing a 1-byte difference in the serialized output. This is a real, if rare, race in the test itself, unrelated to DateType/collection serialization logic.

Reproduced reliably: sweeping ~2,000,000 microsecond-perturbed timestamps through the old code produced 928 mismatches; the same sweep against the fixed code produced 0.

Fix: derive both values from a single canonical millisecond integer, eliminating the possibility of divergent rounding:

now_timestamp_ms = int(time.time() * 1000)
now_datetime = datetime.datetime.fromtimestamp(now_timestamp_ms / 1000.0, tz=datetime.timezone.utc)
now_timestamp = float(now_timestamp_ms)

Verified with 100 repeated subprocess runs, 20,000 in-process iterations using real time.time(), and 500,000 in-process iterations sweeping microsecond-by-microsecond directly across millisecond boundaries (the exact scenario that broke the old code) — all pass, 0 failures. Full tests/unit/ suite: 720 passed, 88 skipped, 0 failed.

linux: tests/unit/io/test_twistedreactor.py::TestTwistedConnection::test_connection_initialization (separate, unrelated, not fixed here)

AssertionError: expected call not found.
Expected: run(installSignalHandlers=False)
  Actual: not called.

This is a different, pre-existing thread-timing race: maybe_start() spawns the Twisted reactor thread asynchronously, and the test asserts reactor.run(...) was called before that thread necessarily got scheduled. It's already been partially addressed once before (commit 94438c6f0, patching reactor.running to False), but that didn't fully close the race, as this CI run shows. Fixing it properly needs a synchronization mechanism (e.g. waiting on an event/condition before asserting) rather than a one-line change, and it's unrelated to this PR's collection-serialization change, so I'm leaving it out of scope here and flagging it separately.

Other review feedback addressed

Also addressed Copilot's three inline review comments in the same amended commit:

  • Updated the _INT32_NULL comment to mention it's used for collection and composite (tuple/UDT) field serialization, not just collections.
  • Added a SetType benchmark section using the previously-unused SetOfInt.
  • Wrapped the benchmark's executable portion in if __name__ == "__main__": so importing the module no longer runs the benchmark as a side effect.

All changes amended into the existing single commit and force-pushed. Still a draft.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 parts list retains every encoded key and value while join allocates the complete result, adding approximately one extra serialized map payload plus two references per fragment at peak. The previous BytesIO loop 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 join has 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 parts keeps the entire fragmented representation alive while join allocates the full result. For large lists/sets this adds roughly another serialized payload plus two list references per element to peak memory, whereas the previous BytesIO path 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_support only 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)

@mykaul mykaul changed the title perf: replace BytesIO with b"".join() in collection serialization (us improvement, x1.25-2.65 speedup) perf: replace BytesIO with b"".join() in collection serialization (us improvement, x1.1-x1.3 speedup) Jul 30, 2026
@mykaul mykaul changed the title perf: replace BytesIO with b"".join() in collection serialization (us improvement, x1.1-x1.3 speedup) perf: replace BytesIO with b"".join() in collection serialization (us improvement, x1.07-1.24 speedup) Jul 30, 2026
@mykaul
mykaul marked this pull request as ready for review August 15, 2026 06:56

@dkropachev dkropachev left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants