Slight refactoring of read table to be able to also handle classes wrapping C++ file interfaces

Change-Id: I4a3d0c4d2a763abb02ca546df35b9556f1060c0e

cpp/src/arrow/parquet/reader.cc

@@ -187,7 +187,7 @@ FileReader::FileReader(
 FileReader::~FileReader() {}
 
 // Static ctor
-Status Open(const std::shared_ptr<io::RandomAccessFile>& file,
+Status OpenFile(const std::shared_ptr<io::RandomAccessFile>& file,
     ParquetAllocator* allocator, std::unique_ptr<FileReader>* reader) {
   std::unique_ptr<ParquetRAS> source(new ParquetReadSource(file, allocator));
 

cpp/src/arrow/parquet/reader.h

@@ -83,11 +83,6 @@ class FlatColumnReader;
 // arrays
 class ARROW_EXPORT FileReader {
  public:
-  // Helper function to create a file reader from an implementation of an Arrow
-  // readable file
-  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);
 
   // Since the distribution of columns amongst a Parquet file's row groups may
@@ -106,7 +101,7 @@ class ARROW_EXPORT FileReader {
   virtual ~FileReader();
 
  private:
-  class PARQUET_NO_EXPORT Impl;
+  class ARROW_NO_EXPORT Impl;
   std::unique_ptr<Impl> impl_;
 };
 
@@ -132,15 +127,20 @@ class ARROW_EXPORT FlatColumnReader {
   Status NextBatch(int batch_size, std::shared_ptr<Array>* out);
 
  private:
-  class Impl;
+  class ARROW_NO_EXPORT Impl;
   std::unique_ptr<Impl> impl_;
   explicit FlatColumnReader(std::unique_ptr<Impl> impl);
 
   friend class FileReader;
 };
 
-}  // namespace parquet
+// Helper function to create a file reader from an implementation of an Arrow
+// readable file
+ARROW_EXPORT
+Status OpenFile(const std::shared_ptr<io::RandomAccessFile>& file,
+    ParquetAllocator* allocator, std::unique_ptr<FileReader>* reader);
 
+}  // namespace parquet
 }  // namespace arrow
 
 #endif  // ARROW_PARQUET_READER_H

cpp/src/arrow/parquet/writer.h

@@ -55,7 +55,7 @@ class ARROW_EXPORT FileWriter {
   MemoryPool* memory_pool() const;
 
  private:
-  class PARQUET_NO_EXPORT Impl;
+  class ARROW_NO_EXPORT Impl;
   std::unique_ptr<Impl> impl_;
 };
 

python/pyarrow/includes/parquet.pxd

@@ -96,6 +96,7 @@ cdef extern from "arrow/parquet/io.h" namespace "arrow::parquet" nogil:
         ParquetAllocator()
         ParquetAllocator(MemoryPool* pool)
         MemoryPool* pool()
