address review comments

apache · Jul 11, 2019 · 0efe91d · 0efe91d
1 parent 17e6e27
commit 0efe91d
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/cpp/src/arrow/array/builder_union.cc b/cpp/src/arrow/array/builder_union.cc
@@ -42,8 +42,9 @@ Status BasicUnionBuilder::FinishInternal(std::shared_ptr<ArrayData>* out) {
   } else {
     type_codes = checked_cast<const UnionType&>(*type_).type_codes();
   }
+  DCHECK_EQ(type_codes.size(), children_.size());
 
-  std::vector<std::shared_ptr<ArrayData>> child_data(type_codes.size());
+  std::vector<std::shared_ptr<ArrayData>> child_data(children_.size());
   for (size_t i = 0; i < children_.size(); ++i) {
     RETURN_NOT_OK(children_[i]->FinishInternal(&child_data[i]));
   }
@@ -74,6 +75,7 @@ BasicUnionBuilder::BasicUnionBuilder(
 
   children_ = children;
   type_id_to_children_.resize(union_type->max_type_code() + 1, nullptr);
+  DCHECK_LT(type_id_to_children_.size(), std::numeric_limits<int8_t>::max());
 
   auto field_it = type->children().begin();
   auto children_it = children.begin();
@@ -93,6 +95,9 @@ int8_t BasicUnionBuilder::AppendChild(const std::shared_ptr<ArrayBuilder>& new_c
   field_names_.push_back(field_name);
   children_.push_back(new_child);
 
+  // Find type_id such that type_id_to_children_[type_id] == nullptr
+  // and use that for the new child. Start searching at dense_type_id_
+  // since type_id_to_children_ is densely packed up at least up to dense_type_id_
   for (; static_cast<size_t>(dense_type_id_) < type_id_to_children_.size();
        ++dense_type_id_) {
     if (type_id_to_children_[dense_type_id_] == nullptr) {
@@ -101,6 +106,9 @@ int8_t BasicUnionBuilder::AppendChild(const std::shared_ptr<ArrayBuilder>& new_c
     }
   }
 
+  DCHECK_LT(type_id_to_children_.size(), std::numeric_limits<int8_t>::max());
+
+  // type_id_to_children_ is already densely packed, so just append the new child
   type_id_to_children_.push_back(new_child);
   return dense_type_id_++;
 }

diff --git a/cpp/src/arrow/array/builder_union.h b/cpp/src/arrow/array/builder_union.h
@@ -63,6 +63,7 @@ class ARROW_EXPORT BasicUnionBuilder : public ArrayBuilder {
 
   UnionMode::type mode_;
   std::vector<std::shared_ptr<ArrayBuilder>> type_id_to_children_;
+  // for all type_id < dense_type_id_, type_id_to_children_[type_id] != nullptr
   int8_t dense_type_id_ = 0;
   TypedBufferBuilder<int8_t> types_builder_;
   std::vector<std::string> field_names_;