Fix the issue that bitmap ops overriding partial leading byte

Author: zanmato1984

cpp/src/arrow/compute/kernels/scalar_if_else_test.cc

@@ -1095,6 +1095,16 @@ TYPED_TEST(TestIfElseDict, DifferentDictionaries) {
   CheckDictionary("if_else", {MakeNullScalar(boolean()), values1, values2});
 }
 
+// GH-47825: BitmapOp overrides partial leading byte for unaligned chunked execution.
+TEST(TestIfElse, ChunkedUnalignedBitmapOp) {
+  auto cond = ArrayFromJSON(boolean(), "[true, true]");
+  auto if_true = ChunkedArrayFromJSON(boolean(), {"[true]", "[true]"});
+  auto if_false = ArrayFromJSON(boolean(), "[true, true]");
+  ASSERT_OK_AND_ASSIGN(auto result, CallFunction("if_else", {cond, if_true, if_false}));
+  ASSERT_OK(result.chunked_array()->ValidateFull());
+  AssertDatumsEqual(ChunkedArrayFromJSON(boolean(), {"[true, true]"}), result);
+}
+
 Datum MakeStruct(const std::vector<Datum>& conds) {
   if (conds.size() == 0) {
     // The tests below want a struct scalar when no condition values passed,

cpp/src/arrow/util/CMakeLists.txt

@@ -49,6 +49,7 @@ add_arrow_test(utility-test
                SOURCES
                align_util_test.cc
                atfork_test.cc
+               bitmap_test.cc
                byte_size_test.cc
                byte_stream_split_test.cc
                cache_test.cc

cpp/src/arrow/util/bitmap_ops.cc

@@ -370,13 +370,25 @@ void UnalignedBitmapOp(const uint8_t* left, int64_t left_offset, const uint8_t*
   }
 }
 
+// XXX: The bits before left/right/out_offset, if unaligned, are untouched. But not for
+// the bits after length. Caller should ensure proper alignment for the tail bits if
+// necessary, or correct the tail bits by subsequent calls.
 template <template <typename> class BitOp>
 void BitmapOp(const uint8_t* left, int64_t left_offset, const uint8_t* right,
               int64_t right_offset, int64_t length, int64_t out_offset, uint8_t* dest) {
-  if ((out_offset % 8 == left_offset % 8) && (out_offset % 8 == right_offset % 8)) {
-    // Fast case: can use bytewise AND
-    AlignedBitmapOp<BitOp>(left, left_offset, right, right_offset, dest, out_offset,
-                           length);
+  if (out_offset % 8 == left_offset % 8 && out_offset % 8 == right_offset % 8) {
+    // Fast case: can use byte-wise BitOp after handling leading unaligned bits.
+    int64_t leading_unaligned_bits = (8 - left_offset % 8) % 8;
+    if (leading_unaligned_bits > 0) {
+      UnalignedBitmapOp<BitOp>(left, left_offset, right, right_offset, dest, out_offset,
+                               leading_unaligned_bits);
+    }
+    if (length > leading_unaligned_bits) {
+      AlignedBitmapOp<BitOp>(left, left_offset + leading_unaligned_bits, right,
+                             right_offset + leading_unaligned_bits, dest,
+                             out_offset + leading_unaligned_bits,
+                             length - leading_unaligned_bits);
+    }
   } else {
     // Unaligned
     UnalignedBitmapOp<BitOp>(left, left_offset, right, right_offset, dest, out_offset,

cpp/src/arrow/util/bitmap_test.cc

@@ -0,0 +1,64 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gtest/gtest.h>
+
+#include "arrow/testing/gtest_util.h"
+#include "arrow/util/bitmap_ops.h"
+
+namespace arrow::internal {
+
+TEST(BitmapOpsTest, PartialLeadingByteSafety) {
+  const int64_t num_bytes = 16;
+  const std::vector<uint8_t> left(num_bytes, 0), right(num_bytes, 0xFF);
+  std::vector<uint8_t> out = right;
+  int64_t num_bits = num_bytes * 8;
+
+  auto left_data = left.data();
+  auto right_data = right.data();
+  auto out_data = out.data();
+
+  for (int64_t out_offset = num_bits - 1; out_offset >= 0; --out_offset) {
+    ARROW_SCOPED_TRACE("out offset = " + std::to_string(out_offset));
+    for (int64_t left_offset : {out_offset, out_offset + 1}) {
+      ARROW_SCOPED_TRACE("left offset = " + std::to_string(left_offset));
+      for (int64_t right_offset : {out_offset, out_offset + 2}) {
+        ARROW_SCOPED_TRACE("right offset = " + std::to_string(right_offset));
+        for (int64_t length = 1;
+             length <= num_bits - out_offset && length <= num_bits - left_offset &&
+             length <= num_bits - right_offset;
+             ++length) {
+          ARROW_SCOPED_TRACE("length = " + std::to_string(length));
+          BitmapAnd(left_data, out_offset, right_data, right_offset, length, out_offset,
+                    out_data);
+          // Bytes before out_offset.
+          for (int64_t i = 0; i < out_offset / 8; ++i) {
+            EXPECT_EQ(out_data[i], 0xFF);
+          }
+          // The byte holding the out_offset bit.
+          EXPECT_EQ(out_data[out_offset / 8], (1 << (out_offset % 8)) - 1);
+          // Bytes after the last bit.
+          for (int64_t i = (out_offset + length + 7) / 8; i < num_bytes; ++i) {
+            EXPECT_EQ(out_data[i], 0);
+          }
+        }
+      }
+    }
+  }
+}
+
+}  // namespace arrow::internal