Merge numbuf.

numbuf/python/src/pynumbuf/adapters/numpy.cc

@@ -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);                                                        \
+  }                                                                          \
+  return Status::OK();
+
+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;
+  }
+  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"); // TODO(pcm): https://github.com/ray-project/numbuf/issues/10
+        }
+        builder.AppendDict(PyDict_Size(result));
+        subdicts.push_back(result);
+      }
+  }
+  Py_XDECREF(contiguous);
+  return Status::OK();
+}
+
+}

numbuf/python/src/pynumbuf/adapters/numpy.h

@@ -0,0 +1,22 @@
+#ifndef PYNUMBUF_NUMPY_H
+#define PYNUMBUF_NUMPY_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 <numbuf/tensor.h>
+#include <numbuf/sequence.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

numbuf/python/src/pynumbuf/adapters/python.cc

@@ -0,0 +1,281 @@
+#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 {
+
+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 = PyString_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");
+    }
+  } else if (PyInt_Check(elem)) {
+    RETURN_NOT_OK(builder.AppendInt64(static_cast<int64_t>(PyInt_AS_LONG(elem))));
+  } else if (PyString_Check(elem)) {
+    auto data = reinterpret_cast<uint8_t*>(PyString_AS_STRING(elem));
+    auto size = PyString_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); // TODO(pcm): Check if this is correct
+    #else
+      PyObject* str = PyUnicode_AsUTF8String(elem);
+      char* data = PyString_AS_STRING(str);
+      size = PyString_GET_SIZE(str);
+    #endif
+    Status s = builder.AppendString(data, size);
+    Py_XDECREF(str);
+    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 " << PyString_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"); // TODO(pcm): https://github.com/ray-project/numbuf/issues/10
+      }
+      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 = PyString_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 = PyString_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"); // TODO(pcm): https://github.com/ray-project/numbuf/issues/10
+    }
+  }
+  *out = result;
+  return Status::OK();
+}
+
+
+}

numbuf/python/src/pynumbuf/adapters/python.h

@@ -0,0 +1,22 @@
+#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

numbuf/python/src/pynumbuf/adapters/scalars.h

@@ -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

numbuf/python/src/pynumbuf/memory.h

@@ -0,0 +1,68 @@
+#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

numbuf/python/src/pynumbuf/numbuf.cc

@@ -0,0 +1,195 @@
+#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 <iostream>
+
+#include <arrow/ipc/metadata.h>
+
+#include "adapters/python.h"
+#include "memory.h"
+
+using namespace arrow;
+using namespace numbuf;
+
+std::shared_ptr<RecordBatch> make_row_batch(std::shared_ptr<Array> data) {
+  auto field = std::make_shared<Field>("list", data->type());
+  std::shared_ptr<Schema> schema(new Schema({field}));
+  return std::shared_ptr<RecordBatch>(new RecordBatch(schema, data->length(), {data}));
+}
+
+extern "C" {
+
+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);
+    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;
+    }
+
+    auto batch = new std::shared_ptr<RecordBatch>();
+    *batch = make_row_batch(array);
+
+    int64_t size = 0;
+    ARROW_CHECK_OK(arrow::ipc::GetRecordBatchSize(batch->get(), &size));
+
+    std::shared_ptr<Buffer> buffer;
+    ARROW_CHECK_OK(ipc::WriteSchema((*batch)->schema().get(), &buffer));
+    auto ptr = reinterpret_cast<const char*>(buffer->data());
+
+    PyObject* r = PyTuple_New(3);
+    PyTuple_SetItem(r, 0, PyByteArray_FromStringAndSize(ptr, buffer->size()));
+    PyTuple_SetItem(r, 1, PyInt_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 PyInt_FromLong(header_end_offset);
+}
+
+/* Documented in doc/numbuf.rst in ray-core */
+static PyObject* read_from_buffer(PyObject* self, PyObject* args) {
+  PyObject* memoryview;
+  PyObject* metadata;
+  int64_t metadata_offset;
+  if (!PyArg_ParseTuple(args, "OOL", &memoryview, &metadata, &metadata_offset)) {
+    return NULL;
+  }
+
+  auto ptr = reinterpret_cast<uint8_t*>(PyByteArray_AsString(metadata));
+  auto schema_buffer = std::make_shared<Buffer>(ptr, PyByteArray_Size(metadata));
+  std::shared_ptr<ipc::Message> message;
+  ARROW_CHECK_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;
+  ARROW_CHECK_OK(schema_msg->GetSchema(&schema));
+
+  Py_buffer* buffer = PyMemoryView_GET_BUFFER(memoryview);
+  auto source = std::make_shared<FixedBufferStream>(reinterpret_cast<uint8_t*>(buffer->buf), buffer->len);
+  std::shared_ptr<arrow::ipc::RecordBatchReader> reader;
+  ARROW_CHECK_OK(arrow::ipc::RecordBatchReader::Open(source.get(), metadata_offset, &reader));
+  auto batch = new std::shared_ptr<arrow::RecordBatch>();
+  ARROW_CHECK_OK(reader->GetRecordBatch(schema, 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);
+  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;
+  }
+  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;
+}
+
+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"},
+ { NULL, NULL, 0, NULL }
+};
+
+PyMODINIT_FUNC initlibnumbuf(void) {
+  PyObject* m;
+  m = Py_InitModule3("libnumbuf", NumbufMethods, "Python C Extension for Numbuf");
+  char numbuf_error[] = "numbuf.error";
+  NumbufError = PyErr_NewException(numbuf_error, NULL, NULL);
+  Py_INCREF(NumbufError);
+  PyModule_AddObject(m, "numbuf_error", NumbufError);
+  import_array();
+}
+
+}

numbuf/python/test/runtest.py

@@ -0,0 +1,114 @@
+import unittest
+import numbuf
+import numpy as np
+from numpy.testing import assert_equal
+
+TEST_OBJECTS = [{(1,2) : 1}, {() : 2}, [1, "hello", 3.0], 42, 43L, "hello world",
+                u"x", u"\u262F", 42.0,
+                1L << 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)]
+
+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.ndarray([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), schema, metadata_offset)
+      result = numbuf.deserialize_list(array)
+      assert_equal(result[0], obj)
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)