Optimize numeric array conversion and serialize directly into Python bytes

capnp/helpers/serialize.h

@@ -0,0 +1,34 @@
+// Copyright (c) 2026 IQ.Lvbs. All rights reserved.
+#pragma once
+
+#include <Python.h>
+#include <capnp/message.h>
+#include <capnp/serialize.h>
+#include <kj/debug.h>
+#include <cstring>
+
+inline PyObject* messageToPythonBytes(capnp::MessageBuilder& builder) {
+  auto segments = builder.getSegmentsForOutput();
+  auto words = capnp::computeSerializedSizeInWords(segments);
+  KJ_REQUIRE(words <= static_cast<size_t>(PY_SSIZE_T_MAX) / sizeof(capnp::word));
+  PyObject* result = PyBytes_FromStringAndSize(nullptr, words * sizeof(capnp::word));
+  if (result == nullptr) return nullptr;
+  char* data = PyBytes_AS_STRING(result);
+  auto write32 = [data](size_t index, uint32_t value) {
+    for (unsigned int byte = 0; byte < 4; ++byte) {
+      data[index * 4 + byte] = static_cast<char>(value >> (byte * 8));
+    }
+  };
+  write32(0, segments.size() - 1);
+  for (size_t index = 0; index < segments.size(); ++index) {
+    write32(index + 1, segments[index].size());
+  }
+  if (segments.size() % 2 == 0) write32(segments.size() + 1, 0);
+  char* output = data + (segments.size() / 2 + 1) * sizeof(capnp::word);
+  for (auto segment : segments) {
+    auto size = segment.size() * sizeof(capnp::word);
+    std::memcpy(output, segment.begin(), size);
+    output += size;
+  }
+  return result;
+}

capnp/lib/capnp.pyx

@@ -13,6 +13,7 @@ from capnp.helpers.helpers cimport init_capnp_api
 from builtins import memoryview as BuiltinsMemoryview
 from cpython cimport Py_buffer, PyObject_CheckBuffer
 from cpython.buffer cimport PyBUF_SIMPLE, PyBUF_CONTIG_RO
+from cpython.bytes cimport PyBytes_FromStringAndSize, PyBytes_AS_STRING
 from cpython.exc cimport PyErr_Clear
 from cython.operator cimport dereference as deref
 from libc.stdlib cimport malloc, free
@@ -28,6 +29,11 @@ import warnings as _warnings
 
 from types import ModuleType as _ModuleType
 
+cdef object _numpy = None
+
+cdef extern from "capnp/helpers/serialize.h":
+    object messageToPythonBytes(schema_cpp.MessageBuilder&) except +reraise_kj_exception
+
 _CAPNP_VERSION_MAJOR = capnp.CAPNP_VERSION_MAJOR
 _CAPNP_VERSION_MINOR = capnp.CAPNP_VERSION_MINOR
 _CAPNP_VERSION_MICRO = capnp.CAPNP_VERSION_MICRO
@@ -243,6 +249,49 @@ cdef class _DynamicListReader:
         ptr = self.thisptr[index]
         return to_python_reader(ptr, self._parent)
 
+    def to_float64_bytes(self):
+        cdef Py_ssize_t size = self.thisptr.size()
+        cdef bytes result = PyBytes_FromStringAndSize(NULL, size * sizeof(double))
+        cdef char* target = PyBytes_AS_STRING(result)
+        cdef Py_ssize_t index
+        cdef double value
+        cdef C_DynamicValue.Reader item
+        cdef int kind
+        if size:
+            item = self.thisptr[0]
+            kind = item.getType()
+            if kind != capnp.TYPE_FLOAT and kind != capnp.TYPE_INT and kind != capnp.TYPE_UINT:
+                raise TypeError("to_float64_bytes requires a numeric list")
+        for index in range(size):
+            item = self.thisptr[index]
+            value = item.asDouble()
+            memcpy(target + index * sizeof(double), &value, sizeof(double))
+        return result
+
+    def __array__(self, dtype=None, copy=None):
+        global _numpy
+        if _numpy is None:
+            _numpy = __import__("numpy")
+        np = _numpy
+        if copy is False:
+            raise ValueError("Cap'n Proto lists require a copy for NumPy conversion")
+        cdef Py_ssize_t size = self.thisptr.size()
+        cdef Py_ssize_t index
+        cdef C_DynamicValue.Reader item
+        cdef int kind
+        cdef double[::1] values64
+        if size:
+            item = self.thisptr[0]
+            kind = item.getType()
+            if kind != capnp.TYPE_FLOAT:
+                return np.array(list(self), dtype=dtype)
+        result = np.empty(size, dtype=np.float64)
+        values64 = result
+        for index in range(size):
+            item = self.thisptr[index]
+            values64[index] = item.asDouble()
+        return result if dtype is None else result.astype(dtype, copy=False)
+
     def __getitem__(self, int64_t index):
         cdef uint size = self.thisptr.size()
         if index >= size:
@@ -289,6 +338,9 @@ cdef class _DynamicListBuilder:
         ptr = self.thisptr[index]
         return to_python_builder(ptr, self._parent)
 
+    def __array__(self, dtype=None, copy=None):
+        return _DynamicListReader()._init(self.thisptr.asReader(), self._parent).__array__(dtype=dtype, copy=copy)
+
     def __getitem__(self, int64_t index):
         cdef uint size = self.thisptr.size()
         if index >= size:
@@ -888,17 +940,13 @@ cdef class _DynamicStructBuilder:
     cpdef to_bytes(_DynamicStructBuilder self):
         """Returns the struct's containing message as a Python bytes object in the unpacked binary format.
 
-        This is inefficient; it makes several copies.
-
         :rtype: bytes
 
         :Raises: :exc:`KjException` if this isn't the message's root struct.
         """
         self._check_write()
         cdef _MessageBuilder builder = self._parent
-        array = schema_cpp.messageToFlatArray(deref(builder.thisptr))
-        cdef const char* ptr = <const char *>array.begin()
-        cdef bytes ret = ptr[:8*array.size()]
+        cdef bytes ret = messageToPythonBytes(deref(builder.thisptr))
         self._is_written = True
         return ret