Get basic Python file read/write working

Change-Id: I9d6e656ea152ef678ac01e22e4e24b3783ccacc0

python/pyarrow/includes/libarrow_io.pxd

@@ -36,6 +36,7 @@ cdef extern from "arrow/io/interfaces.h" namespace "arrow::io" nogil:
         FileMode mode()
 
     cdef cppclass Readable:
+        CStatus ReadB" Read"(int64_t nbytes, shared_ptr[Buffer]* out)
         CStatus Read(int64_t nbytes, int64_t* bytes_read, uint8_t* out)
 
     cdef cppclass Seekable:

python/pyarrow/io.pyx

@@ -406,7 +406,7 @@ cdef class NativeFile:
 
         with nogil:
             check_cstatus(self.rd_file.get()
-                          .Read(nbytes, &out))
+                          .ReadB(nbytes, &out))
 
         result = cp.PyBytes_FromStringAndSize(
             <const char*>out.get().data(), out.get().size())
@@ -425,12 +425,16 @@ cdef class PythonFileInterface(NativeFile):
         self.handle = handle
 
         if mode.startswith('w'):
-            self.wr_file.reset(new PyOutputStream(handle))
+            self.wr_file.reset(new pyarrow.PyOutputStream(handle))
+            self.is_readonly = 0
         elif mode.startswith('r'):
-            self.wr_file.reset(new PyReadableFile(handle))
+            self.rd_file.reset(new pyarrow.PyReadableFile(handle))
+            self.is_readonly = 1
         else:
             raise ValueError('Invalid file mode: {0}'.format(mode))
 
+        self.is_open = True
+
 
 # ----------------------------------------------------------------------
 # Specialization for HDFS

python/pyarrow/tests/test_hdfs.py

@@ -0,0 +1,128 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from io import BytesIO
+from os.path import join as pjoin
+import os
+import random
+
+import pytest
+
+import pyarrow.io as io
+
+# ----------------------------------------------------------------------
+# HDFS tests
+
+
+def hdfs_test_client():
+    host = os.environ.get('ARROW_HDFS_TEST_HOST', 'localhost')
+    user = os.environ['ARROW_HDFS_TEST_USER']
+    try:
+        port = int(os.environ.get('ARROW_HDFS_TEST_PORT', 20500))
+    except ValueError:
+        raise ValueError('Env variable ARROW_HDFS_TEST_PORT was not '
+                         'an integer')
+
+    return io.HdfsClient.connect(host, port, user)
+
+
+libhdfs = pytest.mark.skipif(not io.have_libhdfs(),
+                             reason='No libhdfs available on system')
+
+
+HDFS_TMP_PATH = '/tmp/pyarrow-test-{0}'.format(random.randint(0, 1000))
+
+
+@pytest.fixture(scope='session')
+def hdfs(request):
+    fixture = hdfs_test_client()
+
+    def teardown():
+        fixture.delete(HDFS_TMP_PATH, recursive=True)
+        fixture.close()
+    request.addfinalizer(teardown)
+    return fixture
+
+
+@libhdfs
+def test_hdfs_close():
+    client = hdfs_test_client()
+    assert client.is_open
+    client.close()
+    assert not client.is_open
+
+    with pytest.raises(Exception):
+        client.ls('/')
+
+
+@libhdfs
+def test_hdfs_mkdir(hdfs):
+    path = pjoin(HDFS_TMP_PATH, 'test-dir/test-dir')
+    parent_path = pjoin(HDFS_TMP_PATH, 'test-dir')
+
+    hdfs.mkdir(path)
+    assert hdfs.exists(path)
+
+    hdfs.delete(parent_path, recursive=True)
+    assert not hdfs.exists(path)
+
+
+@libhdfs
+def test_hdfs_ls(hdfs):
+    base_path = pjoin(HDFS_TMP_PATH, 'ls-test')
+    hdfs.mkdir(base_path)
+
+    dir_path = pjoin(base_path, 'a-dir')
+    f1_path = pjoin(base_path, 'a-file-1')
+
+    hdfs.mkdir(dir_path)
+
+    f = hdfs.open(f1_path, 'wb')
+    f.write('a' * 10)
+
+    contents = sorted(hdfs.ls(base_path, False))
+    assert contents == [dir_path, f1_path]
+
+
+@libhdfs
+def test_hdfs_download_upload(hdfs):
+    base_path = pjoin(HDFS_TMP_PATH, 'upload-test')
+
+    data = b'foobarbaz'
+    buf = BytesIO(data)
+    buf.seek(0)
+
+    hdfs.upload(base_path, buf)
+
+    out_buf = BytesIO()
+    hdfs.download(base_path, out_buf)
+    out_buf.seek(0)
+    assert out_buf.getvalue() == data
+
+
+@libhdfs
+def test_hdfs_file_context_manager(hdfs):
+    path = pjoin(HDFS_TMP_PATH, 'ctx-manager')
+
+    data = b'foo'
+    with hdfs.open(path, 'wb') as f:
+        f.write(data)
+
+    with hdfs.open(path, 'rb') as f:
+        assert f.size() == 3
+        result = f.read(10)
+        assert result == data

