Fix overflow bug with child_fields.size() = 128 (kMaxTypeCode + 1)

Author: dahbka-lis

cpp/src/arrow/type.cc

@@ -3270,15 +3270,15 @@ std::shared_ptr<DataType> run_end_encoded(std::shared_ptr<DataType> run_end_type
 std::shared_ptr<DataType> sparse_union(FieldVector child_fields,
                                        std::vector<int8_t> type_codes) {
   if (type_codes.empty()) {
-    type_codes = internal::Iota(static_cast<int8_t>(child_fields.size()));
+    type_codes = internal::Iota<int8_t>(0, child_fields.size());
   }
   return std::make_shared<SparseUnionType>(std::move(child_fields),
                                            std::move(type_codes));
 }
 std::shared_ptr<DataType> dense_union(FieldVector child_fields,
                                       std::vector<int8_t> type_codes) {
   if (type_codes.empty()) {
-    type_codes = internal::Iota(static_cast<int8_t>(child_fields.size()));
+    type_codes = internal::Iota<int8_t>(0, child_fields.size());
   }
   return std::make_shared<DenseUnionType>(std::move(child_fields), std::move(type_codes));
 }
@@ -3310,7 +3310,7 @@ std::shared_ptr<DataType> sparse_union(const ArrayVector& children,
                                        std::vector<std::string> field_names,
                                        std::vector<int8_t> type_codes) {
   if (type_codes.empty()) {
-    type_codes = internal::Iota(static_cast<int8_t>(children.size()));
+    type_codes = internal::Iota<int8_t>(0, children.size());
   }
   auto fields = FieldsFromArraysAndNames(std::move(field_names), children);
   return sparse_union(std::move(fields), std::move(type_codes));
@@ -3320,7 +3320,7 @@ std::shared_ptr<DataType> dense_union(const ArrayVector& children,
                                       std::vector<std::string> field_names,
                                       std::vector<int8_t> type_codes) {
   if (type_codes.empty()) {
-    type_codes = internal::Iota(static_cast<int8_t>(children.size()));
+    type_codes = internal::Iota<int8_t>(0, children.size());
   }
   auto fields = FieldsFromArraysAndNames(std::move(field_names), children);
   return dense_union(std::move(fields), std::move(type_codes));

cpp/src/arrow/util/range.h

@@ -44,6 +44,14 @@ std::vector<T> Iota(T length) {
   return Iota(static_cast<T>(0), length);
 }
 
+/// Create a vector containing the values from start with length elements
+template <typename T>
+std::vector<T> Iota(T start, size_t length) {
+  std::vector<T> result(length);
+  std::iota(result.begin(), result.end(), start);
+  return result;
+}
+
 /// Create a range from a callable which takes a single index parameter
 /// and returns the value of iterator on each call and a length.
 /// Only iterators obtained from the same range should be compared, the