Refactoring, more code sharing with native file interfaces

Change-Id: Ie9f6f13d7e1fb55e948da2a7709864978df7306a

python/pyarrow/__init__.py

@@ -41,6 +41,5 @@ from pyarrow.schema import (null, bool_,
                             list_, struct, field,
                             DataType, Field, Schema, schema)
 
-from pyarrow.array import RowBatch, from_pandas_dataframe
-
-from pyarrow.table import Column, Table
+from pyarrow.array import RowBatch
+from pyarrow.table import Column, Table, from_pandas_dataframe

python/pyarrow/array.pyx

@@ -35,7 +35,6 @@ from pyarrow.scalar import NA
 from pyarrow.schema cimport Schema
 import pyarrow.schema as schema
 
-from pyarrow.table cimport Table
 
 def total_allocated_bytes():
     cdef MemoryPool* pool = pyarrow.GetMemoryPool()
@@ -254,35 +253,6 @@ def from_pandas_series(object series, object mask=None, timestamps_to_ms=False):
     return box_arrow_array(out)
 
 
-def from_pandas_dataframe(object df, name=None, timestamps_to_ms=False):
-    """
-    Convert pandas.DataFrame to an Arrow Table
-
-    Parameters
-    ----------
-    df: pandas.DataFrame
-
-    name: str
-
-    timestamps_to_ms: bool
-        Convert datetime columns to ms resolution. This is needed for
-        compability with other functionality like Parquet I/O which
-        only supports milliseconds.
-    """
-    cdef:
-        list names = []
-        list arrays = []
-
-    for name in df.columns:
-        col = df[name]
-        arr = from_pandas_series(col, timestamps_to_ms=timestamps_to_ms)
-
-        names.append(name)
-        arrays.append(arr)
-
-    return Table.from_arrays(names, arrays, name=name)
-
-
 cdef object series_as_ndarray(object obj):
     import pandas as pd
 
@@ -324,4 +294,3 @@ cdef class RowBatch:
 
     def __getitem__(self, i):
         return self.arrays[i]
-

python/pyarrow/error.pxd

@@ -16,7 +16,7 @@
 # under the License.
 
 from pyarrow.includes.libarrow cimport CStatus
-from pyarrow.includes.pyarrow cimport *
+from pyarrow.includes.pyarrow cimport PyStatus
 
 cdef int check_cstatus(const CStatus& status) nogil except -1
-cdef int check_status(const Status& status) nogil except -1
+cdef int check_status(const PyStatus& status) nogil except -1

python/pyarrow/error.pyx

@@ -30,7 +30,7 @@ cdef int check_cstatus(const CStatus& status) nogil except -1:
     with gil:
         raise ArrowException(frombytes(c_message))
 
-cdef int check_status(const Status& status) nogil except -1:
+cdef int check_status(const PyStatus& status) nogil except -1:
     if status.ok():
         return 0
 

python/pyarrow/includes/libarrow_io.pxd

@@ -18,6 +18,7 @@
 # distutils: language = c++
 
 from pyarrow.includes.common cimport *
+from pyarrow.includes.libarrow cimport MemoryPool
 
 cdef extern from "arrow/io/interfaces.h" namespace "arrow::io" nogil:
     enum FileMode" arrow::io::FileMode::type":
@@ -66,6 +67,24 @@ cdef extern from "arrow/io/interfaces.h" namespace "arrow::io" nogil:
         pass
 
 
+cdef extern from "arrow/io/file.h" namespace "arrow::io" nogil:
+    cdef cppclass FileOutputStream(OutputStream):
+        @staticmethod
+        CStatus Open(const c_string& path, shared_ptr[FileOutputStream]* file)
+
+        int file_descriptor()
+
+    cdef cppclass ReadableFile(ReadableFileInterface):
+        @staticmethod
+        CStatus Open(const c_string& path, shared_ptr[ReadableFile]* file)
+
+        @staticmethod
+        CStatus Open(const c_string& path, MemoryPool* memory_pool,
+                     shared_ptr[ReadableFile]* file)
+
+        int file_descriptor()
+
+
 cdef extern from "arrow/io/hdfs.h" namespace "arrow::io" nogil:
     CStatus ConnectLibHdfs()
 

python/pyarrow/includes/pyarrow.pxd

@@ -18,15 +18,18 @@
 # distutils: language = c++
 
 from pyarrow.includes.common cimport *