python/pyarrow/tests/test_io.py

@@ -24,104 +24,49 @@ import pytest
 
 import pyarrow.io as io
 
-#----------------------------------------------------------------------
-# HDFS tests
+# ----------------------------------------------------------------------
+# Python file-like objects
 
 
-def hdfs_test_client():
-    host = os.environ.get('ARROW_HDFS_TEST_HOST', 'localhost')
-    user = os.environ['ARROW_HDFS_TEST_USER']
-    try:
-        port = int(os.environ.get('ARROW_HDFS_TEST_PORT', 20500))
-    except ValueError:
-        raise ValueError('Env variable ARROW_HDFS_TEST_PORT was not '
-                         'an integer')
+def test_python_file_write():
+    buf = BytesIO()
 
-    return io.HdfsClient.connect(host, port, user)
+    f = io.PythonFileInterface(buf)
+
+    assert f.tell() == 0
+
+    s1 = b'enga\xc3\xb1ado'
+    s2 = b'foobar'
+
+    f.write(s1.decode('utf8'))
+    assert f.tell() == len(s1)
+
+    f.write(s2)
+
+    expected = s1 + s2
+
+    result = buf.getvalue()
+    assert result == expected
+
+    f.close()
 
 
-libhdfs = pytest.mark.skipif(not io.have_libhdfs(),
-                             reason='No libhdfs available on system')
+def test_python_file_read():
+    buf = BytesIO(b'some sample data')
 
+    f = io.PythonFileInterface(buf, mode='r')
 
-HDFS_TMP_PATH = '/tmp/pyarrow-test-{0}'.format(random.randint(0, 1000))
+    assert f.tell() == 0
 
+    assert f.read(4) == b'some'
+    assert f.tell() == 4
 
-@pytest.fixture(scope='session')
-def hdfs(request):
-    fixture = hdfs_test_client()
-    def teardown():
-        fixture.delete(HDFS_TMP_PATH, recursive=True)
-        fixture.close()
-    request.addfinalizer(teardown)
-    return fixture
+    f.seek(0)
+    assert f.tell() == 0
 
+    f.seek(5)
+    assert f.tell() == 5
 
-@libhdfs
-def test_hdfs_close():
-    client = hdfs_test_client()
-    assert client.is_open
-    client.close()
-    assert not client.is_open
+    assert f.read(50) == b'sample data'
 
