<commit_before>/*
 * This file wraps the C++11 Mrpt code to an extension module compatible with Python 3.
 */

#include "Python.h"
#include <cstdio>
#include <cstdlib>
#include <numeric>
#include <string>
#include <vector>
#include <memory>
#include <sys/types.h>
#include <sys/stat.h>

#ifndef _WIN32
#include <sys/mman.h>
#endif

#include "Mrpt.h"
#include "numpy/arrayobject.h"

#include <Eigen/Dense>

using Eigen::MatrixXf;
using Eigen::VectorXf;
using Eigen::VectorXi;

typedef struct {
    PyObject_HEAD
    Mrpt *index;
    float *data;
    int *subset_refs;
    bool mmap;
    int n;
    int dim;
    int k;
} mrptIndex;

static PyObject *Mrpt_new(PyTypeObject *type, PyObject *args, PyObject *kwds) {
    mrptIndex *self = reinterpret_cast<mrptIndex *>(type->tp_alloc(type, 0));

    if (self != NULL) {
        self->index = NULL;
        self->data = NULL;
        self->subset_refs = new int(1);
    }

    return reinterpret_cast<PyObject *>(self);
}

float *read_memory(char *file, int n, int dim) {
    FILE *fd;
    if ((fd = fopen(file, "rb")) == NULL) {
        return NULL;
    }

    float *data = new float[n * dim];

    if (data == NULL) {
        fclose(fd);
        return NULL;
    }

    int read = fread(data, sizeof(float), n * dim, fd);
    fclose(fd);

    if (read != n * dim) {
        delete[] data;
        return NULL;
    }

    return data;
}

#ifndef _WIN32
float *read_mmap(char *file, int n, int dim) {
    FILE *fd;
    if ((fd = fopen(file, "rb")) == NULL)
        return NULL;

    float *data;

    if ((data = reinterpret_cast<float *> (
#ifdef MAP_POPULATE
            mmap(0, n * dim * sizeof(float), PROT_READ,
            MAP_SHARED | MAP_POPULATE, fileno(fd), 0))) == MAP_FAILED) {
#else
            mmap(0, n * dim * sizeof(float), PROT_READ,
            MAP_SHARED, fileno(fd), 0))) == MAP_FAILED) {
#endif
            return NULL;
    }

    fclose(fd);
    return data;
}
#endif

static int Mrpt_init(mrptIndex *self, PyObject *args) {
    PyObject *py_data;
    int n, dim, mmap;

    if (!PyArg_ParseTuple(args, "Oiii", &py_data, &n, &dim, &mmap))
        return -1;

    float *data;
    if (PyUnicode_Check(py_data)) {
        char *file = PyBytes_AsString(py_data);

        struct stat sb;
        if (stat(file, &sb) != 0) {
            PyErr_SetString(PyExc_IOError, strerror(errno));
            return -1;
        }

        if (sb.st_size != static_cast<unsigned>(sizeof(float) * dim * n)) {
            PyErr_SetString(PyExc_ValueError, "Size of the input is not N x dim");
            return -1;
        }

#ifndef _WIN32
        data = mmap ? read_mmap(file, n, dim) : read_memory(file, n, dim);
#else
        data = read_memory(file, n, dim);
#endif

        if (data == NULL) {
            PyErr_SetString(PyExc_IOError, "Unable to read data from file or allocate memory for it");
            return -1;
        }

        self->mmap = mmap;
        self->data = data;
    } else {
        data = reinterpret_cast<float *>(PyArray_DATA(py_data));
    }

    self->n = n;
    self->dim = dim;
    self->index = new Mrpt(data, dim, n);

    return 0;
}

static PyObject *build(mrptIndex *self, PyObject *args) {
    int n_trees, depth;
    float density;

    if (!PyArg_ParseTuple(args, "iif", &n_trees, &depth, &density))
        return NULL;

    try {
        self->index->grow(n_trees, depth, density);
    } catch (const std::exception &e) {
        PyErr_SetString(PyExc_RuntimeError, e.what());
    }

    Py_RETURN_NONE;
}

