wip

Author: Alex-PLACET

include/sparrow-pycapsule/SparrowPythonClass.hpp

@@ -0,0 +1,97 @@
+#pragma once
+
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+
+#include <sparrow/array.hpp>
+
+#include "sparrow-pycapsule/config/config.hpp"
+#include "sparrow-pycapsule/pycapsule.hpp"
+
+namespace sparrow::pycapsule
+{
+    /**
+     * @brief Python object structure for SparrowArray.
+     *
+     * This structure holds a pointer to a sparrow::array. The pointer is used
+     * to avoid issues with C++ objects in C-style Python object structures.
+     */
+    struct SparrowArrayObject
+    {
+        PyObject_HEAD sparrow::array* arr;
+    };
+
+    /**
+     * @brief Deallocator for SparrowArray Python objects.
+     */
+    SPARROW_PYCAPSULE_API void SparrowArray_dealloc(SparrowArrayObject* self);
+
+    /**
+     * @brief Implementation of __arrow_c_array__ method.
+     *
+     * This method exports the wrapped sparrow array as Arrow PyCapsules,
+     * implementing the Arrow PyCapsule Interface (ArrowArrayExportable protocol).
+     *
+     * @param self The SparrowArray object.
+     * @param args Positional arguments (unused).
+     * @param kwargs Keyword arguments (optional requested_schema).
+     * @return A tuple of (schema_capsule, array_capsule).
+     */
+    SPARROW_PYCAPSULE_API PyObject*
+    SparrowArray_arrow_c_array(SparrowArrayObject* self, PyObject* args, PyObject* kwargs);
+
+    /**
+     * @brief Get the size of the wrapped array.
+     *
+     * @param self The SparrowArray object.
+     * @param args Positional arguments (unused).
+     * @return The size of the array as a Python integer.
+     */
+    SPARROW_PYCAPSULE_API PyObject* SparrowArray_size(SparrowArrayObject* self, PyObject* args);
+
+    /**
+     * @brief Get the Python type object for SparrowArray.
+     *
+     * This function returns a pointer to the SparrowArrayType. The type is
+     * initialized on first call if necessary.
+     *
+     * @return Pointer to the SparrowArrayType, or nullptr on error.
+     */
+    SPARROW_PYCAPSULE_API PyTypeObject* get_sparrow_array_type();
+
+    /**
+     * @brief Create a new SparrowArray Python object from a sparrow::array.
+     *
+     * This function creates a new Python object that wraps the given sparrow array.
+     * The array is moved into the Python object, so the caller should not use it
+     * after this call.
+     *
+     * @param arr The sparrow array to wrap (will be moved).
+     * @return A new reference to a SparrowArray Python object, or nullptr on error.
+     */
+    SPARROW_PYCAPSULE_API PyObject* create_sparrow_array_object(sparrow::array&& arr);
+
+    /**
+     * @brief Create a new SparrowArray Python object from PyCapsules.
+     *
+     * This function creates a new Python object by importing from existing
+     * Arrow PyCapsules.
+     *
+     * @param schema_capsule The schema PyCapsule.
+     * @param array_capsule The array PyCapsule.
+     * @return A new reference to a SparrowArray Python object, or nullptr on error.
+     */
+    SPARROW_PYCAPSULE_API PyObject*
+    create_sparrow_array_object_from_capsules(PyObject* schema_capsule, PyObject* array_capsule);
+
+    /**
+     * @brief Register the SparrowArray type with a Python module.
+     *
+     * This function adds the SparrowArray type to the given module.
+     *
+     * @param module The Python module to add the type to.
+     * @return 0 on success, -1 on error.
+     */
+    SPARROW_PYCAPSULE_API int register_sparrow_array_type(PyObject* module);
+
+}  // namespace sparrow::pycapsule

src/SparrowPythonClass.cpp