+        void set_pool(MemoryPool* pool)
 
     cdef cppclass ParquetReadSource:
         ParquetReadSource(const shared_ptr[RandomAccessFile]& file,
@@ -103,19 +104,20 @@ cdef extern from "arrow/parquet/io.h" namespace "arrow::parquet" nogil:
 
 
 cdef extern from "arrow/parquet/reader.h" namespace "arrow::parquet" nogil:
-    cdef cppclass FileReader:
-        @staticmethod
-        CStatus Open(const shared_ptr[RandomAccessFile]& file,
+    CStatus OpenFile(const shared_ptr[RandomAccessFile]& file,
                      ParquetAllocator* allocator,
                      unique_ptr[FileReader]* reader)
 
+    cdef cppclass FileReader:
         FileReader(MemoryPool* pool, unique_ptr[ParquetFileReader] reader)
         CStatus ReadFlatTable(shared_ptr[CTable]* out);
 
 
 cdef extern from "arrow/parquet/schema.h" namespace "arrow::parquet" nogil:
-    CStatus FromParquetSchema(const SchemaDescriptor* parquet_schema, shared_ptr[CSchema]* out)
-    CStatus ToParquetSchema(const CSchema* arrow_schema, shared_ptr[SchemaDescriptor]* out)
+    CStatus FromParquetSchema(const SchemaDescriptor* parquet_schema,
+                              shared_ptr[CSchema]* out)
+    CStatus ToParquetSchema(const CSchema* arrow_schema,
+                            shared_ptr[SchemaDescriptor]* out)
 
 
 cdef extern from "arrow/parquet/writer.h" namespace "arrow::parquet" nogil:

python/pyarrow/io.pyx

@@ -316,10 +316,10 @@ cdef class HdfsClient:
 
 cdef class NativeFileInterface:
 
-    cdef int read_handle(self, shared_ptr[RandomAccessFile]* file) except -1:
+    cdef read_handle(self, shared_ptr[RandomAccessFile]* file):
         raise NotImplementedError
 
-    cdef int write_handle(self, shared_ptr[WriteableFile]* file) except -1:
+    cdef write_handle(self, shared_ptr[WriteableFile]* file):
         raise NotImplementedError
 
 cdef class HdfsFile(NativeFileInterface):
@@ -364,15 +364,13 @@ cdef class HdfsFile(NativeFileInterface):
         if self.is_readonly:
             raise IOError("only valid on writeonly files")
 
-    cdef int read_handle(self, shared_ptr[RandomAccessFile]* file) except -1:
+    cdef read_handle(self, shared_ptr[RandomAccessFile]* file):
         self._assert_readable()
         file[0] = <shared_ptr[RandomAccessFile]> self.rd_file
-        return 0
 
-    cdef int write_handle(self, shared_ptr[WriteableFile]* file) except -1:
+    cdef write_handle(self, shared_ptr[WriteableFile]* file):
         self._assert_writeable()
         file[0] = <shared_ptr[WriteableFile]> self.wr_file
-        return 0
 
     def size(self):
         cdef int64_t size

python/pyarrow/parquet.pyx

@@ -20,35 +20,75 @@
 # cython: embedsignature = True
 
 from pyarrow.includes.libarrow cimport *
-cimport pyarrow.includes.pyarrow as pyarrow
 from pyarrow.includes.parquet cimport *
+from pyarrow.includes.libarrow_io cimport RandomAccessFile, WriteableFile
+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.table cimport Table
 
-def read_table(filename, columns=None):
+from pyarrow.io cimport NativeFileInterface
+
+import six
+
+
+cdef class ParquetReader:
+    cdef:
+        ParquetAllocator allocator
+        unique_ptr[FileReader] reader
+
+    def __cinit__(self):
+        self.allocator.set_pool(default_memory_pool())
+
+    cdef open_local_file(self, file_path):
+        cdef c_string path = tobytes(file_path)
+
+        # Must be in one expression to avoid calling std::move which is not
+        # possible in Cython (due to missing rvalue support)
+
+        # TODO(wesm): ParquetFileReader::OpenFIle can throw?
+        self.reader = unique_ptr[FileReader](
+            new FileReader(default_memory_pool(),
+                           ParquetFileReader.OpenFile(path)))
+
+    cdef open_native_file(self, NativeFileInterface file):
+        cdef shared_ptr[RandomAccessFile] cpp_handle
+        file.read_handle(&cpp_handle)
+
+        check_cstatus(OpenFile(cpp_handle, &self.allocator, &self.reader))
+
+    def read_all(self):
+        cdef:
+            Table table = Table()
+            shared_ptr[CTable] ctable
+
+        with nogil:
+            check_cstatus(self.reader.get()
+                          .ReadFlatTable(&ctable))
+
+        table.init(ctable)
+        return table
+
+
+def read_table(source, columns=None):
     """
     Read a Table from Parquet format
     Returns
     -------
     table: pyarrow.Table
     """
-    cdef unique_ptr[FileReader] reader
-    cdef Table table = Table()
-    cdef shared_ptr[CTable] ctable
+    cdef ParquetReader reader = ParquetReader()
 
-    # 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))))
+    if isinstance(source, six.string_types):
+        reader.open_local_file(source)
+    elif isinstance(source, NativeFileInterface):
+        reader.open_native_file(source)
 
-    with nogil:
-        check_cstatus(reader.get().ReadFlatTable(&ctable))
+    return reader.read_all()
 
-    table.init(ctable)
-    return table
 
 def write_table(table, filename, chunk_size=None, version=None):
     """