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

benchmarks/bench_collection_serialize.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python
+# Copyright ScyllaDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Benchmark: collection serialization - BytesIO vs b"".join(parts).
+
+Measures end-to-end serialize_safe performance for List, Set, Map, Tuple, and
+UserType with varying collection sizes.
+"""
+
+import timeit
+import sys
+import os
+
+# Add project root to path
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from cassandra.cqltypes import (
+    ListType,
+    SetType,
+    MapType,
+    TupleType,
+    UserType,
+    Int32Type,
+    UTF8Type,
+    FloatType,
+)
+
+PROTOCOL_VERSION = 4
+
+# Build parameterized types
+ListOfInt = ListType.apply_parameters([Int32Type])
+SetOfInt = SetType.apply_parameters([Int32Type])
+MapIntToStr = MapType.apply_parameters([Int32Type, UTF8Type])
+
+
+# For TupleType and UserType, we need to set subtypes on a subclass
+class TestTupleType(TupleType):
+    subtypes = (
+        Int32Type,
+        Int32Type,
+        Int32Type,
+        Int32Type,
+        Int32Type,
+        Int32Type,
+        Int32Type,
+        Int32Type,
+        Int32Type,
+        Int32Type,
+    )
+
+
+class TestUserType(UserType):
+    subtypes = (Int32Type, UTF8Type, FloatType, Int32Type, UTF8Type)
+    fieldnames = ("id", "name", "score", "age", "email")
+    typename = "test_udt"
+    keyspace = "test_ks"
+    mapped_class = None
+    tuple_type = None
+
+
+def run_bench(label, fn, args, n):
+    # Warm up
+    for _ in range(min(1000, n)):
+        fn(*args)
+    t = timeit.timeit(lambda: fn(*args), number=n)
+    ns_per_call = t / n * 1e9
+    print(f"  {label:45s} {t:.3f}s  ({ns_per_call:.2f} ns/call)")
+    return t, ns_per_call
+
+
+def main():
+    # Test data
+    list_10 = list(range(10))
+    list_100 = list(range(100))
+    list_1000 = list(range(1000))
+    list_with_nulls = [i if i % 3 != 0 else None for i in range(100)]
+
+    set_10 = set(range(10))
+    set_100 = set(range(100))
+
+    map_10 = {i: f"value_{i}" for i in range(10)}
+    map_100 = {i: f"value_{i}" for i in range(100)}
+
+    tuple_10 = tuple(range(10))
+    udt_val = (1, "test_name", 3.14, 25, "test@example.com")
+
+    N_SMALL = 500_000
+    N_MED = 100_000
+    N_LARGE = 10_000
+
+    print("Collection serialization benchmark")
+    print("=" * 70)
+
+    results = {}
+
+    print("\nListType.serialize (list of int32):")
+    _, r = run_bench(
+        "10 elements", ListOfInt.serialize, (list_10, PROTOCOL_VERSION), N_SMALL
+    )
+    results["list_10"] = r
+    _, r = run_bench(
+        "100 elements", ListOfInt.serialize, (list_100, PROTOCOL_VERSION), N_MED
+    )
+    results["list_100"] = r
+    _, r = run_bench(
+        "1000 elements", ListOfInt.serialize, (list_1000, PROTOCOL_VERSION), N_LARGE
+    )
+    results["list_1000"] = r
+    _, r = run_bench(
+        "100 elements (33% null)",
+        ListOfInt.serialize,
+        (list_with_nulls, PROTOCOL_VERSION),
+        N_MED,
+    )
+    results["list_100_nulls"] = r
+
+    print("\nSetType.serialize (set of int32):")
+    _, r = run_bench(
+        "10 elements", SetOfInt.serialize, (set_10, PROTOCOL_VERSION), N_SMALL
+    )
+    results["set_10"] = r
+    _, r = run_bench(
+        "100 elements", SetOfInt.serialize, (set_100, PROTOCOL_VERSION), N_MED
+    )
+    results["set_100"] = r
+
+    print("\nMapType.serialize (map<int32, text>):")
+    _, r = run_bench(
+        "10 entries", MapIntToStr.serialize, (map_10, PROTOCOL_VERSION), N_SMALL
+    )
+    results["map_10"] = r
+    _, r = run_bench(
+        "100 entries", MapIntToStr.serialize, (map_100, PROTOCOL_VERSION), N_MED
+    )
+    results["map_100"] = r
+
+    print("\nTupleType.serialize (10 x int32):")
+    _, r = run_bench(
+        "10 fields", TestTupleType.serialize, (tuple_10, PROTOCOL_VERSION), N_SMALL
+    )
+    results["tuple_10"] = r
+
+    print("\nUserType.serialize (5 fields: int, text, float, int, text):")
+    _, r = run_bench(
+        "5 fields", TestUserType.serialize, (udt_val, PROTOCOL_VERSION), N_SMALL
+    )
+    results["udt_5"] = r
+
+    # Print summary for easy comparison
+    print(f"\n{'=' * 70}")
+    print("Summary (ns/call):")
+    for k, v in results.items():
+        print(f"  {k:25s}: {v:.2f} ns")
+
+
+if __name__ == "__main__":
+    main()

cassandra/cqltypes.py

@@ -63,6 +63,10 @@
 
 _number_types = frozenset((int, float))
 
+# Pre-computed null sentinel for collection and composite (tuple/UDT) field
+# serialization (int32 -1)
+_INT32_NULL = int32_pack(-1)
+
 
 def _name_from_hex_string(encoded_name):
     bin_str = unhexlify(encoded_name)
@@ -836,17 +840,16 @@ def serialize_safe(cls, items, protocol_version):
             raise TypeError("Received a string for a type that expects a sequence")
 
         subtype, = cls.subtypes
-        buf = io.BytesIO()
-        buf.write(int32_pack(len(items)))
         inner_proto = max(3, protocol_version)
+        parts = [int32_pack(len(items))]
         for item in items:
             if item is None:
-                buf.write(int32_pack(-1))
+                parts.append(_INT32_NULL)
             else:
                 itembytes = subtype.to_binary(item, inner_proto)
-                buf.write(int32_pack(len(itembytes)))
-                buf.write(itembytes)
-        return buf.getvalue()
+                parts.append(int32_pack(len(itembytes)))
+                parts.append(itembytes)
+        return b"".join(parts)
 
 
 class ListType(_SimpleParameterizedType):
@@ -899,27 +902,26 @@ def deserialize_safe(cls, byts, protocol_version):
     @classmethod
     def serialize_safe(cls, themap, protocol_version):
         key_type, value_type = cls.subtypes
-        buf = io.BytesIO()
-        buf.write(int32_pack(len(themap)))
         try:
             items = themap.items()
         except AttributeError:
             raise TypeError("Got a non-map object for a map value")
         inner_proto = max(3, protocol_version)
+        parts = [int32_pack(len(themap))]
         for key, val in items:
             if key is not None:
                 keybytes = key_type.to_binary(key, inner_proto)
-                buf.write(int32_pack(len(keybytes)))
-                buf.write(keybytes)
+                parts.append(int32_pack(len(keybytes)))
+                parts.append(keybytes)
             else:
-                buf.write(int32_pack(-1))
+                parts.append(_INT32_NULL)
             if val is not None:
                 valbytes = value_type.to_binary(val, inner_proto)
-                buf.write(int32_pack(len(valbytes)))
-                buf.write(valbytes)
+                parts.append(int32_pack(len(valbytes)))
+                parts.append(valbytes)
             else:
-                buf.write(int32_pack(-1))
-        return buf.getvalue()
+                parts.append(_INT32_NULL)
+        return b"".join(parts)
 
 
 class TupleType(_ParameterizedType):
@@ -957,15 +959,15 @@ def serialize_safe(cls, val, protocol_version):
                              (len(cls.subtypes), len(val), val))
 
         proto_version = max(3, protocol_version)
-        buf = io.BytesIO()
+        parts = []
         for item, subtype in zip(val, cls.subtypes):
             if item is not None:
                 packed_item = subtype.to_binary(item, proto_version)
-                buf.write(int32_pack(len(packed_item)))
-                buf.write(packed_item)
+                parts.append(int32_pack(len(packed_item)))
+                parts.append(packed_item)
             else:
-                buf.write(int32_pack(-1))
-        return buf.getvalue()
+                parts.append(_INT32_NULL)
+        return b"".join(parts)
 
     @classmethod
     def cql_parameterized_type(cls):
@@ -1026,7 +1028,7 @@ def deserialize_safe(cls, byts, protocol_version):
     @classmethod
     def serialize_safe(cls, val, protocol_version):
         proto_version = max(3, protocol_version)
-        buf = io.BytesIO()
+        parts = []
         for i, (fieldname, subtype) in enumerate(zip(cls.fieldnames, cls.subtypes)):
             # first treat as a tuple, else by custom type
             try:
@@ -1038,11 +1040,11 @@ def serialize_safe(cls, val, protocol_version):
 
             if item is not None:
                 packed_item = subtype.to_binary(item, proto_version)
-                buf.write(int32_pack(len(packed_item)))
-                buf.write(packed_item)
+                parts.append(int32_pack(len(packed_item)))
+                parts.append(packed_item)
             else:
-                buf.write(int32_pack(-1))
-        return buf.getvalue()
+                parts.append(_INT32_NULL)
+        return b"".join(parts)
 
     @classmethod
     def _make_registered_udt_namedtuple(cls, keyspace, name, field_names):

tests/unit/test_types.py

@@ -212,11 +212,23 @@ def test_empty_value(self):
         assert str(EmptyValue()) == 'EMPTY'
 
     def test_datetype(self):
-        now_time_seconds = time.time()
-        now_datetime = datetime.datetime.fromtimestamp(now_time_seconds, tz=datetime.timezone.utc)
+        # Derive now_datetime and now_timestamp from the same integer
+        # millisecond value rather than independently from time.time().
+        # Computing them separately (e.g. now_datetime via
+        # datetime.fromtimestamp(now_time_seconds, ...) and now_timestamp via
+        # now_time_seconds * 1e3) can, in rare cases, round to different
+        # millisecond values: fromtimestamp() rounds the fractional seconds
+        # to the nearest microsecond internally, while a plain
+        # seconds * 1e3 multiplication is truncated towards zero by
+        # DateType.serialize(). When the fractional part sits very close to
+        # a millisecond boundary, those two paths can disagree by 1ms,
+        # making this assertion flaky. Using a single canonical millisecond
+        # source for both eliminates that divergence.
+        now_timestamp_ms = int(time.time() * 1000)
+        now_datetime = datetime.datetime.fromtimestamp(now_timestamp_ms / 1000.0, tz=datetime.timezone.utc)
 
         # Cassandra timestamps in millis
-        now_timestamp = now_time_seconds * 1e3
+        now_timestamp = float(now_timestamp_ms)
 
         # same results serialized
         assert DateType.serialize(now_datetime, 0) == DateType.serialize(now_timestamp, 0)