Skip to content

Commit 5b3693c

Browse files
HuaHuaYHuaHuaY
committed
address review
1 parent ef44df8 commit 5b3693c

2 files changed

Lines changed: 43 additions & 52 deletions

File tree

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

Lines changed: 32 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,10 @@
2525
#include <string>
2626
#include <utility>
2727

28+
#include "arrow/buffer.h"
29+
#include "arrow/memory_pool.h"
2830
#include "arrow/util/bit_util.h"
31+
#include "arrow/util/bitmap_ops.h"
2932
#include "arrow/util/bitmap_reader.h"
3033
#include "arrow/util/endian.h"
3134
#include "arrow/util/logging.h"
@@ -542,7 +545,8 @@ inline void VisitSetBitRunsVoid(const std::shared_ptr<Buffer>& bitmap, int64_t o
542545
template <typename Visit>
543546
inline Status VisitTwoSetBitRuns(const uint8_t* left_bitmap, int64_t left_offset,
544547
const uint8_t* right_bitmap, int64_t right_offset,
545-
int64_t length, Visit&& visit) {
548+
int64_t length, Visit&& visit,
549+
MemoryPool* pool = default_memory_pool()) {
546550
if (length == 0) {
547551
return Status::OK();
548552
}
@@ -554,40 +558,18 @@ inline Status VisitTwoSetBitRuns(const uint8_t* left_bitmap, int64_t left_offset
554558
return VisitSetBitRuns(left_bitmap, left_offset, length, std::forward<Visit>(visit));
555559
}
556560

557-
SetBitRunReader left_reader(left_bitmap, left_offset, length);
558-
SetBitRunReader right_reader(right_bitmap, right_offset, length);
559-
auto left_run = left_reader.NextRun();
560-
auto right_run = right_reader.NextRun();
561-
while (!left_run.AtEnd() && !right_run.AtEnd()) {
562-
const int64_t left_end = left_run.position + left_run.length;
563-
if (left_end <= right_run.position) {
564-
left_run = left_reader.NextRun();
565-
continue;
566-
}
567-
568-
const int64_t right_end = right_run.position + right_run.length;
569-
if (right_end <= left_run.position) {
570-
right_run = right_reader.NextRun();
571-
continue;
572-
}
573-
574-
const int64_t start = std::max(left_run.position, right_run.position);
575-
const int64_t end = std::min(left_end, right_end);
576-
ARROW_RETURN_NOT_OK(visit(start, end - start));
577-
if (end == left_end) {
578-
left_run = left_reader.NextRun();
579-
}
580-
if (end == right_end) {
581-
right_run = right_reader.NextRun();
582-
}
583-
}
584-
return Status::OK();
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));
585566
}
586567

587568
template <typename Visit>
588569
inline Status VisitTwoBitRuns(const uint8_t* left_bitmap, int64_t left_offset,
589570
const uint8_t* right_bitmap, int64_t right_offset,
590-
int64_t length, Visit&& visit) {
571+
int64_t length, Visit&& visit,
572+
MemoryPool* pool = default_memory_pool()) {
591573
int64_t output_position = 0;
592574
ARROW_RETURN_NOT_OK(VisitTwoSetBitRuns(
593575
left_bitmap, left_offset, right_bitmap, right_offset, length,
@@ -598,7 +580,8 @@ inline Status VisitTwoBitRuns(const uint8_t* left_bitmap, int64_t left_offset,
598580
ARROW_RETURN_NOT_OK(visit(position, run_length, true));
599581
output_position = position + run_length;
600582
return Status::OK();
601-
}));
583+
},
584+
pool));
602585
if (output_position < length) {
603586
return visit(output_position, length - output_position, false);
604587
}
@@ -608,25 +591,29 @@ inline Status VisitTwoBitRuns(const uint8_t* left_bitmap, int64_t left_offset,
608591
template <typename Visit>
609592
inline void VisitTwoBitRunsVoid(const uint8_t* left_bitmap, int64_t left_offset,
610593
const uint8_t* right_bitmap, int64_t right_offset,
611-
int64_t length, Visit&& visit) {
612-
ARROW_IGNORE_EXPR(VisitTwoBitRuns(left_bitmap, left_offset, right_bitmap, right_offset,
613-
length,
614-
[&](int64_t position, int64_t length, bool set) {
615-
visit(position, length, set);
616-
return Status::OK();
617-
}));
594+
int64_t length, Visit&& visit,
595+
MemoryPool* pool = default_memory_pool()) {
596+
ARROW_IGNORE_EXPR(VisitTwoBitRuns(
597+
left_bitmap, left_offset, right_bitmap, right_offset, length,
598+
[&](int64_t position, int64_t length, bool set) {
599+
visit(position, length, set);
600+
return Status::OK();
601+
},
602+
pool));
618603
}
619604

620605
template <typename Visit>
621606
inline void VisitTwoSetBitRunsVoid(const uint8_t* left_bitmap, int64_t left_offset,
622607
const uint8_t* right_bitmap, int64_t right_offset,
623-
int64_t length, Visit&& visit) {
624-
ARROW_IGNORE_EXPR(VisitTwoSetBitRuns(left_bitmap, left_offset, right_bitmap,
625-
right_offset, length,
626-
[&](int64_t position, int64_t length) {
627-
visit(position, length);
628-
return Status::OK();
629-
}));
608+
int64_t length, Visit&& visit,
609+
MemoryPool* pool = default_memory_pool()) {
610+
ARROW_IGNORE_EXPR(VisitTwoSetBitRuns(
611+
left_bitmap, left_offset, right_bitmap, right_offset, length,
612+
[&](int64_t position, int64_t length) {
613+
visit(position, length);
614+
return Status::OK();
615+
},
616+
pool));
630617
}
631618

632619
} // namespace internal

