Skip to content

Commit 4f4b336

Browse files
HuaHuaYHuaHuaY
committed
Implement VisitTwoBitRuns and VisitTwoSetBitRuns methods
1 parent f10c93c commit 4f4b336

5 files changed

Lines changed: 547 additions & 15 deletions

File tree

‎cpp/src/arrow/util/CMakeLists.txt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ add_arrow_test(threading-utility-test
118118
test_common.cc
119119
thread_pool_test.cc)
120120

121-
add_arrow_benchmark(bit_block_counter_benchmark)
121+
add_arrow_benchmark(bitmap_traversal_benchmark)
122122
add_arrow_benchmark(bit_util_benchmark)
123123
add_arrow_benchmark(bitmap_reader_benchmark)
124124
add_arrow_benchmark(bpacking_benchmark)
@@ -128,6 +128,7 @@ add_arrow_benchmark(decimal_benchmark)
128128
add_arrow_benchmark(hashing_benchmark)
129129
add_arrow_benchmark(int_util_benchmark)
130130
add_arrow_benchmark(machine_benchmark)
131+
add_arrow_benchmark(nested_bitmap_traversal_benchmark)
131132
add_arrow_benchmark(queue_benchmark)
132133
add_arrow_benchmark(range_benchmark)
133134
add_arrow_benchmark(small_vector_benchmark)

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

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,18 @@
1717

1818
#pragma once
1919

20+
#include <algorithm>
2021
#include <bit>
2122
#include <cassert>
2223
#include <cstdint>
2324
#include <cstring>
2425
#include <string>
26+
#include <utility>
2527

2628
#include "arrow/util/bit_util.h"
2729
#include "arrow/util/bitmap_reader.h"
2830
#include "arrow/util/endian.h"
31+
#include "arrow/util/logging.h"
2932
#include "arrow/util/macros.h"
3033
#include "arrow/util/visibility.h"
3134

@@ -536,5 +539,95 @@ inline void VisitSetBitRunsVoid(const std::shared_ptr<Buffer>& bitmap, int64_t o
536539
std::forward<Visit>(visit));
537540
}
538541

542+
template <typename Visit>
543+
inline Status VisitTwoSetBitRuns(const uint8_t* left_bitmap, int64_t left_offset,
544+
const uint8_t* right_bitmap, int64_t right_offset,
545+
int64_t length, Visit&& visit) {
546+
if (length == 0) {
547+
return Status::OK();
548+
}
549+
if (left_bitmap == NULLPTR) {
550+
return VisitSetBitRuns(right_bitmap, right_offset, length,
551+
std::forward<Visit>(visit));
552+
}
553+
if (right_bitmap == NULLPTR) {
554+
return VisitSetBitRuns(left_bitmap, left_offset, length, std::forward<Visit>(visit));
555+
}
556+
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();
585+
}
586+
587+
template <typename Visit>
588+
inline Status VisitTwoBitRuns(const uint8_t* left_bitmap, int64_t left_offset,
589+
const uint8_t* right_bitmap, int64_t right_offset,
590+
int64_t length, Visit&& visit) {
591+
int64_t output_position = 0;
592+
ARROW_RETURN_NOT_OK(VisitTwoSetBitRuns(
593+
left_bitmap, left_offset, right_bitmap, right_offset, length,
594+
[&](int64_t position, int64_t run_length) {
595+
if (output_position < position) {
596+
ARROW_RETURN_NOT_OK(visit(output_position, position - output_position, false));
597+
}
598+
ARROW_RETURN_NOT_OK(visit(position, run_length, true));
599+
output_position = position + run_length;
600+
return Status::OK();
601+
}));
602+
if (output_position < length) {
603+
return visit(output_position, length - output_position, false);
604+
}
605+
return Status::OK();
606+
}
607+
608+
template <typename Visit>
609+
inline void VisitTwoBitRunsVoid(const uint8_t* left_bitmap, int64_t left_offset,
610+
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+
}));
618+
}
619+
620+
template <typename Visit>
621+
inline void VisitTwoSetBitRunsVoid(const uint8_t* left_bitmap, int64_t left_offset,
622+
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+
}));
630+
}
631+
539632
} // namespace internal
540633
} // namespace arrow

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

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -484,6 +484,150 @@ TEST_F(TestSetBitRunReader, Random) {
484484
}
485485
}
486486

