Skip to content

Commit 710cedd

Browse files
pitroupitrou
committed
EXP: [C++] VisitTwoBitRuns changes
Based on #50281
1 parent 5b3693c commit 710cedd

4 files changed

Lines changed: 70 additions & 89 deletions

File tree

‎cpp/src/arrow/compute/kernels/codegen_internal.h‎

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
#include "arrow/type.h"
3939
#include "arrow/type_traits.h"
4040
#include "arrow/util/bit_block_counter.h"
41+
#include "arrow/util/bit_run_reader.h"
4142
#include "arrow/util/bit_util.h"
4243
#include "arrow/util/bitmap_generate.h"
4344
#include "arrow/util/bitmap_reader.h"
@@ -509,18 +510,24 @@ static void VisitTwoArrayValuesInline(const ArraySpan& arr0, const ArraySpan& ar
509510
ArrayIterator<Arg0Type> arr0_it(arr0);
510511
ArrayIterator<Arg1Type> arr1_it(arr1);
511512

512-
auto visit_valid = [&](int64_t i) {
513-
valid_func(GetViewType<Arg0Type>::LogicalValue(arr0_it()),
514-
GetViewType<Arg1Type>::LogicalValue(arr1_it()));
515-
};
516-
auto visit_null = [&]() {
517-
arr0_it();
518-
arr1_it();
519-
null_func();
513+
auto visit_run = [&](int64_t position, int64_t run_length, bool is_set) {
514+
if (is_set) {
515+
for (int64_t i = 0; i < run_length; ++i) {
516+
valid_func(GetViewType<Arg0Type>::LogicalValue(arr0_it()),
517+
GetViewType<Arg1Type>::LogicalValue(arr1_it()));
518+
}
519+
} else {
520+
for (int64_t i = 0; i < run_length; ++i) {
521+
arr0_it();
522+
arr1_it();
523+
null_func();
524+
}
525+
}
520526
};
521-
VisitTwoBitBlocksVoid(arr0.buffers[0].data, arr0.offset, arr1.buffers[0].data,
522-
arr1.offset, arr0.length, std::move(visit_valid),
523-
std::move(visit_null));
527+
528+
::arrow::internal::VisitTwoBitRunsVoid(arr0.buffers[0].data, arr0.offset,
529+
arr1.buffers[0].data, arr1.offset, arr0.length,
530+
std::move(visit_run));
524531
}
525532

526533
// ----------------------------------------------------------------------

‎cpp/src/arrow/compute/kernels/scalar_if_else.cc‎

Lines changed: 20 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
#include <cstdint>
1819
#include <cstring>
1920
#include <limits>
2021
#include "arrow/array/builder_nested.h"
@@ -1607,26 +1608,14 @@ Status ExecArrayCaseWhen(KernelContext* ctx, const ExecSpan& batch, ExecResult*
16071608

16081609
if (cond_array.GetNullCount() == 0) {
16091610
// If no valid buffer, visit mask & cond bitmap simultaneously
1610-
BinaryBitBlockCounter counter(mask, /*start_offset=*/0, cond_values, cond_offset,
1611-
batch.length);
1612-
while (offset < batch.length) {
1613-
const auto block = counter.NextAndWord();
1614-
if (block.AllSet()) {
1615-
CopyValues<Type>(value, offset, block.length, out_valid, out_values,
1616-
out_offset + offset);
1617-
bit_util::SetBitsTo(mask, offset, block.length, false);
1618-
} else if (block.popcount) {
1619-
for (int64_t j = 0; j < block.length; ++j) {
1620-
if (bit_util::GetBit(mask, offset + j) &&
1621-
bit_util::GetBit(cond_values, cond_offset + offset + j)) {
1622-
CopyValues<Type>(value, offset + j, /*length=*/1, out_valid, out_values,
1623-
out_offset + offset + j);
1624-
bit_util::SetBitTo(mask, offset + j, false);
1625-
}
1626-
}
1627-
}
1628-
offset += block.length;
1629-
}
1611+
auto visit_run = [&](int64_t position, int64_t run_length) {
1612+
CopyValues<Type>(value, position, run_length, out_valid, out_values,
1613+
out_offset + position);
1614+
bit_util::SetBitsTo(mask, position, run_length, false);
1615+
return Status::OK();
1616+
};
1617+
RETURN_NOT_OK(::arrow::internal::VisitTwoSetBitRuns(
1618+
mask, /*start_offset=*/0, cond_values, cond_offset, batch.length, visit_run));
16301619
} else {
16311620
// Visit mask & cond bitmap & cond validity
16321621
const uint8_t* cond_valid = cond_array.buffers[0].data;
@@ -1657,32 +1646,20 @@ Status ExecArrayCaseWhen(KernelContext* ctx, const ExecSpan& batch, ExecResult*
16571646
}
16581647
if (!have_else_arg) {
16591648
// Need to initialize any remaining null slots (uninitialized memory)
1660-
BitBlockCounter counter(mask, /*offset=*/0, batch.length);
1661-
int64_t offset = 0;
16621649
auto bit_width = checked_cast<const FixedWidthType&>(*out->type()).bit_width();
16631650
auto byte_width = bit_util::BytesForBits(bit_width);
1664-
while (offset < batch.length) {
1665-
const auto block = counter.NextWord();
1666-
if (block.AllSet()) {
1667-
if (bit_width == 1) {
1668-
bit_util::SetBitsTo(out_values, out_offset + offset, block.length, false);
1669-
} else {
1670-
std::memset(out_values + (out_offset + offset) * byte_width, 0x00,
1671-
byte_width * block.length);
1672-
}
1673-
} else if (!block.NoneSet()) {
1674-
for (int64_t j = 0; j < block.length; ++j) {
1675-
if (bit_util::GetBit(out_valid, out_offset + offset + j)) continue;
1676-
if (bit_width == 1) {
1677-
bit_util::ClearBit(out_values, out_offset + offset + j);
1678-
} else {
1679-
std::memset(out_values + (out_offset + offset + j) * byte_width, 0x00,
1680-
byte_width);
1681-
}
1682-
}
1651+
1652+
auto visit_run = [&](int64_t position, int64_t run_length) {
1653+
if (bit_width == 1) {
1654+
bit_util::SetBitsTo(out_values, out_offset + position, run_length, false);
1655+
} else {
1656+
std::memset(out_values + (out_offset + position) * byte_width, 0x00,
1657+
byte_width * run_length);
16831658
}
1684-
offset += block.length;
1685-
}
1659+
bit_util::SetBitsTo(mask, position, run_length, false);
1660+
};
1661+
1662+
::arrow::internal::VisitSetBitRunsVoid(mask, /*offset=*/0, batch.length, visit_run);
16861663
}
16871664
return Status::OK();
16881665
}

‎cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc‎

Lines changed: 11 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
#include "arrow/compute/registry.h"
3535
#include "arrow/memory_pool.h"
3636
#include "arrow/record_batch.h"
37+
#include "arrow/status.h"
3738
#include "arrow/table.h"
3839
#include "arrow/type.h"
3940
#include "arrow/type_traits.h"
@@ -135,37 +136,17 @@ Result<std::shared_ptr<ArrayData>> GetTakeIndicesFromBitmapImpl(
135136
// data bitmap together.
136137
DCHECK_EQ(null_selection, FilterOptions::DROP);
137138

138-
// The position relative to the start of the filter
139-
T position = 0;
140-
// The current position taking the filter offset into account
141-
int64_t position_with_offset = filter.offset;
142-
143-
BinaryBitBlockCounter filter_counter(filter_data, filter.offset, filter_is_valid,
144-
filter.offset, filter.length);
145-
while (position < filter.length) {
146-
BitBlockCount and_block = filter_counter.NextAndWord();
147-
RETURN_NOT_OK(builder.Reserve(and_block.popcount));
148-
if (and_block.AllSet()) {
149-
// All the values are selected and non-null
150-
for (int64_t i = 0; i < and_block.length; ++i) {
151-
builder.UnsafeAppend(position++);
152-
}
153-
position_with_offset += and_block.length;
154-
} else if (!and_block.NoneSet()) {
155-
// Some of the values are false or null
156-
for (int64_t i = 0; i < and_block.length; ++i) {
157-
if (bit_util::GetBit(filter_is_valid, position_with_offset) &&
158-
bit_util::GetBit(filter_data, position_with_offset)) {
159-
builder.UnsafeAppend(position);
160-
}
161-
++position;
162-
++position_with_offset;
163-
}
164-
} else {
165-
position += and_block.length;
166-
position_with_offset += and_block.length;
139+
auto visit_run = [&](int64_t position, int64_t run_length) {
140+
RETURN_NOT_OK(builder.Reserve(run_length));
141+
for (int64_t i = 0; i < run_length; ++i) {
142+
builder.UnsafeAppend(static_cast<T>(position + i));
167143
}
168-
}
144+
return Status::OK();
145+
};
146+
147+
RETURN_NOT_OK(::arrow::internal::VisitTwoSetBitRuns(filter_data, filter.offset,
148+
filter_is_valid, filter.offset,
149+
filter.length, visit_run));
169150
} else {
170151
// The filter has no nulls, so we need only look for true values
171152
RETURN_NOT_OK(::arrow::internal::VisitSetBitRuns(

‎cpp/src/arrow/util/bit_run_reader.h‎

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
#pragma once
1919

2020
#include <algorithm>
21+
#include <array>
2122
#include <bit>
2223
#include <cassert>
2324
#include <cstdint>
@@ -27,6 +28,7 @@
2728

2829
#include "arrow/buffer.h"
2930
#include "arrow/memory_pool.h"
31+
#include "arrow/status.h"
3032
#include "arrow/util/bit_util.h"
3133
#include "arrow/util/bitmap_ops.h"
3234
#include "arrow/util/bitmap_reader.h"
@@ -558,11 +560,25 @@ inline Status VisitTwoSetBitRuns(const uint8_t* left_bitmap, int64_t left_offset
558560
return VisitSetBitRuns(left_bitmap, left_offset, length, std::forward<Visit>(visit));
559561
}
560562

561-
ARROW_ASSIGN_OR_RAISE(auto bitmap_and,
562-
BitmapAnd(pool, left_bitmap, left_offset, right_bitmap,
563-
right_offset, length, /*out_offset=*/0));
564-
return VisitSetBitRuns(bitmap_and->data(), /*offset=*/0, length,
565-
std::forward<Visit>(visit));
563+
constexpr int64_t kBitmapChunkLength = 4096; // 512 bytes of scratch space on stack
564+
std::array<uint8_t, bit_util::BytesForBits(kBitmapChunkLength)> bitmap_chunk;
565+
566+
int64_t offset = 0;
567+
while (offset < length) {
568+
const auto chunk_length = std::min(kBitmapChunkLength, length - offset);
569+
BitmapAnd(left_bitmap, left_offset + offset, right_bitmap, right_offset + offset,
570+
chunk_length,
571+
/*out_offset=*/0, bitmap_chunk.data());
572+
573+
auto visit_with_offset = [&](int64_t position_in_chunk, int64_t run_length) {
574+
return visit(position_in_chunk + offset, run_length);
575+
};
576+
577+
ARROW_RETURN_NOT_OK(VisitSetBitRuns(bitmap_chunk.data(), /*offset=*/0, chunk_length,
578+
std::move(visit_with_offset)));
579+
offset += chunk_length;
580+
}
581+
return Status::OK();
566582
}
567583

568584
template <typename Visit>

0 commit comments

Comments
 (0)