-from pyarrow.includes.libarrow cimport (CArray, CColumn, CDataType,
+from pyarrow.includes.libarrow cimport (CArray, CColumn, CDataType, CStatus,
                                         Type, MemoryPool)
 
+cimport pyarrow.includes.libarrow_io as arrow_io
+
+
 cdef extern from "pyarrow/api.h" namespace "pyarrow" nogil:
     # We can later add more of the common status factory methods as needed
-    cdef Status Status_OK "Status::OK"()
+    cdef PyStatus PyStatus_OK "Status::OK"()
 
-    cdef cppclass Status:
-        Status()
+    cdef cppclass PyStatus "pyarrow::Status":
+        PyStatus()
 
         c_string ToString()
 
@@ -40,12 +43,22 @@ cdef extern from "pyarrow/api.h" namespace "pyarrow" nogil:
         c_bool IsArrowError()
 
     shared_ptr[CDataType] GetPrimitiveType(Type type)
-    Status ConvertPySequence(object obj, shared_ptr[CArray]* out)
+    PyStatus ConvertPySequence(object obj, shared_ptr[CArray]* out)
 
-    Status PandasToArrow(MemoryPool* pool, object ao, shared_ptr[CArray]* out)
-    Status PandasMaskedToArrow(MemoryPool* pool, object ao, object mo,
-                               shared_ptr[CArray]* out)
+    PyStatus PandasToArrow(MemoryPool* pool, object ao,
+                           shared_ptr[CArray]* out)
+    PyStatus PandasMaskedToArrow(MemoryPool* pool, object ao, object mo,
+                                 shared_ptr[CArray]* out)
 
-    Status ArrowToPandas(const shared_ptr[CColumn]& arr, object py_ref, PyObject** out)
+    PyStatus ArrowToPandas(const shared_ptr[CColumn]& arr, object py_ref,
+                           PyObject** out)
 
     MemoryPool* GetMemoryPool()
+
+
+cdef extern from "pyarrow/io.h" namespace "pyarrow" nogil:
+    cdef cppclass PyReadableFile(arrow_io.ReadableFileInterface):
+        PyReadableFile(object fo)
+
+    cdef cppclass PyOutputStream(arrow_io.OutputStream):
+        PyOutputStream(object fo)

python/pyarrow/io.pxd

@@ -23,11 +23,16 @@ from pyarrow.includes.libarrow_io cimport (ReadableFileInterface,
                                            OutputStream)
 
 
-cdef class NativeFileInterface:
+cdef class NativeFile:
+    cdef:
+        shared_ptr[ReadableFileInterface] rd_file
+        shared_ptr[OutputStream] wr_file
+        bint is_readonly
+        bint is_open
 
     # By implementing these "virtual" functions (all functions in Cython
-    # extension classes are technically virtual in the C++ sense)m we can
-    # expose the arrow::io abstract file interfaces to other components
-    # throughout the suite of Arrow C++ libraries
+    # extension classes are technically virtual in the C++ sense) we can expose
+    # the arrow::io abstract file interfaces to other components throughout the
+    # suite of Arrow C++ libraries
     cdef read_handle(self, shared_ptr[ReadableFileInterface]* file)
     cdef write_handle(self, shared_ptr[OutputStream]* file)

python/pyarrow/io.pyx

@@ -242,6 +242,9 @@ cdef class HdfsClient:
         cdef int16_t c_replication = replication or 0
         cdef int64_t c_default_block_size = default_block_size or 0
 
+        cdef shared_ptr[HdfsOutputStream] wr_handle
+        cdef shared_ptr[HdfsReadableFile] rd_handle
+
         if mode in ('wb', 'ab'):
             if mode == 'ab':
                 append = True
@@ -251,13 +254,17 @@ cdef class HdfsClient:
                     self.client.get()
                     .OpenWriteable(c_path, append, c_buffer_size,
                                    c_replication, c_default_block_size,
-                                   &out.wr_file))
+                                   &wr_handle))
+
+            out.wr_file = <shared_ptr[OutputStream]> wr_handle
 
             out.is_readonly = False
         else:
             with nogil:
                 check_cstatus(self.client.get()
-                              .OpenReadable(c_path, &out.rd_file))
+                              .OpenReadable(c_path, &rd_handle))
+
+            out.rd_file = <shared_ptr[ReadableFileInterface]> rd_handle
             out.is_readonly = True
 
         if c_buffer_size == 0:
@@ -314,25 +321,8 @@ cdef class HdfsClient:
         f = self.open(path, 'rb', buffer_size=buffer_size)
         f.download(stream)
 
-cdef class NativeFileInterface:
 
