perf: add Cython metadata parser using BytesIOReader (100's-1000's of ns improvements, x1.5-4 speedup)

cassandra/bytesio.pxd

@@ -18,3 +18,4 @@ cdef class BytesIOReader:
     cdef Py_ssize_t pos
     cdef Py_ssize_t size
     cdef char *read(self, Py_ssize_t n = ?) except NULL
+    cdef void seek(self, Py_ssize_t new_pos)

cassandra/bytesio.pyx

@@ -42,3 +42,19 @@ cdef class BytesIOReader:
         cdef char *res = self.buf_ptr + self.pos
         self.pos = newpos
         return res
+
+    cdef void seek(self, Py_ssize_t new_pos):
+        """
+        Reposition the reader to an absolute offset within the buffer, so
+        the next read() starts at `new_pos`.
+
+        `buf_ptr` is the base address of `buf` (the underlying bytes
+        object); it is set once in __init__ and is never reassigned
+        anywhere else -- read() always computes `buf_ptr + pos`, so `pos`
+        is the only mutable cursor state and the only field a caller ever
+        needs to reset. This method exists so callers reposition the
+        reader through a single explicit entry point instead of reaching
+        into `pos` (and, in older code, the always-redundant `buf_ptr`)
+        directly at the call site.
+        """
+        self.pos = new_pos

cassandra/c_shard_info.pyx

@@ -14,9 +14,6 @@
 
 from libc.stdint cimport INT64_MIN, UINT32_MAX, uint64_t, int64_t
 
-cdef extern from *:
-    ctypedef unsigned int __uint128_t
-
 cdef class ShardingInfo():
     cdef readonly int shards_count
     cdef readonly unicode partitioner
@@ -39,5 +36,15 @@ cdef class ShardingInfo():
     def shard_id_from_token(self, int64_t token_input):
         cdef uint64_t biased_token = token_input + (<uint64_t>1 << 63);
         biased_token <<= self.sharding_ignore_msb;
-        cdef int shardId = (<__uint128_t>biased_token * self.shards_count) >> 64;
+        # Compute (biased_token * shards_count) >> 64, i.e. the high 64 bits of the
+        # 64x32-bit product, using only 64-bit arithmetic. This used to rely on the
+        # GCC/Clang-only __uint128_t extension type, which MSVC does not support at
+        # all (no 128-bit integer type), causing a compile error on Windows builds.
+        # The split below is a standard, portable multiply-high decomposition and is
+        # numerically identical to the previous 128-bit computation.
+        cdef uint64_t shards_count = <uint64_t>self.shards_count
+        cdef uint64_t low_product = (biased_token & <uint64_t>UINT32_MAX) * shards_count
+        cdef uint64_t carry = low_product >> 32
+        cdef uint64_t mid = (biased_token >> 32) * shards_count + carry
+        cdef int shardId = <int>(mid >> 32);
         return shardId

cassandra/ioutils.pyx

@@ -15,7 +15,7 @@
 include 'cython_marshal.pyx'
 from cassandra.buffer cimport Buffer, from_ptr_and_size
 
-from libc.stdint cimport int32_t
+from libc.stdint cimport int32_t, uint16_t
 from cassandra.bytesio cimport BytesIOReader
 
 
@@ -45,3 +45,9 @@ cdef inline int32_t read_int(BytesIOReader reader) except ?0xDEAD:
     buf.ptr = reader.read(4)
     buf.size = 4
     return unpack_num[int32_t](&buf)
+
+cdef inline uint16_t read_short(BytesIOReader reader) except ?0xFFFE:
+    cdef Buffer buf
+    buf.ptr = reader.read(2)
+    buf.size = 2
+    return unpack_num[uint16_t](&buf)

cassandra/metadata_parser.pyx

