perf: buffer accumulation in BatchMessage.send_body()

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.

Author: mykaul

benchmarks/bench_batch_send_body.py

@@ -0,0 +1,146 @@
+# Copyright DataStax, 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: BatchMessage.send_body() for vector and scalar workloads.
+
+Measures the actual loaded module's BatchMessage.send_body() method.
+Run this before and after optimization to compare.
+
+Usage:
+  # Build baseline .so, then:
+  python benchmarks/bench_batch_send_body.py
+  # Apply optimization, rebuild .so, then:
+  python benchmarks/bench_batch_send_body.py
+"""
+
+import io
+import struct
+import time
+import timeit
+import sys
+import os
+
+# Ensure the repo root is importable
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+import cassandra.protocol
+from cassandra.protocol import BatchMessage
+from cassandra.query import BatchType
+from cassandra.marshal import int32_pack
+
+
+# ---------------------------------------------------------------------------
+# Scenario builders
+# ---------------------------------------------------------------------------
+
+
+def make_batch_vector_queries(num_queries, dim):
+    """Batch of prepared INSERT with (int32_key, float_vector) params."""
+    vector_bytes = struct.pack(f">{dim}f", *([0.1] * dim))
+    key_bytes = int32_pack(42)
+    return [
+        (True, b"\x01\x02\x03\x04\x05\x06\x07\x08", [key_bytes, vector_bytes])
+        for _ in range(num_queries)
+    ]
+
+
+def make_batch_scalar_queries(num_queries, num_params, param_size=20):
+    """Batch of prepared INSERT with N text columns of param_size bytes."""
+    params = [b"\x41" * param_size for _ in range(num_params)]
+    return [
+        (True, b"\x01\x02\x03\x04\x05\x06\x07\x08", list(params))
+        for _ in range(num_queries)
+    ]
+
+
+def make_batch_unprepared_queries(num_queries, num_params, param_size=20):
+    """Batch of unprepared INSERT statements."""
+    stmt = "INSERT INTO ks.tbl (k, v) VALUES (?, ?)"
+    params = [b"\x41" * param_size for _ in range(num_params)]
+    return [(False, stmt, list(params)) for _ in range(num_queries)]
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+PROTO_VERSION = 4
+ITERATIONS = 50_000
+REPEATS = 3
+
+SCENARIOS = [
+    ("10 queries x 2 params (128D vec)", make_batch_vector_queries(10, 128)),
+    ("10 queries x 2 params (768D vec)", make_batch_vector_queries(10, 768)),
+    ("50 queries x 2 params (128D vec)", make_batch_vector_queries(50, 128)),
+    ("10 queries x 10 text params", make_batch_scalar_queries(10, 10, 20)),
+    ("50 queries x 10 text params", make_batch_scalar_queries(50, 10, 20)),
+    ("10 unprepared x 2 params", make_batch_unprepared_queries(10, 2, 20)),
+]
+
+
+# ---------------------------------------------------------------------------
+# Benchmark
+# ---------------------------------------------------------------------------
+
+
+def bench_batch(queries, iterations, repeats):
+    """Benchmark BatchMessage.send_body(), return best ns/call."""
+    msg = BatchMessage(
+        batch_type=BatchType.LOGGED,
+        queries=queries,
+        consistency_level=1,
+        timestamp=1234567890123456,
+    )
+    f = io.BytesIO()
+
+    def run():
+        f.seek(0)
+        f.truncate()
+        msg.send_body(f, PROTO_VERSION)
+
+    t = timeit.repeat(run, number=iterations, repeat=repeats, timer=time.process_time)
+    return min(t) / iterations * 1e9
+
+
+def main():
+    is_cython = cassandra.protocol.__file__.endswith(".so")
+    print(f"Python:  {sys.version.split()[0]}")
+    print(f"Module:  {cassandra.protocol.__file__}")
+    print(f"Cython:  {'YES (.so loaded)' if is_cython else 'NO (pure Python .py)'}")
+    print(f"Config:  proto v{PROTO_VERSION}, {ITERATIONS:,} iters, best of {REPEATS}")
+    print()
+    print(f"{'Scenario':45s}  {'ns/call':>10s}  {'bytes':>8s}")
+    print(f"{'-' * 45}  {'-' * 10}  {'-' * 8}")
+
+    for label, queries in SCENARIOS:
+        # Measure output size
+        msg = BatchMessage(
+            batch_type=BatchType.LOGGED,
+            queries=queries,
+            consistency_level=1,
+            timestamp=1234567890123456,
+        )
+        f = io.BytesIO()
+        msg.send_body(f, PROTO_VERSION)
+        nbytes = len(f.getvalue())
+
+        ns = bench_batch(queries, ITERATIONS, REPEATS)
+        print(f"{label:45s}  {ns:8.1f}    {nbytes:>6d}")
+
+    print()
+
+
+if __name__ == "__main__":
+    main()

cassandra/protocol.py

@@ -923,21 +923,36 @@ def __init__(self, batch_type, queries, consistency_level,
         self.keyspace = keyspace
 
     def send_body(self, f, protocol_version):
-        write_byte(f, self.batch_type.value)
-        write_short(f, len(self.queries))
+        # Buffer accumulation: collect all bytes and write once.
+        _i32 = int32_pack
+        _u16 = uint16_pack
+        _u8 = uint8_pack
+        parts = [_u8(self.batch_type.value), _u16(len(self.queries))]
+        _p = parts.append
         for prepared, string_or_query_id, params in self.queries:
             if not prepared:
-                write_byte(f, 0)
-                write_longstring(f, string_or_query_id)
+                _p(_u8(0))
+                if isinstance(string_or_query_id, str):
+                    string_or_query_id = string_or_query_id.encode('utf8')
+                _p(_i32(len(string_or_query_id)))
+                _p(string_or_query_id)
             else:
-                write_byte(f, 1)
-                write_short(f, len(string_or_query_id))
-                f.write(string_or_query_id)
-            write_short(f, len(params))
+                _p(_u8(1))
+                _p(_u16(len(string_or_query_id)))
+                _p(string_or_query_id)
+            _p(_u16(len(params)))
             for param in params:
-                write_value(f, param)
+                if param is None:
+                    _p(_i32(-1))
+                elif param is _UNSET_VALUE:
+                    _p(_i32(-2))
+                else:
+                    if isinstance(param, str):
+                        param = param.encode('utf8')
+                    _p(_i32(len(param)))
+                    _p(param)
 
-        write_consistency_level(f, self.consistency_level)
+        _p(_u16(self.consistency_level))
         flags = 0
         if self.serial_consistency_level:
             flags |= _WITH_SERIAL_CONSISTENCY_FLAG
@@ -951,18 +966,24 @@ def send_body(self, f, protocol_version):
                     "Keyspaces may only be set on queries with protocol version "
                     "5 or higher. Consider setting Cluster.protocol_version to 5.")
         if ProtocolVersion.uses_int_query_flags(protocol_version):
-            write_int(f, flags)
+            _p(_i32(flags))
         else:
-            write_byte(f, flags)
+            _p(_u8(flags))
 
         if self.serial_consistency_level:
-            write_consistency_level(f, self.serial_consistency_level)
+            _p(_u16(self.serial_consistency_level))
         if self.timestamp is not None:
-            write_long(f, self.timestamp)
+            _p(uint64_pack(self.timestamp))
 
         if ProtocolVersion.uses_keyspace_flag(protocol_version):
             if self.keyspace is not None:
-                write_string(f, self.keyspace)
+                ks = self.keyspace
+                if isinstance(ks, str):
+                    ks = ks.encode('utf8')
+                _p(_u16(len(ks)))
+                _p(ks)
+
+        f.write(b"".join(parts))
 
 
 known_event_types = frozenset((

tests/unit/test_protocol.py

@@ -170,7 +170,7 @@ def test_keyspace_written_with_length(self):
 
     def test_batch_message_with_keyspace(self):
         self.maxDiff = None
-        io = Mock(name='io')
+        buf = io.BytesIO()
         batch = BatchMessage(
             batch_type=BatchType.LOGGED,
             queries=((False, 'stmt a', ('param a',)),
@@ -180,18 +180,27 @@ def test_batch_message_with_keyspace(self):
             consistency_level=3,
             keyspace='ks'
         )
-        batch.send_body(io, protocol_version=5)
-        self._check_calls(io,
-            ((b'\x00',), (b'\x00\x03',), (b'\x00',),
-             (b'\x00\x00\x00\x06',), (b'stmt a',),
-             (b'\x00\x01',), (b'\x00\x00\x00\x07',), ('param a',),
-             (b'\x00',), (b'\x00\x00\x00\x06',), (b'stmt b',),
-             (b'\x00\x01',), (b'\x00\x00\x00\x07',), ('param b',),
-             (b'\x00',), (b'\x00\x00\x00\x06',), (b'stmt c',),
-             (b'\x00\x01',), (b'\x00\x00\x00\x07',), ('param c',),
-             (b'\x00\x03',),
-             (b'\x00\x00\x00\x80',), (b'\x00\x02',), (b'ks',))
+        batch.send_body(buf, protocol_version=5)
+        expected = (
+            b'\x00'              # batch type LOGGED
+            b'\x00\x03'         # 3 queries
+            b'\x00'              # not prepared
+            b'\x00\x00\x00\x06' b'stmt a'  # longstring 'stmt a'
+            b'\x00\x01'         # 1 param
+            b'\x00\x00\x00\x07' b'param a'  # write_value 'param a'
+            b'\x00'              # not prepared
+            b'\x00\x00\x00\x06' b'stmt b'
+            b'\x00\x01'
+            b'\x00\x00\x00\x07' b'param b'
+            b'\x00'
+            b'\x00\x00\x00\x06' b'stmt c'
+            b'\x00\x01'
+            b'\x00\x00\x00\x07' b'param c'
+            b'\x00\x03'         # consistency level
+            b'\x00\x00\x00\x80'  # flags (keyspace)
+            b'\x00\x02' b'ks'   # keyspace
         )
+        self.assertEqual(buf.getvalue(), expected)
 
 class WriteQueryParamsBufferAccumulationTest(unittest.TestCase):
     """