-    cdef read_handle(self, shared_ptr[ReadableFileInterface]* file):
-        raise NotImplementedError
-
-    cdef write_handle(self, shared_ptr[OutputStream]* file):
-        raise NotImplementedError
-
-cdef class HdfsFile(NativeFileInterface):
-    cdef:
-        shared_ptr[HdfsReadableFile] rd_file
-        shared_ptr[HdfsOutputStream] wr_file
-        bint is_readonly
-        bint is_open
-        object parent
-
-    cdef readonly:
-        int32_t buffer_size
-        object mode
+cdef class NativeFile:
 
     def __cinit__(self):
         self.is_open = False
@@ -356,14 +346,6 @@ cdef class HdfsFile(NativeFileInterface):
                     check_cstatus(self.wr_file.get().Close())
         self.is_open = False
 
-    cdef _assert_readable(self):
-        if not self.is_readonly:
-            raise IOError("only valid on readonly files")
-
-    cdef _assert_writeable(self):
-        if self.is_readonly:
-            raise IOError("only valid on writeonly files")
-
     cdef read_handle(self, shared_ptr[ReadableFileInterface]* file):
         self._assert_readable()
         file[0] = <shared_ptr[ReadableFileInterface]> self.rd_file
@@ -372,6 +354,14 @@ cdef class HdfsFile(NativeFileInterface):
         self._assert_writeable()
         file[0] = <shared_ptr[OutputStream]> self.wr_file
 
+    def _assert_readable(self):
+        if not self.is_readonly:
+            raise IOError("only valid on readonly files")
+
+    def _assert_writeable(self):
+        if self.is_readonly:
+            raise IOError("only valid on writeonly files")
+
     def size(self):
         cdef int64_t size
         self._assert_readable()
@@ -393,6 +383,26 @@ cdef class HdfsFile(NativeFileInterface):
         with nogil:
             check_cstatus(self.rd_file.get().Seek(position))
 