static PyObject *build_autotune(mrptIndex *self, PyObject *args) {
    PyObject *py_data;
    float target_recall, density;
    int n_test, k, trees_max, depth_min, depth_max, votes_max;

    if (!PyArg_ParseTuple(args, "fOiiiiiif", &target_recall, &py_data, &n_test, &k, &trees_max,
                                             &depth_min, &depth_max, &votes_max, &density))
        return NULL;

    bool from_mem = false;

    float *data;
    if (PyUnicode_Check(py_data)) {
        char *file = PyBytes_AsString(py_data);

        struct stat sb;
        if (stat(file, &sb) != 0) {
            PyErr_SetString(PyExc_IOError, strerror(errno));
            return NULL;
        }

        if (sb.st_size != static_cast<unsigned>(sizeof(float) * self->dim * n_test)) {
            PyErr_SetString(PyExc_ValueError, "Size of the input is not n_test x dim");
            return NULL;
        }

        data = read_memory(file, n_test, self->dim);
        if (data == NULL) {
            PyErr_SetString(PyExc_IOError, "Unable to read data from file or allocate memory for it");
            return NULL;
        }

        from_mem = true;
    } else {
        data = reinterpret_cast<float *>(PyArray_DATA(py_data));
    }

    self->k = k;

    try {
        if (target_recall < 0)
            self->index->grow(data, n_test, k, trees_max, depth_max, depth_min, votes_max, density);
        else
            self->index->grow(target_recall, data, n_test, k, trees_max, depth_max, depth_min, votes_max, density);
    } catch (const std::exception &e) {
        PyErr_SetString(PyExc_RuntimeError, e.what());
    }

    if (from_mem)
        delete[] data;

    Py_RETURN_NONE;
}

static PyObject *build_autotune_sample(mrptIndex *self, PyObject *args) {
    float target_recall, density;
    int n_test, k, trees_max, depth_min, depth_max, votes_max;

    if (!PyArg_ParseTuple(args, "fiiiiiif", &target_recall, &n_test, &k, &trees_max,
                                            &depth_min, &depth_max, &votes_max, &density))
        return NULL;

    self->k = k;

    try {
        if (target_recall < 0)
            self->index->grow_autotune(
                k, trees_max, depth_max, depth_min, votes_max, density, 0, n_test);
        else
            self->index->grow_autotune(
                target_recall, k, trees_max, depth_max, depth_min, votes_max, density, 0, n_test);
    } catch (const std::exception &e) {
        PyErr_SetString(PyExc_RuntimeError, e.what());
    }

    Py_RETURN_NONE;
}

static void mrpt_dealloc(mrptIndex *self) {
    (*self->subset_refs)--;
    if (!(*self->subset_refs)) {
        if (self->data) {
#ifndef _WIN32
            if (self->mmap)
                munmap(self->data, self->n * self->dim * sizeof(float));
            else
#endif
                delete[] self->data;

            self->data = NULL;
        }

        delete self->subset_refs;
        self->subset_refs = NULL;
    }

    if (self->index) {
        delete self->index;
        self->index = NULL;
    }

    Py_TYPE(self)->tp_free(reinterpret_cast<PyObject *>(self));
}