@@ -0,0 +1,191 @@
+#include "sparrow-pycapsule/SparrowPythonClass.hpp"
+
+#include <new>
+#include <utility>
+
+namespace sparrow::pycapsule
+{
+    void SparrowArray_dealloc(SparrowArrayObject* self)
+    {
+        delete self->arr;
+        self->arr = nullptr;
+        Py_TYPE(self)->tp_free(reinterpret_cast<PyObject*>(self));
+    }
+
+    PyObject* SparrowArray_arrow_c_array(SparrowArrayObject* self, PyObject* args, PyObject* kwargs)
+    {
+        static const char* kwlist[] = {"requested_schema", nullptr};
+        PyObject* requested_schema = nullptr;
+
+        if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O", const_cast<char**>(kwlist), &requested_schema))
+        {
+            return nullptr;
+        }
+
+        // requested_schema is typically ignored for simple cases
+        // In a full implementation, you might use it to cast to a different type
+        (void) requested_schema;
+
+        if (self->arr == nullptr)
+        {
+            PyErr_SetString(PyExc_ValueError, "SparrowArray contains no data");
+            return nullptr;
+        }
+
+        try
+        {
+            auto [schema_capsule, array_capsule] = export_array_to_capsules(*self->arr);
+
+            if (schema_capsule == nullptr || array_capsule == nullptr)
+            {
+                Py_XDECREF(schema_capsule);
+                Py_XDECREF(array_capsule);
+                PyErr_SetString(PyExc_RuntimeError, "Failed to create Arrow PyCapsules");
+                return nullptr;
+            }
+
+            PyObject* result = PyTuple_Pack(2, schema_capsule, array_capsule);
+            Py_DECREF(schema_capsule);
+            Py_DECREF(array_capsule);
+            return result;
+        }
+        catch (const std::exception& e)
+        {
+            PyErr_SetString(PyExc_RuntimeError, e.what());
+            return nullptr;
+        }
+    }
+
+    PyObject* SparrowArray_size(SparrowArrayObject* self, [[maybe_unused]] PyObject* args)
+    {
+        if (self->arr == nullptr)
+        {
+            PyErr_SetString(PyExc_ValueError, "SparrowArray contains no data");
+            return nullptr;
+        }
+
+        return PyLong_FromSize_t(self->arr->size());
+    }
+
+    static PyMethodDef SparrowArray_methods[] = {
+        {"__arrow_c_array__",
+         reinterpret_cast<PyCFunction>(SparrowArray_arrow_c_array),
+         METH_VARARGS | METH_KEYWORDS,
+         "Export the array via the Arrow PyCapsule interface.\n\n"
+         "Parameters\n"
+         "----------\n"
+         "requested_schema : object, optional\n"
+         "    Requested schema for the output (typically ignored).\n\n"
+         "Returns\n"
+         "-------\n"
+         "tuple[object, object]\n"
+         "    A tuple of (schema_capsule, array_capsule)."},
+        {"size",
+         reinterpret_cast<PyCFunction>(SparrowArray_size),
+         METH_NOARGS,
+         "Get the number of elements in the array.\n\n"
+         "Returns\n"
+         "-------\n"
+         "int\n"
+         "    The size of the array."},
+        {nullptr, nullptr, 0, nullptr}  // Sentinel
+    };
+
+    // The type object - defined as a static variable
+    static PyTypeObject SparrowArrayType = {
+        .ob_base = PyVarObject_HEAD_INIT(nullptr, 0).tp_name = "sparrow.SparrowArray",
+        .tp_basicsize = sizeof(SparrowArrayObject),
+        .tp_itemsize = 0,
+        .tp_dealloc = reinterpret_cast<destructor>(SparrowArray_dealloc),
+        .tp_flags = Py_TPFLAGS_DEFAULT,
+        .tp_doc = PyDoc_STR(
+            "SparrowArray - Arrow array wrapper implementing __arrow_c_array__.\n\n"
+            "This class wraps a sparrow array and implements the Arrow PyCapsule\n"
+            "Interface (ArrowArrayExportable protocol), allowing it to be passed\n"
+            "directly to libraries like Polars via pl.from_arrow()."
+        ),
+        .tp_methods = SparrowArray_methods,
+    };
+
+    static bool type_initialized = false;
+
+    PyTypeObject* get_sparrow_array_type()
+    {
+        if (!type_initialized)
+        {
+            if (PyType_Ready(&SparrowArrayType) < 0)
+            {
+                return nullptr;
+            }
+            type_initialized = true;
+        }
+        return &SparrowArrayType;
+    }
+
+    PyObject* create_sparrow_array_object(sparrow::array&& arr)
+    {
+        PyTypeObject* type = get_sparrow_array_type();
+        if (type == nullptr)
+        {
+            return nullptr;
+        }
+
+        SparrowArrayObject* obj = PyObject_New(SparrowArrayObject, type);
+        if (obj == nullptr)
+        {
+            return nullptr;
+        }
+
+        try
+        {
+            obj->arr = new sparrow::array(std::move(arr));
+        }
+        catch (const std::bad_alloc&)
+        {
+            Py_DECREF(obj);
+            PyErr_NoMemory();
+            return nullptr;
+        }
+        catch (const std::exception& e)
+        {
+            Py_DECREF(obj);
+            PyErr_SetString(PyExc_RuntimeError, e.what());
+            return nullptr;
+        }
+
+        return reinterpret_cast<PyObject*>(obj);
+    }
+
+    PyObject* create_sparrow_array_object_from_capsules(PyObject* schema_capsule, PyObject* array_capsule)
+    {
+        try
+        {
+            sparrow::array arr = import_array_from_capsules(schema_capsule, array_capsule);
+            return create_sparrow_array_object(std::move(arr));
+        }
+        catch (const std::exception& e)
+        {
+            PyErr_SetString(PyExc_RuntimeError, e.what());
+            return nullptr;
+        }
+    }
+
+    int register_sparrow_array_type(PyObject* module)
+    {
+        PyTypeObject* type = get_sparrow_array_type();
+        if (type == nullptr)
+        {
+            return -1;
+        }
+
+        Py_INCREF(type);
+        if (PyModule_AddObject(module, "SparrowArray", reinterpret_cast<PyObject*>(type)) < 0)
+        {
+            Py_DECREF(type);
+            return -1;
+        }
+
+        return 0;
+    }
+
+}  // namespace sparrow::pycapsule