Fix up Python bindings, unit test

Change-Id: I35ab3ace0e4ca7a80fc7d85e55ac55ea222b15dc

cpp/src/parquet/arrow/CMakeLists.txt

@@ -21,7 +21,9 @@ ADD_PARQUET_TEST(arrow-reader-writer-test)
 ADD_BENCHMARK(reader-writer-benchmark
   PREFIX "parquet-arrow"
   EXTRA_LINK_LIBS ${PARQUET_BENCHMARK_LINK_LIBRARIES})
-add_dependencies(parquet parquet-arrow-reader-writer-benchmark)
+if (TARGET parquet-arrow-reader-writer-benchmark)
+  add_dependencies(parquet parquet-arrow-reader-writer-benchmark)
+endif()
 
 # Headers: top level
 install(FILES

python/pyarrow/_parquet.pxd

@@ -19,7 +19,7 @@
 # cython: language_level = 3
 
 from pyarrow.includes.common cimport *
-from pyarrow.includes.libarrow cimport (CArray, CSchema, CStatus,
+from pyarrow.includes.libarrow cimport (CChunkedArray, CSchema, CStatus,
                                         CTable, CMemoryPool,
                                         CKeyValueMetadata,
                                         RandomAccessFile, OutputStream,
@@ -272,8 +272,8 @@ cdef extern from "parquet/arrow/reader.h" namespace "parquet::arrow" nogil:
 
     cdef cppclass FileReader:
         FileReader(CMemoryPool* pool, unique_ptr[ParquetFileReader] reader)
-        CStatus ReadColumn(int i, shared_ptr[CArray]* out)
-        CStatus ReadSchemaField(int i, shared_ptr[CArray]* out)
+        CStatus ReadColumn(int i, shared_ptr[CChunkedArray]* out)
+        CStatus ReadSchemaField(int i, shared_ptr[CChunkedArray]* out)
 
         int num_row_groups()
         CStatus ReadRowGroup(int i, shared_ptr[CTable]* out)

python/pyarrow/_parquet.pyx

@@ -26,6 +26,7 @@ from pyarrow.lib cimport (Array, Schema,
                           check_status,
                           MemoryPool, maybe_unbox_memory_pool,
                           Table,
+                          pyarrow_wrap_chunked_array,
                           pyarrow_wrap_schema,
                           pyarrow_wrap_table,
                           NativeFile, get_reader, get_writer)
@@ -770,28 +771,18 @@ cdef class ParquetReader:
         return self._column_idx_map[tobytes(column_name)]
 
     def read_column(self, int column_index):
-        cdef:
-            Array array = Array()
-            shared_ptr[CArray] carray
-
+        cdef shared_ptr[CChunkedArray] out
         with nogil:
             check_status(self.reader.get()
-                         .ReadColumn(column_index, &carray))
-
-        array.init(carray)
-        return array
+                         .ReadColumn(column_index, &out))
+        return pyarrow_wrap_chunked_array(out)
 
     def read_schema_field(self, int field_index):
-        cdef:
-            Array array = Array()
-            shared_ptr[CArray] carray
-
+        cdef shared_ptr[CChunkedArray] out
         with nogil:
             check_status(self.reader.get()
-                         .ReadSchemaField(field_index, &carray))
-
-        array.init(carray)
-        return array
+                         .ReadSchemaField(field_index, &out))
+        return pyarrow_wrap_chunked_array(out)
 
 
 cdef class ParquetWriter:

python/pyarrow/lib.pxd

@@ -396,6 +396,8 @@ cdef object pyarrow_wrap_metadata(
 #
 
 cdef public object pyarrow_wrap_array(const shared_ptr[CArray]& sp_array)
+cdef public object pyarrow_wrap_chunked_array(
+    const shared_ptr[CChunkedArray]& sp_array)
 # XXX pyarrow.h calls it `wrap_record_batch`
 cdef public object pyarrow_wrap_batch(const shared_ptr[CRecordBatch]& cbatch)
 cdef public object pyarrow_wrap_buffer(const shared_ptr[CBuffer]& buf)

python/pyarrow/tests/test_parquet.py

@@ -1979,7 +1979,9 @@ def test_binary_array_overflow_to_chunked():
 
     col0_data = read_tbl[0].data
     assert isinstance(col0_data, pa.ChunkedArray)
-    assert col0_data.num_chunks == 2
+
+    # Split up into 16MB chunks. 128 * 16 = 2048, so 129
+    assert col0_data.num_chunks == 129
 
     assert tbl.equals(read_tbl)