static PyObject *ann(mrptIndex *self, PyObject *args) {
    PyObject *v;
    int k, elect, dim, n, return_distances;

    if (!PyArg_ParseTuple(args, "Oiii", &v, &k, &elect, &return_distances))
        return NULL;

    float *indata = reinterpret_cast<float *>(PyArray_DATA(v));
    PyObject *nearest;

    if (k < 0)
        k = self->k;
    if (elect < 0) {
        elect = self->index->parameters().votes;
    }

    if (PyArray_NDIM(v) == 1) {
        dim = PyArray_DIM(v, 0);

        npy_intp dims[1] = {k};
        nearest = PyArray_SimpleNew(1, dims, NPY_INT);
        int *outdata = reinterpret_cast<int *>(PyArray_DATA(nearest));

        if (return_distances) {
            PyObject *distances = PyArray_SimpleNew(1, dims, NPY_FLOAT32);
            float *out_distances = reinterpret_cast<float *>(PyArray_DATA(distances));
            self->index->query(indata, k, elect, outdata, out_distances);

            PyObject *out_tuple = PyTuple_New(2);
            PyTuple_SetItem(out_tuple, 0, nearest);
            PyTuple_SetItem(out_tuple, 1, distances);
            return out_tuple;
        } else {
            self->index->query(indata, k, elect, outdata);
            return nearest;
        }
    } else {
        n = PyArray_DIM(v, 0);
        dim = PyArray_DIM(v, 1);

        npy_intp dims[2] = {n, k};
        nearest = PyArray_SimpleNew(2, dims, NPY_INT);
        int *outdata = reinterpret_cast<int *>(PyArray_DATA(nearest));

        if (return_distances) {
            npy_intp dims[2] = {n, k};
            PyObject *distances = PyArray_SimpleNew(2, dims, NPY_FLOAT32);
            float *distances_out = reinterpret_cast<float *>(PyArray_DATA(distances));

            for (int i = 0; i < n; ++i) {
                self->index->query(indata + i * dim, k, elect, outdata + i * k, distances_out + i * k);
            }

            PyObject *out_tuple = PyTuple_New(2);
            PyTuple_SetItem(out_tuple, 0, nearest);
            PyTuple_SetItem(out_tuple, 1, distances);
            return out_tuple;
        } else {
            for (int i = 0; i < n; ++i) {
                self->index->query(indata + i * dim, k, elect, outdata + i * k);
            }
            return nearest;
        }
    }
}

static PyObject *exact_search(mrptIndex *self, PyObject *args) {
    PyObject *v;
    int k, n, dim, return_distances;

    if (!PyArg_ParseTuple(args, "Oii", &v, &k, &return_distances))
        return NULL;

    float *indata = reinterpret_cast<float *>(PyArray_DATA(v));
    PyObject *nearest;

    if (PyArray_NDIM(v) == 1) {
        dim = PyArray_DIM(v, 0);

        npy_intp dims[1] = {k};
        nearest = PyArray_SimpleNew(1, dims, NPY_INT);
        int *outdata = reinterpret_cast<int *>(PyArray_DATA(nearest));

        if (return_distances) {
            PyObject *distances = PyArray_SimpleNew(1, dims, NPY_FLOAT32);
            float *out_distances = reinterpret_cast<float *>(PyArray_DATA(distances));
            self->index->exact_knn(indata, k, outdata, out_distances);

            PyObject *out_tuple = PyTuple_New(2);
            PyTuple_SetItem(out_tuple, 0, nearest);
            PyTuple_SetItem(out_tuple, 1, distances);
            return out_tuple;
        } else {
            self->index->exact_knn(indata, k, outdata);
            return nearest;
        }
    } else {
        n = PyArray_DIM(v, 0);
        dim = PyArray_DIM(v, 1);

        npy_intp dims[2] = {n, k};
        nearest = PyArray_SimpleNew(2, dims, NPY_INT);
        int *outdata = reinterpret_cast<int *>(PyArray_DATA(nearest));

        if (return_distances) {
            npy_intp dims[2] = {n, k};
            PyObject *distances = PyArray_SimpleNew(2, dims, NPY_FLOAT32);
            float *distances_out = reinterpret_cast<float *>(PyArray_DATA(distances));

            for (int i = 0; i < n; ++i) {
                self->index->exact_knn(indata + i * dim, k, outdata + i * k, distances_out + i * k);
            }
            PyObject *out_tuple = PyTuple_New(2);
            PyTuple_SetItem(out_tuple, 0, nearest);
            PyTuple_SetItem(out_tuple, 1, distances);
            return out_tuple;
        } else {
            for (int i = 0; i < n; ++i) {
                self->index->exact_knn(indata + i * dim, k, outdata + i * k);
            }
            return nearest;
        }
    }
}

