pyarrow 表到 PyObject* 通过 pybind11

#include <pybind11/pybind11.h>
#include <Python.h>
#include <iostream>
#include <arrow/python/pyarrow.h>
#include <arrow/table.h>

// Convert pyarrow table to native C++ object and print its contents
void print_table(PyObject* py_table)
{
    // convert pyobject to table
    auto status = arrow::py::unwrap_table(py_table);
    if (!status.ok())
    {
        std::cout << "Error converting pyarrow table to arrow table" << std::endl;
        return;
    }
    std::shared_ptr<arrow::Table> table = status.ValueOrDie();
    std::cout << "Table has " << table->num_rows() << " rows" << std::endl;
    // print table
    for (int i = 0; i < table->num_columns(); i++)
    {
        std::cout << "Column " << i << ": " << table->column(i)->ToString() << std::endl;
    }
}

PYBIND11_MODULE(df_test, m)
{
    arrow::py::import_pyarrow();
    m.doc() = "Pyarrow Extensions";
    m.def("print_table", &print_table);
}

我能够成功构建一个c++图书馆通过pybind11它接受一个PyObject*并希望打印 a 的内容pyarrow表传递给它。然而，在转换我的pandas.DataFrame to a pyarrow.Table如下，

# convert to pyarrow table
table = pa.Table.from_pandas(df)

import df_test
df_test.print_table(table)

该代码抱怨说

TypeError: print_table(): incompatible function arguments. The following argument types are supported:
    1. (arg0: _object) -> None

我如何转换pyarrow.Table to a PyObject*?

归结为编写一个自定义PYBIND11_TYPE_CASTER以便能够正确读取pyarrow.Table从Python传递过来的。这是包括arrow_conversions.hh

#ifndef ARROW_CONVERSIONS
#define ARROW_CONVERSIONS

namespace pybind11
{
    namespace detail
    {
        template <typename TableType>
        struct gen_type_caster
        {
        public:
            PYBIND11_TYPE_CASTER(std::shared_ptr<TableType>, _("pyarrow::Table"));
            // Python -> C++
            bool load(handle src, bool)
            {
                PyObject *source = src.ptr();
                if (!arrow::py::is_table(source))
                    return false;
                arrow::Result<std::shared_ptr<arrow::Table>> result = arrow::py::unwrap_table(source);
                if (!result.ok())
                    return false;
                value = std::static_pointer_cast<TableType>(result.ValueOrDie());
                return true;
            }
            // C++ -> Python
            static handle cast(std::shared_ptr<TableType> src, return_value_policy /* policy */, handle /* parent */)
            {
                return arrow::py::wrap_table(src);
            }
        };
        template <>
        struct type_caster<std::shared_ptr<arrow::Table>> : public gen_type_caster<arrow::Table>
        {
        };
    }
} // namespace pybind11::detail

#endif

And the .cc file

#include <arrow/python/pyarrow.h>
#include <pybind11/pybind11.h>
#include <Python.h>
#include "arrow_conversions.hh"
#include <arrow/table.h>
#include <iostream>

// Convert pyarrow table to native C++ object and print its contents
void print_table(std::shared_ptr<arrow::Table> &table)
{
    // print table
    std::cout << "Table schema: " << std::endl;
    std::cout << table->schema()->ToString() << std::endl;
    std::cout << "Table columns: " << std::endl;
    for (int i = 0; i < table->num_columns(); i++)
    {
        std::cout << "Column " << i << ": " << std::endl;
        std::cout << table->column(i)->ToString() << std::endl;
    }
}

PYBIND11_MODULE(df_test, m)
{
    arrow::py::import_pyarrow();
    m.doc() = "Pyarrow Extensions";
    m.def("print_table", &print_table, pybind11::call_guard<pybind11::gil_scoped_release>());
}