Fix up optional extensions

Change-Id: Icad57e6d5e9ee5302e8623664c3c58ac363bdd69
wesm · May 13, 2017 · 9e6ee24 · 9e6ee24
1 parent cff757d
commit 9e6ee24
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 18 deletions.
diff --git a/python/pyarrow/_jemalloc.pyx b/python/pyarrow/_jemalloc.pyx
@@ -20,7 +20,7 @@
 # cython: embedsignature = True
 
 from pyarrow.includes.libarrow_jemalloc cimport CJemallocMemoryPool
-from pyarrow._memory cimport MemoryPool
+from pyarrow.lib cimport MemoryPool
 
 def default_pool():
     cdef MemoryPool pool = MemoryPool()

diff --git a/python/pyarrow/_parquet.pyx b/python/pyarrow/_parquet.pyx
@@ -22,15 +22,16 @@
 from cython.operator cimport dereference as deref
 from pyarrow.includes.common cimport *
 from pyarrow.includes.libarrow cimport *
-from pyarrow._array cimport Array, Schema, box_schema
-from pyarrow._error cimport check_status
-from pyarrow._memory cimport MemoryPool, maybe_unbox_memory_pool
-from pyarrow._table cimport Table, table_from_ctable
-from pyarrow._io cimport NativeFile, get_reader, get_writer
+from pyarrow.lib cimport (Array, Schema,
+                          check_status,
+                          MemoryPool, maybe_unbox_memory_pool,
+                          Table,
+                          pyarrow_wrap_schema,
+                          pyarrow_wrap_table,
+                          NativeFile, get_reader, get_writer)
 
 from pyarrow.compat import tobytes, frombytes
-from pyarrow._error import ArrowException
-from pyarrow._io import NativeFile
+from pyarrow.lib import ArrowException, NativeFile
 
 import six
 
@@ -212,7 +213,7 @@ cdef class ParquetSchema:
         with nogil:
             check_status(FromParquetSchema(self.schema, &sp_arrow_schema))
 
-        return box_schema(sp_arrow_schema)
+        return pyarrow_wrap_schema(sp_arrow_schema)
 
     def equals(self, ParquetSchema other):
         """
@@ -425,7 +426,7 @@ cdef class ParquetReader:
             with nogil:
                 check_status(self.reader.get()
                              .ReadRowGroup(i, &ctable))
-        return table_from_ctable(ctable)
+        return pyarrow_wrap_table(ctable)
 
     def read_all(self, column_indices=None):
         cdef:
@@ -444,7 +445,7 @@ cdef class ParquetReader:
             with nogil:
                 check_status(self.reader.get()
                              .ReadTable(&ctable))
-        return table_from_ctable(ctable)
+        return pyarrow_wrap_table(ctable)
 
     def column_name_idx(self, column_name):
         """

diff --git a/python/pyarrow/lib.pxd b/python/pyarrow/lib.pxd
@@ -325,3 +325,13 @@ cdef class NativeFile:
 
 cdef get_reader(object source, shared_ptr[RandomAccessFile]* reader)
 cdef get_writer(object source, shared_ptr[OutputStream]* writer)
+
+cdef public object pyarrow_wrap_buffer(const shared_ptr[CBuffer]& buf)
+cdef public object pyarrow_wrap_data_type(const shared_ptr[CDataType]& type)
+cdef public object pyarrow_wrap_field(const shared_ptr[CField]& field)
+cdef public object pyarrow_wrap_schema(const shared_ptr[CSchema]& type)
+cdef public object pyarrow_wrap_array(const shared_ptr[CArray]& sp_array)
+cdef public object pyarrow_wrap_tensor(const shared_ptr[CTensor]& sp_tensor)
+cdef public object pyarrow_wrap_column(const shared_ptr[CColumn]& ccolumn)
+cdef public object pyarrow_wrap_table(const shared_ptr[CTable]& ctable)
+cdef public object pyarrow_wrap_batch(const shared_ptr[CRecordBatch]& cbatch)
diff --git a/python/pyarrow/parquet.py b/python/pyarrow/parquet.py
@@ -24,8 +24,7 @@
                               RowGroupMetaData, ParquetSchema,
                               ParquetWriter)
 import pyarrow._parquet as _parquet  # noqa
-import pyarrow._array as _array
-import pyarrow._table as _table
+import pyarrow.lib as lib
 
 
 # ----------------------------------------------------------------------
@@ -241,8 +240,8 @@ def read(self, columns=None, nthreads=1, partitions=None,
                 # manifest, so ['a', 'b', 'c'] as in our example above.
                 dictionary = partitions.levels[i].dictionary
 
-                arr = _array.DictionaryArray.from_arrays(indices, dictionary)
-                col = _table.Column.from_array(name, arr)
+                arr = lib.DictionaryArray.from_arrays(indices, dictionary)
+                col = lib.Column.from_array(name, arr)
                 table = table.append_column(col)
 
         return table
@@ -298,9 +297,9 @@ def dictionary(self):
         # Only integer and string partition types are supported right now
         try:
             integer_keys = [int(x) for x in self.keys]
-            dictionary = _array.array(integer_keys)
+            dictionary = lib.array(integer_keys)
         except ValueError:
-            dictionary = _array.array(self.keys)
+            dictionary = lib.array(self.keys)
 
         self._dictionary = dictionary
         return dictionary
@@ -539,7 +538,7 @@ def read(self, columns=None, nthreads=1):
                                open_file_func=open_file)
             tables.append(table)
 
-        all_data = _table.concat_tables(tables)
+        all_data = lib.concat_tables(tables)
         return all_data
 
     def _get_open_file_func(self):