Draft PyReadableFile implementation, not yet tested

Change-Id: I2feee0fd6a6faa5c486ab4699a7cdc8c042d5726

python/CMakeLists.txt

@@ -432,6 +432,7 @@ set(PYARROW_SRCS
   src/pyarrow/common.cc
   src/pyarrow/config.cc
   src/pyarrow/helpers.cc
+  src/pyarrow/io.cc
   src/pyarrow/status.cc
 
   src/pyarrow/adapters/builtin.cc

python/src/pyarrow/common.cc

@@ -68,4 +68,19 @@ arrow::MemoryPool* GetMemoryPool() {
   return &memory_pool;
 }
 
+// ----------------------------------------------------------------------
+// PyBytesBuffer
+
+PyBytesBuffer::PyBytesBuffer(PyObject* obj)
+    : Buffer(reinterpret_cast<const uint8_t*>(PyBytes_AS_STRING(obj_)),
+        PyBytes_GET_SIZE(obj_)),
+      obj_(obj) {
+  Py_INCREF(obj_);
+}
+
+PyBytesBuffer::~PyBytesBuffer() {
+  PyAcquireGIL_RAII lock;
+  Py_DECREF(obj_);
+}
+
 } // namespace pyarrow

python/src/pyarrow/common.h

@@ -20,6 +20,7 @@
 
 #include "pyarrow/config.h"
 #include "arrow/util/buffer.h"
+#include "arrow/util/macros.h"
 #include "pyarrow/visibility.h"
 
 namespace arrow { class MemoryPool; }
@@ -81,6 +82,20 @@ struct PyObjectStringify {
   }
 };
 
