Fix small bug in BinaryArray::Equals, add rudimentary StreamReader/Writer interface and tests

Change-Id: Ic578b6c747e8033d19d7b0f93c5d065bf8d96e70

cpp/src/arrow/array.cc

@@ -359,7 +359,7 @@ bool BinaryArray::EqualsExact(const BinaryArray& other) const {
 
   if (!data_buffer_ && !(other.data_buffer_)) { return true; }
 
-  return data_buffer_->Equals(*other.data_buffer_, data_buffer_->size());
+  return data_buffer_->Equals(*other.data_buffer_, raw_offsets()[length_]);
 }
 
 bool BinaryArray::Equals(const std::shared_ptr<Array>& arr) const {

cpp/src/arrow/ipc/stream.cc

@@ -117,9 +117,9 @@ Status StreamWriter::WriteRecordBatch(const RecordBatch& batch) {
 }
 
 Status StreamWriter::Close() {
-  // Close the stream
-  RETURN_NOT_OK(CheckStarted());
-  return sink_->Close();
+  // Write the schema if not already written
+  // User is responsible for closing the OutputStream
+  return CheckStarted();
 }
 
 // ----------------------------------------------------------------------

cpp/src/arrow/ipc/stream.h

@@ -54,6 +54,9 @@ class ARROW_EXPORT StreamWriter {
       std::shared_ptr<StreamWriter>* out);
 
   virtual Status WriteRecordBatch(const RecordBatch& batch);
+
+  /// Perform any logic necessary to finish the stream. User is responsible for
+  /// closing the actual OutputStream
   virtual Status Close();
 
  protected:

python/pyarrow/io.pyx

@@ -905,7 +905,7 @@ cdef class _StreamReader:
     def __cinit__(self):
         pass
 
-    def _open(self, source, footer_offset=None):
+    def _open(self, source):
         cdef:
             shared_ptr[ReadableFileInterface] reader
             shared_ptr[InputStream] in_stream
@@ -920,6 +920,10 @@ cdef class _StreamReader:
         schema.init_schema(self.reader.get().schema())
 
     def get_next_batch(self):
+        """
+        Read next RecordBatch from the stream. Raises StopIteration at end of
+        stream
+        """
         cdef shared_ptr[CRecordBatch] batch
 
         with nogil:

python/pyarrow/ipc.py

@@ -21,11 +21,35 @@ import pyarrow.io as io
 
 
 class StreamReader(io._StreamReader):
-    pass
+    """
+    Reader for the Arrow streaming binary format
+
+    Parameters
+    ----------
+    source : str, pyarrow.NativeFile, or file-like Python object
+        Either a file path, or a readable file object
+    """
+    def __init__(self, source):
+        self._open(source)
+
+    def __iter__(self):
+        while True:
+            yield self.get_next_batch()
 
 
 class StreamWriter(io._StreamWriter):
-    pass
+    """
+    Writer for the Arrow streaming binary format
+
+    Parameters
+    ----------
+    sink : str, pyarrow.NativeFile, or file-like Python object
+        Either a file path, or a writeable file object
+    schema : pyarrow.Schema
+        The Arrow schema for data to be written to the file
+    """
+    def __init__(self, sink, schema):
+        self._open(sink, schema)
 
 
 class FileReader(io._FileReader):

python/pyarrow/schema.pyx

@@ -112,6 +112,7 @@ cdef class Field:
         def __get__(self):
             return frombytes(self.field.name)
 
+
 cdef class Schema:
 
     def __cinit__(self):

python/pyarrow/tests/test_ipc.py

@@ -16,21 +16,20 @@
 # under the License.
 
 import io
+import pytest
 
 import numpy as np
 
 from pandas.util.testing import assert_frame_equal
 import pandas as pd
 
-import pyarrow as A
-import pyarrow.io as aio
-import pyarrow.ipc as ipc
+from pyarrow.compat import unittest
+import pyarrow as pa
 
 
-class RoundtripTest(object):
-    # Also tests writing zero-copy NumPy array with additional padding
+class MessagingTest(object):
 
-    def __init__(self):
+    def setUp(self):
         self.sink = self._get_sink()
 
     def _get_sink(self):
@@ -39,14 +38,15 @@ class RoundtripTest(object):
     def _get_source(self):
         return self.sink.getvalue()
 
-    def run(self):
+    def write_batches(self):
         nrows = 5
         df = pd.DataFrame({
             'one': np.random.randn(nrows),
             'two': ['foo', np.nan, 'bar', 'bazbaz', 'qux']})
 
-        batch = A.RecordBatch.from_pandas(df)
-        writer = ipc.FileWriter(self.sink, batch.schema)
+        batch = pa.RecordBatch.from_pandas(df)
+
+        writer = self._get_writer(self.sink, batch.schema)
 
         num_batches = 5
         frames = []
@@ -55,46 +55,73 @@ class RoundtripTest(object):
             unique_df = df.copy()
             unique_df['one'] = np.random.randn(nrows)
 
-            batch = A.RecordBatch.from_pandas(unique_df)
+            batch = pa.RecordBatch.from_pandas(unique_df)
             writer.write_batch(batch)
             frames.append(unique_df)
             batches.append(batch)
 
         writer.close()
+        return batches
 
+
+class TestFile(MessagingTest, unittest.TestCase):
+    # Also tests writing zero-copy NumPy array with additional padding
+
+    def _get_writer(self, sink, schema):
+        return pa.FileWriter(sink, schema)
+
+    def test_simple_roundtrip(self):
+        batches = self.write_batches()
         file_contents = self._get_source()
-        reader = ipc.FileReader(aio.BufferReader(file_contents))
 
-        assert reader.num_record_batches == num_batches
+        reader = pa.FileReader(pa.BufferReader(file_contents))
 
-        for i in range(num_batches):
+        assert reader.num_record_batches == len(batches)
+
+        for i, batch in enumerate(batches):
             # it works. Must convert back to DataFrame
             batch = reader.get_batch(i)
             assert batches[i].equals(batch)
 
 
-class InMemoryStreamTest(RoundtripTest):
+class TestStream(MessagingTest, unittest.TestCase):
+
+    def _get_writer(self, sink, schema):
+        return pa.StreamWriter(sink, schema)
+
+    def test_simple_roundtrip(self):
+        batches = self.write_batches()
+        file_contents = self._get_source()
+        reader = pa.StreamReader(pa.BufferReader(file_contents))
+
+        total = 0
+        for i, next_batch in enumerate(reader):
+            assert next_batch.equals(batches[i])
+            total += 1
+
+        assert total == len(batches)
+
+        with pytest.raises(StopIteration):
+            reader.get_next_batch()
+
+
+class TestInMemoryFile(TestFile):
 
     def _get_sink(self):
-        return aio.InMemoryOutputStream()
+        return pa.InMemoryOutputStream()
 
     def _get_source(self):
         return self.sink.get_result()
 
 
-def test_ipc_file_simple_roundtrip():
-    helper = RoundtripTest()
-    helper.run()
-
-
 def test_ipc_zero_copy_numpy():
     df = pd.DataFrame({'foo': [1.5]})
 
-    batch = A.RecordBatch.from_pandas(df)
-    sink = aio.InMemoryOutputStream()
+    batch = pa.RecordBatch.from_pandas(df)
+    sink = pa.InMemoryOutputStream()
     write_file(batch, sink)
     buffer = sink.get_result()
-    reader = aio.BufferReader(buffer)
+    reader = pa.BufferReader(buffer)
 
     batches = read_file(reader)
 
@@ -104,12 +131,12 @@ def test_ipc_zero_copy_numpy():
 
 
 def write_file(batch, sink):
-    writer = ipc.FileWriter(sink, batch.schema)
+    writer = pa.FileWriter(sink, batch.schema)
     writer.write_batch(batch)
     writer.close()
 
 
 def read_file(source):
-    reader = ipc.FileReader(source)
+    reader = pa.FileReader(source)
     return [reader.get_batch(i)
             for i in range(reader.num_record_batches)]