ARROW-60: Struct type builder API

Implement the basic classes, StructArray and StructBuilder, meanwhile,

add the perspective test cases for them.

Other necessary methods will be added subsequetly.

Author: fengguangyuan

cpp/src/arrow/type.h

@@ -161,6 +161,7 @@ struct Field {
 
   std::string ToString() const;
 };
+typedef std::shared_ptr<Field> FieldPtr;
 
 template <typename Derived>
 struct PrimitiveType : public DataType {

cpp/src/arrow/types/construct.cc

@@ -23,6 +23,7 @@
 #include "arrow/types/list.h"
 #include "arrow/types/primitive.h"
 #include "arrow/types/string.h"
+#include "arrow/types/struct.h"
 #include "arrow/util/buffer.h"
 #include "arrow/util/status.h"
 
@@ -71,6 +72,20 @@ Status MakeBuilder(MemoryPool* pool, const std::shared_ptr<DataType>& type,
   }
 }
 
+Status MakeStructBuilder(MemoryPool* pool, const std::shared_ptr<DataType>& type,
+    const std::vector<std::shared_ptr<Field>>& fields,
+    std::shared_ptr<ArrayBuilder>* out) {
+  std::vector<std::shared_ptr<ArrayBuilder>> values_builder;
+
+  for (auto it = fields.cbegin(); it != fields.cend(); it++) {
+    std::shared_ptr<ArrayBuilder> builder;
+    RETURN_NOT_OK(MakeBuilder(pool, it->get()->type, &builder));
+    values_builder.push_back(builder);
+  }
+  out->reset(new StructBuilder(pool, type, fields, values_builder));
+  return Status::OK();
+}
+
 #define MAKE_PRIMITIVE_ARRAY_CASE(ENUM, ArrayType)                          \
   case Type::ENUM:                                                          \
     out->reset(new ArrayType(type, length, data, null_count, null_bitmap)); \

cpp/src/arrow/types/construct.h

@@ -20,19 +20,24 @@
 
 #include <cstdint>
 #include <memory>
