]> git.ipfire.org Git - thirdparty/psycopg.git/commitdiff
Added Transformer.load_rows()
authorDaniele Varrazzo <daniele.varrazzo@gmail.com>
Sat, 26 Dec 2020 23:28:29 +0000 (00:28 +0100)
committerDaniele Varrazzo <daniele.varrazzo@gmail.com>
Sun, 27 Dec 2020 16:14:47 +0000 (17:14 +0100)
There is some good optimisation to be had moving in C the two loops.

psycopg3/psycopg3/_transform.py
psycopg3/psycopg3/cursor.py
psycopg3/psycopg3/proto.py
psycopg3_c/psycopg3_c/_psycopg3.pyi
psycopg3_c/psycopg3_c/transform.pyx

index 385ad8e81e7177262787751c267d7876100666ad..11a70c0c8afecbe695f5855530e9daafc5d94a34 100644 (file)
@@ -130,6 +130,28 @@ class Transformer(AdaptContext):
             f" to format {Format(format).name}"
         )
 
+    def load_rows(self, row0: int, row1: int) -> Sequence[Tuple[Any, ...]]:
+        res = self._pgresult
+        if not res:
+            raise e.InterfaceError("result not set")
+
+        if not (0 <= row0 <= self._ntuples and 0 <= row1 <= self._ntuples):
+            raise e.InterfaceError(
+                f"rows must be included between 0 and {self._ntuples}"
+            )
+
+        records: List[Tuple[Any, ...]]
+        records = [None] * (row1 - row0)  # type: ignore[list-item]
+        for row in range(row0, row1):
+            record: List[Any] = [None] * self._nfields
+            for col in range(self._nfields):
+                val = res.get_value(row, col)
+                if val is not None:
+                    record[col] = self._row_loaders[col](val)
+            records[row - row0] = tuple(record)
+
+        return records
+
     def load_row(self, row: int) -> Optional[Tuple[Any, ...]]:
         res = self._pgresult
         if not res:
@@ -138,13 +160,13 @@ class Transformer(AdaptContext):
         if not 0 <= row < self._ntuples:
             return None
 
-        rv: List[Any] = [None] * self._nfields
+        record: List[Any] = [None] * self._nfields
         for col in range(self._nfields):
             val = res.get_value(row, col)
             if val is not None:
-                rv[col] = self._row_loaders[col](val)
+                record[col] = self._row_loaders[col](val)
 
