Provide a means to expose abstract native file handles

Change-Id: I86e43b42582276302332eb3c61afffd6f7187c40

cpp/src/arrow/io/interfaces.h

@@ -19,6 +19,7 @@
 #define ARROW_IO_INTERFACES_H
 
 #include <cstdint>
+#include <memory>
 
 namespace arrow {
 

cpp/src/arrow/parquet/io.cc

@@ -56,10 +56,8 @@ void ParquetAllocator::Free(uint8_t* buffer, int64_t size) {
 // ParquetReadSource
 
 ParquetReadSource::ParquetReadSource(
-    const std::shared_ptr<ArrowROFile>& file, MemoryPool* pool)
-    : file_(file) {
-  allocator_.set_pool(pool);
-}
+    const std::shared_ptr<ArrowROFile>& file, ParquetAllocator* allocator)
+    : file_(file), allocator_(allocator) {}
 
 void ParquetReadSource::Close() {
   PARQUET_THROW_NOT_OK(file_->Close());
@@ -84,7 +82,7 @@ int64_t ParquetReadSource::Read(int64_t nbytes, uint8_t* out) {
 std::shared_ptr<::parquet::Buffer> ParquetReadSource::Read(int64_t nbytes) {
   // TODO(wesm): This code is duplicated from parquet/util/input.cc; suggests
   // that there should be more code sharing amongst file-like sources
-  auto result = std::make_shared<::parquet::OwnedMutableBuffer>(0, &allocator_);
+  auto result = std::make_shared<::parquet::OwnedMutableBuffer>(0, allocator_);
   result->Resize(nbytes);
 
   int64_t bytes_read = Read(nbytes, result->mutable_data());

cpp/src/arrow/parquet/io.h

@@ -51,7 +51,7 @@ class ARROW_EXPORT ParquetAllocator : public ::parquet::MemoryAllocator {
 
   void set_pool(MemoryPool* pool) { pool_ = pool; }
 
-  MemoryPool* pool() { return pool_; }
+  MemoryPool* pool() const { return pool_; }
 
  private:
   MemoryPool* pool_;
@@ -59,7 +59,8 @@ class ARROW_EXPORT ParquetAllocator : public ::parquet::MemoryAllocator {
 
 class ARROW_EXPORT ParquetReadSource : public ::parquet::RandomAccessSource {
  public:
-  ParquetReadSource(const std::shared_ptr<io::RandomAccessFile>& file, MemoryPool* pool);
+  ParquetReadSource(
+      const std::shared_ptr<io::RandomAccessFile>& file, ParquetAllocator* allocator);
 
   void Close() override;
   int64_t Tell() const override;
@@ -72,7 +73,7 @@ class ARROW_EXPORT ParquetReadSource : public ::parquet::RandomAccessSource {
   std::shared_ptr<io::RandomAccessFile> file_;
 
   // The allocator is required for creating managed buffers
-  ParquetAllocator allocator_;
+  ParquetAllocator* allocator_;
 };
 
 }  // namespace parquet

cpp/src/arrow/parquet/parquet-io-test.cc

@@ -147,8 +147,10 @@ TEST(TestParquetReadSource, Basics) {
   std::string data = "this is the data";
   auto data_buffer = reinterpret_cast<const uint8_t*>(data.c_str());
 
+  ParquetAllocator allocator(default_memory_pool());
+
   auto file = std::make_shared<BufferReader>(data_buffer, data.size());
-  auto source = std::make_shared<ParquetReadSource>(file, default_memory_pool());
+  auto source = std::make_shared<ParquetReadSource>(file, &allocator);
 
   ASSERT_EQ(0, source->Tell());
   ASSERT_NO_THROW(source->Seek(5));

cpp/src/arrow/parquet/reader.cc

@@ -187,14 +187,16 @@ FileReader::FileReader(
 FileReader::~FileReader() {}
 
 // Static ctor
-Status Open(MemoryPool* pool, const std::shared_ptr<io::RandomAccessFile>& file,
-    std::unique_ptr<FileReader>* reader) {
-  std::unique_ptr<ParquetRAS> source(new ParquetReadSource(file, pool));
+Status Open(const std::shared_ptr<io::RandomAccessFile>& file,
+    ParquetAllocator* allocator, std::unique_ptr<FileReader>* reader) {
+  std::unique_ptr<ParquetRAS> source(new ParquetReadSource(file, allocator));
 
   // TODO(wesm): reader properties
   std::unique_ptr<ParquetReader> pq_reader;
   PARQUET_CATCH_NOT_OK(pq_reader = ParquetReader::Open(std::move(source)));
-  reader->reset(new FileReader(pool, std::move(pq_reader)));
+
+  // Use the same memory pool as the ParquetAllocator
+  reader->reset(new FileReader(allocator->pool(), std::move(pq_reader)));
   return Status::OK();
 }
 

cpp/src/arrow/parquet/reader.h

@@ -85,8 +85,8 @@ class ARROW_EXPORT FileReader {
  public:
   // Helper function to create a file reader from an implementation of an Arrow
   // readable file
-  static Status Open(MemoryPool* pool, const std::shared_ptr<io::RandomAccessFile>& file,
-      std::unique_ptr<FileReader>* reader);
+  static Status Open(const std::shared_ptr<io::RandomAccessFile>& file,
+      ParquetAllocator* allocator, std::unique_ptr<FileReader>* reader);
 
   FileReader(MemoryPool* pool, std::unique_ptr<::parquet::ParquetFileReader> reader);
 

python/pyarrow/includes/libarrow_io.pxd

@@ -19,11 +19,37 @@
 
 from pyarrow.includes.common cimport *
 
-cdef extern from "arrow/io/interfaces.h" nogil:
+cdef extern from "arrow/io/interfaces.h" namespace "arrow::io" nogil:
+    enum FileMode" arrow::io::FileMode::type":
+        FileMode_READ" arrow::io::FileMode::READ"
+        FileMode_WRITE" arrow::io::FileMode::WRITE"
+        FileMode_READWRITE" arrow::io::FileMode::READWRITE"
+
     enum ObjectType" arrow::io::ObjectType::type":
         ObjectType_FILE" arrow::io::ObjectType::FILE"
         ObjectType_DIRECTORY" arrow::io::ObjectType::DIRECTORY"
 
+    cdef cppclass FileBase:
+        CStatus Close()
+        CStatus Tell(int64_t* position)
+
+    cdef cppclass ReadableFile(FileBase):
+        CStatus GetSize(int64_t* size)
+        CStatus Read(int64_t nbytes, int64_t* bytes_read,
+                     uint8_t* buffer)
+
+        CStatus ReadAt(int64_t position, int64_t nbytes,
+                       int64_t* bytes_read, uint8_t* buffer)
+
+    cdef cppclass RandomAccessFile(ReadableFile):
+        CStatus Seek(int64_t position)
+
+    cdef cppclass WriteableFile(FileBase):
+        CStatus Write(const uint8_t* buffer, int64_t nbytes)
+        # CStatus Write(const uint8_t* buffer, int64_t nbytes,
+        #               int64_t* bytes_written)
+
+
 cdef extern from "arrow/io/hdfs.h" namespace "arrow::io" nogil:
     CStatus ConnectLibHdfs()
 
@@ -44,24 +70,11 @@ cdef extern from "arrow/io/hdfs.h" namespace "arrow::io" nogil:
         int64_t block_size
         int16_t permissions
 
-    cdef cppclass CHdfsFile:
-        CStatus Close()
-        CStatus Seek(int64_t position)
-        CStatus Tell(int64_t* position)
+    cdef cppclass HdfsReadableFile(RandomAccessFile):
+        pass
 
-    cdef cppclass HdfsReadableFile(CHdfsFile):
-        CStatus GetSize(int64_t* size)
-        CStatus Read(int64_t nbytes, int64_t* bytes_read,
-                     uint8_t* buffer)
-
-        CStatus ReadAt(int64_t position, int64_t nbytes,
-                       int64_t* bytes_read, uint8_t* buffer)
-
-    cdef cppclass HdfsWriteableFile(CHdfsFile):
-        CStatus Write(const uint8_t* buffer, int64_t nbytes)
-
-        CStatus Write(const uint8_t* buffer, int64_t nbytes,
-                      int64_t* bytes_written)
+    cdef cppclass HdfsWriteableFile(WriteableFile):
+        pass
 
     cdef cppclass CHdfsClient" arrow::io::HdfsClient":
         @staticmethod

python/pyarrow/includes/parquet.pxd

@@ -19,6 +19,7 @@
 
 from pyarrow.includes.common cimport *
 from pyarrow.includes.libarrow cimport CSchema, CStatus, CTable, MemoryPool
+from pyarrow.includes.libarrow_io cimport RandomAccessFile
 
 
 cdef extern from "parquet/api/schema.h" namespace "parquet::schema" nogil:
@@ -90,8 +91,24 @@ cdef extern from "parquet/api/writer.h" namespace "parquet" nogil:
             shared_ptr[WriterProperties] build()
 
 
+cdef extern from "arrow/parquet/io.h" namespace "arrow::parquet" nogil:
+    cdef cppclass ParquetAllocator:
+        ParquetAllocator()
+        ParquetAllocator(MemoryPool* pool)
+        MemoryPool* pool()
+
+    cdef cppclass ParquetReadSource:
+        ParquetReadSource(const shared_ptr[RandomAccessFile]& file,
+                          ParquetAllocator* allocator)
+
+
 cdef extern from "arrow/parquet/reader.h" namespace "arrow::parquet" nogil:
     cdef cppclass FileReader:
+        @staticmethod
+        CStatus Open(const shared_ptr[RandomAccessFile]& file,
+                     ParquetAllocator* allocator,
+                     unique_ptr[FileReader]* reader)
+
         FileReader(MemoryPool* pool, unique_ptr[ParquetFileReader] reader)
         CStatus ReadFlatTable(shared_ptr[CTable]* out);
 
@@ -105,4 +122,3 @@ cdef extern from "arrow/parquet/writer.h" namespace "arrow::parquet" nogil:
     cdef CStatus WriteFlatTable(const CTable* table, MemoryPool* pool,
             const shared_ptr[OutputStream]& sink, int64_t chunk_size,
             const shared_ptr[WriterProperties]& properties)
-

python/pyarrow/io.pyx

@@ -164,7 +164,7 @@ cdef class HdfsClient:
                           .ListDirectory(c_path, &listing))
 
         cdef const HdfsPathInfo* info
-        for i in range(listing.size()):
+        for i in range(<int> listing.size()):
             info = &listing[i]
 
             # Try to trim off the hdfs://HOST:PORT piece
@@ -314,8 +314,15 @@ cdef class HdfsClient:
         f = self.open(path, 'rb', buffer_size=buffer_size)
         f.download(stream)
 
+cdef class NativeFileInterface:
 
-cdef class HdfsFile:
+    cdef int read_handle(self, shared_ptr[RandomAccessFile]* file) except -1:
+        raise NotImplementedError
+
+    cdef int write_handle(self, shared_ptr[WriteableFile]* file) except -1:
+        raise NotImplementedError
+
+cdef class HdfsFile(NativeFileInterface):
     cdef:
         shared_ptr[HdfsReadableFile] rd_file
         shared_ptr[HdfsWriteableFile] wr_file
@@ -357,6 +364,16 @@ cdef class HdfsFile:
         if self.is_readonly:
             raise IOError("only valid on writeonly files")
 
+    cdef int read_handle(self, shared_ptr[RandomAccessFile]* file) except -1:
+        self._assert_readable()
+        file[0] = <shared_ptr[RandomAccessFile]> self.rd_file
+        return 0
+
+    cdef int write_handle(self, shared_ptr[WriteableFile]* file) except -1:
+        self._assert_writeable()
+        file[0] = <shared_ptr[WriteableFile]> self.wr_file
+        return 0
+
     def size(self):
         cdef int64_t size
         self._assert_readable()

python/pyarrow/parquet.pyx

@@ -39,10 +39,11 @@ def read_table(filename, columns=None):
     cdef Table table = Table()
     cdef shared_ptr[CTable] ctable
 
-    # Must be in one expression to avoid calling std::move which is not possible
-    # in Cython (due to missing rvalue support)
+    # Must be in one expression to avoid calling std::move which is not
+    # possible in Cython (due to missing rvalue support)
     reader = unique_ptr[FileReader](new FileReader(default_memory_pool(),
         ParquetFileReader.OpenFile(tobytes(filename))))
+
     with nogil:
         check_cstatus(reader.get().ReadFlatTable(&ctable))
 
@@ -84,4 +85,3 @@ def write_table(table, filename, chunk_size=None, version=None):
     with nogil:
         check_cstatus(WriteFlatTable(ctable_, default_memory_pool(), sink,
             chunk_size_, properties_builder.build()))
-