Use Cython-based deserializers whenever available

markflorisson · markflorisson · commit 7ce8e3a3c174 · 2015-08-06T21:45:49.000+01:00
diff --git a/cassandra/cython_deps.py b/cassandra/cython_deps.py
@@ -0,0 +1,5 @@
+try:
+    from cassandra.rowparser import make_recv_results_rows
+    HAVE_CYTHON = True
+except ImportError:
+    HAVE_CYTHON = False
diff --git a/cassandra/protocol.py b/cassandra/protocol.py
@@ -40,6 +40,7 @@
                                 TupleType, lookup_casstype, SimpleDateType,
                                 TimeType, ByteType, ShortType)
 from cassandra.policies import WriteType
+from cassandra.cython_deps import HAVE_CYTHON
 from cassandra import util
 
 log = logging.getLogger(__name__)
@@ -69,10 +70,16 @@ class InternalError(Exception):
 
 _UNSET_VALUE = object()
 
+def register_class(cls):
+    _message_types_by_opcode[cls.opcode] = cls
+
+def get_registered_classes():
+    return _message_types_by_opcode.copy()
+
 class _RegisterMessageType(type):
     def __init__(cls, name, bases, dct):
         if not name.startswith('_'):
-            _message_types_by_opcode[cls.opcode] = cls
+            register_class(cls)
 
 
 @six.add_metaclass(_RegisterMessageType)
@@ -987,6 +994,62 @@ def decode_message(cls, protocol_version, user_type_map, stream_id, flags, opcod
         return msg
 
 
+def cython_protocol_handler(colparser):
+    """
+    Given a column parser to deserialize ResultMessages, return a suitable
+    Cython-based protocol handler.
+
+    There are three Cython-based protocol handlers (least to most performant):
+
+        1. objparser.ListParser
+            this parser decodes result messages into a list of tuples
+
+        2. objparser.LazyParser
+            this parser decodes result messages lazily by returning an iterator
+
+        3. numpyparser.NumPyParser
+            this parser decodes result messages into NumPy arrays
+
+    The default is to use objparser.ListParser
+    """
+    # TODO: It may be cleaner to turn ProtocolHandler and ResultMessage into
+    # TODO:     instances and use methods instead of class methods
+    from cassandra.rowparser import make_recv_results_rows
+
+    class FastResultMessage(ResultMessage):
+        """
+        Cython version of Result Message that has a faster implementation of
+        recv_results_row.
+        """
+        # type_codes = ResultMessage.type_codes.copy()
+        code_to_type = dict((v, k) for k, v in ResultMessage.type_codes.items())
+        recv_results_rows = classmethod(make_recv_results_rows(colparser))
+
+    class CythonProtocolHandler(ProtocolHandler):
+        """
+        Use FastResultMessage to decode query result message messages.
+        """
+
+        my_opcodes = ProtocolHandler.message_types_by_opcode.copy()
+        my_opcodes[FastResultMessage.opcode] = FastResultMessage
+        message_types_by_opcode = my_opcodes
+
+    return CythonProtocolHandler
+
+
+if HAVE_CYTHON:
+    from cassandra.objparser import ListParser, LazyParser
+    from cassandra.numpyparser import NumpyParser
+
+    ProtocolHandler = cython_protocol_handler(ListParser())
+    LazyProtocolHandler = cython_protocol_handler(LazyParser())
+    NumpyProtocolHandler = cython_protocol_handler(NumpyParser())
+else:
+    # Use Python-based ProtocolHandler
+    LazyProtocolHandler = None
+    NumpyProtocolHandler = None
+
+
 def read_byte(f):
     return int8_unpack(f.read(1))
 
diff --git a/cassandra/rowparser.pyx b/cassandra/rowparser.pyx
@@ -0,0 +1,26 @@
+# -- cython: profile=True
+
+from cassandra.parsing cimport ParseDesc, ColumnParser
+from cassandra.deserializers import make_deserializers
+
+include "ioutils.pyx"
+
+def make_recv_results_rows(ColumnParser colparser):
+    def recv_results_rows(cls, f, int protocol_version, user_type_map):
+        """
+        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
+        """
+        paging_state, column_metadata = cls.recv_results_metadata(f, user_type_map)
+
+        colnames = [c[2] for c in column_metadata]
+        coltypes = [c[3] for c in column_metadata]
+
+        desc = ParseDesc(colnames, coltypes, make_deserializers(coltypes),
+                         protocol_version)
+        reader = BytesIOReader(f.read())
+        parsed_rows = colparser.parse_rows(reader, desc)
+
+        return (paging_state, (colnames, parsed_rows))
+
+    return recv_results_rows
diff --git a/tests/integration/standard/test_cython_protocol_handlers.py b/tests/integration/standard/test_cython_protocol_handlers.py
@@ -8,19 +8,15 @@
     import unittest
 
 from cassandra.cluster import Cluster
+from cassandra.protocol import ProtocolHandler, LazyProtocolHandler, NumpyProtocolHandler
 from tests.integration import use_singledc, PROTOCOL_VERSION
 from tests.integration.datatype_utils import update_datatypes
 from tests.integration.standard.utils import create_table_with_all_types, get_all_primitive_params
-from six import next
 
-try:
-    from cassandra.cython_protocol_handler import make_protocol_handler
-except ImportError as e:
+from cassandra.cython_deps import HAVE_CYTHON
+if not HAVE_CYTHON:
     raise unittest.SkipTest("Skipping test, not compiled with Cython enabled")
 
-from cassandra.numpyparser import NumpyParser
-from cassandra.objparser import ListParser, LazyParser
-
 
 def setup_module():
     use_singledc()
@@ -47,20 +43,20 @@ def test_cython_parser(self):
         """
         Test Cython-based parser that returns a list of tuples
         """
-        self.cython_parser(ListParser())
+        self.cython_parser(ProtocolHandler)
 
     def test_cython_lazy_parser(self):
         """
         Test Cython-based parser that returns a list of tuples
         """
-        self.cython_parser(LazyParser())
+        self.cython_parser(LazyProtocolHandler)
 
-    def cython_parser(self, colparser):
+    def cython_parser(self, protocol_handler):
         cluster = Cluster(protocol_version=PROTOCOL_VERSION)
         session = cluster.connect(keyspace="testspace")
 
         # use our custom protocol handler
-        session.client_protocol_handler = make_protocol_handler(colparser)
+        session.client_protocol_handler = protocol_handler
         # session.row_factory = tuple_factory
 
         # verify data
