Optimize column_encryption_policy checks in Cython parsers

Split TupleRowParser.unpack_row() into unpack_plain_row() and
unpack_col_encrypted_row(), branching once in the callers (ListParser,
LazyParser, row_parser error path) rather than per value.

- Add boundscheck/wraparound Cython optimizations on both methods
- Extract try/except outside the per-column loop
- NumpyParser.parse_rows() raises NotImplementedError when column
  encryption is configured, since it has no decryption support
- Update parsing.pxd declarations to match

Author: mykaul

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

@@ -39,8 +39,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/row_parser.pyx

@@ -44,7 +44,11 @@ def make_recv_results_rows(ColumnParser colparser):
             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