Implement RecordBatch::Equals, compare in Python ipc file writes

Change-Id: I84cdb9a07f3853f06436c53b65fa082e3b843ed2

cpp/src/arrow/table-test.cc

@@ -123,4 +123,31 @@ TEST_F(TestTable, InvalidColumns) {
   ASSERT_RAISES(Invalid, table_->ValidateColumns());
 }
 
+class TestRecordBatch : public TestBase {};
+
+TEST_F(TestRecordBatch, Equals) {
+  const int length = 10;
+
+  auto f0 = std::make_shared<Field>("f0", INT32);
+  auto f1 = std::make_shared<Field>("f1", UINT8);
+  auto f2 = std::make_shared<Field>("f2", INT16);
+
+  vector<shared_ptr<Field>> fields = {f0, f1, f2};
+  auto schema = std::make_shared<Schema>(fields);
+
+  auto a0 = MakePrimitive<Int32Array>(length);
+  auto a1 = MakePrimitive<UInt8Array>(length);
+  auto a2 = MakePrimitive<Int16Array>(length);
+
+  RecordBatch b1(schema, length, {a0, a1, a2});
+  RecordBatch b2(schema, 5, {a0, a1, a2});
+  RecordBatch b3(schema, length, {a0, a1});
+  RecordBatch b4(schema, length, {a0, a1, a1});
+
+  ASSERT_TRUE(b1.Equals(b1));
+  ASSERT_FALSE(b1.Equals(b2));
+  ASSERT_FALSE(b1.Equals(b3));
+  ASSERT_FALSE(b1.Equals(b4));
+}
+
 }  // namespace arrow

cpp/src/arrow/table.cc

@@ -21,6 +21,7 @@
 #include <memory>
 #include <sstream>
 
+#include "arrow/array.h"
 #include "arrow/column.h"
 #include "arrow/schema.h"
 #include "arrow/util/status.h"
@@ -35,6 +36,21 @@ const std::string& RecordBatch::column_name(int i) const {
   return schema_->field(i)->name;
 }
 
+bool RecordBatch::Equals(const RecordBatch& other) const {
+  if (num_columns() != other.num_columns() || num_rows_ != other.num_rows()) {
+    return false;
+  }
+
+  for (int i = 0; i < num_columns(); ++i) {
+    if (!column(i)->Equals(other.column(i))) { return false; }
+  }
+
+  return true;
+}
+
+// ----------------------------------------------------------------------
+// Table methods
+
 Table::Table(const std::string& name, const std::shared_ptr<Schema>& schema,
     const std::vector<std::shared_ptr<Column>>& columns)
     : name_(name), schema_(schema), columns_(columns) {

cpp/src/arrow/table.h

@@ -43,6 +43,8 @@ class ARROW_EXPORT RecordBatch {
   RecordBatch(const std::shared_ptr<Schema>& schema, int32_t num_rows,
       const std::vector<std::shared_ptr<Array>>& columns);
 
+  bool Equals(const RecordBatch& other) const;
+
   // @returns: the table's schema
   const std::shared_ptr<Schema>& schema() const { return schema_; }
 

cpp/src/arrow/types/primitive-test.cc

@@ -238,8 +238,7 @@ void TestPrimitiveBuilder<PBoolean>::Check(
 }
 
 typedef ::testing::Types<PBoolean, PUInt8, PUInt16, PUInt32, PUInt64, PInt8, PInt16,
-    PInt32, PInt64, PFloat, PDouble>
-    Primitives;
+    PInt32, PInt64, PFloat, PDouble> Primitives;
 
 TYPED_TEST_CASE(TestPrimitiveBuilder, Primitives);
 

cpp/src/arrow/util/logging.h

@@ -118,9 +118,9 @@ class CerrLog {
 class FatalLog : public CerrLog {
  public:
   explicit FatalLog(int /* severity */)  // NOLINT
-      : CerrLog(ARROW_FATAL){}           // NOLINT
+      : CerrLog(ARROW_FATAL) {}          // NOLINT
 
-            [[noreturn]] ~FatalLog() {
+  [[noreturn]] ~FatalLog() {
     if (has_logged_) { std::cerr << std::endl; }
     std::exit(1);
   }

python/pyarrow/includes/libarrow.pxd

@@ -153,6 +153,8 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         CRecordBatch(const shared_ptr[CSchema]& schema, int32_t num_rows,
                      const vector[shared_ptr[CArray]]& columns)
 
+        c_bool Equals(const CRecordBatch& other)
+
         const shared_ptr[CSchema]& schema()
         const shared_ptr[CArray]& column(int i)
         const c_string& column_name(int i)

python/pyarrow/table.pyx

@@ -19,6 +19,8 @@
 # distutils: language = c++
 # cython: embedsignature = True
 
+from cython.operator cimport dereference as deref
+
 from pyarrow.includes.libarrow cimport *
 cimport pyarrow.includes.pyarrow as pyarrow
 
@@ -45,8 +47,8 @@ cdef class ChunkedArray:
 
     cdef _check_nullptr(self):
         if self.chunked_array == NULL:
-            raise ReferenceError("ChunkedArray object references a NULL pointer."
-                    "Not initialized.")
+            raise ReferenceError("ChunkedArray object references a NULL "
+                                 "pointer. Not initialized.")
 
     def length(self):
         self._check_nullptr()
@@ -225,6 +227,12 @@ cdef class RecordBatch:
         arr.init(self.batch.column(i))
         return arr
 
+    def equals(self, RecordBatch other):
+        self._check_nullptr()
+        other._check_nullptr()
+
+        return self.batch.Equals(deref(other.batch))
+
     @classmethod
     def from_pandas(cls, df):
         """

python/pyarrow/tests/test_ipc.py

@@ -17,7 +17,6 @@
 
 import io
 
-from pandas.util.testing import assert_frame_equal
 import numpy as np
 import pandas as pd
 
@@ -25,7 +24,7 @@ import pyarrow as A
 import pyarrow.ipc as ipc
 
 
-def test_ipc_simple_roundtrip():
+def test_ipc_file_simple_roundtrip():
     # Also tests writing zero-copy NumPy array with additional padding
 
     nrows = 5
@@ -40,6 +39,7 @@ def test_ipc_simple_roundtrip():
 
     num_batches = 5
     frames = []
+    batches = []
     for i in range(num_batches):
         unique_df = df.copy()
         unique_df['one'] = np.random.randn(nrows)
@@ -47,6 +47,7 @@ def test_ipc_simple_roundtrip():
         batch = A.RecordBatch.from_pandas(unique_df)
         writer.write_record_batch(batch)
         frames.append(unique_df)
+        batches.append(batch)
 
     writer.close()
 
@@ -59,6 +60,8 @@ def test_ipc_simple_roundtrip():
         # it works. Must convert back to DataFrame
         batch = reader.get_record_batch(i)
 
+        assert batches[i].equals(batch)
+
 
 # XXX: For benchmarking