-        return tuple(rv)
+        return tuple(record)
 
     def load_sequence(
         self, record: Sequence[Optional[bytes]]
index 100fa513e2a5a99157d89d675bb132e042e103e3..bfb1a0086e0f77f3eb0f5d5a05704f17f38153d0 100644 (file)
@@ -427,10 +427,10 @@ class Cursor(BaseCursor["Connection"]):
         Return `!None` the recordset is finished.
         """
         self._check_result()
-        rv = self._transformer.load_row(self._pos)
-        if rv is not None:
+        record = self._transformer.load_row(self._pos)
+        if record is not None:
             self._pos += 1
-        return rv
+        return record
 
     def fetchmany(self, size: int = 0) -> List[Sequence[Any]]:
         """
@@ -439,34 +439,25 @@ class Cursor(BaseCursor["Connection"]):
         *size* default to `!self.arraysize` if not specified.
         """
         self._check_result()
-        if not size:
-            size = self.arraysize
-
         assert self.pgresult
-        load = self._transformer.load_row
-        rv: List[Any] = [None] * (min(size, self.pgresult.ntuples - self._pos))
-
-        for i in range(len(rv)):
-            rv[i] = load(i + self._pos)
 
-        self._pos += len(rv)
-        return rv
+        if not size:
+            size = self.arraysize
+        records = self._transformer.load_rows(
+            self._pos, min(self._pos + size, self.pgresult.ntuples)
+        )
+        self._pos += len(records)
+        return records  # type: ignore[return-value]
 
     def fetchall(self) -> List[Sequence[Any]]:
         """
         Return all the remaining records from the current recordset.
         """
         self._check_result()
-
         assert self.pgresult
-        load = self._transformer.load_row
-
-        rv: List[Any] = [None] * (self.pgresult.ntuples - self._pos)
-        for i in range(len(rv)):
-            rv[i] = load(i + self._pos)
-
-        self._pos += len(rv)
-        return rv
+        records = self._transformer.load_rows(self._pos, self.pgresult.ntuples)
+        self._pos += self.pgresult.ntuples
+        return records  # type: ignore[return-value]
 
     def __iter__(self) -> Iterator[Sequence[Any]]:
         self._check_result()
@@ -537,31 +528,22 @@ class AsyncCursor(BaseCursor["AsyncConnection"]):
 
     async def fetchmany(self, size: int = 0) -> List[Sequence[Any]]:
         self._check_result()
-        if not size:
-            size = self.arraysize
-
         assert self.pgresult
-        load = self._transformer.load_row
-        rv: List[Any] = [None] * (min(size, self.pgresult.ntuples - self._pos))
-
-        for i in range(len(rv)):
-            rv[i] = load(i + self._pos)
 
-        self._pos += len(rv)
-        return rv
+        if not size:
+            size = self.arraysize
+        records = self._transformer.load_rows(
+            self._pos, min(self._pos + size, self.pgresult.ntuples)
+        )
+        self._pos += len(records)
+        return records  # type: ignore[return-value]
 
     async def fetchall(self) -> List[Sequence[Any]]:
         self._check_result()
-
         assert self.pgresult
-        load = self._transformer.load_row
-
-        rv: List[Any] = [None] * (self.pgresult.ntuples - self._pos)
-        for i in range(len(rv)):
-            rv[i] = load(i + self._pos)
-
-        self._pos += len(rv)
-        return rv
+        records = self._transformer.load_rows(self._pos, self.pgresult.ntuples)
+        self._pos += self.pgresult.ntuples
+        return records  # type: ignore[return-value]
 
     async def __aiter__(self) -> AsyncIterator[Sequence[Any]]:
         self._check_result()
index 29a29cd5e69dadc164d08703bf1d079cf2fb65c1..8cb92e5889a26aa9df68dfbb88dbb377476b5ee0 100644 (file)
@@ -98,6 +98,9 @@ class Transformer(Protocol):
     def get_dumper(self, obj: Any, format: Format) -> "Dumper":
         ...
 
+    def load_rows(self, row0: int, row1: int) -> Sequence[Tuple[Any, ...]]:
+        ...
+
     def load_row(self, row: int) -> Optional[Tuple[Any, ...]]:
         ...
 
index 146436bca6e28657992b7512fc22678653e53ea8..d7f0840d78ee82894175e38ad8d5ada2567a8734 100644 (file)
@@ -29,6 +29,7 @@ class Transformer(AdaptContext):
         self, types: Sequence[int], formats: Sequence[Format]
     ) -> None: ...
     def get_dumper(self, obj: Any, format: Format) -> Dumper: ...
+    def load_rows(self, row0: int, row1: int) -> Sequence[Tuple[Any, ...]]: ...
     def load_row(self, row: int) -> Optional[Tuple[Any, ...]]: ...
     def load_sequence(
         self, record: Sequence[Optional[bytes]]
index 7671dc926b709bee5b9d31a0084ad691c98cb744..d46a46a0275cd0323cede81cfc2cea75a0b7a337 100644 (file)
@@ -180,6 +180,58 @@ cdef class Transformer:
             f" to format {Format(format).name}"
         )
 
+    def load_rows(self, row0: int, row1: int) -> Sequence[Tuple[Any, ...]]:
+        if self._pgresult is None:
+            raise e.InterfaceError("result not set")
+
+        cdef int crow0 = row0
+        cdef int crow1 = row1
+        if not (0 <= row0 <= self._ntuples and 0 <= row1 <= self._ntuples):
+            raise e.InterfaceError(
+                f"rows must be included between 0 and {self._ntuples}"
+            )
+
+        cdef libpq.PGresult *res = self._pgresult.pgresult_ptr
+        # cheeky access to the internal PGresult structure
+        cdef pg_result_int *ires = <pg_result_int*>res
+
+        cdef RowLoader loader
+        cdef int row
+        cdef int col
+        cdef int length
+        cdef PGresAttValue *attval
+        cdef const char *val
+        cdef tuple record
+        cdef list records = [None] * (row1 - row0)
+
+        for row in range(row0, row1):
+            record = PyTuple_New(self._nfields)
+            for col in range(self._nfields):
+                attval = &(ires.tuples[row][col])
+                length = attval.len
+                if length == -1:  # NULL_LEN
+                    Py_INCREF(None)
+                    PyTuple_SET_ITEM(record, col, None)
+                    continue
+
+                # TODO: the is some visible python churn around this lookup.
+                # replace with a C array of borrowed references pointing to
+                # the cloader.cload function pointer
+                loader = self._row_loaders[col]
+                val = attval.value
+                if loader.cloader is not None:
+                    pyval = loader.cloader.cload(val, length)
+                else:
+                    # TODO: no copy
+                    pyval = loader.pyloader(val[:length])
+
+                Py_INCREF(pyval)
+                PyTuple_SET_ITEM(record, col, pyval)
+
+            records[row - row0] = record
+
+        return records
+
     def load_row(self, row: int) -> Optional[Tuple[Any, ...]]:
         if self._pgresult is None:
             return None
@@ -189,23 +241,23 @@ cdef class Transformer:
             return None
 
         cdef libpq.PGresult *res = self._pgresult.pgresult_ptr
+        # cheeky access to the internal PGresult structure
+        cdef pg_result_int *ires = <pg_result_int*>res
 
         cdef RowLoader loader
         cdef int col
         cdef int length
-        cdef const char *val
-
-        # cheeky access to the internal PGresult structure
-        cdef pg_result_int *ires = <pg_result_int*>res
         cdef PGresAttValue *attval
+        cdef const char *val
+        cdef tuple record
 
-        rv = PyTuple_New(self._nfields)
+        record = PyTuple_New(self._nfields)
         for col in range(self._nfields):
             attval = &(ires.tuples[crow][col])
             length = attval.len
             if length == -1:  # NULL_LEN
                 Py_INCREF(None)
-                PyTuple_SET_ITEM(rv, col, None)
+                PyTuple_SET_ITEM(record, col, None)
                 continue
 
             # TODO: the is some visible python churn around this lookup.
@@ -220,9 +272,9 @@ cdef class Transformer:
                 pyval = loader.pyloader(val[:length])
 
             Py_INCREF(pyval)
-            PyTuple_SET_ITEM(rv, col, pyval)
+            PyTuple_SET_ITEM(record, col, pyval)
 
-        return rv
+        return record
 
     def load_sequence(self, record: Sequence[Optional[bytes]]) -> Tuple[Any, ...]:
         cdef int length = len(record)