Passing C++ vector to Numpy through Cython without copying and taking care of memory management automatically

Dealing with processing large matrices (NxM with 1K <= N <= 20K & 10K <= M <= 200K), I often need to pass Numpy matrices to C++ through Cython to get the job done and this works as expected & without copying.

However, there are times when I need to initiate and preprocess a matrix in C++ and pass it to Numpy (Python 3.6). Let's assume the matrices are linearized (so the size is N*M and it's a 1D matrix - col/row major doesn't matter here). Following the information in here: exposing C-computed arrays in Python without data copies & modifying it for C++ compatibility, I'm able to pass C++ array.

The problem is if I want to use std vector instead of initiating array, I'd get Segmentation fault. For example, considering the following files:

fast.h

#include <iostream>
#include <vector>

using std::cout; using std::endl; using std::vector;
int* doit(int length);

fast.cpp

#include "fast.h"
int* doit(int length) {
    // Something really heavy
    cout << "C++: doing it fast " << endl; 

    vector<int> WhyNot;

    // Heavy stuff - like reading a big file and preprocessing it
    for(int i=0; i<length; ++i)
        WhyNot.push_back(i); // heavy stuff

    cout << "C++: did it really fast" << endl;
    return &WhyNot[0]; // or WhyNot.data()
}

faster.pyx

cimport numpy as np
import numpy as np
from libc.stdlib cimport free
from cpython cimport PyObject, Py_INCREF

np.import_array()

cdef extern from "fast.h":
    int* doit(int length)

cdef class ArrayWrapper:
    cdef void* data_ptr
    cdef int size

    cdef set_data(self, int size, void* data_ptr):
        self.data_ptr = data_ptr
        self.size = size

    def __array__(self):
        print ("Cython: __array__ called")
        cdef np.npy_intp shape[1]
        shape[0] = <np.npy_intp> self.size
        ndarray = np.PyArray_SimpleNewFromData(1, shape,
                                               np.NPY_INT, self.data_ptr)
        print ("Cython: __array__ done")
        return ndarray

    def __dealloc__(self):
        print("Cython: __dealloc__ called")
        free(<void*>self.data_ptr)
        print("Cython: __dealloc__ done")


def faster(length):
    print("Cython: calling C++ function to do it")
    cdef int *array = doit(length)
    print("Cython: back from C++")
    cdef np.ndarray ndarray
    array_wrapper = ArrayWrapper()
    array_wrapper.set_data(length, <void*> array)
    print("Ctyhon: array wrapper set")
    ndarray = np.array(array_wrapper, copy=False)
    ndarray.base = <PyObject*> array_wrapper
    Py_INCREF(array_wrapper)
    print("Cython: all done - returning")
    return ndarray 

setup.py

from distutils.core import setup
from distutils.extension import Extension
from Cython.Distutils import build_ext
import numpy 

ext_modules = [Extension(
    "faster", 
    ["faster.pyx", "fast.cpp"], 
    language='c++',
    extra_compile_args=["-std=c++11"],
    extra_link_args=["-std=c++11"]
)]

setup(
    cmdclass = {'build_ext': build_ext}, 
    ext_modules = ext_modules,
    include_dirs=[numpy.get_include()]
)

If you build this with

python setup.py build_ext --inplace

and run Python 3.6 interpreter, if you enter the following you'd get seg fault after a couple of tries.

>>> from faster import faster
>>> a = faster(1000000)
Cython: calling C++ function to do it
C++: doing it fast
C++: did it really fast
Cython: back from C++
Ctyhon: array wrapper set
Cython: __array__ called
Cython: __array__ done
Cython: all done - returning
>>> a = faster(1000000)
Cython: calling C++ function to do it
C++: doing it fast
C++: did it really fast
Cython: back from C++
Ctyhon: array wrapper set
Cython: __array__ called
Cython: __array__ done
Cython: all done - returning
Cython: __dealloc__ called
Segmentation fault (core dumped)

Couple of things to note:

• If you use array instead of vector (in fast.cpp) this would work like a charm!
• If you call faster(1000000) and put the result into something other than variable a this would work.

If you enter smaller number like faster(10) you'd get a more detailed info like:

Cython: calling C++ function to do it
C++: doing it fast
C++: did it really fast
Cython: back from C++
Ctyhon: array wrapper set
Cython: __array__ called
Cython: __array__ done
Cython: all done - returning
Cython: __dealloc__ called <--- Perhaps this happened too early or late?
*** Error in 'python': double free or corruption (fasttop): 0x0000000001365570 ***
======= Backtrace: =========
More info here ....

It's really puzzling that why this doesn't happen with arrays? No matter what!

I make use of vectors a lot and would love to be able to use them in these scenarios.