‎cpp/src/arrow/util/nested_bitmap_traversal_benchmark.cc‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
#include "arrow/array/array_nested.h"
2525
#include "arrow/array/array_primitive.h"
2626
#include "arrow/buffer.h"
27+
#include "arrow/memory_pool.h"
2728
#include "arrow/testing/gtest_util.h"
2829
#include "arrow/testing/random.h"
2930
#include "arrow/util/bit_block_counter.h"
@@ -35,13 +36,14 @@ namespace arrow::internal {
3536

3637
struct NestedBitmapTraversalBenchmark {
3738
benchmark::State& state;
39+
MemoryPool* memory_pool;
3840
int64_t parent_length;
3941
std::shared_ptr<ListArray> list_array;
4042
std::shared_ptr<Buffer> outer_bitmap;
4143
int64_t expected;
4244

4345
explicit NestedBitmapTraversalBenchmark(benchmark::State& state)
44-
: state(state), parent_length(1 << 20) {
46+
: state(state), memory_pool(default_memory_pool()), parent_length(1 << 20) {
4547
constexpr double kValuesNullProbability = 0.1;
4648
random::RandomArrayGenerator rng(/*seed=*/0);
4749
const auto parent_null_probability = 1. / static_cast<double>(state.range(0));
@@ -82,9 +84,8 @@ struct NestedBitmapTraversalBenchmark {
8284
const auto& values = static_cast<const Int8Array&>(*list_array->values());
8385
for (auto _ : state) {
8486
std::shared_ptr<Buffer> visible_bitmap;
85-
ABORT_NOT_OK(BitmapAnd(default_memory_pool(), outer_bitmap_data,
86-
/*left_offset=*/0, list_bitmap, /*right_offset=*/0,
87-
parent_length,
87+
ABORT_NOT_OK(BitmapAnd(memory_pool, outer_bitmap_data, /*left_offset=*/0,
88+
list_bitmap, /*right_offset=*/0, parent_length,
8889
/*out_offset=*/0)
8990
.Value(&visible_bitmap));
9091

@@ -136,7 +137,8 @@ struct NestedBitmapTraversalBenchmark {
136137
const int64_t child_end = list_array->value_offset(position + length);
137138
result += CountValues(values, child_start, child_end - child_start);
138139
}
139-
});
140+
},
141+
memory_pool);
140142
CheckResult(result);
141143
}
142144
state.SetItemsProcessed(state.iterations() * parent_length);
@@ -150,11 +152,13 @@ struct NestedBitmapTraversalBenchmark {
150152
int64_t result = 0;
151153
VisitTwoSetBitRunsVoid(
152154
outer_bitmap_data, /*left_offset=*/0, list_bitmap,
153-
/*right_offset=*/0, parent_length, [&](int64_t position, int64_t length) {
155+
/*right_offset=*/0, parent_length,
156+
[&](int64_t position, int64_t length) {
154157
const int64_t child_start = list_array->value_offset(position);
155158
const int64_t child_end = list_array->value_offset(position + length);
156159
result += CountValues(values, child_start, child_end - child_start);
157-
});
160+
},
161+
memory_pool);
158162
CheckResult(result);
159163
}
160164
state.SetItemsProcessed(state.iterations() * parent_length);

0 commit comments

Comments
 (0)