This is an automated email from the ASF dual-hosted git repository.
pitrou pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new aa78eb89733 GH-50280: [C++] Implement VisitTwoBitRuns and
VisitTwoSetBitRuns methods (#50281)
aa78eb89733 is described below
commit aa78eb89733134c4a8db8b5840960676f96df103
Author: Zehua Zou <[email protected]>
AuthorDate: Fri Jul 24 00:34:31 2026 +0800
GH-50280: [C++] Implement VisitTwoBitRuns and VisitTwoSetBitRuns methods
(#50281)
### Rationale for this change
We already have `VisitTwoBitBlocks` which can iterate over two bitmaps
simultaneously. But there is no mirror code about `BitRunReader` and
`SetBitRunReader`.
### What changes are included in this PR?
Implement `VisitTwoBitRuns` and `VisitTwoSetBitRuns` methods and refactor
`VisitTwoBitBlocks` in `GroupedPivotAccumulator` as examples.
### Are these changes tested?
Yes.
### Are there any user-facing changes?
No.
* GitHub Issue: #50280
Lead-authored-by: Zehua Zou <[email protected]>
Co-authored-by: Antoine Pitrou <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>
---
cpp/src/arrow/compute/kernels/codegen_internal.h | 29 ++--
cpp/src/arrow/compute/kernels/scalar_if_else.cc | 63 +++-----
.../kernels/vector_selection_take_internal.cc | 41 ++---
cpp/src/arrow/util/bit_block_counter_benchmark.cc | 168 +++++++++++++++++++--
cpp/src/arrow/util/bit_run_reader.h | 85 +++++++++++
cpp/src/arrow/util/bitmap_test.cc | 123 +++++++++++++++
6 files changed, 411 insertions(+), 98 deletions(-)
diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.h
b/cpp/src/arrow/compute/kernels/codegen_internal.h
index d33ffacb2f5..9d12d62d95f 100644
--- a/cpp/src/arrow/compute/kernels/codegen_internal.h
+++ b/cpp/src/arrow/compute/kernels/codegen_internal.h
@@ -38,6 +38,7 @@
#include "arrow/type.h"
#include "arrow/type_traits.h"
#include "arrow/util/bit_block_counter.h"
+#include "arrow/util/bit_run_reader.h"
#include "arrow/util/bit_util.h"
#include "arrow/util/bitmap_generate.h"
#include "arrow/util/bitmap_reader.h"
@@ -509,18 +510,24 @@ static void VisitTwoArrayValuesInline(const ArraySpan&
arr0, const ArraySpan& ar
ArrayIterator<Arg0Type> arr0_it(arr0);
ArrayIterator<Arg1Type> arr1_it(arr1);
- auto visit_valid = [&](int64_t i) {
- valid_func(GetViewType<Arg0Type>::LogicalValue(arr0_it()),
- GetViewType<Arg1Type>::LogicalValue(arr1_it()));
- };
- auto visit_null = [&]() {
- arr0_it();
- arr1_it();
- null_func();
+ auto visit_run = [&](int64_t position, int64_t run_length, bool is_set) {
+ if (is_set) {
+ for (int64_t i = 0; i < run_length; ++i) {
+ valid_func(GetViewType<Arg0Type>::LogicalValue(arr0_it()),
+ GetViewType<Arg1Type>::LogicalValue(arr1_it()));
+ }
+ } else {
+ for (int64_t i = 0; i < run_length; ++i) {
+ arr0_it();
+ arr1_it();
+ null_func();
+ }
+ }
};
- VisitTwoBitBlocksVoid(arr0.buffers[0].data, arr0.offset,
arr1.buffers[0].data,
- arr1.offset, arr0.length, std::move(visit_valid),
- std::move(visit_null));
+
+ ::arrow::internal::VisitTwoBitRunsVoid(arr0.buffers[0].data, arr0.offset,
+ arr1.buffers[0].data, arr1.offset,
arr0.length,
+ std::move(visit_run));
}
// ----------------------------------------------------------------------
diff --git a/cpp/src/arrow/compute/kernels/scalar_if_else.cc
b/cpp/src/arrow/compute/kernels/scalar_if_else.cc
index 0c5cfe1c903..5618f24e5ce 100644
--- a/cpp/src/arrow/compute/kernels/scalar_if_else.cc
+++ b/cpp/src/arrow/compute/kernels/scalar_if_else.cc
@@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.
+#include <cstdint>
#include <cstring>
#include <limits>
#include "arrow/array/builder_nested.h"
@@ -1607,26 +1608,14 @@ Status ExecArrayCaseWhen(KernelContext* ctx, const
ExecSpan& batch, ExecResult*
if (cond_array.GetNullCount() == 0) {
// If no valid buffer, visit mask & cond bitmap simultaneously
- BinaryBitBlockCounter counter(mask, /*start_offset=*/0, cond_values,
cond_offset,
- batch.length);
- while (offset < batch.length) {
- const auto block = counter.NextAndWord();
- if (block.AllSet()) {
- CopyValues<Type>(value, offset, block.length, out_valid, out_values,
- out_offset + offset);
- bit_util::SetBitsTo(mask, offset, block.length, false);
- } else if (block.popcount) {
- for (int64_t j = 0; j < block.length; ++j) {
- if (bit_util::GetBit(mask, offset + j) &&
- bit_util::GetBit(cond_values, cond_offset + offset + j)) {
- CopyValues<Type>(value, offset + j, /*length=*/1, out_valid,
out_values,
- out_offset + offset + j);
- bit_util::SetBitTo(mask, offset + j, false);
- }
- }
- }
- offset += block.length;
- }
+ auto visit_run = [&](int64_t position, int64_t run_length) {
+ CopyValues<Type>(value, position, run_length, out_valid, out_values,
+ out_offset + position);
+ bit_util::SetBitsTo(mask, position, run_length, false);
+ return Status::OK();
+ };
+ RETURN_NOT_OK(::arrow::internal::VisitTwoSetBitRuns(
+ mask, /*start_offset=*/0, cond_values, cond_offset, batch.length,
visit_run));
} else {
// Visit mask & cond bitmap & cond validity
const uint8_t* cond_valid = cond_array.buffers[0].data;
@@ -1657,32 +1646,20 @@ Status ExecArrayCaseWhen(KernelContext* ctx, const
ExecSpan& batch, ExecResult*
}
if (!have_else_arg) {
// Need to initialize any remaining null slots (uninitialized memory)
- BitBlockCounter counter(mask, /*offset=*/0, batch.length);
- int64_t offset = 0;
auto bit_width = checked_cast<const
FixedWidthType&>(*out->type()).bit_width();
auto byte_width = bit_util::BytesForBits(bit_width);
- while (offset < batch.length) {
- const auto block = counter.NextWord();
- if (block.AllSet()) {
- if (bit_width == 1) {
- bit_util::SetBitsTo(out_values, out_offset + offset, block.length,
false);
- } else {
- std::memset(out_values + (out_offset + offset) * byte_width, 0x00,
- byte_width * block.length);
- }
- } else if (!block.NoneSet()) {
- for (int64_t j = 0; j < block.length; ++j) {
- if (bit_util::GetBit(out_valid, out_offset + offset + j)) continue;
- if (bit_width == 1) {
- bit_util::ClearBit(out_values, out_offset + offset + j);
- } else {
- std::memset(out_values + (out_offset + offset + j) * byte_width,
0x00,
- byte_width);
- }
- }
+
+ auto visit_run = [&](int64_t position, int64_t run_length) {
+ if (bit_width == 1) {
+ bit_util::SetBitsTo(out_values, out_offset + position, run_length,
false);
+ } else {
+ std::memset(out_values + (out_offset + position) * byte_width, 0x00,
+ byte_width * run_length);
}
- offset += block.length;
- }
+ bit_util::SetBitsTo(mask, position, run_length, false);
+ };
+
+ ::arrow::internal::VisitSetBitRunsVoid(mask, /*offset=*/0, batch.length,
visit_run);
}
return Status::OK();
}
diff --git a/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc
b/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc
index fedafeb5bea..1237aae35c6 100644
--- a/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc
+++ b/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc
@@ -34,6 +34,7 @@
#include "arrow/compute/registry.h"
#include "arrow/memory_pool.h"
#include "arrow/record_batch.h"
+#include "arrow/status.h"
#include "arrow/table.h"
#include "arrow/type.h"
#include "arrow/type_traits.h"
@@ -135,37 +136,17 @@ Result<std::shared_ptr<ArrayData>>
GetTakeIndicesFromBitmapImpl(
// data bitmap together.
DCHECK_EQ(null_selection, FilterOptions::DROP);
- // The position relative to the start of the filter
- T position = 0;
- // The current position taking the filter offset into account
- int64_t position_with_offset = filter.offset;
-
- BinaryBitBlockCounter filter_counter(filter_data, filter.offset,
filter_is_valid,
- filter.offset, filter.length);
- while (position < filter.length) {
- BitBlockCount and_block = filter_counter.NextAndWord();
- RETURN_NOT_OK(builder.Reserve(and_block.popcount));
- if (and_block.AllSet()) {
- // All the values are selected and non-null
- for (int64_t i = 0; i < and_block.length; ++i) {
- builder.UnsafeAppend(position++);
- }
- position_with_offset += and_block.length;
- } else if (!and_block.NoneSet()) {
- // Some of the values are false or null
- for (int64_t i = 0; i < and_block.length; ++i) {
- if (bit_util::GetBit(filter_is_valid, position_with_offset) &&
- bit_util::GetBit(filter_data, position_with_offset)) {
- builder.UnsafeAppend(position);
- }
- ++position;
- ++position_with_offset;
- }
- } else {
- position += and_block.length;
- position_with_offset += and_block.length;
+ auto visit_run = [&](int64_t position, int64_t run_length) {
+ RETURN_NOT_OK(builder.Reserve(run_length));
+ for (int64_t i = 0; i < run_length; ++i) {
+ builder.UnsafeAppend(static_cast<T>(position + i));
}
- }
+ return Status::OK();
+ };
+
+ RETURN_NOT_OK(::arrow::internal::VisitTwoSetBitRuns(filter_data,
filter.offset,
+ filter_is_valid,
filter.offset,
+ filter.length,
visit_run));
} else {
// The filter has no nulls, so we need only look for true values
RETURN_NOT_OK(::arrow::internal::VisitSetBitRuns(
diff --git a/cpp/src/arrow/util/bit_block_counter_benchmark.cc
b/cpp/src/arrow/util/bit_block_counter_benchmark.cc
index c08fcfdb604..7be82b24b86 100644
--- a/cpp/src/arrow/util/bit_block_counter_benchmark.cc
+++ b/cpp/src/arrow/util/bit_block_counter_benchmark.cc
@@ -23,22 +23,24 @@
#include "arrow/array/array_base.h"
#include "arrow/array/array_primitive.h"
+#include "arrow/testing/gtest_util.h"
#include "arrow/testing/random.h"
#include "arrow/util/bit_block_counter.h"
+#include "arrow/util/bit_run_reader.h"
#include "arrow/util/bit_util.h"
#include "arrow/util/bitmap_reader.h"
namespace arrow {
namespace internal {
-struct UnaryBitBlockBenchmark {
+struct UnaryBitmapTraversalBenchmark {
benchmark::State& state;
int64_t offset;
int64_t bitmap_length;
std::shared_ptr<Array> arr;
int64_t expected;
- explicit UnaryBitBlockBenchmark(benchmark::State& state, int64_t offset = 0)
+ explicit UnaryBitmapTraversalBenchmark(benchmark::State& state, int64_t
offset = 0)
: state(state), offset(offset), bitmap_length(1 << 20) {
random::RandomArrayGenerator rng(/*seed=*/0);
// State parameter is the average number of total values for each null
@@ -111,9 +113,57 @@ struct UnaryBitBlockBenchmark {
}
state.SetItemsProcessed(state.iterations() * bitmap_length);
}
+
+ void BenchVisitBitRuns() {
+ const auto& int8_arr = static_cast<const Int8Array&>(*arr);
+ const uint8_t* bitmap = arr->null_bitmap_data();
+ for (auto _ : state) {
+ int64_t result = 0;
+ ABORT_NOT_OK(VisitBitRuns(bitmap, this->offset, bitmap_length -
this->offset,
+ [&](int64_t position, int64_t length, bool
set) {
+ if (set) {
+ int64_t run_sum = 0;
+ const int64_t end = position + length;
+ for (int64_t i = position; i < end; ++i) {
+ run_sum += int8_arr.Value(this->offset +
i);
+ }
+ result += run_sum;
+ }
+ return Status::OK();
+ }));
+ // Sanity check
+ if (result != expected) {
+ std::abort();
+ }
+ }
+ state.SetItemsProcessed(state.iterations() * bitmap_length);
+ }
+
+ void BenchVisitSetBitRuns() {
+ const auto& int8_arr = static_cast<const Int8Array&>(*arr);
+ const uint8_t* bitmap = arr->null_bitmap_data();
+ for (auto _ : state) {
+ int64_t result = 0;
+ ABORT_NOT_OK(VisitSetBitRuns(bitmap, this->offset, bitmap_length -
this->offset,
+ [&](int64_t position, int64_t length) {
+ int64_t run_sum = 0;
+ const int64_t end = position + length;
+ for (int64_t i = position; i < end; ++i) {
+ run_sum += int8_arr.Value(this->offset
+ i);
+ }
+ result += run_sum;
+ return Status::OK();
+ }));
+ // Sanity check
+ if (result != expected) {
+ std::abort();
+ }
+ }
+ state.SetItemsProcessed(state.iterations() * bitmap_length);
+ }
};
-struct BinaryBitBlockBenchmark {
+struct BinaryBitmapTraversalBenchmark {
benchmark::State& state;
int64_t offset;
int64_t bitmap_length;
@@ -123,7 +173,7 @@ struct BinaryBitBlockBenchmark {
const Int8Array* left_int8;
const Int8Array* right_int8;
- explicit BinaryBitBlockBenchmark(benchmark::State& state, int64_t offset = 0)
+ explicit BinaryBitmapTraversalBenchmark(benchmark::State& state, int64_t
offset = 0)
: state(state), offset(offset), bitmap_length(1 << 20) {
random::RandomArrayGenerator rng(/*seed=*/0);
@@ -202,52 +252,134 @@ struct BinaryBitBlockBenchmark {
}
state.SetItemsProcessed(state.iterations() * bitmap_length);
}
+
+ void BenchVisitTwoBitRuns() {
+ const uint8_t* left_bitmap = left->null_bitmap_data();
+ const uint8_t* right_bitmap = right->null_bitmap_data();
+ for (auto _ : state) {
+ int64_t result = 0;
+ VisitTwoBitRunsVoid(left_bitmap, this->offset, right_bitmap,
this->offset,
+ bitmap_length - this->offset,
+ [&](int64_t position, int64_t length, bool set) {
+ if (set) {
+ int64_t run_sum = 0;
+ const int64_t end = position + length;
+ for (int64_t i = position; i < end; ++i) {
+ run_sum += left_int8->Value(this->offset + i) +
+ right_int8->Value(this->offset + i);
+ }
+ result += run_sum;
+ }
+ });
+ // Sanity check
+ if (result != expected) {
+ std::abort();
+ }
+ }
+ state.SetItemsProcessed(state.iterations() * bitmap_length);
+ }
+
+ void BenchVisitTwoSetBitRuns() {
+ const uint8_t* left_bitmap = left->null_bitmap_data();
+ const uint8_t* right_bitmap = right->null_bitmap_data();
+ for (auto _ : state) {
+ int64_t result = 0;
+ VisitTwoSetBitRunsVoid(left_bitmap, this->offset, right_bitmap,
this->offset,
+ bitmap_length - this->offset,
+ [&](int64_t position, int64_t length) {
+ int64_t run_sum = 0;
+ const int64_t end = position + length;
+ for (int64_t i = position; i < end; ++i) {
+ run_sum += left_int8->Value(this->offset + i)
+
+ right_int8->Value(this->offset +
i);
+ }
+ result += run_sum;
+ });
+ // Sanity check
+ if (result != expected) {
+ std::abort();
+ }
+ }
+ state.SetItemsProcessed(state.iterations() * bitmap_length);
+ }
};
static void BitBlockCounterSum(benchmark::State& state) {
- UnaryBitBlockBenchmark(state, /*offset=*/0)
+ UnaryBitmapTraversalBenchmark(state, /*offset=*/0)
.BenchBitBlockCounter([](BitBlockCounter* counter) { return
counter->NextWord(); });
}
static void BitBlockCounterSumWithOffset(benchmark::State& state) {
- UnaryBitBlockBenchmark(state, /*offset=*/4)
+ UnaryBitmapTraversalBenchmark(state, /*offset=*/4)
.BenchBitBlockCounter([](BitBlockCounter* counter) { return
counter->NextWord(); });
}
static void BitBlockCounterFourWordsSum(benchmark::State& state) {
- UnaryBitBlockBenchmark(state, /*offset=*/0)
+ UnaryBitmapTraversalBenchmark(state, /*offset=*/0)
.BenchBitBlockCounter(
[](BitBlockCounter* counter) { return counter->NextFourWords(); });
}
static void BitBlockCounterFourWordsSumWithOffset(benchmark::State& state) {
- UnaryBitBlockBenchmark(state, /*offset=*/4)
+ UnaryBitmapTraversalBenchmark(state, /*offset=*/4)
.BenchBitBlockCounter(
[](BitBlockCounter* counter) { return counter->NextFourWords(); });
}
static void BitmapReaderSum(benchmark::State& state) {
- UnaryBitBlockBenchmark(state, /*offset=*/0).BenchBitmapReader();
+ UnaryBitmapTraversalBenchmark(state, /*offset=*/0).BenchBitmapReader();
}
static void BitmapReaderSumWithOffset(benchmark::State& state) {
- UnaryBitBlockBenchmark(state, /*offset=*/4).BenchBitmapReader();
+ UnaryBitmapTraversalBenchmark(state, /*offset=*/4).BenchBitmapReader();
+}
+
+static void VisitBitRunsSum(benchmark::State& state) {
+ UnaryBitmapTraversalBenchmark(state, /*offset=*/0).BenchVisitBitRuns();
+}
+
+static void VisitBitRunsSumWithOffset(benchmark::State& state) {
+ UnaryBitmapTraversalBenchmark(state, /*offset=*/4).BenchVisitBitRuns();
+}
+
+static void VisitSetBitRunsSum(benchmark::State& state) {
+ UnaryBitmapTraversalBenchmark(state, /*offset=*/0).BenchVisitSetBitRuns();
+}
+
+static void VisitSetBitRunsSumWithOffset(benchmark::State& state) {
+ UnaryBitmapTraversalBenchmark(state, /*offset=*/4).BenchVisitSetBitRuns();
}
static void BinaryBitBlockCounterSum(benchmark::State& state) {
- BinaryBitBlockBenchmark(state, /*offset=*/0).BenchBitBlockCounter();
+ BinaryBitmapTraversalBenchmark(state, /*offset=*/0).BenchBitBlockCounter();
}
static void BinaryBitBlockCounterSumWithOffset(benchmark::State& state) {
- BinaryBitBlockBenchmark(state, /*offset=*/4).BenchBitBlockCounter();
+ BinaryBitmapTraversalBenchmark(state, /*offset=*/4).BenchBitBlockCounter();
}
static void BinaryBitmapReaderSum(benchmark::State& state) {
- BinaryBitBlockBenchmark(state, /*offset=*/0).BenchBitmapReader();
+ BinaryBitmapTraversalBenchmark(state, /*offset=*/0).BenchBitmapReader();
}
static void BinaryBitmapReaderSumWithOffset(benchmark::State& state) {
- BinaryBitBlockBenchmark(state, /*offset=*/4).BenchBitmapReader();
+ BinaryBitmapTraversalBenchmark(state, /*offset=*/4).BenchBitmapReader();
+}
+
+static void BinaryVisitTwoBitRunsSum(benchmark::State& state) {
+ BinaryBitmapTraversalBenchmark(state, /*offset=*/0).BenchVisitTwoBitRuns();
+}
+
+static void BinaryVisitTwoBitRunsSumWithOffset(benchmark::State& state) {
+ BinaryBitmapTraversalBenchmark(state, /*offset=*/4).BenchVisitTwoBitRuns();
+}
+
+static void BinaryVisitTwoSetBitRunsSum(benchmark::State& state) {
+ BinaryBitmapTraversalBenchmark(state,
/*offset=*/0).BenchVisitTwoSetBitRuns();
+}
+
+static void BinaryVisitTwoSetBitRunsSumWithOffset(benchmark::State& state) {
+ BinaryBitmapTraversalBenchmark(state,
/*offset=*/4).BenchVisitTwoSetBitRuns();
}
// Range value: average number of total values per null
@@ -257,10 +389,18 @@ BENCHMARK(BitBlockCounterFourWordsSum)->Range(2, 1 << 16);
BENCHMARK(BitBlockCounterFourWordsSumWithOffset)->Range(2, 1 << 16);
BENCHMARK(BitmapReaderSum)->Range(2, 1 << 16);
BENCHMARK(BitmapReaderSumWithOffset)->Range(2, 1 << 16);
+BENCHMARK(VisitBitRunsSum)->Range(2, 1 << 16);
+BENCHMARK(VisitBitRunsSumWithOffset)->Range(2, 1 << 16);
+BENCHMARK(VisitSetBitRunsSum)->Range(2, 1 << 16);
+BENCHMARK(VisitSetBitRunsSumWithOffset)->Range(2, 1 << 16);
BENCHMARK(BinaryBitBlockCounterSum)->Range(2, 1 << 16);
BENCHMARK(BinaryBitBlockCounterSumWithOffset)->Range(2, 1 << 16);
BENCHMARK(BinaryBitmapReaderSum)->Range(2, 1 << 16);
BENCHMARK(BinaryBitmapReaderSumWithOffset)->Range(2, 1 << 16);
+BENCHMARK(BinaryVisitTwoBitRunsSum)->Range(2, 1 << 16);
+BENCHMARK(BinaryVisitTwoBitRunsSumWithOffset)->Range(2, 1 << 16);
+BENCHMARK(BinaryVisitTwoSetBitRunsSum)->Range(2, 1 << 16);
+BENCHMARK(BinaryVisitTwoSetBitRunsSumWithOffset)->Range(2, 1 << 16);
} // namespace internal
} // namespace arrow
diff --git a/cpp/src/arrow/util/bit_run_reader.h
b/cpp/src/arrow/util/bit_run_reader.h
index 1a9638880c5..84366e2789c 100644
--- a/cpp/src/arrow/util/bit_run_reader.h
+++ b/cpp/src/arrow/util/bit_run_reader.h
@@ -17,15 +17,21 @@
#pragma once
+#include <algorithm>
+#include <array>
#include <bit>
#include <cassert>
#include <cstdint>
#include <cstring>
#include <string>
+#include <utility>
+#include "arrow/status.h"
#include "arrow/util/bit_util.h"
+#include "arrow/util/bitmap_ops.h"
#include "arrow/util/bitmap_reader.h"
#include "arrow/util/endian.h"
+#include "arrow/util/logging.h"
#include "arrow/util/macros.h"
#include "arrow/util/visibility.h"
@@ -536,5 +542,84 @@ inline void VisitSetBitRunsVoid(const
std::shared_ptr<Buffer>& bitmap, int64_t o
std::forward<Visit>(visit));
}
+template <typename Visit>
+inline Status VisitTwoSetBitRuns(const uint8_t* left_bitmap, int64_t
left_offset,
+ const uint8_t* right_bitmap, int64_t
right_offset,
+ int64_t length, Visit&& visit) {
+ if (length == 0) {
+ return Status::OK();
+ }
+ if (left_bitmap == NULLPTR) {
+ return VisitSetBitRuns(right_bitmap, right_offset, length,
+ std::forward<Visit>(visit));
+ }
+ if (right_bitmap == NULLPTR) {
+ return VisitSetBitRuns(left_bitmap, left_offset, length,
std::forward<Visit>(visit));
+ }
+
+ constexpr int64_t kBitmapChunkLength = 4096; // 512 bytes of scratch space
on stack
+ std::array<uint8_t, bit_util::BytesForBits(kBitmapChunkLength)> bitmap_chunk;
+
+ int64_t offset = 0;
+ while (offset < length) {
+ const auto chunk_length = std::min(kBitmapChunkLength, length - offset);
+ BitmapAnd(left_bitmap, left_offset + offset, right_bitmap, right_offset +
offset,
+ chunk_length,
+ /*out_offset=*/0, bitmap_chunk.data());
+
+ auto visit_with_offset = [&](int64_t position_in_chunk, int64_t
run_length) {
+ return visit(position_in_chunk + offset, run_length);
+ };
+
+ ARROW_RETURN_NOT_OK(VisitSetBitRuns(bitmap_chunk.data(), /*offset=*/0,
chunk_length,
+ std::move(visit_with_offset)));
+ offset += chunk_length;
+ }
+ return Status::OK();
+}
+
+template <typename Visit>
+inline Status VisitTwoBitRuns(const uint8_t* left_bitmap, int64_t left_offset,
+ const uint8_t* right_bitmap, int64_t
right_offset,
+ int64_t length, Visit&& visit) {
+ int64_t output_position = 0;
+ ARROW_RETURN_NOT_OK(VisitTwoSetBitRuns(
+ left_bitmap, left_offset, right_bitmap, right_offset, length,
+ [&](int64_t position, int64_t run_length) {
+ if (output_position < position) {
+ ARROW_RETURN_NOT_OK(visit(output_position, position -
output_position, false));
+ }
+ ARROW_RETURN_NOT_OK(visit(position, run_length, true));
+ output_position = position + run_length;
+ return Status::OK();
+ }));
+ if (output_position < length) {
+ return visit(output_position, length - output_position, false);
+ }
+ return Status::OK();
+}
+
+template <typename Visit>
+inline void VisitTwoBitRunsVoid(const uint8_t* left_bitmap, int64_t
left_offset,
+ const uint8_t* right_bitmap, int64_t
right_offset,
+ int64_t length, Visit&& visit) {
+ ARROW_CHECK_OK(VisitTwoBitRuns(left_bitmap, left_offset, right_bitmap,
right_offset,
+ length, [&](int64_t position, int64_t length,
bool set) {
+ visit(position, length, set);
+ return Status::OK();
+ }));
+}
+
+template <typename Visit>
+inline void VisitTwoSetBitRunsVoid(const uint8_t* left_bitmap, int64_t
left_offset,
+ const uint8_t* right_bitmap, int64_t
right_offset,
+ int64_t length, Visit&& visit) {
+ ARROW_CHECK_OK(VisitTwoSetBitRuns(left_bitmap, left_offset, right_bitmap,
right_offset,
+ length, [&](int64_t position, int64_t
length) {
+ visit(position, length);
+ return Status::OK();
+ }));
+}
+
} // namespace internal
} // namespace arrow
diff --git a/cpp/src/arrow/util/bitmap_test.cc
b/cpp/src/arrow/util/bitmap_test.cc
index 0a71f4e6038..4ede199e626 100644
--- a/cpp/src/arrow/util/bitmap_test.cc
+++ b/cpp/src/arrow/util/bitmap_test.cc
@@ -332,6 +332,49 @@ class TestSetBitRunReader : public ::testing::Test {
AssertBitRuns(buffer.data(), start_offset, length, expected);
}
+ void AssertVisitTwoBitRuns(const uint8_t* left_data, int64_t left_offset,
+ const uint8_t* right_data, int64_t right_offset,
+ int64_t length,
+ const std::vector<int64_t>& expected_positions,
+ const std::vector<BitRun>& expected_runs,
+ const std::vector<SetBitRun>& expected_set_runs) {
+ std::vector<int64_t> positions;
+ std::vector<BitRun> runs;
+ ASSERT_OK(VisitTwoBitRuns(left_data, left_offset, right_data,
right_offset, length,
+ [&](int64_t position, int64_t run_length, bool
set) {
+ positions.push_back(position);
+ runs.push_back({run_length, set});
+ return Status::OK();
+ }));
+ ASSERT_THAT(positions, ElementsAreArray(expected_positions));
+ ASSERT_THAT(runs, ElementsAreArray(expected_runs));
+
+ positions.clear();
+ runs.clear();
+ VisitTwoBitRunsVoid(left_data, left_offset, right_data, right_offset,
length,
+ [&](int64_t position, int64_t run_length, bool set) {
+ positions.push_back(position);
+ runs.push_back({run_length, set});
+ });
+ ASSERT_THAT(positions, ElementsAreArray(expected_positions));
+ ASSERT_THAT(runs, ElementsAreArray(expected_runs));
+
+ std::vector<SetBitRun> set_runs;
+ ASSERT_OK(VisitTwoSetBitRuns(left_data, left_offset, right_data,
right_offset, length,
+ [&](int64_t position, int64_t run_length) {
+ set_runs.push_back({position, run_length});
+ return Status::OK();
+ }));
+ ASSERT_THAT(set_runs, ElementsAreArray(expected_set_runs));
+
+ set_runs.clear();
+ VisitTwoSetBitRunsVoid(left_data, left_offset, right_data, right_offset,
length,
+ [&](int64_t position, int64_t run_length) {
+ set_runs.push_back({position, run_length});
+ });
+ ASSERT_THAT(set_runs, ElementsAreArray(expected_set_runs));
+ }
+
struct Range {
int64_t offset;
int64_t length;
@@ -484,6 +527,86 @@ TEST_F(TestSetBitRunReader, Random) {
}
}
+TEST_F(TestSetBitRunReader, VisitTwoBitRuns) {
+ auto left = BitmapFromString("11110000 11111100");
+ auto right = BitmapFromString("11001100 00111111");
+ const int64_t left_offset = 1;
+ const int64_t right_offset = 2;
+ const int64_t length = 12;
+
+ // With left_offset = 1 and right_offset = 2:
+ // position: 0 1 2 3 4 5 6 7 8 9 10 11
+ // left[+1]: 1 1 1 0 0 0 0 1 1 1 1 1
+ // right[+2]: 0 0 1 1 0 0 0 0 1 1 1 1
+ // AND: 0 0 1 0 0 0 0 0 1 1 1 1
+ const std::vector<int64_t> expected_positions = {0, 2, 3, 8};
+ const std::vector<BitRun> expected_runs = {{.length = 2, .set = false},
+ {.length = 1, .set = true},
+ {.length = 5, .set = false},
+ {.length = 4, .set = true}};
+ const std::vector<SetBitRun> expected_set_runs = {{.position = 2, .length =
1},
+ {.position = 8, .length =
4}};
+
+ AssertVisitTwoBitRuns(left->data(), left_offset, right->data(),
right_offset, length,
+ expected_positions, expected_runs, expected_set_runs);
+}
+
+TEST_F(TestSetBitRunReader, VisitTwoBitRunsWithNullBitmap) {
+ auto bitmap = BitmapFromString("01101101");
+ const int64_t left_offset = 0;
+ const int64_t right_offset = 1;
+ const int64_t length = 6;
+
+ // With a null left bitmap and right_offset = 1:
+ // left all set: 1 1 1 1 1 1
+ // right[1..6]: 1 1 0 1 1 0
+ // AND: 1 1 0 1 1 0
+ const std::vector<int64_t> expected_positions = {0, 2, 3, 5};
+ const std::vector<BitRun> expected_runs = {{.length = 2, .set = true},
+ {.length = 1, .set = false},
+ {.length = 2, .set = true},
+ {.length = 1, .set = false}};
+ const std::vector<SetBitRun> expected_set_runs = {{.position = 0, .length =
2},
+ {.position = 3, .length =
2}};
+
+ AssertVisitTwoBitRuns(nullptr, left_offset, bitmap->data(), right_offset,
length,
+ expected_positions, expected_runs, expected_set_runs);
+}
+
+TEST_F(TestSetBitRunReader, VisitTwoBitRunsFalseTail) {
+ auto left = BitmapFromString("11111111");
+ auto right = BitmapFromString("11110000");
+
+ std::vector<int64_t> positions;
+ std::vector<BitRun> runs;
+ ASSERT_OK(VisitTwoBitRuns(left->data(), /*left_offset=*/0, right->data(),
+ /*right_offset=*/0, /*length=*/8,
+ [&](int64_t position, int64_t length, bool set) {
+ positions.push_back(position);
+ runs.push_back({length, set});
+ return Status::OK();
+ }));
+ ASSERT_THAT(positions, ElementsAreArray({0, 4}));
+ ASSERT_THAT(runs, ElementsAreArray({BitRun{4, true}, BitRun{4, false}}));
+}
+
+TEST_F(TestSetBitRunReader, VisitTwoBitRunsNoOverlap) {
+ auto left = BitmapFromString("11110000");
+ auto right = BitmapFromString("00001111");
+
+ std::vector<int64_t> positions;
+ std::vector<BitRun> runs;
+ ASSERT_OK(VisitTwoBitRuns(left->data(), /*left_offset=*/0, right->data(),
+ /*right_offset=*/0, /*length=*/8,
+ [&](int64_t position, int64_t length, bool set) {
+ positions.push_back(position);
+ runs.push_back({length, set});
+ return Status::OK();
+ }));
+ ASSERT_THAT(positions, ElementsAreArray({0}));
+ ASSERT_THAT(runs, ElementsAreArray({BitRun{8, false}}));
+}
+
// Tests for BitRunReader.
TEST(BitRunReader, ZeroLength) {