+class PyAcquireGIL_RAII {
+ public:
+  PyAcquireGIL_RAII() {
+    state_ = PyGILState_Ensure();
+  }
+
+  ~PyAcquireGIL_RAII() {
+    PyGILState_Release(state_);
+  }
+ private:
+  PyGILState_STATE state_;
+  DISALLOW_COPY_AND_ASSIGN(PyAcquireGIL_RAII);
+};
+
 // TODO(wesm): We can just let errors pass through. To be explored later
 #define RETURN_IF_PYERROR()                         \
   if (PyErr_Occurred()) {                           \
@@ -98,8 +113,8 @@ PYARROW_EXPORT arrow::MemoryPool* GetMemoryPool();
 
 class PYARROW_EXPORT NumPyBuffer : public arrow::Buffer {
  public:
-  NumPyBuffer(PyArrayObject* arr) :
-      Buffer(nullptr, 0) {
+  NumPyBuffer(PyArrayObject* arr)
+    : Buffer(nullptr, 0) {
     arr_ = arr;
     Py_INCREF(arr);
 
@@ -115,6 +130,15 @@ class PYARROW_EXPORT NumPyBuffer : public arrow::Buffer {
   PyArrayObject* arr_;
 };
 
+class PYARROW_EXPORT PyBytesBuffer : public arrow::Buffer {
+ public:
+  PyBytesBuffer(PyObject* obj);
+  virtual ~PyBytesBuffer();
+
+ private:
+  PyObject* obj_;
+};
+
 } // namespace pyarrow
 
 #endif // PYARROW_COMMON_H

python/src/pyarrow/io.cc

@@ -15,13 +15,15 @@
 // specific language governing permissions and limitations
 // under the License.
 
-#include "pyarrow/common.h"
+#include "pyarrow/io.h"
 
 #include <cstdlib>
 #include <sstream>
 
 #include <arrow/util/memory-pool.h>
 #include <arrow/util/status.h>
+
+#include "pyarrow/common.h"
 #include "pyarrow/status.h"
 
 namespace pyarrow {
@@ -31,28 +33,105 @@ namespace pyarrow {
 
 PyReadableFile::PyReadableFile(PyObject* file)
     : file_(file) {
-  Py_INCREF(file);
+  Py_INCREF(file_);
 }
 
 PyReadableFile::~PyReadableFile() {
-  Py_DECREF();
+  Py_DECREF(file_);
+}
+
+#define ARROW_RETURN_IF_PYERROR()                   \
+  if (PyErr_Occurred()) {                           \
+    PyObject *exc_type, *exc_value, *traceback;     \
+    PyErr_Fetch(&exc_type, &exc_value, &traceback); \
+    PyObjectStringify stringified(exc_value);       \
+    std::string message(stringified.bytes);         \
+    Py_DECREF(exc_type);                            \
+    Py_DECREF(exc_value);                           \
+    Py_DECREF(traceback);                           \
+    PyErr_Clear();                                  \
+    return arrow::Status::IOError(message);    \
+  }
+
+static arrow::Status SeekNoGIL(PyObject* file, int64_t position, int whence) {
+  // whence: 0 for relative to start of file, 2 for end of file
+  PyObject* result = PyObject_CallMethod(file, "seek", "(i)", position);
+  Py_XDECREF(result);
+  ARROW_RETURN_IF_PYERROR();
+  return arrow::Status::OK();
+}
+
+static arrow::Status ReadNoGIL(PyObject* file, int64_t nbytes, PyObject** out) {
+  PyObject* result = PyObject_CallMethod(file, "read", "(i)", nbytes);
+  ARROW_RETURN_IF_PYERROR();
+  *out = result;
+  return arrow::Status::OK();
+}
+
+static arrow::Status TellNoGIL(PyObject* file, int64_t* position) {
+  PyObject* result = PyObject_CallMethod(file, "tell", "()");
+  ARROW_RETURN_IF_PYERROR();
+
+  *position = PyLong_AsLongLong(result);
+  Py_DECREF(result);
+
+  // PyLong_AsLongLong can raise OverflowError
+  ARROW_RETURN_IF_PYERROR();
+
+ return arrow::Status::OK();
+}
+
+arrow::Status PyReadableFile::Seek(int64_t position) {
+  PyAcquireGIL_RAII lock;
+  return SeekNoGIL(file_, position, 0);
 }
 
 arrow::Status PyReadableFile::ReadAt(
-    int64_t position, int64_t nbytes, int64_t* bytes_read, uint8_t* out) override {
+    int64_t position, int64_t nbytes, int64_t* bytes_read, uint8_t* out) {
+  PyAcquireGIL_RAII lock;
+  ARROW_RETURN_NOT_OK(SeekNoGIL(file_, position, 0));
 
-  return Status::OK();
+  PyObject* bytes_obj;
+  ARROW_RETURN_NOT_OK(ReadNoGIL(file_, nbytes, &bytes_obj));
+
+  *bytes_read = PyBytes_GET_SIZE(bytes_obj);
+  std::memcpy(out, PyBytes_AS_STRING(bytes_obj), *bytes_read);
+  Py_DECREF(bytes_obj);
+
+  return arrow::Status::OK();
 }
 
-arrow::Status PyReadableFile::GetSize(int64_t* size) override {
-  return Status::OK();
+arrow::Status PyReadableFile::GetSize(int64_t* size) {
+  PyAcquireGIL_RAII lock;
+
+  int64_t current_position;;
+  ARROW_RETURN_NOT_OK(TellNoGIL(file_, &current_position));
+
+  ARROW_RETURN_NOT_OK(SeekNoGIL(file_, 0, 2));
+
+  int64_t file_size;
+  ARROW_RETURN_NOT_OK(TellNoGIL(file_, &file_size));
+
+  // Restore previous file position
+  ARROW_RETURN_NOT_OK(SeekNoGIL(file_, current_position, 0));
+
+  *size = file_size;
+  return arrow::Status::OK();
 }
 
-  // Does not copy if not necessary
-Status PyReadableFile::ReadAt(
-    int64_t position, int64_t nbytes, std::shared_ptr<Buffer>* out) override {
+// Does not copy if not necessary
+arrow::Status PyReadableFile::ReadAt(
+    int64_t position, int64_t nbytes, std::shared_ptr<arrow::Buffer>* out) {
+  PyAcquireGIL_RAII lock;
+  ARROW_RETURN_NOT_OK(SeekNoGIL(file_, position, 0));
 
-  return Status::OK();
+  PyObject* bytes_obj;
+  ARROW_RETURN_NOT_OK(ReadNoGIL(file_, nbytes, &bytes_obj));
+
+  *out = std::make_shared<PyBytesBuffer>(bytes_obj);
+  Py_DECREF(bytes_obj);
+
+  return arrow::Status::OK();
 }
 
 bool PyReadableFile::supports_zero_copy() const {
@@ -71,20 +150,16 @@ PyOutputStream::~PyOutputStream() {
   Py_DECREF(file_);
 }
 
-Status PyOutputStream::Close() override {
+arrow::Status PyOutputStream::Close() {
   return arrow::Status::OK();
 }
 
-Status PyOutputStream::Tell(int64_t* position) {
+arrow::Status PyOutputStream::Tell(int64_t* position) {
   return arrow::Status::OK();
 }
 
-Status PyOutputStream::Write(const uint8_t* data, int64_t nbytes) {
+arrow::Status PyOutputStream::Write(const uint8_t* data, int64_t nbytes) {
   return arrow::Status::OK();
 }
 
-private:
-  PyObject* file_;
-};
-
 } // namespace pyarrow

python/src/pyarrow/io.h

@@ -15,8 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
-#ifndef PYARROW_COMMON_H
-#define PYARROW_COMMON_H
+#ifndef PYARROW_IO_H
+#define PYARROW_IO_H
 
 #include "arrow/io/interfaces.h"
 
@@ -39,7 +39,9 @@ public:
 
   // Does not copy if not necessary
   arrow::Status ReadAt(
-      int64_t position, int64_t nbytes, std::shared_ptr<Buffer>* out) override;
+      int64_t position, int64_t nbytes, std::shared_ptr<arrow::Buffer>* out) override;
+
+  arrow::Status Seek(int64_t position) override;
 
   bool supports_zero_copy() const override;
 
@@ -61,3 +63,5 @@ private:
 };
 
 } // namespace pyarrow
+
+#endif  // PYARROW_IO_H