+    def write(self, data):
+        """
+        Write bytes-like (unicode, encoded to UTF-8) to file
+        """
+        self._assert_writeable()
+
+        data = tobytes(data)
+
+        cdef const uint8_t* buf = <const uint8_t*> cp.PyBytes_AS_STRING(data)
+        cdef int64_t bufsize = len(data)
+        with nogil:
+            check_cstatus(self.wr_file.get().Write(buf, bufsize))
+
+
+cdef class HdfsFile(NativeFile):
+    cdef readonly:
+        int32_t buffer_size
+        object mode
+        object parent
+
     def read(self, int nbytes):
         """
         Read indicated number of bytes from the file, up to EOF
@@ -504,16 +514,3 @@ cdef class HdfsFile(NativeFileInterface):
         writer_thread.join()
         if exc_info is not None:
             raise exc_info[0], exc_info[1], exc_info[2]
-
-    def write(self, data):
-        """
-        Write bytes-like (unicode, encoded to UTF-8) to file
-        """
-        self._assert_writeable()
-
-        data = tobytes(data)
-
-        cdef const uint8_t* buf = <const uint8_t*> cp.PyBytes_AS_STRING(data)
-        cdef int64_t bufsize = len(data)
-        with nogil:
-            check_cstatus(self.wr_file.get().Write(buf, bufsize))

python/pyarrow/parquet.pyx

@@ -27,10 +27,10 @@ cimport pyarrow.includes.pyarrow as pyarrow
 from pyarrow.compat import tobytes
 from pyarrow.error import ArrowException
 from pyarrow.error cimport check_cstatus
-from pyarrow.io import NativeFileInterface
+from pyarrow.io import NativeFile
 from pyarrow.table cimport Table
 
-from pyarrow.io cimport NativeFileInterface
+from pyarrow.io cimport NativeFile
 
 import six
 
@@ -54,7 +54,7 @@ cdef class ParquetReader:
             new FileReader(default_memory_pool(),
                            ParquetFileReader.OpenFile(path)))
 
-    cdef open_native_file(self, NativeFileInterface file):
+    cdef open_native_file(self, NativeFile file):
         cdef shared_ptr[ReadableFileInterface] cpp_handle
         file.read_handle(&cpp_handle)
 
@@ -84,7 +84,7 @@ def read_table(source, columns=None):
 
     if isinstance(source, six.string_types):
         reader.open_local_file(source)
-    elif isinstance(source, NativeFileInterface):
+    elif isinstance(source, NativeFile):
         reader.open_native_file(source)
 
     return reader.read_all()

python/pyarrow/table.pyx

@@ -25,10 +25,12 @@ cimport pyarrow.includes.pyarrow as pyarrow
 import pyarrow.config
 
 from pyarrow.array cimport Array, box_arrow_array
-from pyarrow.compat import frombytes, tobytes
 from pyarrow.error cimport check_status
 from pyarrow.schema cimport box_data_type, box_schema
 
+from pyarrow.compat import frombytes, tobytes
+
+
 cdef class ChunkedArray:
     '''
     Do not call this class's constructor directly.
@@ -161,7 +163,7 @@ cdef class Table:
 
     @staticmethod
     def from_pandas(df, name=None):
-        pass
+        return from_pandas_dataframe(df, name=name)
 
     @staticmethod
     def from_arrays(names, arrays, name=None):
@@ -264,3 +266,34 @@ cdef class Table:
         def __get__(self):
             return (self.num_rows, self.num_columns)
 
+
+
+def from_pandas_dataframe(object df, name=None, timestamps_to_ms=False):
+    """
+    Convert pandas.DataFrame to an Arrow Table
+
+    Parameters
+    ----------
+    df: pandas.DataFrame
+
+    name: str
+
+    timestamps_to_ms: bool
+        Convert datetime columns to ms resolution. This is needed for
+        compability with other functionality like Parquet I/O which
+        only supports milliseconds.
+    """
+    from pyarrow.array import from_pandas_series
+
+    cdef:
+        list names = []
+        list arrays = []
+
+    for name in df.columns:
+        col = df[name]
+        arr = from_pandas_series(col, timestamps_to_ms=timestamps_to_ms)
+
+        names.append(name)
+        arrays.append(arr)
+
+    return Table.from_arrays(names, arrays, name=name)

python/src/pyarrow/io.cc

@@ -29,70 +29,104 @@
 namespace pyarrow {
 
 // ----------------------------------------------------------------------
-// Seekable input stream
+// Python file
 
-PyReadableFile::PyReadableFile(PyObject* file)
-    : file_(file) {
+PythonFile::PythonFile(PyObject* file) {
   Py_INCREF(file_);
 }
 
-PyReadableFile::~PyReadableFile() {
+PythonFile::~PythonFile() {
   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 CheckPyError() {
+  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();
+arrow::Status PythonFile::Close() {
+  // whence: 0 for relative to start of file, 2 for end of file
+  PyObject* result = PyObject_CallMethod(file_, "close", "()");
+  Py_XDECREF(result);
+  ARROW_RETURN_NOT_OK(CheckPyError());
+  return arrow::Status::OK();
+}
+
+arrow::Status PythonFile::Seek(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_NOT_OK(CheckPyError());
+  return arrow::Status::OK();
+}
+
+arrow::Status PythonFile::Read(int64_t nbytes, PyObject** out) {
+  PyObject* result = PyObject_CallMethod(file_, "read", "(i)", nbytes);
+  ARROW_RETURN_NOT_OK(CheckPyError());
   *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();
+arrow::Status PythonFile::Write(const uint8_t* data, int64_t nbytes) {
+  PyObject* py_data = PyBytes_FromStringAndSize(
+      reinterpret_cast<const char*>(data), nbytes);
+  ARROW_RETURN_NOT_OK(CheckPyError());
+
+  PyObject* result = PyObject_CallMethod(file_, "write", "(O)", py_data);
+  Py_DECREF(py_data);
+  ARROW_RETURN_NOT_OK(CheckPyError());
+  return arrow::Status::OK();
+}
+
+arrow::Status PythonFile::Tell(int64_t* position) {
+  PyObject* result = PyObject_CallMethod(file_, "tell", "()");
+  ARROW_RETURN_NOT_OK(CheckPyError());
 
   *position = PyLong_AsLongLong(result);
   Py_DECREF(result);
 
   // PyLong_AsLongLong can raise OverflowError
-  ARROW_RETURN_IF_PYERROR();
+  ARROW_RETURN_NOT_OK(CheckPyError());
 
  return arrow::Status::OK();
 }
 
+// ----------------------------------------------------------------------
+// Seekable input stream
+
+PyReadableFile::PyReadableFile(PyObject* file) {
+  file_.reset(new PythonFile(file));
+}
+
+PyReadableFile::~PyReadableFile() {}
+
+arrow::Status PyReadableFile::Close() {
+  PyAcquireGIL_RAII lock;
+  return file_->Close();
+}
+
 arrow::Status PyReadableFile::Seek(int64_t position) {
   PyAcquireGIL_RAII lock;
-  return SeekNoGIL(file_, position, 0);
+  return file_->Seek(position, 0);
 }
 
 arrow::Status PyReadableFile::ReadAt(
     int64_t position, int64_t nbytes, int64_t* bytes_read, uint8_t* out) {
   PyAcquireGIL_RAII lock;
-  ARROW_RETURN_NOT_OK(SeekNoGIL(file_, position, 0));
+  ARROW_RETURN_NOT_OK(file_->Seek(position, 0));
 
   PyObject* bytes_obj;
-  ARROW_RETURN_NOT_OK(ReadNoGIL(file_, nbytes, &bytes_obj));
+  ARROW_RETURN_NOT_OK(file_->Read(nbytes, &bytes_obj));
 
   *bytes_read = PyBytes_GET_SIZE(bytes_obj);
   std::memcpy(out, PyBytes_AS_STRING(bytes_obj), *bytes_read);
@@ -105,15 +139,15 @@ 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(file_->Tell(&current_position));
 
-  ARROW_RETURN_NOT_OK(SeekNoGIL(file_, 0, 2));
+  ARROW_RETURN_NOT_OK(file_->Seek(0, 2));
 
   int64_t file_size;
-  ARROW_RETURN_NOT_OK(TellNoGIL(file_, &file_size));
+  ARROW_RETURN_NOT_OK(file_->Tell(&file_size));
 
   // Restore previous file position
-  ARROW_RETURN_NOT_OK(SeekNoGIL(file_, current_position, 0));
+  ARROW_RETURN_NOT_OK(file_->Seek(current_position, 0));
 
   *size = file_size;
   return arrow::Status::OK();
@@ -123,10 +157,10 @@ arrow::Status PyReadableFile::GetSize(int64_t* size) {
 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));
+  ARROW_RETURN_NOT_OK(file_->Seek(position, 0));
 
   PyObject* bytes_obj;
-  ARROW_RETURN_NOT_OK(ReadNoGIL(file_, nbytes, &bytes_obj));
+  ARROW_RETURN_NOT_OK(file_->Read(nbytes, &bytes_obj));
 
   *out = std::make_shared<PyBytesBuffer>(bytes_obj);
   Py_DECREF(bytes_obj);
@@ -141,14 +175,11 @@ bool PyReadableFile::supports_zero_copy() const {
 // ----------------------------------------------------------------------
 // Output stream
 
-PyOutputStream::PyOutputStream(PyObject* file)
-    : file_(file) {
-  Py_INCREF(file);
+PyOutputStream::PyOutputStream(PyObject* file) {
+  file_.reset(new PythonFile(file));
 }
 
-PyOutputStream::~PyOutputStream() {
-  Py_DECREF(file_);
-}
+PyOutputStream::~PyOutputStream() {}
 
 arrow::Status PyOutputStream::Close() {
   return arrow::Status::OK();

python/src/pyarrow/io.h

@@ -27,11 +27,30 @@ namespace arrow { class MemoryPool; }
 
 namespace pyarrow {
 
+// A common interface to a Python file-like object. Must acquire GIL before
+// calling any methods
+class PythonFile {
+ public:
+  PythonFile(PyObject* file);
+  ~PythonFile();
+
+  arrow::Status Close();
+  arrow::Status Seek(int64_t position, int whence);
+  arrow::Status Read(int64_t nbytes, PyObject** out);
+  arrow::Status Tell(int64_t* position);
+  arrow::Status Write(const uint8_t* data, int64_t nbytes);
+
+ private:
+  PyObject* file_;
+};
+
 class PYARROW_EXPORT PyReadableFile : public arrow::io::ReadableFileInterface {
 public:
   explicit PyReadableFile(PyObject* file);
   ~PyReadableFile();
 
+  arrow::Status Close() override;
+
   arrow::Status ReadAt(
       int64_t position, int64_t nbytes, int64_t* bytes_read, uint8_t* out) override;
 
@@ -46,7 +65,7 @@ public:
   bool supports_zero_copy() const override;
 
 private:
-  PyObject* file_;
+  std::unique_ptr<PythonFile> file_;
 };
 
 class PYARROW_EXPORT PyOutputStream : public arrow::io::OutputStream {
@@ -59,9 +78,11 @@ public:
   arrow::Status Write(const uint8_t* data, int64_t nbytes) override;
 
 private:
-  PyObject* file_;
+  std::unique_ptr<PythonFile> file_;
 };
 
+// TODO(wesm): seekable output files
+
 } // namespace pyarrow
 
 #endif  // PYARROW_IO_H