static PyObject *subset(mrptIndex *self, PyObject *args) {
    float target_recall;

    if (!PyArg_ParseTuple(args, "f", &target_recall))
        return NULL;

    mrptIndex *new_idx = (mrptIndex *) PyObject_New(mrptIndex, Py_TYPE(self));
    new_idx = (mrptIndex *) PyObject_Init((PyObject *) new_idx, Py_TYPE(self));

    new_idx->data = self->data;
    new_idx->subset_refs = self->subset_refs;
    (*new_idx->subset_refs)++;
    new_idx->mmap = self->mmap;
    new_idx->n = self->n;
    new_idx->dim = self->dim;
    new_idx->k = self->k;

    new_idx->index = self->index->subset_pointer(target_recall);
    return reinterpret_cast<PyObject *>(new_idx);
}

static PyObject *parameters(mrptIndex *self, PyObject *args) {
    Mrpt_Parameters par = self->index->parameters();

    PyObject *tup = PyTuple_New(6);
    PyTuple_SetItem(tup, 0, PyLong_FromLong(par.n_trees));
    PyTuple_SetItem(tup, 1, PyLong_FromLong(par.depth));
    PyTuple_SetItem(tup, 2, PyLong_FromLong(par.votes));
    PyTuple_SetItem(tup, 3, PyLong_FromLong(par.k));
    PyTuple_SetItem(tup, 4, PyFloat_FromDouble(par.estimated_qtime));<commit_msg>Fix ref counts for data in Python bindings<commit_after>/*
 * This file wraps the C++11 Mrpt code to an extension module compatible with Python 3.
 */

#include "Python.h"
#include <cstdio>
#include <cstdlib>
#include <numeric>
#include <string>
#include <vector>
#include <memory>
#include <sys/types.h>
#include <sys/stat.h>

#ifndef _WIN32
#include <sys/mman.h>
#endif

#include "Mrpt.h"
#include "numpy/arrayobject.h"

#include <Eigen/Dense>

using Eigen::MatrixXf;
using Eigen::VectorXf;
using Eigen::VectorXi;

typedef struct {
    PyObject_HEAD
    Mrpt *index;
    PyObject *py_data;
    float *data;
    int *subset_refs;
    bool mmap;
    int n;
    int dim;
    int k;
} mrptIndex;

static PyObject *Mrpt_new(PyTypeObject *type, PyObject *args, PyObject *kwds) {
    mrptIndex *self = reinterpret_cast<mrptIndex *>(type->tp_alloc(type, 0));

    if (self != NULL) {
        self->index = NULL;
        self->data = NULL;
        self->py_data = NULL;
        self->subset_refs = new int(1);
    }

    return reinterpret_cast<PyObject *>(self);
}

float *read_memory(char *file, int n, int dim) {
    FILE *fd;
    if ((fd = fopen(file, "rb")) == NULL) {
        return NULL;
    }

    float *data = new float[n * dim];

    if (data == NULL) {
        fclose(fd);
        return NULL;
    }

    int read = fread(data, sizeof(float), n * dim, fd);
    fclose(fd);

    if (read != n * dim) {
        delete[] data;
        return NULL;
    }

    return data;
}

#ifndef _WIN32
float *read_mmap(char *file, int n, int dim) {
    FILE *fd;
    if ((fd = fopen(file, "rb")) == NULL)
        return NULL;

    float *data;

    if ((data = reinterpret_cast<float *> (
#ifdef MAP_POPULATE
            mmap(0, n * dim * sizeof(float), PROT_READ,
            MAP_SHARED | MAP_POPULATE, fileno(fd), 0))) == MAP_FAILED) {
#else
            mmap(0, n * dim * sizeof(float), PROT_READ,
            MAP_SHARED, fileno(fd), 0))) == MAP_FAILED) {
#endif
            return NULL;
    }

    fclose(fd);
    return data;
}
#endif