487+
TEST_F(TestSetBitRunReader, VisitTwoBitRuns) {
488+
auto left = BitmapFromString("11110000 11111100");
489+
auto right = BitmapFromString("11001100 00111111");
490+
491+
std::vector<int64_t> positions;
492+
std::vector<BitRun> runs;
493+
ASSERT_OK(VisitTwoBitRuns(left->data(), /*left_offset=*/1, right->data(),
494+
/*right_offset=*/2, /*length=*/12,
495+
[&](int64_t position, int64_t length, bool set) {
496+
positions.push_back(position);
497+
runs.push_back({length, set});
498+
return Status::OK();
499+
}));
500+
ASSERT_THAT(positions, ElementsAreArray({0, 2, 3, 8}));
501+
ASSERT_THAT(runs, ElementsAreArray({BitRun{2, false}, BitRun{1, true}, BitRun{5, false},
502+
BitRun{4, true}}));
503+
}
504+
505+
TEST_F(TestSetBitRunReader, VisitTwoBitRunsFalseTail) {
506+
auto left = BitmapFromString("11111111");
507+
auto right = BitmapFromString("11110000");
508+
509+
std::vector<int64_t> positions;
510+
std::vector<BitRun> runs;
511+
ASSERT_OK(VisitTwoBitRuns(left->data(), /*left_offset=*/0, right->data(),
512+
/*right_offset=*/0, /*length=*/8,
513+
[&](int64_t position, int64_t length, bool set) {
514+
positions.push_back(position);
515+
runs.push_back({length, set});
516+
return Status::OK();
517+
}));
518+
ASSERT_THAT(positions, ElementsAreArray({0, 4}));
519+
ASSERT_THAT(runs, ElementsAreArray({BitRun{4, true}, BitRun{4, false}}));
520+
}
521+
522+
TEST_F(TestSetBitRunReader, VisitTwoBitRunsNoOverlap) {
523+
auto left = BitmapFromString("11110000");
524+
auto right = BitmapFromString("00001111");
525+
526+
std::vector<int64_t> positions;
527+
std::vector<BitRun> runs;
528+
ASSERT_OK(VisitTwoBitRuns(left->data(), /*left_offset=*/0, right->data(),
529+
/*right_offset=*/0, /*length=*/8,
530+
[&](int64_t position, int64_t length, bool set) {
531+
positions.push_back(position);
532+
runs.push_back({length, set});
533+
return Status::OK();
534+
}));
535+
ASSERT_THAT(positions, ElementsAreArray({0}));
536+
ASSERT_THAT(runs, ElementsAreArray({BitRun{8, false}}));
537+
}
538+
539+
TEST_F(TestSetBitRunReader, VisitTwoBitRunsWithNullBitmap) {
540+
auto bitmap = BitmapFromString("11110000");
541+
542+
std::vector<int64_t> positions;
543+
std::vector<BitRun> runs;
544+
ASSERT_OK(VisitTwoBitRuns(nullptr, /*left_offset=*/0, bitmap->data(),
545+
/*right_offset=*/0, /*length=*/8,
546+
[&](int64_t position, int64_t length, bool set) {
547+
positions.push_back(position);
548+
runs.push_back({length, set});
549+
return Status::OK();
550+
}));
551+
ASSERT_THAT(positions, ElementsAreArray({0, 4}));
552+
ASSERT_THAT(runs, ElementsAreArray({BitRun{4, true}, BitRun{4, false}}));
553+
}
554+
555+
TEST_F(TestSetBitRunReader, VisitTwoSetBitRuns) {
556+
auto left = BitmapFromString("11110000 11111100");
557+
auto right = BitmapFromString("11001100 00111111");
558+
559+
std::vector<SetBitRun> runs;
560+
ASSERT_OK(VisitTwoSetBitRuns(left->data(), /*left_offset=*/1, right->data(),
561+
/*right_offset=*/2, /*length=*/12,
562+
[&](int64_t position, int64_t length) {
563+
runs.push_back({position, length});
564+
return Status::OK();
565+
}));
566+
ASSERT_THAT(runs, ElementsAreArray({SetBitRun{2, 1}, SetBitRun{8, 4}}));
567+
}
568+
569+
TEST_F(TestSetBitRunReader, VisitTwoSetBitRunsWithNullBitmap) {
570+
auto bitmap = BitmapFromString("01101101");
571+
572+
std::vector<SetBitRun> runs;
573+
ASSERT_OK(VisitTwoSetBitRuns(nullptr, /*left_offset=*/0, bitmap->data(),
574+
/*right_offset=*/1, /*length=*/6,
575+
[&](int64_t position, int64_t length) {
576+
runs.push_back({position, length});
577+
return Status::OK();
578+
}));
579+
ASSERT_THAT(runs, ElementsAreArray({SetBitRun{0, 2}, SetBitRun{3, 2}}));
580+
}
581+
582+
TEST_F(TestSetBitRunReader, VisitTwoBitRunsVoid) {
583+
auto left = BitmapFromString("11110000 11111100");
584+
auto right = BitmapFromString("11001100 00111111");
585+
586+
std::vector<int64_t> positions;
587+
std::vector<BitRun> runs;
588+
VisitTwoBitRunsVoid(left->data(), /*left_offset=*/1, right->data(),
589+
/*right_offset=*/2, /*length=*/12,
590+
[&](int64_t position, int64_t length, bool set) {
591+
positions.push_back(position);
592+
runs.push_back({length, set});
593+
});
594+
// With left_offset = 1 and right_offset = 2:
595+
// position: 0 1 2 3 4 5 6 7 8 9 10 11
596+
// left[+1]: 1 1 1 0 0 0 0 1 1 1 1 1
597+
// right[+2]: 0 0 1 1 0 0 0 0 1 1 1 1
598+
// AND: 0 0 1 0 0 0 0 0 1 1 1 1
599+
ASSERT_THAT(positions, ElementsAreArray({0, 2, 3, 8}));
600+
ASSERT_THAT(runs, ElementsAreArray({BitRun{2, false}, BitRun{1, true}, BitRun{5, false},
601+
BitRun{4, true}}));
602+
}
603+
604+
TEST_F(TestSetBitRunReader, VisitTwoSetBitRunsVoid) {
605+
auto left = BitmapFromString("11110000 11111100");
606+
auto right = BitmapFromString("11001100 00111111");
607+
608+
std::vector<SetBitRun> runs;
609+
VisitTwoSetBitRunsVoid(
610+
left->data(), /*left_offset=*/1, right->data(),
611+
/*right_offset=*/2, /*length=*/12,
612+
[&](int64_t position, int64_t length) { runs.push_back({position, length}); });
613+
ASSERT_THAT(runs, ElementsAreArray({SetBitRun{2, 1}, SetBitRun{8, 4}}));
614+
}
615+
616+
TEST_F(TestSetBitRunReader, VisitTwoSetBitRunsVoidWithNullBitmap) {
617+
auto bitmap = BitmapFromString("01101101");
618+
619+
std::vector<SetBitRun> runs;
620+
VisitTwoSetBitRunsVoid(
621+
nullptr, /*left_offset=*/0, bitmap->data(),
622+
/*right_offset=*/1, /*length=*/6,
623+
[&](int64_t position, int64_t length) { runs.push_back({position, length}); });
624+
// With a null left bitmap and right_offset = 1:
625+
// left all set: 1 1 1 1 1 1
626+
// right[1..6]: 1 1 0 1 1 0
627+
// AND: 1 1 0 1 1 0
628+
ASSERT_THAT(runs, ElementsAreArray({SetBitRun{0, 2}, SetBitRun{3, 2}}));
629+
}
630+
487631
// Tests for BitRunReader.
488632

489633
TEST(BitRunReader, ZeroLength) {

0 commit comments

Comments
 (0)