-
+#include <vector>
 namespace arrow {
 
 class Array;
 class ArrayBuilder;
 class Buffer;
 struct DataType;
+struct Field;
 class MemoryPool;
 class Status;
 
 Status MakeBuilder(MemoryPool* pool, const std::shared_ptr<DataType>& type,
     std::shared_ptr<ArrayBuilder>* out);
 
+Status MakeStructBuilder(MemoryPool* pool, const std::shared_ptr<DataType>& type,
+    const std::vector<std::shared_ptr<Field>>& fields,
+    std::shared_ptr<ArrayBuilder>* out);
+
 // Create new arrays for logical types that are backed by primitive arrays.
 Status MakePrimitiveArray(const std::shared_ptr<DataType>& type, int32_t length,
     const std::shared_ptr<Buffer>& data, int32_t null_count,

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

@@ -22,6 +22,16 @@
 #include "gtest/gtest.h"
 
 #include "arrow/type.h"
+#include "arrow/array.h"
+#include "arrow/builder.h"
+#include "arrow/test-util.h"
+#include "arrow/types/struct.h"
+#include "arrow/types/construct.h"
+#include "arrow/types/list.h"
+#include "arrow/types/primitive.h"
+#include "arrow/types/test-common.h"
+#include "arrow/util/status.h"
+
 
 using std::shared_ptr;
 using std::string;
@@ -52,4 +62,132 @@ TEST(TestStructType, Basics) {
   // TODO(wesm): out of bounds for field(...)
 }
 
-}  // namespace arrow
+// .............................................................................
+// Struct test
+class TestStructBuilder : public TestBuilder {
+ public:
+  void SetUp() {
+    TestBuilder::SetUp();
+
+    auto value_type = TypePtr(new Int32Type());
+    auto char_type = TypePtr(new Int8Type());
+    auto list_type = TypePtr(new ListType(char_type));
+
+    std::vector<TypePtr> types = {list_type, value_type};
+    std::vector<FieldPtr> fields;
+    fields.push_back(FieldPtr(new Field("list", list_type)));
+    fields.push_back(FieldPtr(new Field("int", value_type)));
+
+    type_ = TypePtr(new StructType(fields));
+    value_fields_ = fields;
+
+    std::shared_ptr<ArrayBuilder> tmp;
+    ASSERT_OK(MakeStructBuilder(pool_, type_, fields, &tmp));
+
+    builder_ = std::dynamic_pointer_cast<StructBuilder>(tmp);
+  }
+
+  void Done() {
+    result_ = std::dynamic_pointer_cast<StructArray>(builder_->Finish());
+  }
+
+ protected:
+  std::vector<FieldPtr> value_fields_;
+  TypePtr type_;
+
+  std::shared_ptr<StructBuilder> builder_;
+  std::shared_ptr<StructArray> result_;
+};
+
+TEST_F(TestStructBuilder, TestAppendNull) {
+  ASSERT_OK(builder_->AppendNull());
+  ASSERT_OK(builder_->AppendNull());
+  ASSERT_EQ(2, builder_->value_builder().size());
+
+  Done();
+
+  ASSERT_EQ(2, result_->values().size());
+  ASSERT_TRUE(result_->IsNull(0));
+  ASSERT_TRUE(result_->IsNull(1));
+
+
+  auto list_char = static_cast<ListArray*>(result_->values(0).get());
+  auto chars = static_cast<Int8Array*>(list_char->values().get());
+  auto int32 = static_cast<Int32Array*>(result_->values(1).get());
+  ASSERT_EQ(0, list_char->length());
+  ASSERT_EQ(0, chars->length());
+  ASSERT_EQ(0, int32->length());
+
+  ASSERT_EQ(Type::LIST, list_char->type_enum());
+  ASSERT_EQ(Type::INT8, list_char->values()->type_enum());
+  ASSERT_EQ(Type::INT32, int32->type_enum());
+}
+
+TEST_F(TestStructBuilder, TestBasics) {
+  vector<int32_t> int_values = {1, 2, 3, 4};
+  vector<char> list_values = {'j', 'o', 'e', 'b', 'o', 'b', 'm', 'a', 'r', 'k'};
+  vector<int> list_lengths = {3, 0, 3, 4};
+  vector<int> list_offsets = {0, 3, 3, 6, 10};
+  vector<uint8_t> list_is_not_null = {1, 0, 1, 1};
+  vector<uint8_t> struct_is_not_null = {1, 1, 1, 1};
+
+  ListBuilder* list_vb = static_cast<ListBuilder*>(
+      builder_->value_builder().at(0).get());
+  Int8Builder* char_vb = static_cast<Int8Builder*>(
+      list_vb->value_builder().get());
+  Int32Builder* int_vb = static_cast<Int32Builder*>(
+      builder_->value_builder().at(1).get());
+  ASSERT_EQ(2, builder_->value_builder().size());
+
+  //EXPECT_OK(builder_->Reserve(list_lengths.size()));
+  EXPECT_OK(char_vb->Reserve(list_values.size()));
+  EXPECT_OK(int_vb->Reserve(int_values.size()));
+
+  int pos = 0;
+  for (size_t i = 0; i < list_lengths.size(); ++i) {
+    ASSERT_OK(list_vb->Append(list_is_not_null[i] > 0));
+    int_vb->Append(int_values[i]);
+    for (int j = 0; j < list_lengths[i]; ++j) {
+     char_vb->Append(list_values[pos++]);
+    }
+  }
+
+  for (size_t i = 0; i < struct_is_not_null.size(); ++i) {
+    ASSERT_OK(builder_->Append(struct_is_not_null[i] > 0));
+  }
+
+  Done();
+
+  ASSERT_EQ(2, result_->length());
+
+  auto list_char = static_cast<ListArray*>(result_->values(0).get());
+  auto chars = static_cast<Int8Array*>(list_char->values().get());
+  auto int32 = static_cast<Int32Array*>(result_->values(1).get());
+
+  ASSERT_EQ(0, result_->null_count());
+  ASSERT_EQ(1, list_char->null_count());
+  ASSERT_EQ(0, int32->null_count());
+
+
+  for (int i = 0; i < result_->length(); ++i) {
+    ASSERT_EQ(!static_cast<bool>(struct_is_not_null[i]), result_->IsNull(i));
+    ASSERT_EQ(!static_cast<bool>(list_is_not_null[i]), list_char->IsNull(i));
+  }
+
+  // List<char>
+  ASSERT_EQ(4, list_char->length());
+  ASSERT_EQ(10, list_char->values()->length());
+  for (size_t i = 0; i < list_offsets.size(); ++i) {
+    ASSERT_EQ(list_offsets[i], list_char->offsets()[i]);
+  }
+  for (size_t i = 0; i < list_values.size(); ++i) {
+    ASSERT_EQ(list_values[i], chars->Value(i));
+  }
+
+  // Int32
+  ASSERT_EQ(4, int32->length());
+  for (size_t i = 0; i < int_values.size(); ++i) {
+    ASSERT_EQ(int_values[i], int32->Value(i));
+  }
+}
+} // namespace arrow