static int Mrpt_init(mrptIndex *self, PyObject *args) {
    PyObject *py_data;
    int n, dim, mmap;

    if (!PyArg_ParseTuple(args, "Oiii", &py_data, &n, &dim, &mmap))
        return -1;

    float *data;
    if (PyUnicode_Check(py_data)) {
        char *file = PyBytes_AsString(py_data);

        struct stat sb;
        if (stat(file, &sb) != 0) {
            PyErr_SetString(PyExc_IOError, strerror(errno));
            return -1;
        }

        if (sb.st_size != static_cast<unsigned>(sizeof(float) * dim * n)) {
            PyErr_SetString(PyExc_ValueError, "Size of the input is not N x dim");
            return -1;
        }

#ifndef _WIN32
        data = mmap ? read_mmap(file, n, dim) : read_memory(file, n, dim);
#else
        data = read_memory(file, n, dim);
#endif

        if (data == NULL) {
            PyErr_SetString(PyExc_IOError, "Unable to read data from file or allocate memory for it");
            return -1;
        }

        self->mmap = mmap;
        self->data = data;
    } else {
        data = reinterpret_cast<float *>(PyArray_DATA(py_data));
        self->py_data = py_data;
        Py_XINCREF(self->py_data);
    }

    self->n = n;
    self->dim = dim;
    self->index = new Mrpt(data, dim, n);

    return 0;
}

static PyObject *build(mrptIndex *self, PyObject *args) {
    int n_trees, depth;
    float density;

    if (!PyArg_ParseTuple(args, "iif", &n_trees, &depth, &density))
        return NULL;

    try {
        self->index->grow(n_trees, depth, density);
    } catch (const std::exception &e) {
        PyErr_SetString(PyExc_RuntimeError, e.what());
    }

    Py_RETURN_NONE;
}

static PyObject *build_autotune(mrptIndex *self, PyObject *args) {
    PyObject *py_data;
    float target_recall, density;
    int n_test, k, trees_max, depth_min, depth_max, votes_max;

    if (!PyArg_ParseTuple(args, "fOiiiiiif", &target_recall, &py_data, &n_test, &k, &trees_max,
                                             &depth_min, &depth_max, &votes_max, &density))
        return NULL;

    bool from_mem = false;

    float *data;
    if (PyUnicode_Check(py_data)) {
        char *file = PyBytes_AsString(py_data);

        struct stat sb;
        if (stat(file, &sb) != 0) {
            PyErr_SetString(PyExc_IOError, strerror(errno));
            return NULL;
        }

        if (sb.st_size != static_cast<unsigned>(sizeof(float) * self->dim * n_test)) {
            PyErr_SetString(PyExc_ValueError, "Size of the input is not n_test x dim");
            return NULL;
        }

        data = read_memory(file, n_test, self->dim);
        if (data == NULL) {
            PyErr_SetString(PyExc_IOError, "Unable to read data from file or allocate memory for it");
            return NULL;
        }

        from_mem = true;
    } else {
        data = reinterpret_cast<float *>(PyArray_DATA(py_data));
    }

    self->k = k;

    try {
        if (target_recall < 0)
            self->index->grow(data, n_test, k, trees_max, depth_max, depth_min, votes_max, density);
        else
            self->index->grow(target_recall, data, n_test, k, trees_max, depth_max, depth_min, votes_max, density);
    } catch (const std::exception &e) {
        PyErr_SetString(PyExc_RuntimeError, e.what());
    }

    if (from_mem)
        delete[] data;

    Py_RETURN_NONE;
}