-    with pytest.raises(Exception):
-        client.ls('/')
-
-
-@libhdfs
-def test_hdfs_mkdir(hdfs):
-    path = pjoin(HDFS_TMP_PATH, 'test-dir/test-dir')
-    parent_path = pjoin(HDFS_TMP_PATH, 'test-dir')
-
-    hdfs.mkdir(path)
-    assert hdfs.exists(path)
-
-    hdfs.delete(parent_path, recursive=True)
-    assert not hdfs.exists(path)
-
-
-@libhdfs
-def test_hdfs_ls(hdfs):
-    base_path = pjoin(HDFS_TMP_PATH, 'ls-test')
-    hdfs.mkdir(base_path)
-
-    dir_path = pjoin(base_path, 'a-dir')
-    f1_path = pjoin(base_path, 'a-file-1')
-
-    hdfs.mkdir(dir_path)
-
-    f = hdfs.open(f1_path, 'wb')
-    f.write('a' * 10)
-
-    contents = sorted(hdfs.ls(base_path, False))
-    assert contents == [dir_path, f1_path]
-
-
-@libhdfs
-def test_hdfs_download_upload(hdfs):
-    base_path = pjoin(HDFS_TMP_PATH, 'upload-test')
-
-    data = b'foobarbaz'
-    buf = BytesIO(data)
-    buf.seek(0)
-
-    hdfs.upload(base_path, buf)
-
-    out_buf = BytesIO()
-    hdfs.download(base_path, out_buf)
-    out_buf.seek(0)
-    assert out_buf.getvalue() == data
-
-
-@libhdfs
-def test_hdfs_file_context_manager(hdfs):
-    path = pjoin(HDFS_TMP_PATH, 'ctx-manager')
-
-    data = b'foo'
-    with hdfs.open(path, 'wb') as f:
-        f.write(data)
-
-    with hdfs.open(path, 'rb') as f:
-        assert f.size() == 3
-        result = f.read(10)
-        assert result == data
+    f.close()

python/src/pyarrow/common.cc

