Implement Table::FromRecordBatches

Change-Id: I740bc111d0a3e81f8a5076a6607926a8d164ddc1

cpp/CMakeLists.txt

@@ -76,7 +76,7 @@ if("${CMAKE_SOURCE_DIR}" STREQUAL "${CMAKE_CURRENT_SOURCE_DIR}")
 
   option(ARROW_JEMALLOC
     "Build the Arrow jemalloc-based allocator"
-    ON)
+    OFF)
 
   option(ARROW_BOOST_USE_SHARED
     "Rely on boost shared libraries where relevant"

cpp/src/arrow/table-test.cc

@@ -44,16 +44,23 @@ class TestTable : public TestBase {
     vector<shared_ptr<Field>> fields = {f0, f1, f2};
     schema_ = std::make_shared<Schema>(fields);
 
+    arrays_ = {
+      MakePrimitive<Int32Array>(length),
+      MakePrimitive<UInt8Array>(length),
+      MakePrimitive<Int16Array>(length)};
+
     columns_ = {
-        std::make_shared<Column>(schema_->field(0), MakePrimitive<Int32Array>(length)),
-        std::make_shared<Column>(schema_->field(1), MakePrimitive<UInt8Array>(length)),
-        std::make_shared<Column>(schema_->field(2), MakePrimitive<Int16Array>(length))};
+        std::make_shared<Column>(schema_->field(0), arrays_[0]),
+        std::make_shared<Column>(schema_->field(1), arrays_[1]),
+        std::make_shared<Column>(schema_->field(2), arrays_[2])};
   }
 
  protected:
   std::shared_ptr<Table> table_;
   shared_ptr<Schema> schema_;
-  vector<std::shared_ptr<Column>> columns_;
+
+  std::vector<std::shared_ptr<Array>> arrays_;
+  std::vector<std::shared_ptr<Column>> columns_;
 };
 
 TEST_F(TestTable, EmptySchema) {
@@ -65,7 +72,7 @@ TEST_F(TestTable, EmptySchema) {
 }
 
 TEST_F(TestTable, Ctors) {
-  int length = 100;
+  const int length = 100;
   MakeExample1(length);
 
   std::string name = "data";
@@ -83,7 +90,7 @@ TEST_F(TestTable, Ctors) {
 }
 
 TEST_F(TestTable, Metadata) {
-  int length = 100;
+  const int length = 100;
   MakeExample1(length);
 
   std::string name = "data";
@@ -98,7 +105,7 @@ TEST_F(TestTable, Metadata) {
 
 TEST_F(TestTable, InvalidColumns) {
   // Check that columns are all the same length
-  int length = 100;
+  const int length = 100;
   MakeExample1(length);
 
   table_.reset(new Table("data", schema_, columns_, length - 1));
@@ -120,7 +127,7 @@ TEST_F(TestTable, InvalidColumns) {
 }
 
 TEST_F(TestTable, Equals) {
-  int length = 100;
+  const int length = 100;
   MakeExample1(length);
 
   std::string name = "data";
@@ -145,6 +152,40 @@ TEST_F(TestTable, Equals) {
   ASSERT_FALSE(table_->Equals(std::make_shared<Table>(name, schema_, other_columns)));
 }
 
+TEST_F(TestTable, FromRecordBatches) {
+  const int32_t length = 10;
+  MakeExample1(length);
+
+  auto batch1 = std::make_shared<RecordBatch>(schema_, length, arrays_);
+
+  std::shared_ptr<Table> result, expected;
+  ASSERT_OK(Table::FromRecordBatches("foo", {batch1}, &result));
+
+  expected = std::make_shared<Table>("foo", schema_, columns_);
+  ASSERT_TRUE(result->Equals(expected));
+
+  std::vector<std::shared_ptr<Column>> other_columns;
+  for (int i = 0; i < schema_->num_fields(); ++i) {
+    std::vector<std::shared_ptr<Array>> col_arrays = {arrays_[i], arrays_[i]};
+    other_columns.push_back(std::make_shared<Column>(schema_->field(i), col_arrays));
+  }
+
+  ASSERT_OK(Table::FromRecordBatches("foo", {batch1, batch1}, &result));
+  expected = std::make_shared<Table>("foo", schema_, other_columns);
+  ASSERT_TRUE(result->Equals(expected));
+
+  // Error states
+  std::vector<std::shared_ptr<RecordBatch>> empty_batches;
+  ASSERT_RAISES(Invalid, Table::FromRecordBatches("", empty_batches, &result));
+
+  std::vector<std::shared_ptr<Field>> fields = {schema_->field(0), schema_->field(1)};
+  auto other_schema = std::make_shared<Schema>(fields);
+
+  std::vector<std::shared_ptr<Array>> other_arrays = {arrays_[0], arrays_[1]};
+  auto batch2 = std::make_shared<RecordBatch>(other_schema, length, other_arrays);
+  ASSERT_RAISES(Invalid, Table::FromRecordBatches("", {batch1, batch2}, &result));
+}
+
 class TestRecordBatch : public TestBase {};
 
 TEST_F(TestRecordBatch, Equals) {

cpp/src/arrow/table.cc

@@ -77,7 +77,45 @@ Table::Table(const std::string& name, const std::shared_ptr<Schema>& schema,
     const std::vector<std::shared_ptr<Column>>& columns, int64_t num_rows)
     : name_(name), schema_(schema), columns_(columns), num_rows_(num_rows) {}
 
+Status Table::FromRecordBatches(const std::string& name,
+    const std::vector<std::shared_ptr<RecordBatch>>& batches,
+    std::shared_ptr<Table>* table) {
+  if (batches.size() == 0) {
+    return Status::Invalid("Must pass at least one record batch");
+  }
+
+  std::shared_ptr<Schema> schema = batches[0]->schema();
+
+  const int nbatches = static_cast<int>(batches.size());
+  const int ncolumns = static_cast<int>(schema->num_fields());
+
+  for (int i = 1; i < nbatches; ++i) {
+    if (!batches[i]->schema()->Equals(schema)) {
+      std::stringstream ss;
+      ss << "Schema at index " << static_cast<int>(i)
+         << " was different: \n"
+         << schema->ToString() << "\nvs\n"
+         << batches[i]->schema()->ToString();
+      return Status::Invalid(ss.str());
+    }
+  }
+
+  std::vector<std::shared_ptr<Column>> columns(ncolumns);
+  std::vector<std::shared_ptr<Array>> column_arrays(nbatches);
+
+  for (int i = 0; i < ncolumns; ++i) {
+    for (int j = 0; j < nbatches; ++j) {
+      column_arrays[j] = batches[j]->column(i);
+    }
+    columns[i] = std::make_shared<Column>(schema->field(i), column_arrays);
+  }
+
+  *table = std::make_shared<Table>(name, schema, columns);
+  return Status::OK();
+}
+
 bool Table::Equals(const Table& other) const {
+
   if (name_ != other.name()) { return false; }
   if (!schema_->Equals(other.schema())) { return false; }
   if (static_cast<int64_t>(columns_.size()) != other.num_columns()) { return false; }

cpp/src/arrow/table.h

@@ -82,7 +82,18 @@ class ARROW_EXPORT Table {
   // same length as num_rows -- you can validate this using
   // Table::ValidateColumns
   Table(const std::string& name, const std::shared_ptr<Schema>& schema,
-      const std::vector<std::shared_ptr<Column>>& columns, int64_t num_rows);
+      const std::vector<std::shared_ptr<Column>>& columns, int64_t nubm_rows);
+
+  // Construct table from RecordBatch, but only if all of the batch schemas are
+  // equal. Returns Status::Invalid if there is some problem
+  static Status FromRecordBatches(const std::string& name,
+      const std::vector<std::shared_ptr<RecordBatch>>& batches,
+      std::shared_ptr<Table>* table);
+
+  // Construct table from multiple input tables
+  Status Concatenate(const std::string& output_name,
+      const std::vector<std::shared_ptr<Table>>& others,
+      std::shared_ptr<Table>* table);
 
   // @returns: the table's name, if any (may be length 0)
   const std::string& name() const { return name_; }

python/pyarrow/includes/libarrow.pxd

@@ -207,6 +207,12 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         CTable(const c_string& name, const shared_ptr[CSchema]& schema,
                const vector[shared_ptr[CColumn]]& columns)
 
+        @staticmethod
+        CStatus FromRecordBatches(
+            const c_string& name,
+            const vector[shared_ptr[CRecordBatch]]& batches,
+            shared_ptr[CTable]* table)
+
         int num_columns()
         int num_rows()
 

python/pyarrow/table.pyx

@@ -584,7 +584,7 @@ cdef class Table:
         return result
 
     @staticmethod
-    def from_batches(batches):
+    def from_batches(batches, name=None):
         """
         Construct a Table from a list of Arrow RecordBatches
 
@@ -596,37 +596,22 @@ cdef class Table:
         """
 
         cdef:
-            vector[shared_ptr[CArray]] c_array_chunks
-            vector[shared_ptr[CColumn]] c_columns
+            vector[shared_ptr[CRecordBatch]] c_batches
             shared_ptr[CTable] c_table
-            Array arr
-            Schema schema
+            RecordBatch batch
+            Table table
 
-        import pandas as pd
+        if name is None:
+            name = b''
+        else:
+            name = tobytes(name)
 
-        schema = batches[0].schema
+        for batch in batches:
+            c_batches.push_back(batch.sp_batch)
 
-        # check schemas are equal
-        for other in batches[1:]:
-            if not schema.equals(other.schema):
-                raise ArrowException("Error converting list of RecordBatches "
-                        "to DataFrame, not all schemas are equal: {%s} != {%s}"
-                        % (str(schema), str(other.schema)))
+        with nogil:
+            check_status(CTable.FromRecordBatches(name, c_batches, &c_table))
 
-        cdef int K = batches[0].num_columns
-
-        # create chunked columns from the batches
-        c_columns.resize(K)
-        for i in range(K):
-            for batch in batches:
-                arr = batch[i]
-                c_array_chunks.push_back(arr.sp_array)
-            c_columns[i].reset(new CColumn(schema.sp_schema.get().field(i),
-                               c_array_chunks))
-            c_array_chunks.clear()
-
-        # create a Table from columns and convert to DataFrame
-        c_table.reset(new CTable('', schema.sp_schema, c_columns))
         table = Table()
         table.init(c_table)
         return table