Speed up Cursor.fetchall by draining batches directly

pynuodb/cursor.py

@@ -184,49 +184,66 @@ def executemany(self, operation, seq_of_parameters):
     def fetchone(self):
         # type: () -> Optional[result_set.Row]
         """Return the next row of results from the previous SQL operation."""
-        self._check_closed()
+        # Inline _check_closed to avoid per-row function-call overhead.
+        if self.closed:
+            raise Error("cursor is closed")
+        if self.session.closed:
+            raise Error("connection is closed")
         if self._result_set is None:
             raise Error("Previous execute did not produce any results or no call was issued yet")
         self.rownumber += 1
         if not self._result_set.is_complete():
             self.session.fetch_result_set_next(self._result_set)
         return self._result_set.fetchone()
 
+    def _drain(self, limit):
+        # type: (Optional[int]) -> List[result_set.Row]
+        """Drain rows from the current result set.
+
+        :param limit: maximum rows to return; None means "no cap".
+        """
+        self._check_closed()
+        if self._result_set is None:
+            raise Error("Previous execute did not produce any results or no call was issued yet")
+
+        fetched_rows = []  # type: List[result_set.Row]
+        remaining = limit  # None == unlimited
+        while remaining is None or remaining > 0:
+            # Drain the current in-memory batch in one shot instead of
+            # calling fetchone() per row.
+            idx = self._result_set.results_idx
+            batch = self._result_set.results
+            avail = len(batch) - idx
+            if avail > 0:
+                if remaining is None:
+                    take = avail
+                else:
+                    take = avail if avail < remaining else remaining
+                    remaining -= take
+                fetched_rows.extend(batch[idx:idx + take])
+                self._result_set.results_idx = idx + take
+                if remaining == 0:
+                    break
+            if self._result_set.complete:
+                break
+            self.session.fetch_result_set_next(self._result_set)
+        self.rownumber += len(fetched_rows)
+        return fetched_rows
+
     def fetchmany(self, size=None):
         # type: (Optional[int]) -> List[result_set.Row]
         """Return SIZE rows from the previous SQL operation.
 
         If size is None, uses the default size for this Cursor.
         """
-        self._check_closed()
-
         if size is None:
             size = self.arraysize
-
-        fetched_rows = []
-        num_fetched_rows = 0
-        while num_fetched_rows < size:
-            row = self.fetchone()
-            if row is None:
-                break
-            else:
-                fetched_rows.append(row)
-                num_fetched_rows += 1
-        return fetched_rows
+        return self._drain(size)
 
     def fetchall(self):
         # type: () -> List[result_set.Row]
         """Return all rows generated by the previous SQL operation."""
-        self._check_closed()
-
-        fetched_rows = []
-        while True:
-            row = self.fetchone()
-            if row is None:
-                break
-            else:
-                fetched_rows.append(row)
-        return fetched_rows
+        return self._drain(None)
 
     def nextset(self):  # pylint: disable=no-self-use
         # type: () -> None