@@ -72,14 +72,14 @@ arrow::MemoryPool* GetMemoryPool() {
 // PyBytesBuffer
 
 PyBytesBuffer::PyBytesBuffer(PyObject* obj)
-    : Buffer(reinterpret_cast<const uint8_t*>(PyBytes_AS_STRING(obj_)),
-        PyBytes_GET_SIZE(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;
+  PyGILGuard lock;
   Py_DECREF(obj_);
 }
 

python/src/pyarrow/common.h

@@ -82,18 +82,18 @@ struct PyObjectStringify {
   }
 };
 
-class PyAcquireGIL_RAII {
+class PyGILGuard {
  public:
-  PyAcquireGIL_RAII() {
+  PyGILGuard() {
     state_ = PyGILState_Ensure();
   }
 
-  ~PyAcquireGIL_RAII() {
+  ~PyGILGuard() {
     PyGILState_Release(state_);
   }
  private:
   PyGILState_STATE state_;
-  DISALLOW_COPY_AND_ASSIGN(PyAcquireGIL_RAII);
+  DISALLOW_COPY_AND_ASSIGN(PyGILGuard);
 };
 
 // TODO(wesm): We can just let errors pass through. To be explored later

python/src/pyarrow/io.cc

@@ -31,7 +31,8 @@ namespace pyarrow {
 // ----------------------------------------------------------------------
 // Python file
 
-PythonFile::PythonFile(PyObject* file) {
+PythonFile::PythonFile(PyObject* file)
+    : file_(file) {
   Py_INCREF(file_);
 }
 
@@ -84,6 +85,7 @@ arrow::Status PythonFile::Write(const uint8_t* data, int64_t nbytes) {
 
   PyObject* result = PyObject_CallMethod(file_, "write", "(O)", py_data);
   Py_DECREF(py_data);
+  Py_XDECREF(result);
   ARROW_RETURN_NOT_OK(CheckPyError());
   return arrow::Status::OK();
 }
@@ -111,20 +113,22 @@ PyReadableFile::PyReadableFile(PyObject* file) {
 PyReadableFile::~PyReadableFile() {}
 
 arrow::Status PyReadableFile::Close() {
-  PyAcquireGIL_RAII lock;
+  PyGILGuard lock;
   return file_->Close();
 }
 
 arrow::Status PyReadableFile::Seek(int64_t position) {
-  PyAcquireGIL_RAII lock;
+  PyGILGuard lock;
   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(file_->Seek(position, 0));
+arrow::Status PyReadableFile::Tell(int64_t* position) {
+  PyGILGuard lock;
+  return file_->Tell(position);
+}
 
+arrow::Status PyReadableFile::Read(int64_t nbytes, int64_t* bytes_read, uint8_t* out) {
+  PyGILGuard lock;
   PyObject* bytes_obj;
   ARROW_RETURN_NOT_OK(file_->Read(nbytes, &bytes_obj));
 
@@ -135,8 +139,20 @@ arrow::Status PyReadableFile::ReadAt(
   return arrow::Status::OK();
 }
 
+arrow::Status PyReadableFile::Read(int64_t nbytes, std::shared_ptr<arrow::Buffer>* out) {
+  PyGILGuard lock;
+
+  PyObject* bytes_obj;
+  ARROW_RETURN_NOT_OK(file_->Read(nbytes, &bytes_obj));
+
+  *out = std::make_shared<PyBytesBuffer>(bytes_obj);
+  Py_DECREF(bytes_obj);
+
+  return arrow::Status::OK();
+}
+
 arrow::Status PyReadableFile::GetSize(int64_t* size) {
-  PyAcquireGIL_RAII lock;
+  PyGILGuard lock;
 
   int64_t current_position;;
   ARROW_RETURN_NOT_OK(file_->Tell(&current_position));
@@ -153,21 +169,6 @@ arrow::Status PyReadableFile::GetSize(int64_t* size) {
   return arrow::Status::OK();
 }
 
-// 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(file_->Seek(position, 0));
-
-  PyObject* bytes_obj;
-  ARROW_RETURN_NOT_OK(file_->Read(nbytes, &bytes_obj));
-
-  *out = std::make_shared<PyBytesBuffer>(bytes_obj);
-  Py_DECREF(bytes_obj);
-
-  return arrow::Status::OK();
-}
-
 bool PyReadableFile::supports_zero_copy() const {
   return false;
 }
@@ -182,15 +183,18 @@ PyOutputStream::PyOutputStream(PyObject* file) {
 PyOutputStream::~PyOutputStream() {}
 
 arrow::Status PyOutputStream::Close() {
-  return arrow::Status::OK();
+  PyGILGuard lock;
+  return file_->Close();
 }
 
 arrow::Status PyOutputStream::Tell(int64_t* position) {
-  return arrow::Status::OK();
+  PyGILGuard lock;
+  return file_->Tell(position);
 }
 
 arrow::Status PyOutputStream::Write(const uint8_t* data, int64_t nbytes) {
-  return arrow::Status::OK();
+  PyGILGuard lock;
+  return file_->Write(data, nbytes);
 }
 
 } // namespace pyarrow

python/src/pyarrow/io.h

@@ -51,17 +51,15 @@ public:
 
   arrow::Status Close() override;
 
-  arrow::Status ReadAt(
-      int64_t position, int64_t nbytes, int64_t* bytes_read, uint8_t* out) override;
+  arrow::Status Read(int64_t nbytes, int64_t* bytes_read, uint8_t* out) override;
+  arrow::Status Read(int64_t nbytes, std::shared_ptr<arrow::Buffer>* out) override;
 
   arrow::Status GetSize(int64_t* size) override;
 
-  // Does not copy if not necessary
-  arrow::Status ReadAt(
-      int64_t position, int64_t nbytes, std::shared_ptr<arrow::Buffer>* out) override;
-
   arrow::Status Seek(int64_t position) override;
 
+  arrow::Status Tell(int64_t* position) override;
+
   bool supports_zero_copy() const override;
 
 private: