mirror of
https://github.com/wassname/ray.git
synced 2026-08-06 13:31:10 +08:00
Switch build system to use CMake completely. (#200)
* switch to CMake completely ... * cleanup * Run C tests, update installation instructions.
This commit is contained in:
Philipp Moritz
committed by
Robert Nishihara
Robert Nishihara
parent
ba8933e10f
commit
a708e36225
@@ -0,0 +1,132 @@
|
||||
#include "numpy.h"
|
||||
#include "python.h"
|
||||
|
||||
#include <sstream>
|
||||
|
||||
#include <numbuf/tensor.h>
|
||||
|
||||
using namespace arrow;
|
||||
|
||||
extern "C" {
|
||||
extern PyObject* numbuf_serialize_callback;
|
||||
extern PyObject* numbuf_deserialize_callback;
|
||||
}
|
||||
|
||||
namespace numbuf {
|
||||
|
||||
#define ARROW_TYPE_TO_NUMPY_CASE(TYPE) \
|
||||
case Type::TYPE: \
|
||||
return NPY_##TYPE;
|
||||
|
||||
#define DESERIALIZE_ARRAY_CASE(TYPE, ArrayType, type) \
|
||||
case Type::TYPE: { \
|
||||
auto values = std::dynamic_pointer_cast<ArrayType>(content->values()); \
|
||||
DCHECK(values); \
|
||||
type* data = const_cast<type*>(values->raw_data()) + content->offset(offset); \
|
||||
*out = PyArray_SimpleNewFromData( \
|
||||
num_dims, dim.data(), NPY_##TYPE, reinterpret_cast<void*>(data)); \
|
||||
if (base != Py_None) { PyArray_SetBaseObject((PyArrayObject*)*out, base); } \
|
||||
Py_XINCREF(base); \
|
||||
} break;
|
||||
|
||||
Status DeserializeArray(
|
||||
std::shared_ptr<Array> array, int32_t offset, PyObject* base, PyObject** out) {
|
||||
DCHECK(array);
|
||||
auto tensor = std::dynamic_pointer_cast<StructArray>(array);
|
||||
DCHECK(tensor);
|
||||
auto dims = std::dynamic_pointer_cast<ListArray>(tensor->field(0));
|
||||
auto content = std::dynamic_pointer_cast<ListArray>(tensor->field(1));
|
||||
npy_intp num_dims = dims->value_length(offset);
|
||||
std::vector<npy_intp> dim(num_dims);
|
||||
for (int i = dims->offset(offset); i < dims->offset(offset + 1); ++i) {
|
||||
dim[i - dims->offset(offset)] =
|
||||
std::dynamic_pointer_cast<Int64Array>(dims->values())->Value(i);
|
||||
}
|
||||
switch (content->value_type()->type) {
|
||||
DESERIALIZE_ARRAY_CASE(INT8, Int8Array, int8_t)
|
||||
DESERIALIZE_ARRAY_CASE(INT16, Int16Array, int16_t)
|
||||
DESERIALIZE_ARRAY_CASE(INT32, Int32Array, int32_t)
|
||||
DESERIALIZE_ARRAY_CASE(INT64, Int64Array, int64_t)
|
||||
DESERIALIZE_ARRAY_CASE(UINT8, UInt8Array, uint8_t)
|
||||
DESERIALIZE_ARRAY_CASE(UINT16, UInt16Array, uint16_t)
|
||||
DESERIALIZE_ARRAY_CASE(UINT32, UInt32Array, uint32_t)
|
||||
DESERIALIZE_ARRAY_CASE(UINT64, UInt64Array, uint64_t)
|
||||
DESERIALIZE_ARRAY_CASE(FLOAT, FloatArray, float)
|
||||
DESERIALIZE_ARRAY_CASE(DOUBLE, DoubleArray, double)
|
||||
default:
|
||||
DCHECK(false) << "arrow type not recognized: " << content->value_type()->type;
|
||||
}
|
||||
/* Mark the array as immutable. */
|
||||
PyObject* flags = PyObject_GetAttrString(*out, "flags");
|
||||
DCHECK(flags != NULL) << "Could not mark Numpy array immutable";
|
||||
int flag_set = PyObject_SetAttrString(flags, "writeable", Py_False);
|
||||
DCHECK(flag_set == 0) << "Could not mark Numpy array immutable";
|
||||
Py_XDECREF(flags);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status SerializeArray(
|
||||
PyArrayObject* array, SequenceBuilder& builder, std::vector<PyObject*>& subdicts) {
|
||||
size_t ndim = PyArray_NDIM(array);
|
||||
int dtype = PyArray_TYPE(array);
|
||||
std::vector<int64_t> dims(ndim);
|
||||
for (int i = 0; i < ndim; ++i) {
|
||||
dims[i] = PyArray_DIM(array, i);
|
||||
}
|
||||
// TODO(pcm): Once we don't use builders any more below and directly share
|
||||
// the memory buffer, we need to be more careful about this and not
|
||||
// decrease the reference count of "contiguous" before the serialization
|
||||
// is finished
|
||||
auto contiguous = PyArray_GETCONTIGUOUS(array);
|
||||
auto data = PyArray_DATA(contiguous);
|
||||
switch (dtype) {
|
||||
case NPY_UINT8:
|
||||
RETURN_NOT_OK(builder.AppendTensor(dims, reinterpret_cast<uint8_t*>(data)));
|
||||
break;
|
||||
case NPY_INT8:
|
||||
RETURN_NOT_OK(builder.AppendTensor(dims, reinterpret_cast<int8_t*>(data)));
|
||||
break;
|
||||
case NPY_UINT16:
|
||||
RETURN_NOT_OK(builder.AppendTensor(dims, reinterpret_cast<uint16_t*>(data)));
|
||||
break;
|
||||
case NPY_INT16:
|
||||
RETURN_NOT_OK(builder.AppendTensor(dims, reinterpret_cast<int16_t*>(data)));
|
||||
break;
|
||||
case NPY_UINT32:
|
||||
RETURN_NOT_OK(builder.AppendTensor(dims, reinterpret_cast<uint32_t*>(data)));
|
||||
break;
|
||||
case NPY_INT32:
|
||||
RETURN_NOT_OK(builder.AppendTensor(dims, reinterpret_cast<int32_t*>(data)));
|
||||
break;
|
||||
case NPY_UINT64:
|
||||
RETURN_NOT_OK(builder.AppendTensor(dims, reinterpret_cast<uint64_t*>(data)));
|
||||
break;
|
||||
case NPY_INT64:
|
||||
RETURN_NOT_OK(builder.AppendTensor(dims, reinterpret_cast<int64_t*>(data)));
|
||||
break;
|
||||
case NPY_FLOAT:
|
||||
RETURN_NOT_OK(builder.AppendTensor(dims, reinterpret_cast<float*>(data)));
|
||||
break;
|
||||
case NPY_DOUBLE:
|
||||
RETURN_NOT_OK(builder.AppendTensor(dims, reinterpret_cast<double*>(data)));
|
||||
break;
|
||||
default:
|
||||
if (!numbuf_serialize_callback) {
|
||||
std::stringstream stream;
|
||||
stream << "numpy data type not recognized: " << dtype;
|
||||
return Status::NotImplemented(stream.str());
|
||||
} else {
|
||||
PyObject* arglist = Py_BuildValue("(O)", array);
|
||||
// The reference count of the result of the call to PyObject_CallObject
|
||||
// must be decremented. This is done in SerializeDict in python.cc.
|
||||
PyObject* result = PyObject_CallObject(numbuf_serialize_callback, arglist);
|
||||
Py_XDECREF(arglist);
|
||||
if (!result) { return Status::NotImplemented("python error"); }
|
||||
builder.AppendDict(PyDict_Size(result));
|
||||
subdicts.push_back(result);
|
||||
}
|
||||
}
|
||||
Py_XDECREF(contiguous);
|
||||
return Status::OK();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
#ifndef PYNUMBUF_NUMPY_H
|
||||
#define PYNUMBUF_NUMPY_H
|
||||
|
||||
#include <Python.h>
|
||||
#include <arrow/api.h>
|
||||
|
||||
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
|
||||
#define NO_IMPORT_ARRAY
|
||||
#define PY_ARRAY_UNIQUE_SYMBOL NUMBUF_ARRAY_API
|
||||
#include <numpy/arrayobject.h>
|
||||
|
||||
#include <numbuf/sequence.h>
|
||||
#include <numbuf/tensor.h>
|
||||
|
||||
namespace numbuf {
|
||||
|
||||
arrow::Status SerializeArray(
|
||||
PyArrayObject* array, SequenceBuilder& builder, std::vector<PyObject*>& subdicts);
|
||||
arrow::Status DeserializeArray(
|
||||
std::shared_ptr<arrow::Array> array, int32_t offset, PyObject* base, PyObject** out);
|
||||
}
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,295 @@
|
||||
#include "python.h"
|
||||
|
||||
#include <sstream>
|
||||
|
||||
#include "scalars.h"
|
||||
|
||||
using namespace arrow;
|
||||
|
||||
int32_t MAX_RECURSION_DEPTH = 100;
|
||||
|
||||
extern "C" {
|
||||
|
||||
extern PyObject* numbuf_serialize_callback;
|
||||
extern PyObject* numbuf_deserialize_callback;
|
||||
}
|
||||
|
||||
namespace numbuf {
|
||||
|
||||
#if PY_MAJOR_VERSION >= 3
|
||||
#define PyInt_FromLong PyLong_FromLong
|
||||
#endif
|
||||
|
||||
Status get_value(
|
||||
ArrayPtr arr, int32_t index, int32_t type, PyObject* base, PyObject** result) {
|
||||
switch (arr->type()->type) {
|
||||
case Type::BOOL:
|
||||
*result =
|
||||
PyBool_FromLong(std::static_pointer_cast<BooleanArray>(arr)->Value(index));
|
||||
return Status::OK();
|
||||
case Type::INT64:
|
||||
*result = PyInt_FromLong(std::static_pointer_cast<Int64Array>(arr)->Value(index));
|
||||
return Status::OK();
|
||||
case Type::BINARY: {
|
||||
int32_t nchars;
|
||||
const uint8_t* str =
|
||||
std::static_pointer_cast<BinaryArray>(arr)->GetValue(index, &nchars);
|
||||
*result = PyBytes_FromStringAndSize(reinterpret_cast<const char*>(str), nchars);
|
||||
return Status::OK();
|
||||
}
|
||||
case Type::STRING: {
|
||||
int32_t nchars;
|
||||
const uint8_t* str =
|
||||
std::static_pointer_cast<StringArray>(arr)->GetValue(index, &nchars);
|
||||
*result = PyUnicode_FromStringAndSize(reinterpret_cast<const char*>(str), nchars);
|
||||
return Status::OK();
|
||||
}
|
||||
case Type::FLOAT:
|
||||
*result =
|
||||
PyFloat_FromDouble(std::static_pointer_cast<FloatArray>(arr)->Value(index));
|
||||
return Status::OK();
|
||||
case Type::DOUBLE:
|
||||
*result =
|
||||
PyFloat_FromDouble(std::static_pointer_cast<DoubleArray>(arr)->Value(index));
|
||||
return Status::OK();
|
||||
case Type::STRUCT: {
|
||||
auto s = std::static_pointer_cast<StructArray>(arr);
|
||||
auto l = std::static_pointer_cast<ListArray>(s->field(0));
|
||||
if (s->type()->child(0)->name == "list") {
|
||||
return DeserializeList(l->values(), l->value_offset(index),
|
||||
l->value_offset(index + 1), base, result);
|
||||
} else if (s->type()->child(0)->name == "tuple") {
|
||||
return DeserializeTuple(l->values(), l->value_offset(index),
|
||||
l->value_offset(index + 1), base, result);
|
||||
} else if (s->type()->child(0)->name == "dict") {
|
||||
return DeserializeDict(l->values(), l->value_offset(index),
|
||||
l->value_offset(index + 1), base, result);
|
||||
} else {
|
||||
return DeserializeArray(arr, index, base, result);
|
||||
}
|
||||
}
|
||||
default:
|
||||
DCHECK(false) << "union tag not recognized " << type;
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status append(PyObject* elem, SequenceBuilder& builder, std::vector<PyObject*>& sublists,
|
||||
std::vector<PyObject*>& subtuples, std::vector<PyObject*>& subdicts) {
|
||||
// The bool case must precede the int case (PyInt_Check passes for bools)
|
||||
if (PyBool_Check(elem)) {
|
||||
RETURN_NOT_OK(builder.AppendBool(elem == Py_True));
|
||||
} else if (PyFloat_Check(elem)) {
|
||||
RETURN_NOT_OK(builder.AppendDouble(PyFloat_AS_DOUBLE(elem)));
|
||||
} else if (PyLong_Check(elem)) {
|
||||
int overflow = 0;
|
||||
int64_t data = PyLong_AsLongLongAndOverflow(elem, &overflow);
|
||||
RETURN_NOT_OK(builder.AppendInt64(data));
|
||||
if (overflow) { return Status::NotImplemented("long overflow"); }
|
||||
#if PY_MAJOR_VERSION < 3
|
||||
} else if (PyInt_Check(elem)) {
|
||||
RETURN_NOT_OK(builder.AppendInt64(static_cast<int64_t>(PyInt_AS_LONG(elem))));
|
||||
#endif
|
||||
} else if (PyBytes_Check(elem)) {
|
||||
auto data = reinterpret_cast<uint8_t*>(PyBytes_AS_STRING(elem));
|
||||
auto size = PyBytes_GET_SIZE(elem);
|
||||
RETURN_NOT_OK(builder.AppendBytes(data, size));
|
||||
} else if (PyUnicode_Check(elem)) {
|
||||
Py_ssize_t size;
|
||||
#if PY_MAJOR_VERSION >= 3
|
||||
char* data = PyUnicode_AsUTF8AndSize(elem, &size);
|
||||
#else
|
||||
PyObject* str = PyUnicode_AsUTF8String(elem);
|
||||
char* data = PyString_AS_STRING(str);
|
||||
size = PyString_GET_SIZE(str);
|
||||
Py_XDECREF(str);
|
||||
#endif
|
||||
Status s = builder.AppendString(data, size);
|
||||
RETURN_NOT_OK(s);
|
||||
} else if (PyList_Check(elem)) {
|
||||
builder.AppendList(PyList_Size(elem));
|
||||
sublists.push_back(elem);
|
||||
} else if (PyDict_Check(elem)) {
|
||||
builder.AppendDict(PyDict_Size(elem));
|
||||
subdicts.push_back(elem);
|
||||
} else if (PyTuple_CheckExact(elem)) {
|
||||
builder.AppendTuple(PyTuple_Size(elem));
|
||||
subtuples.push_back(elem);
|
||||
} else if (PyArray_IsScalar(elem, Generic)) {
|
||||
RETURN_NOT_OK(AppendScalar(elem, builder));
|
||||
} else if (PyArray_Check(elem)) {
|
||||
RETURN_NOT_OK(SerializeArray((PyArrayObject*)elem, builder, subdicts));
|
||||
} else if (elem == Py_None) {
|
||||
RETURN_NOT_OK(builder.AppendNone());
|
||||
} else {
|
||||
if (!numbuf_serialize_callback) {
|
||||
std::stringstream ss;
|
||||
ss << "data type of " << PyBytes_AS_STRING(PyObject_Repr(elem))
|
||||
<< " not recognized and custom serialization handler not registered";
|
||||
return Status::NotImplemented(ss.str());
|
||||
} else {
|
||||
PyObject* arglist = Py_BuildValue("(O)", elem);
|
||||
// The reference count of the result of the call to PyObject_CallObject
|
||||
// must be decremented. This is done in SerializeDict in this file.
|
||||
PyObject* result = PyObject_CallObject(numbuf_serialize_callback, arglist);
|
||||
Py_XDECREF(arglist);
|
||||
if (!result) { return Status::NotImplemented("python error"); }
|
||||
builder.AppendDict(PyDict_Size(result));
|
||||
subdicts.push_back(result);
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status SerializeSequences(std::vector<PyObject*> sequences, int32_t recursion_depth,
|
||||
std::shared_ptr<Array>* out) {
|
||||
DCHECK(out);
|
||||
if (recursion_depth >= MAX_RECURSION_DEPTH) {
|
||||
return Status::NotImplemented(
|
||||
"This object exceeds the maximum recursion depth. It may contain itself "
|
||||
"recursively.");
|
||||
}
|
||||
SequenceBuilder builder(nullptr);
|
||||
std::vector<PyObject *> sublists, subtuples, subdicts;
|
||||
for (const auto& sequence : sequences) {
|
||||
PyObject* item;
|
||||
PyObject* iterator = PyObject_GetIter(sequence);
|
||||
while ((item = PyIter_Next(iterator))) {
|
||||
Status s = append(item, builder, sublists, subtuples, subdicts);
|
||||
Py_DECREF(item);
|
||||
// if an error occurs, we need to decrement the reference counts before returning
|
||||
if (!s.ok()) {
|
||||
Py_DECREF(iterator);
|
||||
return s;
|
||||
}
|
||||
}
|
||||
Py_DECREF(iterator);
|
||||
}
|
||||
std::shared_ptr<Array> list;
|
||||
if (sublists.size() > 0) {
|
||||
RETURN_NOT_OK(SerializeSequences(sublists, recursion_depth + 1, &list));
|
||||
}
|
||||
std::shared_ptr<Array> tuple;
|
||||
if (subtuples.size() > 0) {
|
||||
RETURN_NOT_OK(SerializeSequences(subtuples, recursion_depth + 1, &tuple));
|
||||
}
|
||||
std::shared_ptr<Array> dict;
|
||||
if (subdicts.size() > 0) {
|
||||
RETURN_NOT_OK(SerializeDict(subdicts, recursion_depth + 1, &dict));
|
||||
}
|
||||
return builder.Finish(list, tuple, dict, out);
|
||||
}
|
||||
|
||||
#define DESERIALIZE_SEQUENCE(CREATE, SET_ITEM) \
|
||||
auto data = std::dynamic_pointer_cast<UnionArray>(array); \
|
||||
int32_t size = array->length(); \
|
||||
PyObject* result = CREATE(stop_idx - start_idx); \
|
||||
auto types = std::make_shared<Int8Array>(size, data->types()); \
|
||||
auto offsets = std::make_shared<Int32Array>(size, data->offset_buf()); \
|
||||
for (size_t i = start_idx; i < stop_idx; ++i) { \
|
||||
if (data->IsNull(i)) { \
|
||||
Py_INCREF(Py_None); \
|
||||
SET_ITEM(result, i - start_idx, Py_None); \
|
||||
} else { \
|
||||
int32_t offset = offsets->Value(i); \
|
||||
int8_t type = types->Value(i); \
|
||||
ArrayPtr arr = data->child(type); \
|
||||
PyObject* value; \
|
||||
RETURN_NOT_OK(get_value(arr, offset, type, base, &value)); \
|
||||
SET_ITEM(result, i - start_idx, value); \
|
||||
} \
|
||||
} \
|
||||
*out = result; \
|
||||
return Status::OK();
|
||||
|
||||
Status DeserializeList(std::shared_ptr<Array> array, int32_t start_idx, int32_t stop_idx,
|
||||
PyObject* base, PyObject** out) {
|
||||
DESERIALIZE_SEQUENCE(PyList_New, PyList_SetItem)
|
||||
}
|
||||
|
||||
Status DeserializeTuple(std::shared_ptr<Array> array, int32_t start_idx, int32_t stop_idx,
|
||||
PyObject* base, PyObject** out) {
|
||||
DESERIALIZE_SEQUENCE(PyTuple_New, PyTuple_SetItem)
|
||||
}
|
||||
|
||||
Status SerializeDict(
|
||||
std::vector<PyObject*> dicts, int32_t recursion_depth, std::shared_ptr<Array>* out) {
|
||||
DictBuilder result;
|
||||
if (recursion_depth >= MAX_RECURSION_DEPTH) {
|
||||
return Status::NotImplemented(
|
||||
"This object exceeds the maximum recursion depth. It may contain itself "
|
||||
"recursively.");
|
||||
}
|
||||
std::vector<PyObject *> key_tuples, val_lists, val_tuples, val_dicts, dummy;
|
||||
for (const auto& dict : dicts) {
|
||||
PyObject *key, *value;
|
||||
Py_ssize_t pos = 0;
|
||||
while (PyDict_Next(dict, &pos, &key, &value)) {
|
||||
RETURN_NOT_OK(append(key, result.keys(), dummy, key_tuples, dummy));
|
||||
DCHECK(dummy.size() == 0);
|
||||
RETURN_NOT_OK(append(value, result.vals(), val_lists, val_tuples, val_dicts));
|
||||
}
|
||||
}
|
||||
std::shared_ptr<Array> key_tuples_arr;
|
||||
if (key_tuples.size() > 0) {
|
||||
RETURN_NOT_OK(SerializeSequences(key_tuples, recursion_depth + 1, &key_tuples_arr));
|
||||
}
|
||||
std::shared_ptr<Array> val_list_arr;
|
||||
if (val_lists.size() > 0) {
|
||||
RETURN_NOT_OK(SerializeSequences(val_lists, recursion_depth + 1, &val_list_arr));
|
||||
}
|
||||
std::shared_ptr<Array> val_tuples_arr;
|
||||
if (val_tuples.size() > 0) {
|
||||
RETURN_NOT_OK(SerializeSequences(val_tuples, recursion_depth + 1, &val_tuples_arr));
|
||||
}
|
||||
std::shared_ptr<Array> val_dict_arr;
|
||||
if (val_dicts.size() > 0) {
|
||||
RETURN_NOT_OK(SerializeDict(val_dicts, recursion_depth + 1, &val_dict_arr));
|
||||
}
|
||||
result.Finish(key_tuples_arr, val_list_arr, val_tuples_arr, val_dict_arr, out);
|
||||
|
||||
// This block is used to decrement the reference counts of the results
|
||||
// returned by the serialization callback, which is called in SerializeArray
|
||||
// in numpy.cc as well as in DeserializeDict and in append in this file.
|
||||
static PyObject* py_type = PyUnicode_FromString("_pytype_");
|
||||
for (const auto& dict : dicts) {
|
||||
if (PyDict_Contains(dict, py_type)) {
|
||||
// If the dictionary contains the key "_pytype_", then the user has to
|
||||
// have registered a callback.
|
||||
ARROW_CHECK(numbuf_serialize_callback);
|
||||
Py_XDECREF(dict);
|
||||
}
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status DeserializeDict(std::shared_ptr<Array> array, int32_t start_idx, int32_t stop_idx,
|
||||
PyObject* base, PyObject** out) {
|
||||
auto data = std::dynamic_pointer_cast<StructArray>(array);
|
||||
// TODO(pcm): error handling, get rid of the temporary copy of the list
|
||||
PyObject *keys, *vals;
|
||||
PyObject* result = PyDict_New();
|
||||
ARROW_RETURN_NOT_OK(DeserializeList(data->field(0), start_idx, stop_idx, base, &keys));
|
||||
ARROW_RETURN_NOT_OK(DeserializeList(data->field(1), start_idx, stop_idx, base, &vals));
|
||||
for (size_t i = start_idx; i < stop_idx; ++i) {
|
||||
PyDict_SetItem(
|
||||
result, PyList_GetItem(keys, i - start_idx), PyList_GetItem(vals, i - start_idx));
|
||||
}
|
||||
Py_XDECREF(keys); // PyList_GetItem(keys, ...) incremented the reference count
|
||||
Py_XDECREF(vals); // PyList_GetItem(vals, ...) incremented the reference count
|
||||
static PyObject* py_type = PyUnicode_FromString("_pytype_");
|
||||
if (PyDict_Contains(result, py_type) && numbuf_deserialize_callback) {
|
||||
PyObject* arglist = Py_BuildValue("(O)", result);
|
||||
// The result of the call to PyObject_CallObject will be passed to Python
|
||||
// and its reference count will be decremented by the interpreter.
|
||||
PyObject* callback_result = PyObject_CallObject(numbuf_deserialize_callback, arglist);
|
||||
Py_XDECREF(arglist);
|
||||
Py_XDECREF(result);
|
||||
result = callback_result;
|
||||
if (!callback_result) { return Status::NotImplemented("python error"); }
|
||||
}
|
||||
*out = result;
|
||||
return Status::OK();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
#ifndef PYNUMBUF_PYTHON_H
|
||||
#define PYNUMBUF_PYTHON_H
|
||||
|
||||
#include <Python.h>
|
||||
|
||||
#include <arrow/api.h>
|
||||
#include <numbuf/dict.h>
|
||||
#include <numbuf/sequence.h>
|
||||
|
||||
#include "numpy.h"
|
||||
|
||||
namespace numbuf {
|
||||
|
||||
arrow::Status SerializeSequences(std::vector<PyObject*> sequences,
|
||||
int32_t recursion_depth, std::shared_ptr<arrow::Array>* out);
|
||||
arrow::Status SerializeDict(std::vector<PyObject*> dicts, int32_t recursion_depth,
|
||||
std::shared_ptr<arrow::Array>* out);
|
||||
arrow::Status DeserializeList(std::shared_ptr<arrow::Array> array, int32_t start_idx,
|
||||
int32_t stop_idx, PyObject* base, PyObject** out);
|
||||
arrow::Status DeserializeTuple(std::shared_ptr<arrow::Array> array, int32_t start_idx,
|
||||
int32_t stop_idx, PyObject* base, PyObject** out);
|
||||
arrow::Status DeserializeDict(std::shared_ptr<arrow::Array> array, int32_t start_idx,
|
||||
int32_t stop_idx, PyObject* base, PyObject** out);
|
||||
}
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,54 @@
|
||||
#ifndef PYNUMBUF_SCALARS_H
|
||||
#define PYNUMBUF_SCALARS_H
|
||||
|
||||
#include <arrow/api.h>
|
||||
|
||||
#include <Python.h>
|
||||
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
|
||||
#define NO_IMPORT_ARRAY
|
||||
#define PY_ARRAY_UNIQUE_SYMBOL NUMBUF_ARRAY_API
|
||||
#include <numpy/arrayobject.h>
|
||||
#include <numpy/arrayscalars.h>
|
||||
|
||||
#include <numbuf/sequence.h>
|
||||
|
||||
namespace numbuf {
|
||||
|
||||
arrow::Status AppendScalar(PyObject* obj, SequenceBuilder& builder) {
|
||||
if (PyArray_IsScalar(obj, Bool)) {
|
||||
return builder.AppendBool(((PyBoolScalarObject*)obj)->obval != 0);
|
||||
} else if (PyArray_IsScalar(obj, Float)) {
|
||||
return builder.AppendFloat(((PyFloatScalarObject*)obj)->obval);
|
||||
} else if (PyArray_IsScalar(obj, Double)) {
|
||||
return builder.AppendDouble(((PyDoubleScalarObject*)obj)->obval);
|
||||
}
|
||||
int64_t value = 0;
|
||||
if (PyArray_IsScalar(obj, Byte)) {
|
||||
value = ((PyByteScalarObject*)obj)->obval;
|
||||
} else if (PyArray_IsScalar(obj, UByte)) {
|
||||
value = ((PyUByteScalarObject*)obj)->obval;
|
||||
} else if (PyArray_IsScalar(obj, Short)) {
|
||||
value = ((PyShortScalarObject*)obj)->obval;
|
||||
} else if (PyArray_IsScalar(obj, UShort)) {
|
||||
value = ((PyUShortScalarObject*)obj)->obval;
|
||||
} else if (PyArray_IsScalar(obj, Int)) {
|
||||
value = ((PyIntScalarObject*)obj)->obval;
|
||||
} else if (PyArray_IsScalar(obj, UInt)) {
|
||||
value = ((PyUIntScalarObject*)obj)->obval;
|
||||
} else if (PyArray_IsScalar(obj, Long)) {
|
||||
value = ((PyLongScalarObject*)obj)->obval;
|
||||
} else if (PyArray_IsScalar(obj, ULong)) {
|
||||
value = ((PyULongScalarObject*)obj)->obval;
|
||||
} else if (PyArray_IsScalar(obj, LongLong)) {
|
||||
value = ((PyLongLongScalarObject*)obj)->obval;
|
||||
} else if (PyArray_IsScalar(obj, ULongLong)) {
|
||||
value = ((PyULongLongScalarObject*)obj)->obval;
|
||||
} else {
|
||||
DCHECK(false) << "scalar type not recognized";
|
||||
}
|
||||
return builder.AppendInt64(value);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
#endif // PYNUMBUF_SCALARS_H
|
||||
@@ -0,0 +1,67 @@
|
||||
#ifndef PYNUMBUF_MEMORY_H
|
||||
#define PYNUMBUF_MEMORY_H
|
||||
|
||||
#include <arrow/io/interfaces.h>
|
||||
|
||||
namespace numbuf {
|
||||
|
||||
class FixedBufferStream : public arrow::io::OutputStream,
|
||||
public arrow::io::ReadableFileInterface {
|
||||
public:
|
||||
virtual ~FixedBufferStream() {}
|
||||
|
||||
explicit FixedBufferStream(uint8_t* data, int64_t nbytes)
|
||||
: data_(data), position_(0), size_(nbytes) {}
|
||||
|
||||
arrow::Status Read(int64_t nbytes, std::shared_ptr<arrow::Buffer>* out) override {
|
||||
DCHECK(out);
|
||||
DCHECK(position_ + nbytes <= size_) << "position: " << position_
|
||||
<< " nbytes: " << nbytes << "size: " << size_;
|
||||
*out = std::make_shared<arrow::Buffer>(data_ + position_, nbytes);
|
||||
position_ += nbytes;
|
||||
return arrow::Status::OK();
|
||||
}
|
||||
|
||||
arrow::Status Read(int64_t nbytes, int64_t* bytes_read, uint8_t* out) {
|
||||
assert(0);
|
||||
return arrow::Status::OK();
|
||||
}
|
||||
|
||||
arrow::Status Seek(int64_t position) override {
|
||||
position_ = position;
|
||||
return arrow::Status::OK();
|
||||
}
|
||||
|
||||
arrow::Status Close() override { return arrow::Status::OK(); }
|
||||
|
||||
arrow::Status Tell(int64_t* position) override {
|
||||
*position = position_;
|
||||
return arrow::Status::OK();
|
||||
}
|
||||
|
||||
arrow::Status Write(const uint8_t* data, int64_t nbytes) override {
|
||||
DCHECK(position_ >= 0 && position_ < size_);
|
||||
DCHECK(position_ + nbytes <= size_) << "position: " << position_
|
||||
<< " nbytes: " << nbytes << "size: " << size_;
|
||||
uint8_t* dst = data_ + position_;
|
||||
memcpy(dst, data, nbytes);
|
||||
position_ += nbytes;
|
||||
return arrow::Status::OK();
|
||||
}
|
||||
|
||||
arrow::Status GetSize(int64_t* size) override {
|
||||
*size = size_;
|
||||
return arrow::Status::OK();
|
||||
}
|
||||
|
||||
bool supports_zero_copy() const override { return true; }
|
||||
|
||||
private:
|
||||
uint8_t* data_;
|
||||
int64_t position_;
|
||||
int64_t size_;
|
||||
};
|
||||
|
||||
} // namespace numbuf
|
||||
|
||||
#endif // PYNUMBUF_MEMORY_H
|
||||
@@ -0,0 +1,421 @@
|
||||
#include <Python.h>
|
||||
#include <arrow/api.h>
|
||||
#include <arrow/ipc/adapter.h>
|
||||
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
|
||||
#define PY_ARRAY_UNIQUE_SYMBOL NUMBUF_ARRAY_API
|
||||
#include <numpy/arrayobject.h>
|
||||
|
||||
#include "bytesobject.h"
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include <arrow/ipc/metadata.h>
|
||||
|
||||
#include "adapters/python.h"
|
||||
#include "memory.h"
|
||||
|
||||
#ifdef HAS_PLASMA
|
||||
extern "C" {
|
||||
#include "format/plasma_reader.h"
|
||||
#include "plasma_client.h"
|
||||
}
|
||||
|
||||
PyObject* NumbufPlasmaOutOfMemoryError;
|
||||
PyObject* NumbufPlasmaObjectExistsError;
|
||||
#endif
|
||||
|
||||
using namespace arrow;
|
||||
using namespace numbuf;
|
||||
|
||||
int64_t make_schema_and_batch(std::shared_ptr<Array> data,
|
||||
std::shared_ptr<Buffer>* metadata_out, std::shared_ptr<RecordBatch>* batch_out) {
|
||||
auto field = std::make_shared<Field>("list", data->type());
|
||||
std::shared_ptr<Schema> schema(new Schema({field}));
|
||||
*batch_out =
|
||||
std::shared_ptr<RecordBatch>(new RecordBatch(schema, data->length(), {data}));
|
||||
int64_t size = 0;
|
||||
ARROW_CHECK_OK(ipc::GetRecordBatchSize(batch_out->get(), &size));
|
||||
ARROW_CHECK_OK(ipc::WriteSchema((*batch_out)->schema().get(), metadata_out));
|
||||
return size;
|
||||
}
|
||||
|
||||
Status read_batch(std::shared_ptr<Buffer> schema_buffer, int64_t header_end_offset,
|
||||
uint8_t* data, int64_t size, std::shared_ptr<RecordBatch>* batch_out) {
|
||||
std::shared_ptr<ipc::Message> message;
|
||||
RETURN_NOT_OK(ipc::Message::Open(schema_buffer, &message));
|
||||
DCHECK_EQ(ipc::Message::SCHEMA, message->type());
|
||||
std::shared_ptr<ipc::SchemaMessage> schema_msg = message->GetSchema();
|
||||
std::shared_ptr<Schema> schema;
|
||||
RETURN_NOT_OK(schema_msg->GetSchema(&schema));
|
||||
auto source = std::make_shared<FixedBufferStream>(data, size);
|
||||
std::shared_ptr<arrow::ipc::RecordBatchReader> reader;
|
||||
RETURN_NOT_OK(ipc::RecordBatchReader::Open(source.get(), header_end_offset, &reader));
|
||||
RETURN_NOT_OK(reader->GetRecordBatch(schema, batch_out));
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
extern "C" {
|
||||
|
||||
#define CHECK_SERIALIZATION_ERROR(STATUS) \
|
||||
do { \
|
||||
Status _s = (STATUS); \
|
||||
if (!_s.ok()) { \
|
||||
/* If this condition is true, there was an error in the callback that \
|
||||
* needs to be passed through */ \
|
||||
if (!PyErr_Occurred()) { PyErr_SetString(NumbufError, _s.ToString().c_str()); } \
|
||||
return NULL; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
static PyObject* NumbufError;
|
||||
|
||||
PyObject* numbuf_serialize_callback = NULL;
|
||||
PyObject* numbuf_deserialize_callback = NULL;
|
||||
|
||||
int PyObjectToArrow(PyObject* object, std::shared_ptr<RecordBatch>** result) {
|
||||
if (PyCapsule_IsValid(object, "arrow")) {
|
||||
*result = reinterpret_cast<std::shared_ptr<RecordBatch>*>(
|
||||
PyCapsule_GetPointer(object, "arrow"));
|
||||
return 1;
|
||||
} else {
|
||||
PyErr_SetString(PyExc_TypeError, "must be an 'arrow' capsule");
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
|
||||
static void ArrowCapsule_Destructor(PyObject* capsule) {
|
||||
delete reinterpret_cast<std::shared_ptr<RecordBatch>*>(
|
||||
PyCapsule_GetPointer(capsule, "arrow"));
|
||||
}
|
||||
|
||||
/* Documented in doc/numbuf.rst in ray-core */
|
||||
static PyObject* serialize_list(PyObject* self, PyObject* args) {
|
||||
PyObject* value;
|
||||
if (!PyArg_ParseTuple(args, "O", &value)) { return NULL; }
|
||||
std::shared_ptr<Array> array;
|
||||
if (PyList_Check(value)) {
|
||||
int32_t recursion_depth = 0;
|
||||
Status s =
|
||||
SerializeSequences(std::vector<PyObject*>({value}), recursion_depth, &array);
|
||||
CHECK_SERIALIZATION_ERROR(s);
|
||||
|
||||
auto batch = new std::shared_ptr<RecordBatch>();
|
||||
std::shared_ptr<Buffer> metadata;
|
||||
int64_t size = make_schema_and_batch(array, &metadata, batch);
|
||||
|
||||
auto ptr = reinterpret_cast<const char*>(metadata->data());
|
||||
PyObject* r = PyTuple_New(3);
|
||||
PyTuple_SetItem(r, 0, PyByteArray_FromStringAndSize(ptr, metadata->size()));
|
||||
PyTuple_SetItem(r, 1, PyLong_FromLong(size));
|
||||
PyTuple_SetItem(r, 2,
|
||||
PyCapsule_New(reinterpret_cast<void*>(batch), "arrow", &ArrowCapsule_Destructor));
|
||||
return r;
|
||||
}
|
||||
return NULL;
|
||||
}
|
||||
|
||||
/* Documented in doc/numbuf.rst in ray-core */
|
||||
static PyObject* write_to_buffer(PyObject* self, PyObject* args) {
|
||||
std::shared_ptr<RecordBatch>* batch;
|
||||
PyObject* memoryview;
|
||||
if (!PyArg_ParseTuple(args, "O&O", &PyObjectToArrow, &batch, &memoryview)) {
|
||||
return NULL;
|
||||
}
|
||||
if (!PyMemoryView_Check(memoryview)) { return NULL; }
|
||||
Py_buffer* buffer = PyMemoryView_GET_BUFFER(memoryview);
|
||||
auto target = std::make_shared<FixedBufferStream>(
|
||||
reinterpret_cast<uint8_t*>(buffer->buf), buffer->len);
|
||||
int64_t body_end_offset;
|
||||
int64_t header_end_offset;
|
||||
ARROW_CHECK_OK(ipc::WriteRecordBatch((*batch)->columns(), (*batch)->num_rows(),
|
||||
target.get(), &body_end_offset, &header_end_offset));
|
||||
return PyLong_FromLong(header_end_offset);
|
||||
}
|
||||
|
||||
/* Documented in doc/numbuf.rst in ray-core */
|
||||
static PyObject* read_from_buffer(PyObject* self, PyObject* args) {
|
||||
PyObject* data_memoryview;
|
||||
PyObject* metadata_memoryview;
|
||||
int64_t header_end_offset;
|
||||
if (!PyArg_ParseTuple(
|
||||
args, "OOL", &data_memoryview, &metadata_memoryview, &header_end_offset)) {
|
||||
return NULL;
|
||||
}
|
||||
|
||||
Py_buffer* metadata_buffer = PyMemoryView_GET_BUFFER(metadata_memoryview);
|
||||
Py_buffer* data_buffer = PyMemoryView_GET_BUFFER(data_memoryview);
|
||||
auto ptr = reinterpret_cast<uint8_t*>(metadata_buffer->buf);
|
||||
auto schema_buffer = std::make_shared<Buffer>(ptr, metadata_buffer->len);
|
||||
|
||||
auto batch = new std::shared_ptr<arrow::RecordBatch>();
|
||||
ARROW_CHECK_OK(read_batch(schema_buffer, header_end_offset,
|
||||
reinterpret_cast<uint8_t*>(data_buffer->buf), data_buffer->len, batch));
|
||||
|
||||
return PyCapsule_New(reinterpret_cast<void*>(batch), "arrow", &ArrowCapsule_Destructor);
|
||||
}
|
||||
|
||||
/* Documented in doc/numbuf.rst in ray-core */
|
||||
static PyObject* deserialize_list(PyObject* self, PyObject* args) {
|
||||
std::shared_ptr<RecordBatch>* data;
|
||||
PyObject* base = Py_None;
|
||||
if (!PyArg_ParseTuple(args, "O&|O", &PyObjectToArrow, &data, &base)) { return NULL; }
|
||||
PyObject* result;
|
||||
Status s = DeserializeList((*data)->column(0), 0, (*data)->num_rows(), base, &result);
|
||||
CHECK_SERIALIZATION_ERROR(s);
|
||||
return result;
|
||||
}
|
||||
|
||||
static PyObject* register_callbacks(PyObject* self, PyObject* args) {
|
||||
PyObject* result = NULL;
|
||||
PyObject* serialize_callback;
|
||||
PyObject* deserialize_callback;
|
||||
if (PyArg_ParseTuple(
|
||||
args, "OO:register_callbacks", &serialize_callback, &deserialize_callback)) {
|
||||
if (!PyCallable_Check(serialize_callback)) {
|
||||
PyErr_SetString(PyExc_TypeError, "serialize_callback must be callable");
|
||||
return NULL;
|
||||
}
|
||||
if (!PyCallable_Check(deserialize_callback)) {
|
||||
PyErr_SetString(PyExc_TypeError, "deserialize_callback must be callable");
|
||||
return NULL;
|
||||
}
|
||||
Py_XINCREF(serialize_callback); // Add a reference to new serialization callback
|
||||
Py_XINCREF(deserialize_callback); // Add a reference to new deserialization callback
|
||||
Py_XDECREF(numbuf_serialize_callback); // Dispose of old serialization callback
|
||||
Py_XDECREF(numbuf_deserialize_callback); // Dispose of old deserialization callback
|
||||
numbuf_serialize_callback = serialize_callback;
|
||||
numbuf_deserialize_callback = deserialize_callback;
|
||||
Py_INCREF(Py_None);
|
||||
result = Py_None;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
#ifdef HAS_PLASMA
|
||||
|
||||
#include "plasma_extension.h"
|
||||
|
||||
/**
|
||||
* Release the object when its associated PyCapsule goes out of scope.
|
||||
*
|
||||
* The PyCapsule is used as the base object for the Python object that
|
||||
* is stored with store_list and retrieved with retrieve_list. The base
|
||||
* object ensures that the reference count of the capsule is non-zero
|
||||
* during the lifetime of the Python object returned by retrieve_list.
|
||||
*
|
||||
* @param capsule The capsule that went out of scope.
|
||||
* @return Void.
|
||||
*/
|
||||
static void BufferCapsule_Destructor(PyObject* capsule) {
|
||||
object_id* id = reinterpret_cast<object_id*>(PyCapsule_GetPointer(capsule, "buffer"));
|
||||
auto context = reinterpret_cast<PyObject*>(PyCapsule_GetContext(capsule));
|
||||
/* We use the context of the connection capsule to indicate if the connection
|
||||
* is still active (if the context is NULL) or if it is closed (if the context
|
||||
* is (void*) 0x1). This is neccessary because the primary pointer of the
|
||||
* capsule cannot be NULL. */
|
||||
if (PyCapsule_GetContext(context) == NULL) {
|
||||
plasma_connection* conn;
|
||||
CHECK(PyObjectToPlasmaConnection(context, &conn));
|
||||
plasma_release(conn, *id);
|
||||
}
|
||||
Py_XDECREF(context);
|
||||
delete id;
|
||||
}
|
||||
|
||||
/**
|
||||
* Store a PyList in the plasma store.
|
||||
*
|
||||
* This function converts the PyList into an arrow RecordBatch, constructs the
|
||||
* metadata (schema) of the PyList, creates a new plasma object, puts the data
|
||||
* into the plasma buffer and the schema into the plasma metadata. This raises
|
||||
*
|
||||
*
|
||||
* @param args Contains the object ID the list is stored under, the
|
||||
* connection to the plasma store and the PyList we want to store.
|
||||
* @return None.
|
||||
*/
|
||||
static PyObject* store_list(PyObject* self, PyObject* args) {
|
||||
object_id obj_id;
|
||||
plasma_connection* conn;
|
||||
PyObject* value;
|
||||
if (!PyArg_ParseTuple(args, "O&O&O", PyStringToUniqueID, &obj_id,
|
||||
PyObjectToPlasmaConnection, &conn, &value)) {
|
||||
return NULL;
|
||||
}
|
||||
if (!PyList_Check(value)) { return NULL; }
|
||||
|
||||
std::shared_ptr<Array> array;
|
||||
int32_t recursion_depth = 0;
|
||||
Status s = SerializeSequences(std::vector<PyObject*>({value}), recursion_depth, &array);
|
||||
CHECK_SERIALIZATION_ERROR(s);
|
||||
|
||||
std::shared_ptr<RecordBatch> batch;
|
||||
std::shared_ptr<Buffer> metadata;
|
||||
int64_t size = make_schema_and_batch(array, &metadata, &batch);
|
||||
|
||||
uint8_t* data;
|
||||
/* The arrow schema is stored as the metadata of the plasma object and
|
||||
* both the arrow data and the header end offset are
|
||||
* stored in the plasma data buffer. The header end offset is stored in
|
||||
* the first sizeof(int64_t) bytes of the data buffer. The RecordBatch
|
||||
* data is stored after that. */
|
||||
int error_code = plasma_create(conn, obj_id, sizeof(size) + size,
|
||||
(uint8_t*)metadata->data(), metadata->size(), &data);
|
||||
if (error_code == PlasmaError_ObjectExists) {
|
||||
PyErr_SetString(NumbufPlasmaObjectExistsError,
|
||||
"An object with this ID already exists in the plasma "
|
||||
"store.");
|
||||
return NULL;
|
||||
}
|
||||
if (error_code == PlasmaError_OutOfMemory) {
|
||||
PyErr_SetString(NumbufPlasmaOutOfMemoryError,
|
||||
"The plasma store ran out of memory and could not create "
|
||||
"this object.");
|
||||
return NULL;
|
||||
}
|
||||
CHECK(error_code == PlasmaError_OK);
|
||||
|
||||
auto target = std::make_shared<FixedBufferStream>(sizeof(size) + data, size);
|
||||
int64_t body_end_offset;
|
||||
int64_t header_end_offset;
|
||||
ARROW_CHECK_OK(ipc::WriteRecordBatch(batch->columns(), batch->num_rows(), target.get(),
|
||||
&body_end_offset, &header_end_offset));
|
||||
|
||||
/* Save the header end offset at the beginning of the plasma data buffer. */
|
||||
*((int64_t*)data) = header_end_offset;
|
||||
/* Do the plasma_release corresponding to the call to plasma_create. */
|
||||
plasma_release(conn, obj_id);
|
||||
/* Seal the object. */
|
||||
plasma_seal(conn, obj_id);
|
||||
Py_RETURN_NONE;
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieve a PyList from the plasma store.
|
||||
*
|
||||
* This reads the arrow schema from the plasma metadata, constructs
|
||||
* Python objects from the plasma data according to the schema and
|
||||
* returns the object.
|
||||
*
|
||||
* @param args Object ID of the PyList to be retrieved and connection to the
|
||||
* plasma store.
|
||||
* @return The PyList.
|
||||
*/
|
||||
static PyObject* retrieve_list(PyObject* self, PyObject* args) {
|
||||
object_id obj_id;
|
||||
PyObject* plasma_conn;
|
||||
if (!PyArg_ParseTuple(args, "O&O", PyStringToUniqueID, &obj_id, &plasma_conn)) {
|
||||
return NULL;
|
||||
}
|
||||
plasma_connection* conn;
|
||||
if (!PyObjectToPlasmaConnection(plasma_conn, &conn)) { return NULL; }
|
||||
object_id* buffer_obj_id = new object_id(obj_id);
|
||||
/* This keeps a Plasma buffer in scope as long as an object that is backed by that
|
||||
* buffer is in scope. This prevents memory in the object store from getting
|
||||
* released while it is still being used to back a Python object. */
|
||||
PyObject* base = PyCapsule_New(buffer_obj_id, "buffer", BufferCapsule_Destructor);
|
||||
PyCapsule_SetContext(base, plasma_conn);
|
||||
Py_XINCREF(plasma_conn);
|
||||
|
||||
int64_t size, metadata_size;
|
||||
uint8_t *data, *metadata;
|
||||
plasma_get(conn, obj_id, &size, &data, &metadata_size, &metadata);
|
||||
|
||||
/* Remember: The metadata offset was written at the beginning of the plasma buffer. */
|
||||
int64_t header_end_offset = *((int64_t*)data);
|
||||
auto schema_buffer = std::make_shared<Buffer>(metadata, metadata_size);
|
||||
auto batch = std::shared_ptr<RecordBatch>();
|
||||
ARROW_CHECK_OK(read_batch(schema_buffer, header_end_offset, data + sizeof(size),
|
||||
size - sizeof(size), &batch));
|
||||
|
||||
PyObject* result;
|
||||
Status s = DeserializeList(batch->column(0), 0, batch->num_rows(), base, &result);
|
||||
CHECK_SERIALIZATION_ERROR(s);
|
||||
Py_XDECREF(base);
|
||||
return result;
|
||||
}
|
||||
|
||||
#endif // HAS_PLASMA
|
||||
|
||||
static PyMethodDef NumbufMethods[] = {
|
||||
{"serialize_list", serialize_list, METH_VARARGS, "serialize a Python list"},
|
||||
{"deserialize_list", deserialize_list, METH_VARARGS, "deserialize a Python list"},
|
||||
{"write_to_buffer", write_to_buffer, METH_VARARGS, "write serialized data to buffer"},
|
||||
{"read_from_buffer", read_from_buffer, METH_VARARGS,
|
||||
"read serialized data from buffer"},
|
||||
{"register_callbacks", register_callbacks, METH_VARARGS,
|
||||
"set serialization and deserialization callbacks"},
|
||||
#ifdef HAS_PLASMA
|
||||
{"store_list", store_list, METH_VARARGS, "store a Python list in plasma"},
|
||||
{"retrieve_list", retrieve_list, METH_VARARGS, "retrieve a Python list from plasma"},
|
||||
#endif
|
||||
{NULL, NULL, 0, NULL}};
|
||||
|
||||
// clang-format off
|
||||
#if PY_MAJOR_VERSION >= 3
|
||||
static struct PyModuleDef moduledef = {
|
||||
PyModuleDef_HEAD_INIT,
|
||||
"libnumbuf", /* m_name */
|
||||
"Python C Extension for Numbuf", /* m_doc */
|
||||
0, /* m_size */
|
||||
NumbufMethods, /* m_methods */
|
||||
NULL, /* m_reload */
|
||||
NULL, /* m_traverse */
|
||||
NULL, /* m_clear */
|
||||
NULL, /* m_free */
|
||||
};
|
||||
#endif
|
||||
// clang-format on
|
||||
|
||||
#if PY_MAJOR_VERSION >= 3
|
||||
#define INITERROR return NULL
|
||||
#else
|
||||
#define INITERROR return
|
||||
#endif
|
||||
|
||||
#ifndef PyMODINIT_FUNC /* declarations for DLL import/export */
|
||||
#define PyMODINIT_FUNC void
|
||||
#endif
|
||||
|
||||
#if PY_MAJOR_VERSION >= 3
|
||||
#define MOD_INIT(name) PyMODINIT_FUNC PyInit_##name(void)
|
||||
#else
|
||||
#define MOD_INIT(name) PyMODINIT_FUNC init##name(void)
|
||||
#endif
|
||||
|
||||
MOD_INIT(libnumbuf) {
|
||||
#if PY_MAJOR_VERSION >= 3
|
||||
PyObject* m = PyModule_Create(&moduledef);
|
||||
#else
|
||||
PyObject* m =
|
||||
Py_InitModule3("libnumbuf", NumbufMethods, "Python C Extension for Numbuf");
|
||||
#endif
|
||||
|
||||
#if HAS_PLASMA
|
||||
/* Create a custom exception for when an object ID is reused. */
|
||||
char numbuf_plasma_object_exists_error[] = "numbuf_plasma_object_exists.error";
|
||||
NumbufPlasmaObjectExistsError =
|
||||
PyErr_NewException(numbuf_plasma_object_exists_error, NULL, NULL);
|
||||
Py_INCREF(NumbufPlasmaObjectExistsError);
|
||||
PyModule_AddObject(
|
||||
m, "pnumbuf_lasma_object_exists_error", NumbufPlasmaObjectExistsError);
|
||||
/* Create a custom exception for when the plasma store is out of memory. */
|
||||
char numbuf_plasma_out_of_memory_error[] = "numbuf_plasma_out_of_memory.error";
|
||||
NumbufPlasmaOutOfMemoryError =
|
||||
PyErr_NewException(numbuf_plasma_out_of_memory_error, NULL, NULL);
|
||||
Py_INCREF(NumbufPlasmaOutOfMemoryError);
|
||||
PyModule_AddObject(
|
||||
m, "numbuf_plasma_out_of_memory_error", NumbufPlasmaOutOfMemoryError);
|
||||
#endif
|
||||
|
||||
char numbuf_error[] = "numbuf.error";
|
||||
NumbufError = PyErr_NewException(numbuf_error, NULL, NULL);
|
||||
Py_INCREF(NumbufError);
|
||||
PyModule_AddObject(m, "numbuf_error", NumbufError);
|
||||
import_array();
|
||||
|
||||
#if PY_MAJOR_VERSION >= 3
|
||||
return m;
|
||||
#endif
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,130 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import unittest
|
||||
import numbuf
|
||||
import numpy as np
|
||||
from numpy.testing import assert_equal
|
||||
import sys
|
||||
|
||||
TEST_OBJECTS = [{(1,2) : 1}, {() : 2}, [1, "hello", 3.0], 42, 43, "hello world",
|
||||
u"x", u"\u262F", 42.0,
|
||||
1 << 62, (1.0, "hi"),
|
||||
None, (None, None), ("hello", None),
|
||||
True, False, (True, False), "hello",
|
||||
{True: "hello", False: "world"},
|
||||
{"hello" : "world", 1: 42, 1.0: 45}, {},
|
||||
np.int8(3), np.int32(4), np.int64(5),
|
||||
np.uint8(3), np.uint32(4), np.uint64(5),
|
||||
np.float32(1.0), np.float64(1.0)]
|
||||
|
||||
if sys.version_info < (3, 0):
|
||||
TEST_OBJECTS += [long(42), long(1 << 62)]
|
||||
|
||||
class SerializationTests(unittest.TestCase):
|
||||
|
||||
def roundTripTest(self, data):
|
||||
schema, size, serialized = numbuf.serialize_list(data)
|
||||
result = numbuf.deserialize_list(serialized)
|
||||
assert_equal(data, result)
|
||||
|
||||
def testSimple(self):
|
||||
self.roundTripTest([1, 2, 3])
|
||||
self.roundTripTest([1.0, 2.0, 3.0])
|
||||
self.roundTripTest(['hello', 'world'])
|
||||
self.roundTripTest([1, 'hello', 1.0])
|
||||
self.roundTripTest([{'hello': 1.0, 'world': 42}])
|
||||
self.roundTripTest([True, False])
|
||||
|
||||
def testNone(self):
|
||||
self.roundTripTest([1, 2, None, 3])
|
||||
|
||||
def testNested(self):
|
||||
self.roundTripTest([{"hello": {"world": (1, 2, 3)}}])
|
||||
self.roundTripTest([((1,), (1, 2, 3, (4, 5, 6), "string"))])
|
||||
self.roundTripTest([{"hello": [1, 2, 3]}])
|
||||
self.roundTripTest([{"hello": [1, [2, 3]]}])
|
||||
self.roundTripTest([{"hello": (None, 2, [3, 4])}])
|
||||
self.roundTripTest([{"hello": (None, 2, [3, 4], np.array([1.0, 2.0, 3.0]))}])
|
||||
|
||||
def numpyTest(self, t):
|
||||
a = np.random.randint(0, 10, size=(100, 100)).astype(t)
|
||||
self.roundTripTest([a])
|
||||
|
||||
def testArrays(self):
|
||||
for t in ["int8", "uint8", "int16", "uint16", "int32", "uint32", "float32", "float64"]:
|
||||
self.numpyTest(t)
|
||||
|
||||
def testRay(self):
|
||||
for obj in TEST_OBJECTS:
|
||||
self.roundTripTest([obj])
|
||||
|
||||
def testCallback(self):
|
||||
|
||||
class Foo(object):
|
||||
def __init__(self):
|
||||
self.x = 1
|
||||
|
||||
class Bar(object):
|
||||
def __init__(self):
|
||||
self.foo = Foo()
|
||||
|
||||
def serialize(obj):
|
||||
return dict(obj.__dict__, **{"_pytype_": type(obj).__name__})
|
||||
|
||||
def deserialize(obj):
|
||||
if obj["_pytype_"] == "Foo":
|
||||
result = Foo()
|
||||
elif obj["_pytype_"] == "Bar":
|
||||
result = Bar()
|
||||
|
||||
obj.pop("_pytype_", None)
|
||||
result.__dict__ = obj
|
||||
return result
|
||||
|
||||
bar = Bar()
|
||||
bar.foo.x = 42
|
||||
|
||||
numbuf.register_callbacks(serialize, deserialize)
|
||||
|
||||
metadata, size, serialized = numbuf.serialize_list([bar])
|
||||
self.assertEqual(numbuf.deserialize_list(serialized)[0].foo.x, 42)
|
||||
|
||||
def testObjectArray(self):
|
||||
x = np.array([1, 2, "hello"], dtype=object)
|
||||
y = np.array([[1, 2], [3, 4]], dtype=object)
|
||||
|
||||
def myserialize(obj):
|
||||
return {"_pytype_": "numpy.array", "data": obj.tolist()}
|
||||
|
||||
def mydeserialize(obj):
|
||||
if obj["_pytype_"] == "numpy.array":
|
||||
return np.array(obj["data"], dtype=object)
|
||||
|
||||
numbuf.register_callbacks(myserialize, mydeserialize)
|
||||
|
||||
metadata, size, serialized = numbuf.serialize_list([x, y])
|
||||
|
||||
assert_equal(numbuf.deserialize_list(serialized), [x, y])
|
||||
|
||||
def testBuffer(self):
|
||||
for (i, obj) in enumerate(TEST_OBJECTS):
|
||||
schema, size, batch = numbuf.serialize_list([obj])
|
||||
size = size + 4096 # INITIAL_METADATA_SIZE in arrow.
|
||||
buff = np.zeros(size, dtype="uint8")
|
||||
metadata_offset = numbuf.write_to_buffer(batch, memoryview(buff))
|
||||
array = numbuf.read_from_buffer(memoryview(buff), memoryview(schema), metadata_offset)
|
||||
result = numbuf.deserialize_list(array)
|
||||
assert_equal(result[0], obj)
|
||||
|
||||
def testObjectArrayImmutable(self):
|
||||
obj = np.zeros([10])
|
||||
schema, size, serialized = numbuf.serialize_list([obj])
|
||||
result = numbuf.deserialize_list(serialized)
|
||||
assert_equal(result[0], obj)
|
||||
with self.assertRaises(ValueError):
|
||||
result[0][0] = 1
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
Reference in New Issue
Block a user