pyarrow Table to PyObject* via pybind11

user_4555763 0
Viewed 197
#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);
}

I'm able to successfully build a c++ library via pybind11 which accepts a PyObject* and hopefully prints the contents of a pyarrow table passed to it. However, after converting my pandas.DataFrame to a pyarrow.Table as follows,

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

import df_test
df_test.print_table(table)

the code complains that

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

How do I convert the pyarrow.Table to a PyObject*?

pybind11 apache-arrow pyarrow c++
1 Answers

It boiled down to writing a custom PYBIND11_TYPE_CASTER to be able to correctly read the pyarrow.Table passed from python. Here's the include 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>());
}
user_45557630
Related