@@ -373,3 +382,90 @@ def test_single_unset_param(self):
         raw = self._execute_msg_bytes(msg, protocol_version=4)
         self.assertIn(expected, raw)
 
+    # -- BatchMessage buffer accumulation tests ---------------------------
+
+    @staticmethod
+    def _batch_msg_bytes(queries, protocol_version=4, **kwargs):
+        """Serialize a BatchMessage and return the raw bytes."""
+        msg = BatchMessage(batch_type=BatchType.LOGGED, queries=queries,
+                           consistency_level=1, **kwargs)
+        buf = io.BytesIO()
+        msg.send_body(buf, protocol_version)
+        return buf.getvalue()
+
+    def test_batch_prepared_queries_with_params(self):
+        """Batch of prepared queries with byte params serializes correctly."""
+        queries = [
+            (True, b'\x01\x02\x03\x04', [b'val1', b'val2']),
+            (True, b'\x01\x02\x03\x04', [b'val3', None]),
+        ]
+        raw = self._batch_msg_bytes(queries)
+        self.assertIn(b'val1', raw)
+        self.assertIn(b'val2', raw)
+        self.assertIn(b'val3', raw)
+        self.assertIn(int32_pack(-1), raw)  # NULL
+
+    def test_batch_unprepared_queries(self):
+        """Batch of unprepared (string) queries serializes correctly."""
+        queries = [
+            (False, 'INSERT INTO t (k) VALUES (?)', [b'\x01']),
+            (False, 'INSERT INTO t (k) VALUES (?)', [b'\x02']),
+        ]
+        raw = self._batch_msg_bytes(queries)
+        self.assertIn(b'INSERT INTO t (k) VALUES (?)', raw)
+
+    def test_batch_mixed_prepared_unprepared(self):
+        """Batch mixing prepared and unprepared queries."""
+        queries = [
+            (False, 'SELECT 1', []),
+            (True, b'\xab\xcd', [b'data']),
+        ]
+        raw = self._batch_msg_bytes(queries)
+        self.assertIn(b'SELECT 1', raw)
+        self.assertIn(b'data', raw)
+
+    def test_batch_empty_queries(self):
+        """Batch with zero queries."""
+        raw = self._batch_msg_bytes([])
+        self.assertIn(uint16_pack(0), raw)
+
+    def test_batch_many_queries(self):
+        """Batch with 50 queries to exercise accumulation at scale."""
+        queries = [
+            (True, b'\x01\x02', [b'param_%03d' % i])
+            for i in range(50)
+        ]
+        raw = self._batch_msg_bytes(queries)
+        self.assertIn(uint16_pack(50), raw)
+        for i in range(50):
+            self.assertIn(b'param_%03d' % i, raw)
+
+    def test_batch_null_and_unset_params(self):
+        """Batch params with NULL and UNSET values."""
+        queries = [
+            (True, b'\x01', [None, _UNSET_VALUE, b'ok']),
+        ]
+        raw = self._batch_msg_bytes(queries, protocol_version=4)
+        self.assertIn(int32_pack(-1), raw)   # NULL
+        self.assertIn(int32_pack(-2), raw)   # UNSET
+        self.assertIn(b'ok', raw)
+
+    def test_batch_vector_params(self):
+        """Batch with large vector params (simulating bulk vector INSERT)."""
+        vector = struct.pack('128f', *([0.5] * 128))
+        queries = [
+            (True, b'\x01\x02', [int32_pack(i), vector])
+            for i in range(10)
+        ]
+        raw = self._batch_msg_bytes(queries)
+        # 10 copies of the vector should appear
+        count = 0
+        start = 0
+        while True:
+            idx = raw.find(vector, start)
+            if idx == -1:
+                break
+            count += 1
+            start = idx + 1
+        self.assertEqual(count, 10)
+