@@ -0,0 +1,218 @@
+# 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.
+
+"""
+Cython-optimized metadata parsing for CQL protocol ResultMessage.
+
+Uses BytesIOReader for zero-copy reads, eliminating per-read bytes allocation
+that dominates recv_results_metadata cost.
+"""
+
+include "ioutils.pyx"
+
+
+# ---------- low-level readers on BytesIOReader ----------
+# read_int(BytesIOReader) and read_short(BytesIOReader) are provided by ioutils.pyx
+
+cdef inline str read_string_br(BytesIOReader reader):
+    """Read a [string]: a [short] n, followed by n bytes of UTF-8."""
+    cdef uint16_t size = read_short(reader)
+    cdef char *ptr = reader.read(size)
+    return ptr[:size].decode('utf8')
+
+cdef inline bytes read_binary_string_br(BytesIOReader reader):
+    """Read a [short bytes]: a [short] n, followed by n raw bytes."""
+    cdef uint16_t size = read_short(reader)
+    cdef char *ptr = reader.read(size)
+    return ptr[:size]
+
+cdef inline bytes read_binary_longstring_br(BytesIOReader reader):
+    """Read a [bytes]: an [int] n, followed by n raw bytes.
+
+    Unlike row values (where a negative length is a valid NULL/"not set"
+    sentinel handled by get_buf() in ioutils.pyx), the [bytes] values read
+    here (paging_state, result_metadata_id) are only read when the protocol
+    says they are present, so a negative length here means a malformed or
+    corrupted frame. Reject it explicitly instead of passing it through:
+    BytesIOReader.read() treats any negative n as "read to the end of the
+    buffer" rather than raising, and slicing a char* with a negative stop
+    bound is not a validated length -- on CPython that currently surfaces
+    as an opaque SystemError from PyBytes_FromStringAndSize instead of a
+    clean, catchable protocol error.
+    """
+    cdef int32_t size = read_int(reader)
+    if size < 0:
+        raise ValueError(
+            "Invalid negative length %d for [bytes] value" % size)
+    cdef char *ptr = reader.read(size)
+    return ptr[:size]
+
+
+# ---------- flag constants (mirrored from ResultMessage) ----------
+# These MUST stay in sync with the class attributes in ResultMessage (protocol.py).
+# They are duplicated here as compile-time DEF constants for Cython performance.
+
+DEF _FLAGS_GLOBAL_TABLES_SPEC = 0x0001
+DEF _HAS_MORE_PAGES_FLAG      = 0x0002
+DEF _NO_METADATA_FLAG          = 0x0004
+DEF _METADATA_ID_FLAG          = 0x0008
+DEF _CONTINUOUS_PAGING_FLAG      = 0x40000000
+DEF _CONTINUOUS_PAGING_LAST_FLAG = 0x80000000
+
+
+# ---------- read_type using BytesIOReader ----------
+
+cdef object _read_type_br(BytesIOReader reader, dict type_codes_map, object user_type_map,
+                           object ListType, object SetType, object MapType,
+                           object TupleType, object UserType, object CUSTOM_TYPE,
+                           object lookup_casstype, object NotSupportedError):
+    """
+    Cython version of ResultMessage.read_type() operating on BytesIOReader.
+
+    Parameters are passed in to avoid module-level imports from protocol.py
+    (which would create circular dependencies). They are captured once in the
+    closure created by make_recv_results_metadata().
+    """
+    cdef uint16_t optid = read_short(reader)
+
+    typeclass = type_codes_map.get(optid)
+    if typeclass is None:
+        raise NotSupportedError(
+            "Unknown data type code 0x%04x. Have to skip entire result set." % (optid,))
+
+    if typeclass is ListType or typeclass is SetType:
+        subtype = _read_type_br(reader, type_codes_map, user_type_map,
+                                ListType, SetType, MapType, TupleType, UserType,
+                                CUSTOM_TYPE, lookup_casstype, NotSupportedError)
+        typeclass = typeclass.apply_parameters((subtype,))
+    elif typeclass is MapType:
+        keysubtype = _read_type_br(reader, type_codes_map, user_type_map,
+                                   ListType, SetType, MapType, TupleType, UserType,
+                                   CUSTOM_TYPE, lookup_casstype, NotSupportedError)
+        valsubtype = _read_type_br(reader, type_codes_map, user_type_map,
+                                   ListType, SetType, MapType, TupleType, UserType,
+                                   CUSTOM_TYPE, lookup_casstype, NotSupportedError)
+        typeclass = typeclass.apply_parameters((keysubtype, valsubtype))
+    elif typeclass is TupleType:
+        num_items = read_short(reader)
+        types = tuple(_read_type_br(reader, type_codes_map, user_type_map,
+                                    ListType, SetType, MapType, TupleType, UserType,
+                                    CUSTOM_TYPE, lookup_casstype, NotSupportedError)
+                      for _ in range(num_items))
+        typeclass = typeclass.apply_parameters(types)
+    elif typeclass is UserType:
+        ks = read_string_br(reader)
+        udt_name = read_string_br(reader)
+        num_fields = read_short(reader)
+        names_and_types = tuple(
+            (read_string_br(reader),
+             _read_type_br(reader, type_codes_map, user_type_map,
+                           ListType, SetType, MapType, TupleType, UserType,
+                           CUSTOM_TYPE, lookup_casstype, NotSupportedError))
+            for _ in range(num_fields))
+        # zip(*()) with num_fields == 0 raises ValueError (not enough values
+        # to unpack), matching ResultMessage.read_type()'s pure-Python
+        # behavior for the same malformed/degenerate case: a real UDT
+        # always has at least one field, so a UserType with num_fields == 0
+        # on the wire indicates a corrupted or malformed frame and should
+        # raise rather than silently producing an empty-field UDT.
+        names, types = zip(*names_and_types)
+        specialized_type = typeclass.make_udt_class(ks, udt_name, names, types)
+        specialized_type.mapped_class = user_type_map.get(ks, {}).get(udt_name)
+        typeclass = specialized_type
+    elif typeclass is CUSTOM_TYPE:
+        classname = read_string_br(reader)
+        typeclass = lookup_casstype(classname)
+
+    return typeclass
+
+
+# ---------- public factory: creates closures that capture type objects ----------
+
+def make_recv_results_metadata(dict type_codes_map, object CUSTOM_TYPE,
+                                object ListType, object SetType, object MapType,
+                                object TupleType, object UserType,
+                                object lookup_casstype, object NotSupportedError):
+    """
+    Factory that returns a recv_results_metadata function suitable for use
+    as an unbound method replacement on FastResultMessage.
+
+    The closure captures the type-code map and type objects once, so they
+    don't have to be looked up on every call.
+
+    All type objects -- including the NotSupportedError exception class
+    raised for an unrecognized type code -- are passed in by the caller
+    (cython_protocol_handler() in protocol.py, which already has them as
+    module-level names) rather than imported here. cassandra.metadata_parser
+    is itself imported from inside cassandra.protocol's own module body
+    (cython_protocol_handler() runs while "import cassandra.protocol" is
+    still executing), so a module-level
+    "from cassandra.protocol import ResultMessage, CUSTOM_TYPE" here would be a
+    self-referential import back into the partially-initialized protocol
+    module -- the same class of import-order fragility that used to make
+    cassandra.cython_deps.HAVE_CYTHON detection depend on which cassandra
+    module a process happened to import first. It happens to work today only
+    because ResultMessage/CUSTOM_TYPE are defined earlier in protocol.py than
+    cython_protocol_handler() is called; passing them in avoids relying on
+    that ordering at all.
+    """
+    def read_type_br_closure(BytesIOReader reader, user_type_map):
+        return _read_type_br(reader, type_codes_map, user_type_map,
+                             ListType, SetType, MapType, TupleType, UserType,
+                             CUSTOM_TYPE, lookup_casstype, NotSupportedError)
+
+    def recv_results_metadata(self, BytesIOReader reader, user_type_map):
+        """
+        Cython-optimized recv_results_metadata operating on BytesIOReader.
+        Replaces ResultMessage.recv_results_metadata.
+        """
+        cdef int32_t flags = read_int(reader)
+        cdef int32_t colcount = read_int(reader)
+
+        if flags & _HAS_MORE_PAGES_FLAG:
+            self.paging_state = read_binary_longstring_br(reader)
+
+        if flags & _NO_METADATA_FLAG:
+            return
+
+        if flags & _CONTINUOUS_PAGING_FLAG:
+            self.continuous_paging_seq = read_int(reader)
+            self.continuous_paging_last = flags & _CONTINUOUS_PAGING_LAST_FLAG
+
+        if flags & _METADATA_ID_FLAG:
+            self.result_metadata_id = read_binary_string_br(reader)
+
+        cdef str ksname, cfname, colname
+        cdef object coltype
+        cdef int i
+        cdef list column_metadata = [None] * colcount
+
+        if flags & _FLAGS_GLOBAL_TABLES_SPEC:
+            ksname = read_string_br(reader)
+            cfname = read_string_br(reader)
+            for i in range(colcount):
+                colname = read_string_br(reader)
+                coltype = read_type_br_closure(reader, user_type_map)
+                column_metadata[i] = (ksname, cfname, colname, coltype)
+        else:
+            for i in range(colcount):
+                ksname = read_string_br(reader)
+                cfname = read_string_br(reader)
+                colname = read_string_br(reader)
+                coltype = read_type_br_closure(reader, user_type_map)
+                column_metadata[i] = (ksname, cfname, colname, coltype)
+
+        self.column_metadata = column_metadata
+
+    return recv_results_metadata