cpp/src/arrow/types/struct.h

@@ -23,7 +23,135 @@
 #include <vector>
 
 #include "arrow/type.h"
+#include "arrow/types/primitive.h"
+#include "arrow/types/list.h"
 
-namespace arrow {}  // namespace arrow
+namespace arrow {
 
-#endif  // ARROW_TYPES_STRUCT_H
+class StructArray : public Array {
+ public:
+  StructArray(const TypePtr& type, int32_t length,
+      std::vector<ArrayPtr>& values,
+      int32_t null_count = 0,
+      std::shared_ptr<Buffer> null_bitmap = nullptr) :
+      Array(type, length, null_count, null_bitmap) {
+    type_ = type;
+    values_ = values;
+  }
+
+  virtual ~StructArray() {}
+
+  // Return a shared pointer in case the requestor desires to share ownership
+  // with this array.
+  const std::shared_ptr<Array>& values(int32_t pos) const {return values_.at(pos);}
+  const std::vector<ArrayPtr>& values() const {return values_;}
+
+  const std::shared_ptr<DataType>& value_type(int32_t pos) const {
+    return values_.at(pos)->type();
+  }
+
+  // return the count of Array in values_
+  int32_t length() const { return values_.size(); }
+
+  bool EqualsExact(const ListArray& other) const {
+    return true;
+  }
+  bool Equals(const std::shared_ptr<Array>& arr) const override {
+    return true;
+  }
+
+ protected:
+  // Contains kinds of Arrays.
+  std::vector<ArrayPtr> values_;
+};
+
+// ............................................................................
+// StrcutArray builder
+class StructBuilder : public ArrayBuilder {
+ public:
+  StructBuilder(MemoryPool* pool, const std::shared_ptr<DataType>& type,
+    const std::vector<FieldPtr>& fields,
+    std::vector<std::shared_ptr<ArrayBuilder>>& value_builder)
+    : ArrayBuilder(pool, type) {
+    fields_ = fields;
+    value_builder_ = value_builder;
+  }
+
+  Status Init(int32_t elements) {
+    return ArrayBuilder::Init(elements);
+  }
+
+  Status Resize(int32_t capacity) {
+    // Need space for the end offset
+    if (capacity < MIN_BUILDER_CAPACITY) {
+      capacity = MIN_BUILDER_CAPACITY;
+    }
+
+    if (capacity_ == 0) {
+      RETURN_NOT_OK(ArrayBuilder::Init(capacity));
+    } else {
+      RETURN_NOT_OK(ArrayBuilder::Resize(capacity));
+    }
+    capacity_ = capacity;
+
+    return Status::OK();
+  }
+
+  // Vector append
+  //
+  // If passed, valid_bytes is of equal length to values, and any zero byte
+  // will be considered as a null for that slot
+  Status Append(const uint8_t* null_bitmap, int32_t length) {
+    RETURN_NOT_OK(Reserve(length));
+    UnsafeAppendToBitmap(null_bitmap, length);
+    return Status::OK();
+  }
+
+  template <typename Container>
+  std::shared_ptr<Array> Transfer() {
+    DCHECK(value_builder_.size());
+
+    std::vector<std::shared_ptr<Array>> items;
+    for (auto it = value_builder_.cbegin(); it != value_builder_.cend(); it++) {
+      items.push_back(it->get()->Finish());
+    }
+    // Here, for ListArray, offsets_ is needed, but StructArray dont need it.
+    auto result = std::make_shared<StructArray>(type_, length_, items,
+        null_count_, null_bitmap_);
+
+    null_bitmap_ = nullptr;
+    capacity_ = length_ = null_count_ = 0;
+
+    return result;
+  }
+
+  std::shared_ptr<Array> Finish() override {
+    return Transfer<StructArray>();
+  }
+
+  // Start a new variable-length list slot
+  //
+  // This function should be called before beginning to append elements to the
+  // value builder
+  Status Append(bool is_valid = true) {
+    RETURN_NOT_OK(Reserve(1));
+    UnsafeAppendToBitmap(is_valid);
+    return Status::OK();
+  }
+
+  Status AppendNull() {
+    return Append(false);
+  }
+
+  const std::vector<std::shared_ptr<ArrayBuilder>>& value_builder() const {
+    return value_builder_;
+  }
+
+ protected:
+  std::vector<std::shared_ptr<ArrayBuilder>> value_builder_;
+  std::vector<FieldPtr> fields_;
+};
+
+} // namespace arrow
+
+#endif // ARROW_TYPES_STRUCT_H