static PyObject *build_autotune_sample(mrptIndex *self, PyObject *args) {
    float target_recall, density;
    int n_test, k, trees_max, depth_min, depth_max, votes_max;

    if (!PyArg_ParseTuple(args, "fiiiiiif", &target_recall, &n_test, &k, &trees_max,
                                            &depth_min, &depth_max, &votes_max, &density))
        return NULL;

    self->k = k;

    try {
        if (target_recall < 0)
            self->index->grow_autotune(
                k, trees_max, depth_max, depth_min, votes_max, density, 0, n_test);
        else
            self->index->grow_autotune(
                target_recall, k, trees_max, depth_max, depth_min, votes_max, density, 0, n_test);
    } catch (const std::exception &e) {
        PyErr_SetString(PyExc_RuntimeError, e.what());
    }

    Py_RETURN_NONE;
}

static void mrpt_dealloc(mrptIndex *self) {
    (*self->subset_refs)--;
    if (!(*self->subset_refs)) {
        if (self->data) {
#ifndef _WIN32
            if (self->mmap)
                munmap(self->data, self->n * self->dim * sizeof(float));
            else
#endif
                delete[] self->data;

            self->data = NULL;
        }

        delete self->subset_refs;
        self->subset_refs = NULL;
    }

    if (self->index) {
        delete self->index;
        self->index = NULL;
    }

    Py_XDECREF(self->py_data);
    self->py_data = NULL;

    Py_TYPE(self)->tp_free(reinterpret_cast<PyObject *>(self));
}

static PyObject *ann(mrptIndex *self, PyObject *args) {
    PyObject *v;
    int k, elect, dim, n, return_distances;

    if (!PyArg_ParseTuple(args, "Oiii", &v, &k, &elect, &return_distances))
        return NULL;

    float *indata = reinterpret_cast<float *>(PyArray_DATA(v));
    PyObject *nearest;

    if (k < 0)
        k = self->k;
    if (elect < 0) {
        elect = self->index->parameters().votes;
    }

    if (PyArray_NDIM(v) == 1) {
        dim = PyArray_DIM(v, 0);

        npy_intp dims[1] = {k};
        nearest = PyArray_SimpleNew(1, dims, NPY_INT);
        int *outdata = reinterpret_cast<int *>(PyArray_DATA(nearest));

        if (return_distances) {
            PyObject *distances = PyArray_SimpleNew(1, dims, NPY_FLOAT32);
            float *out_distances = reinterpret_cast<float *>(PyArray_DATA(distances));
            self->index->query(indata, k, elect, outdata, out_distances);

            PyObject *out_tuple = PyTuple_New(2);
            PyTuple_SetItem(out_tuple, 0, nearest);
            PyTuple_SetItem(out_tuple, 1, distances);
            return out_tuple;
        } else {
            self->index->query(indata, k, elect, outdata);
            return nearest;
        }
    } else {
        n = PyArray_DIM(v, 0);
        dim = PyArray_DIM(v, 1);

        npy_intp dims[2] = {n, k};
        nearest = PyArray_SimpleNew(2, dims, NPY_INT);
        int *outdata = reinterpret_cast<int *>(PyArray_DATA(nearest));

        if (return_distances) {
            npy_intp dims[2] = {n, k};
            PyObject *distances = PyArray_SimpleNew(2, dims, NPY_FLOAT32);
            float *distances_out = reinterpret_cast<float *>(PyArray_DATA(distances));

            for (int i = 0; i < n; ++i) {
                self->index->query(indata + i * dim, k, elect, outdata + i * k, distances_out + i * k);
            }

            PyObject *out_tuple = PyTuple_New(2);
            PyTuple_SetItem(out_tuple, 0, nearest);
            PyTuple_SetItem(out_tuple, 1, distances);
            return out_tuple;
        } else {
            for (int i = 0; i < n; ++i) {
                self->index->query(indata + i * dim, k, elect, outdata + i * k);
            }
            return nearest;
        }
    }
}