cassandra/protocol.py

@@ -1264,15 +1264,30 @@ def cython_protocol_handler(colparser):
     The default is to use obj_parser.ListParser
     """
     from cassandra.row_parser import make_recv_results_rows
+    from cassandra.metadata_parser import make_recv_results_metadata
+
+    # Pass already-available module-level names in explicitly instead of
+    # letting metadata_parser import them itself: cassandra.metadata_parser is
+    # imported here, from inside cassandra.protocol's own module body, so a
+    # "from cassandra.protocol import ResultMessage, CUSTOM_TYPE" inside
+    # metadata_parser would be a self-referential import back into this
+    # (still-initializing) module. NotSupportedError is injected the same
+    # way, so metadata_parser never needs to import it from protocol.py at
+    # runtime either.
+    recv_results_metadata_br = make_recv_results_metadata(
+        ResultMessage.type_codes, CUSTOM_TYPE,
+        ListType, SetType, MapType, TupleType, UserType, lookup_casstype,
+        NotSupportedError,
+    )
 
     class FastResultMessage(ResultMessage):
         """
         Cython version of Result Message that has a faster implementation of
-        recv_results_row.
+        recv_results_rows using BytesIOReader for zero-copy metadata + row parsing.
         """
         # type_codes = ResultMessage.type_codes.copy()
         code_to_type = dict((v, k) for k, v in ResultMessage.type_codes.items())
-        recv_results_rows = make_recv_results_rows(colparser)
+        recv_results_rows = make_recv_results_rows(colparser, recv_results_metadata_br)
 
     class CythonProtocolHandler(_ProtocolHandler):
         """

