ARROW-962: [Python] Add schema attribute to RecordBatchFileReader

Author: Wes McKinney <[email protected]> Closes apache#843 from wesm/ARROW-962 and squashes the following commits: b7144aa [Wes McKinney] Add schema attribute to RecordBatchFileReader
luffyhwl · Jul 15, 2017 · bfe3959 · bfe3959
1 parent cb31b8b
commit bfe3959
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/python/pyarrow/io.pxi b/python/pyarrow/io.pxi
@@ -990,8 +990,7 @@ cdef class _RecordBatchReader:
             check_status(CRecordBatchStreamReader.Open(in_stream, &reader))
 
         self.reader = <shared_ptr[CRecordBatchReader]> reader
-        self.schema = Schema()
-        self.schema.init_schema(self.reader.get().schema())
+        self.schema = pyarrow_wrap_schema(self.reader.get().schema())
 
     def get_next_batch(self):
         """
@@ -1049,6 +1048,9 @@ cdef class _RecordBatchFileReader:
     cdef:
         shared_ptr[CRecordBatchFileReader] reader
 
+    cdef readonly:
+        Schema schema
+
     def __cinit__(self):
         pass
 
@@ -1067,6 +1069,8 @@ cdef class _RecordBatchFileReader:
             else:
                 check_status(CRecordBatchFileReader.Open(reader, &self.reader))
 
+        self.schema = pyarrow_wrap_schema(self.reader.get().schema())
+
     property num_record_batches:
 
         def __get__(self):

diff --git a/python/pyarrow/tests/test_ipc.py b/python/pyarrow/tests/test_ipc.py
@@ -89,6 +89,7 @@ def test_simple_roundtrip(self):
             # it works. Must convert back to DataFrame
             batch = reader.get_batch(i)
             assert batches[i].equals(batch)
+            assert reader.schema.equals(batches[0].schema)
 
     def test_read_all(self):
         batches = self.write_batches()