static PyObject *exact_search(mrptIndex *self, PyObject *args) {
    PyObject *v;
    int k, n, dim, return_distances;

    if (!PyArg_ParseTuple(args, "Oii", &v, &k, &return_distances))
        return NULL;

    float *indata = reinterpret_cast<float *>(PyArray_DATA(v));
    PyObject *nearest;

    if (PyArray_NDIM(v) == 1) {
        dim = PyArray_DIM(v, 0);

        npy_intp dims[1] = {k};
        nearest = PyArray_SimpleNew(1, dims, NPY_INT);
        int *outdata = reinterpret_cast<int *>(PyArray_DATA(nearest));

        if (return_distances) {
            PyObject *distances = PyArray_SimpleNew(1, dims, NPY_FLOAT32);
            float *out_distances = reinterpret_cast<float *>(PyArray_DATA(distances));
            self->index->exact_knn(indata, k, outdata, out_distances);

            PyObject *out_tuple = PyTuple_New(2);
            PyTuple_SetItem(out_tuple, 0, nearest);
            PyTuple_SetItem(out_tuple, 1, distances);
            return out_tuple;
        } else {
            self->index->exact_knn(indata, k, outdata);
            return nearest;
        }
    } else {
        n = PyArray_DIM(v, 0);
        dim = PyArray_DIM(v, 1);

        npy_intp dims[2] = {n, k};
        nearest = PyArray_SimpleNew(2, dims, NPY_INT);
        int *outdata = reinterpret_cast<int *>(PyArray_DATA(nearest));

        if (return_distances) {
            npy_intp dims[2] = {n, k};
            PyObject *distances = PyArray_SimpleNew(2, dims, NPY_FLOAT32);
            float *distances_out = reinterpret_cast<float *>(PyArray_DATA(distances));

            for (int i = 0; i < n; ++i) {
                self->index->exact_knn(indata + i * dim, k, outdata + i * k, distances_out + i * k);
            }
            PyObject *out_tuple = PyTuple_New(2);
            PyTuple_SetItem(out_tuple, 0, nearest);
            PyTuple_SetItem(out_tuple, 1, distances);
            return out_tuple;
        } else {
            for (int i = 0; i < n; ++i) {
                self->index->exact_knn(indata + i * dim, k, outdata + i * k);
            }
            return nearest;
        }
    }
}

static PyObject *subset(mrptIndex *self, PyObject *args) {
    float target_recall;

    if (!PyArg_ParseTuple(args, "f", &target_recall))
        return NULL;

    mrptIndex *new_idx = (mrptIndex *) PyObject_New(mrptIndex, Py_TYPE(self));
    new_idx = (mrptIndex *) PyObject_Init((PyObject *) new_idx, Py_TYPE(self));

    new_idx->data = self->data;
    new_idx->subset_refs = self->subset_refs;
    (*new_idx->subset_refs)++;
    new_idx->mmap = self->mmap;
    new_idx->n = self->n;
    new_idx->dim = self->dim;
    new_idx->k = self->k;
    new_idx->py_data = self->py_data;
    Py_XINCREF(new_idx->py_data);

    new_idx->index = self->index->subset_pointer(target_recall);
    return reinterpret_cast<PyObject *>(new_idx);
}

static PyObject *parameters(mrptIndex *self, PyObject *args) {
    Mrpt_Parameters par = self->index->parameters();

    PyObject *tup = PyTuple_New(6);
    PyTuple_SetItem(tup, 0, PyLong_FromLong(par.n_trees));
    PyTuple_SetItem(tup, 1, PyLong_FromLong(par.depth));
    PyTuple_SetItem(tup, 2, PyLong_FromLong(par.votes));
    PyTuple_SetItem(tup, 3, PyLong_FromLong(par.k));
    PyTuple_SetItem(tup, 4, PyFloat_FromDouble(par.estimated_qtime));
{'fraction_non_alphanumeric': 0.13203445447087778, 'fraction_numerical': 0.006726825266611977, 'mean_word_length': 2.3806156405990015, 'pattern_counts': {'":': 0, '<': 77, '<?xml version=': 0, '>': 219, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7203951', 'n_tokens_mistral': 8634, 'n_tokens_neox': 8409, 'n_words': 1927}