feat(cpp): Add UnsafeAppend methods to StructBuilder

Author: Arkur745

cpp/src/arrow/array/array_struct_test.cc

@@ -748,4 +748,92 @@ TEST(TestFieldRef, GetChildren) {
   AssertArraysEqual(*a, *expected_a);
 }
 
-}  // namespace arrow
+TEST(TestStructBuilder, UnsafeAppend) {
+  
+  auto int_type = int32();
+  auto str_type = utf8();
+  auto struct_type = struct_({field("a", int_type), field("b", str_type)});
+  auto pool = default_memory_pool();
+  std::shared_ptr<Array> final_array;
+  auto int_builder = std::make_shared<Int32Builder>(pool);
+  auto str_builder = std::make_shared<StringBuilder>(pool);
+  StructBuilder builder(struct_type, pool, {int_builder, str_builder});
+  
+  builder.UnsafeAppend();
+  ASSERT_OK(int_builder->Append(1));
+  ASSERT_OK(str_builder->Append("hello"));
+
+  builder.UnsafeAppend();
+  ASSERT_OK(int_builder->Append(2));
+  ASSERT_OK(str_builder->Append("arrow"));
+  
+  ASSERT_OK(builder.Finish(&final_array));
+  ASSERT_EQ(2, final_array->length());
+  ASSERT_EQ(0, final_array->null_count());
+  auto expected_json = R"([{"a": 1, "b": "hello"}, {"a": 2, "b": "arrow"}])";
+  auto expected_array = ArrayFromJSON(struct_type, expected_json);
+  ASSERT_TRUE(final_array->Equals(expected_array));
+}
+
+TEST(TestStructBuilder, UnsafeAppendNull) {
+  
+  auto int_type = int32();
+  auto str_type = utf8();
+  auto struct_type = struct_({field("a", int_type), field("b", str_type)});
+  auto pool = default_memory_pool();
+  std::shared_ptr<Array> final_array;
+  auto int_builder = std::make_shared<Int32Builder>(pool);
+  auto str_builder = std::make_shared<StringBuilder>(pool);
+  StructBuilder builder(struct_type, pool, {int_builder, str_builder});
+
+  builder.UnsafeAppend();
+  ASSERT_OK(int_builder->Append(1));
+  ASSERT_OK(str_builder->Append("hello"));
+
+  ASSERT_OK(builder.UnsafeAppendNull());
+
+  builder.UnsafeAppend();
+  ASSERT_OK(int_builder->Append(2));
+  ASSERT_OK(str_builder->Append("arrow"));
+
+  ASSERT_OK(builder.Finish(&final_array));
+  ASSERT_EQ(3, final_array->length());
+  ASSERT_EQ(1, final_array->null_count());
+  ASSERT_TRUE(final_array->IsNull(1));
+  auto expected_json = R"([{"a": 1, "b": "hello"}, null, {"a": 2, "b": "arrow"}])";
+  auto expected_array = ArrayFromJSON(struct_type, expected_json);
+  ASSERT_TRUE(final_array->Equals(expected_array));
+}
+
+TEST(TestStructBuilder, UnsafeAppendNulls) {
+  
+  auto int_type = int32();
+  auto str_type = utf8();
+  auto struct_type = struct_({field("a", int_type), field("b", str_type)});
+  auto pool = default_memory_pool();
+  std::shared_ptr<Array> final_array;
+  auto int_builder = std::make_shared<Int32Builder>(pool);
+  auto str_builder = std::make_shared<StringBuilder>(pool);
+  StructBuilder builder(struct_type, pool, {int_builder, str_builder});
+
+  builder.UnsafeAppend();
+  ASSERT_OK(int_builder->Append(1));
+  ASSERT_OK(str_builder->Append("hello"));
+
+  ASSERT_OK(builder.UnsafeAppendNulls(2));
+
+  builder.UnsafeAppend();
+  ASSERT_OK(int_builder->Append(2));
+  ASSERT_OK(str_builder->Append("arrow"));
+
+  ASSERT_OK(builder.Finish(&final_array));
+  ASSERT_EQ(4, final_array->length());
+  ASSERT_EQ(2, final_array->null_count());
+  ASSERT_TRUE(final_array->IsNull(1));
+  ASSERT_TRUE(final_array->IsNull(2));
+  auto expected_json = R"([{"a": 1, "b": "hello"}, null, null, {"a": 2, "b": "arrow"}])";
+  auto expected_array = ArrayFromJSON(struct_type, expected_json);
+  ASSERT_TRUE(final_array->Equals(expected_array));
+}
+
+} // namespace arrow

cpp/src/arrow/array/builder_nested.h

@@ -771,6 +771,13 @@ class ARROW_EXPORT StructBuilder : public ArrayBuilder {
     return Status::OK();
   }
 
+  /// This method is "unsafe" because it does not update the null bitmap
+  /// It doesn't check the capacity. Caller is responsible for calling Reserve()
+  /// beforehand.
+  void UnsafeAppend(){
+    ++length_;
+  }
+
   /// \brief Append a null value. Automatically appends an empty value to each child
   /// builder.
   Status AppendNull() final {
@@ -780,6 +787,19 @@ class ARROW_EXPORT StructBuilder : public ArrayBuilder {
     return Append(false);
   }
 
+  /// This method is "unsafe" because it does not check for capacity.
+  /// The caller is responsible for calling Reserve() beforehand to ensure
+  /// there is enough space. This method will append nulls to all child
+  /// builders to maintain a consistent state.
+  Status UnsafeAppendNull(){
+    null_bitmap_builder_.UnsafeAppend(false);
+    for(const auto& child: children_){
+      ARROW_RETURN_NOT_OK(child->AppendNull());
+    }
+    length_++;
+    return Status::OK();
+  }
+
   /// \brief Append multiple null values. Automatically appends empty values to each
   /// child builder.
   Status AppendNulls(int64_t length) final {
@@ -791,6 +811,20 @@ class ARROW_EXPORT StructBuilder : public ArrayBuilder {
     return Status::OK();
   }
 
+  /// This method is "unsafe" because it does not check for capacity.
+  /// The caller is responsible for calling Reserve() beforehand to ensure
+  /// there is enough space. This method will append nulls to all child
+  /// builders to maintain a consistent state.
+  /// param length The number of null slots to append.
+  Status UnsafeAppendNulls(int64_t length){
+    null_bitmap_builder_.UnsafeAppend(length, false);
+    for(const auto& child:children_){
+      ARROW_RETURN_NOT_OK(child->AppendNulls(length));
+    }
+    ++length_;
+    return Status::OK();
+  }
+
   Status AppendEmptyValue() final {
     for (const auto& field : children_) {
       ARROW_RETURN_NOT_OK(field->AppendEmptyValue());