perf: Optimize column_encryption_policy checks in recv_results_rows (1.7x speedup for Cython, 1.2x for Python)

cassandra/numpy_parser.pyx

@@ -81,6 +81,10 @@ cdef class NumpyParser(ColumnParser):
         cdef ArrDesc[::1] array_descs
         cdef ArrDesc *arrs
 
+        if desc.column_encryption_policy:
+            raise NotImplementedError(
+                "NumpyParser does not support column encryption")
+
         rowcount = read_int(reader)
         array_descs, arrays = make_arrays(desc, rowcount)
         arrs = &array_descs[0]
@@ -97,7 +101,7 @@ cdef _parse_rows(BytesIOReader reader, ParseDesc desc,
     cdef Py_ssize_t i
 
     for i in range(rowcount):
-        unpack_row(reader, desc, arrs)
+        unpack_plain_row(reader, desc, arrs)
 
 
 ### Helper functions to create NumPy arrays and array descriptors
@@ -144,7 +148,7 @@ def make_array(coltype, array_size):
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-cdef inline int unpack_row(
+cdef inline int unpack_plain_row(
         BytesIOReader reader, ParseDesc desc, ArrDesc *arrays) except -1:
     cdef Buffer buf
     cdef Py_ssize_t i, rowsize = desc.rowsize

cassandra/obj_parser.pyx

@@ -22,6 +22,7 @@ from cassandra.parsing cimport ParseDesc, ColumnParser, RowParser
 from cassandra.tuple cimport tuple_new, tuple_set
 
 from cpython.bytes cimport PyBytes_AsStringAndSize
+cimport cython
 
 
 cdef class ListParser(ColumnParser):
@@ -31,7 +32,10 @@ cdef class ListParser(ColumnParser):
         cdef Py_ssize_t i, rowcount
         rowcount = read_int(reader)
         cdef RowParser rowparser = TupleRowParser()
-        return [rowparser.unpack_row(reader, desc) for i in range(rowcount)]
+        if desc.column_encryption_policy:
+            return [rowparser.unpack_col_encrypted_row(reader, desc) for i in range(rowcount)]
+        else:
+            return [rowparser.unpack_plain_row(reader, desc) for i in range(rowcount)]
 
 
 cdef class LazyParser(ColumnParser):
@@ -47,48 +51,78 @@ def parse_rows_lazy(BytesIOReader reader, ParseDesc desc):
     cdef Py_ssize_t i, rowcount
     rowcount = read_int(reader)
     cdef RowParser rowparser = TupleRowParser()
-    return (rowparser.unpack_row(reader, desc) for i in range(rowcount))
+    if desc.column_encryption_policy:
+        return (rowparser.unpack_col_encrypted_row(reader, desc) for i in range(rowcount))
+    else:
+        return (rowparser.unpack_plain_row(reader, desc) for i in range(rowcount))
 
 
 cdef class TupleRowParser(RowParser):
     """
     Parse a single returned row into a tuple of objects:
 
         (obj1, ..., objN)
+    If CE (Column encryption) policy is enabled - use unpack_col_encrypted_row(),
+    otherwise use unpack_plain_row()
     """
 
-    cpdef unpack_row(self, BytesIOReader reader, ParseDesc desc):
-        assert desc.rowsize >= 0
-
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    cpdef unpack_col_encrypted_row(self, BytesIOReader reader, ParseDesc desc):
         cdef Buffer buf
         cdef Buffer newbuf
         cdef Py_ssize_t i, rowsize = desc.rowsize
         cdef Deserializer deserializer
         cdef tuple res = tuple_new(desc.rowsize)
 
         ce_policy = desc.column_encryption_policy
-        for i in range(rowsize):
-            # Read the next few bytes
-            get_buf(reader, &buf)
-
-            # Deserialize bytes to python object
-            deserializer = desc.deserializers[i]
-            coldesc = desc.coldescs[i]
-            uses_ce = ce_policy and ce_policy.contains_column(coldesc)
-            try:
+        try:
+            for i in range(rowsize):
+                # Read the next few bytes
+                get_buf(reader, &buf)
+
+                # Deserialize bytes to python object
+                deserializer = desc.deserializers[i]
+                coldesc = desc.coldescs[i]
+                uses_ce = ce_policy.contains_column(coldesc)
                 if uses_ce:
                     col_type = ce_policy.column_type(coldesc)
                     decrypted_bytes = ce_policy.decrypt(coldesc, to_bytes(&buf))
                     PyBytes_AsStringAndSize(decrypted_bytes, &newbuf.ptr, &newbuf.size)
-                    deserializer = find_deserializer(ce_policy.column_type(coldesc))
+                    deserializer = find_deserializer(col_type)
                     val = from_binary(deserializer, &newbuf, desc.protocol_version)
                 else:
                     val = from_binary(deserializer, &buf, desc.protocol_version)
-            except Exception as e:
-                raise DriverException('Failed decoding result column "%s" of type %s: %s' % (desc.colnames[i],
-                                                                                             desc.coltypes[i].cql_parameterized_type(),
-                                                                                             str(e)))
-            # Insert new object into tuple
-            tuple_set(res, i, val)
+                # Insert new object into tuple
+                tuple_set(res, i, val)
+        except Exception as e:
+            raise DriverException('Failed decoding result column "%s" of type %s: %s' % (desc.colnames[i],
+                                                                                         desc.coltypes[i].cql_parameterized_type(),
+                                                                                         str(e)))
+
+        return res
+
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    cpdef unpack_plain_row(self, BytesIOReader reader, ParseDesc desc):
+        cdef Buffer buf
+        cdef Py_ssize_t i, rowsize = desc.rowsize
+        cdef Deserializer deserializer
+        cdef tuple res = tuple_new(desc.rowsize)
+
+        try:
+            for i in range(rowsize):
+                # Read the next few bytes
+                get_buf(reader, &buf)
+
+                # Deserialize bytes to python object
+                deserializer = desc.deserializers[i]
+                val = from_binary(deserializer, &buf, desc.protocol_version)
+                # Insert new object into tuple
+                tuple_set(res, i, val)
+        except Exception as e:
+            raise DriverException('Failed decoding result column "%s" of type %s: %s' % (desc.colnames[i],
+                                                                                         desc.coltypes[i].cql_parameterized_type(),
+                                                                                         str(e)))
 
         return res

cassandra/parsing.pxd

@@ -28,5 +28,6 @@ cdef class ColumnParser:
     cpdef parse_rows(self, BytesIOReader reader, ParseDesc desc)
 
 cdef class RowParser:
-    cpdef unpack_row(self, BytesIOReader reader, ParseDesc desc)
+    cpdef unpack_plain_row(self, BytesIOReader reader, ParseDesc desc)
+    cpdef unpack_col_encrypted_row(self, BytesIOReader reader, ParseDesc desc)
 

cassandra/parsing.pyx

@@ -20,6 +20,18 @@ cdef class ParseDesc:
     """Description of what structure to parse"""
 
     def __init__(self, colnames, coltypes, column_encryption_policy, coldescs, deserializers, protocol_version):
+        if len(deserializers) != len(colnames):
+            # The row parsers (obj_parser.pyx TupleRowParser.unpack_plain_row /
+            # unpack_col_encrypted_row) index into `deserializers` with
+            # @cython.boundscheck(False), bounded by rowsize == len(colnames).
+            # A length mismatch here would turn into an out-of-bounds memory
+            # read at parse time instead of a clean, immediate error, so this
+            # invariant is validated once at construction time. Use a real
+            # exception (not `assert`) so the guard cannot be stripped by
+            # running Python with optimizations enabled (-O).
+            raise ValueError(
+                "deserializers must have the same length as colnames "
+                "(got %d deserializers for %d columns)" % (len(deserializers), len(colnames)))
         self.colnames = colnames
         self.coltypes = coltypes
         self.column_encryption_policy = column_encryption_policy
@@ -39,8 +51,14 @@ cdef class ColumnParser:
 cdef class RowParser:
     """Parser for a single row"""
 
-    cpdef unpack_row(self, BytesIOReader reader, ParseDesc desc):
+    cpdef unpack_plain_row(self, BytesIOReader reader, ParseDesc desc):
         """
         Unpack a single row of data in a ResultMessage.
         """
         raise NotImplementedError
+
+    cpdef unpack_col_encrypted_row(self, BytesIOReader reader, ParseDesc desc):
+        """
+        Unpack a single row of data in a ResultMessage, with column encryption support.
+        """
+        raise NotImplementedError

cassandra/protocol.py

@@ -766,28 +766,46 @@ def recv_results_rows(self, f, protocol_version, user_type_map, result_metadata,
         rows = [self.recv_row(f, len(column_metadata)) for _ in range(rowcount)]
         self.column_names = [c[2] for c in column_metadata]
         self.column_types = [c[3] for c in column_metadata]
-        col_descs = [ColDesc(md[0], md[1], md[2]) for md in column_metadata]
 
-        def decode_val(val, col_md, col_desc):
-            uses_ce = column_encryption_policy and column_encryption_policy.contains_column(col_desc)
-            col_type = column_encryption_policy.column_type(col_desc) if uses_ce else col_md[3]
-            raw_bytes = column_encryption_policy.decrypt(col_desc, val) if uses_ce else val
-            return col_type.from_binary(raw_bytes, protocol_version)
+        if column_encryption_policy:
+            col_descs = [ColDesc(md[0], md[1], md[2]) for md in column_metadata]
 
-        def decode_row(row):
-            return tuple(decode_val(val, col_md, col_desc) for val, col_md, col_desc in zip(row, column_metadata, col_descs))
+            def decode_val(val, col_md, col_desc):
+                uses_ce = column_encryption_policy.contains_column(col_desc)
+                if uses_ce:
+                    col_type = column_encryption_policy.column_type(col_desc)
+                    raw_bytes = column_encryption_policy.decrypt(col_desc, val)
+                    return col_type.from_binary(raw_bytes, protocol_version)
+                else:
+                    return col_md[3].from_binary(val, protocol_version)
+
+            def decode_row(row):
+                return tuple(decode_val(val, col_md, col_desc) for val, col_md, col_desc in zip(row, column_metadata, col_descs))
+        else:
+            def decode_row(row):
+                return tuple(col_md[3].from_binary(val, protocol_version) for val, col_md in zip(row, column_metadata))
 
         try:
             self.parsed_rows = [decode_row(row) for row in rows]
         except Exception:
-            for row in rows:
-                for val, col_md, col_desc in zip(row, column_metadata, col_descs):
-                    try:
-                        decode_val(val, col_md, col_desc)
-                    except Exception as e:
-                        raise DriverException('Failed decoding result column "%s" of type %s: %s' % (col_md[2],
-                                                                                                     col_md[3].cql_parameterized_type(),
-                                                                                                     str(e)))
+            if column_encryption_policy:
+                for row in rows:
+                    for val, col_md, col_desc in zip(row, column_metadata, col_descs):
+                        try:
+                            decode_val(val, col_md, col_desc)
+                        except Exception as e:
+                            raise DriverException('Failed decoding result column "%s" of type %s: %s' % (col_md[2],
+                                                                                                         col_md[3].cql_parameterized_type(),
+                                                                                                         str(e)))
+            else:
+                for row in rows:
+                    for val, col_md in zip(row, column_metadata):
+                        try:
+                            col_md[3].from_binary(val, protocol_version)
+                        except Exception as e:
+                            raise DriverException('Failed decoding result column "%s" of type %s: %s' % (col_md[2],
+                                                                                                         col_md[3].cql_parameterized_type(),
+                                                                                                         str(e)))
 
     def recv_results_prepared(self, f, protocol_version, protocol_features, user_type_map):
         self.query_id = read_binary_string(f)

cassandra/query.py

@@ -684,6 +684,23 @@ def bind(self, values):
 
         self.raw_values = values
         self.values = []
+
+        # The policy existence check happens once here (not per-value below)
+        # to select how column bytes get serialized; the per-value loop
+        # itself stays identical for both branches.
+        if ce_policy:
+            def _serialize(value, col_spec):
+                col_desc = ColDesc(col_spec.keyspace_name, col_spec.table_name, col_spec.name)
+                if ce_policy.contains_column(col_desc):
+                    col_type = ce_policy.column_type(col_desc)
+                    col_bytes = col_type.serialize(value, proto_version)
+                    return ce_policy.encrypt(col_desc, col_bytes)
+                else:
+                    return col_spec.type.serialize(value, proto_version)
+        else:
+            def _serialize(value, col_spec):
+                return col_spec.type.serialize(value, proto_version)
+
         for value, col_spec in zip(values, col_meta):
             if value is None:
                 self.values.append(None)
@@ -694,13 +711,7 @@ def bind(self, values):
                     raise ValueError("Attempt to bind UNSET_VALUE while using unsuitable protocol version (%d < 4)" % proto_version)
             else:
                 try:
-                    col_desc = ColDesc(col_spec.keyspace_name, col_spec.table_name, col_spec.name)
-                    uses_ce = ce_policy and ce_policy.contains_column(col_desc)
-                    col_type = ce_policy.column_type(col_desc) if uses_ce else col_spec.type
-                    col_bytes = col_type.serialize(value, proto_version)
-                    if uses_ce:
-                        col_bytes = ce_policy.encrypt(col_desc, col_bytes)
-                    self.values.append(col_bytes)
+                    self.values.append(_serialize(value, col_spec))
                 except (TypeError, struct.error) as exc:
                     actual_type = type(value)
                     message = ('Received an argument of invalid type for column "%s". '

cassandra/row_parser.pyx

@@ -32,19 +32,30 @@ def make_recv_results_rows(ColumnParser colparser):
         self.column_names = [md[2] for md in column_metadata]
         self.column_types = [md[3] for md in column_metadata]
 
+        coldescs = [ColDesc(md[0], md[1], md[2]) for md in column_metadata] if column_encryption_policy else None
         desc = ParseDesc(self.column_names, self.column_types, column_encryption_policy,
-                        [ColDesc(md[0], md[1], md[2]) for md in column_metadata],
+                        coldescs,
                         make_deserializers(self.column_types), protocol_version)
         reader = BytesIOReader(f.read())
         try:
             self.parsed_rows = colparser.parse_rows(reader, desc)
+        except NotImplementedError:
+            # e.g. NumpyParser does not support column encryption. This
+            # signals an unsupported configuration, not a decoding failure,
+            # so it must propagate to the caller instead of being silently
+            # swallowed by the TupleRowParser fallback below.
+            raise
         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
             rowcount = read_int(reader)
-            for i in range(rowcount):
-                rowparser.unpack_row(reader, desc)
+            if desc.column_encryption_policy:
+                for i in range(rowcount):
+                    rowparser.unpack_col_encrypted_row(reader, desc)
+            else:
+                for i in range(rowcount):
+                    rowparser.unpack_plain_row(reader, desc)
 
     return recv_results_rows