cassandra/row_parser.pyx

@@ -19,13 +19,27 @@ from cassandra.deserializers import make_deserializers
 
 include "ioutils.pyx"
 
-def make_recv_results_rows(ColumnParser colparser):
+def make_recv_results_rows(ColumnParser colparser, recv_results_metadata_br):
+    """
+    Create a recv_results_rows closure that uses:
+      - recv_results_metadata_br: Cython metadata parser operating on BytesIOReader
+      - colparser: Cython column parser for row data
+
+    A single BytesIOReader is created from the full remaining buffer and used
+    for both metadata parsing and row parsing, eliminating the per-read bytes
+    allocation overhead of Python BytesIO.
+    """
     def recv_results_rows(self, f, int protocol_version, user_type_map, result_metadata, column_encryption_policy):
         """
         Parse protocol data given as a BytesIO f into a set of columns (e.g. list of tuples)
         This is used as the recv_results_rows method of (Fast)ResultMessage
         """
-        self.recv_results_metadata(f, user_type_map)
+        # Create ONE BytesIOReader for the entire remaining buffer.
+        # This is used for both metadata parsing and row parsing.
+        reader = BytesIOReader(f.read())
+
+        # Use Cython-optimized metadata parsing on BytesIOReader
+        recv_results_metadata_br(self, reader, user_type_map)
 
         column_metadata = self.column_metadata or result_metadata
 
@@ -35,14 +49,16 @@ def make_recv_results_rows(ColumnParser colparser):
         desc = ParseDesc(self.column_names, self.column_types, column_encryption_policy,
                         [ColDesc(md[0], md[1], md[2]) for md in column_metadata],
                         make_deserializers(self.column_types), protocol_version)
-        reader = BytesIOReader(f.read())
+        # The reader's position is now right after the metadata;
+        # row data follows immediately — no need to create a second reader.
+        # Save position so we can rewind to the start of row data on error.
+        cdef Py_ssize_t rows_start_pos = reader.pos
         try:
             self.parsed_rows = colparser.parse_rows(reader, desc)
         except Exception as e:
             # Use explicitly the TupleRowParser to display better error messages for column decoding failures
             rowparser = TupleRowParser()
-            reader.buf_ptr = reader.buf
-            reader.pos = 0
+            reader.seek(rows_start_pos)
             rowcount = read_int(reader)
             for i in range(rowcount):
                 rowparser.unpack_row(reader, desc)

tests/unit/cython/bytesio_testhelper.pyx

@@ -37,3 +37,21 @@ def test_read_eof():
     with pytest.raises(EOFError):
         reader.read(2)
     reader.read(1) # see that we can still read this
+
+def test_seek():
+    # seek() must reposition the reader so that a subsequent read() lands
+    # at the correct byte offset and returns the correct bytes -- this is
+    # the same repositioning row_parser.pyx relies on when it rewinds the
+    # reader on the recv_results_rows exception-recovery path.
+    cdef BytesIOReader reader = BytesIOReader(b'abcdef')
+    assert reader.read(4)[:4] == b'abcd'  # pos -> 4
+
+    reader.seek(1)
+    assert reader.read(3)[:3] == b'bcd'  # pos -> 4 again, via a different route
+
+    reader.seek(0)
+    assert reader.read(6)[:6] == b'abcdef'  # full rewind still reads correctly
+
+    reader.seek(6)
+    with pytest.raises(EOFError):
+        reader.read(1)