replace use of builder with SetBit

apache · Nov 27, 2024 · 943a047 · 943a047
1 parent 7ee000a
commit 943a047
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 34 deletions.
diff --git a/cpp/src/arrow/array/array_test.cc b/cpp/src/arrow/array/array_test.cc
@@ -874,24 +874,9 @@ TEST_F(TestArray, TestMakeMaskArray) {
   ASSERT_OK(array->ValidateFull());
   ASSERT_EQ(array->length(), 10);
 
-  ASSERT_OK_AND_ASSIGN(auto true_scalar, MakeScalar(boolean(), true));
-  ASSERT_OK_AND_ASSIGN(auto false_scalar, MakeScalar(boolean(), false));
-
   // Only values at index 5 and 8 should be true.
-  ASSERT_OK_AND_ASSIGN(auto scalar, array->GetScalar(1));
-  AssertScalarsEqual(*scalar, *false_scalar);
-
-  ASSERT_OK_AND_ASSIGN(scalar, array->GetScalar(5));
-  AssertScalarsEqual(*scalar, *true_scalar);
-
-  ASSERT_OK_AND_ASSIGN(scalar, array->GetScalar(6));
-  AssertScalarsEqual(*scalar, *false_scalar);
-
-  ASSERT_OK_AND_ASSIGN(scalar, array->GetScalar(8));
-  AssertScalarsEqual(*scalar, *true_scalar);
-
-  ASSERT_OK_AND_ASSIGN(scalar, array->GetScalar(9));
-  AssertScalarsEqual(*scalar, *false_scalar);
+  auto expected = ArrayFromJSON(boolean(), "[false, false, false, false, false, true, false, false, true, false]");
+  AssertArraysEqual(*array, *expected);
 }
 
 TEST_F(TestArray, ExtensionSpanRoundTrip) {

diff --git a/cpp/src/arrow/array/util.cc b/cpp/src/arrow/array/util.cc
@@ -918,25 +918,15 @@ Result<std::shared_ptr<Array>> MakeEmptyArray(std::shared_ptr<DataType> type,
 
 Result<std::shared_ptr<Array>> MakeMaskArray(const std::vector<int64_t>& indices,
                                              int64_t length, MemoryPool* pool) {
-  BooleanBuilder builder(pool);
-  RETURN_NOT_OK(builder.Resize(length));
-
-  auto indices_end = indices.end();
-  auto i = indices.begin();
-  for (int64_t builder_i = 0; builder_i < length; builder_i++) {
-    if (i == indices_end) {
-      RETURN_NOT_OK(
-          builder.AppendValues(static_cast<int64_t>(length - builder.length()), false));
-      break;
-    } else if (builder_i == *i) {
-      builder.UnsafeAppend(true);
-      i++;
-    } else {
-      builder.UnsafeAppend(false);
+  ARROW_ASSIGN_OR_RAISE(auto buffer, AllocateBitmap(length, pool));
+  bit_util::SetBitsTo(buffer->mutable_data(), 0, length, false);
+  for (int64_t index : indices) {
+    if (index < 0 || index >= length) {
+      return Status::IndexError("Index out of bounds: ", index);
     }
+    bit_util::SetBit(buffer->mutable_data(), index);
   }
-
-  return builder.Finish();
+  return std::make_shared<BooleanArray>(length, buffer);
 }
 
 namespace internal {