Skip to content

Commit 36c7481

Browse files
HuaHuaYHuaHuaY
committed
Implement VisitTwoBitRuns and VisitTwoSetBitRuns methods
1 parent 388e95b commit 36c7481

3 files changed

Lines changed: 146 additions & 15 deletions

File tree

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

Lines changed: 17 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,8 @@
3131
#include "arrow/compute/kernels/pivot_internal.h"
3232
#include "arrow/compute/registry_internal.h"
3333
#include "arrow/compute/row/grouper.h"
34-
#include "arrow/util/bit_block_counter.h"
34+
#include "arrow/util/bit_run_reader.h"
35+
#include "arrow/util/bit_util.h"
3536
#include "arrow/util/checked_cast.h"
3637
#include "arrow/util/logging_internal.h"
3738
#include "arrow/visit_type_inline.h"
@@ -125,23 +126,24 @@ struct GroupedPivotAccumulator {
125126
const uint32_t* key_values = keys->GetValues<uint32_t>(1);
126127
const uint8_t* values_null_bitmap =
127128
(values.GetNullCount() != 0) ? values.GetValues<uint8_t>(0, 0) : nullptr;
128-
return ::arrow::internal::VisitTwoBitBlocks(
129+
return ::arrow::internal::VisitTwoSetBitRuns(
129130
keys_null_bitmap, keys->offset, values_null_bitmap, values.offset,
130-
values.length,
131-
[&](int64_t i) {
132-
// Non-null key, non-null value
133-
const uint32_t group = groups[i];
134-
const uint32_t key = key_values[i];
135-
DCHECK_LT(static_cast<int>(key), num_keys_);
136-
if (ARROW_PREDICT_FALSE(bit_util::GetBit(take_bitmap_data[key], group))) {
137-
return DuplicateValue();
131+
values.length, [&](int64_t position, int64_t length) {
132+
const int64_t end = position + length;
133+
for (int64_t i = position; i < end; ++i) {
134+
const uint32_t group = groups[i];
135+
const uint32_t key = key_values[i];
136+
DCHECK_LT(static_cast<int>(key), num_keys_);
137+
if (ARROW_PREDICT_FALSE(bit_util::GetBit(take_bitmap_data[key], group))) {
138+
return DuplicateValue();
139+
}
140+
// For row #group in column #key,
141+
// we are going to take the value at index #i
142+
bit_util::SetBit(take_bitmap_data[key], group);
143+
take_indices_data[key][group] = static_cast<TakeIndex>(i);
138144
}
139-
// For row #group in column #key, we are going to take the value at index #i
140-
bit_util::SetBit(take_bitmap_data[key], group);
141-
take_indices_data[key][group] = static_cast<TakeIndex>(i);
142145
return Status::OK();
143-
},
144-
[] { return Status::OK(); });
146+
});
145147
};
146148

147149
// Call compute_take_indices with the optimal integer width

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

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
#pragma once
1919

20+
#include <algorithm>
2021
#include <bit>
2122
#include <cassert>
2223
#include <cstdint>
@@ -504,6 +505,89 @@ inline Status VisitSetBitRuns(const uint8_t* bitmap, int64_t offset, int64_t len
504505
return Status::OK();
505506
}
506507

508+
template <typename Visit>
509+
inline Status VisitTwoBitRuns(const uint8_t* left_bitmap, int64_t left_offset,
510+
const uint8_t* right_bitmap, int64_t right_offset,
511+
int64_t length, Visit&& visit) {
512+
if (length == 0) {
513+
return Status::OK();
514+
}
515+
if (left_bitmap == NULLPTR) {
516+
return VisitBitRuns(right_bitmap, right_offset, length, std::forward<Visit>(visit));
517+
}
518+
if (right_bitmap == NULLPTR) {
519+
return VisitBitRuns(left_bitmap, left_offset, length, std::forward<Visit>(visit));
520+
}
521+
522+
BitRunReader left_reader(left_bitmap, left_offset, length);
523+
BitRunReader right_reader(right_bitmap, right_offset, length);
524+
auto left_run = left_reader.NextRun();
525+
auto right_run = right_reader.NextRun();
526+
int64_t left_end = left_run.length;
527+
int64_t right_end = right_run.length;
528+
int64_t position = 0;
529+
int64_t output_start = 0;
530+
bool output_set = left_run.set && right_run.set;
531+
532+
while (position < length) {
533+
const int64_t segment_end = std::min(left_end, right_end);
534+
const bool segment_set = left_run.set && right_run.set;
535+
if (segment_set != output_set) {
536+
ARROW_RETURN_NOT_OK(visit(output_start, position - output_start, output_set));
537+
output_start = position;
538+
output_set = segment_set;
539+
}
540+
541+
position = segment_end;
542+
if (position == left_end) {
543+
left_run = left_reader.NextRun();
544+
left_end += left_run.length;
545+
}
546+
if (position == right_end) {
547+
right_run = right_reader.NextRun();
548+
right_end += right_run.length;
549+
}
550+
}
551+
return visit(output_start, length - output_start, output_set);
552+
}
553+
554+
template <typename Visit>
555+
inline Status VisitTwoSetBitRuns(const uint8_t* left_bitmap, int64_t left_offset,
556+
const uint8_t* right_bitmap, int64_t right_offset,
557+
int64_t length, Visit&& visit) {
558+
if (length == 0) {
559+
return Status::OK();
560+
}
561+
if (left_bitmap == NULLPTR) {
562+
return VisitSetBitRuns(right_bitmap, right_offset, length,
563+
std::forward<Visit>(visit));
564+
}
565+
if (right_bitmap == NULLPTR) {
566+
return VisitSetBitRuns(left_bitmap, left_offset, length, std::forward<Visit>(visit));
567+
}
568+
569+
SetBitRunReader left_reader(left_bitmap, left_offset, length);
570+
SetBitRunReader right_reader(right_bitmap, right_offset, length);
571+
auto left_run = left_reader.NextRun();
572+
auto right_run = right_reader.NextRun();
573+
while (!left_run.AtEnd() && !right_run.AtEnd()) {
574+
const int64_t start = std::max(left_run.position, right_run.position);
575+
const int64_t left_end = left_run.position + left_run.length;
576+
const int64_t right_end = right_run.position + right_run.length;
577+
const int64_t end = std::min(left_end, right_end);
578+
if (start < end) {
579+
ARROW_RETURN_NOT_OK(visit(start, end - start));
580+
}
581+
if (left_end <= right_end) {
582+
left_run = left_reader.NextRun();
583+
}
584+
if (right_end <= left_end) {
585+
right_run = right_reader.NextRun();
586+
}
587+
}
588+
return Status::OK();
589+
}
590+
507591
template <typename Visit>
508592
inline void VisitSetBitRunsVoid(const uint8_t* bitmap, int64_t offset, int64_t length,
509593
Visit&& visit) {

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

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -484,6 +484,51 @@ 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, VisitTwoSetBitRuns) {
506+
auto left = BitmapFromString("11110000 11111100");
507+
auto right = BitmapFromString("11001100 00111111");
508+
509+
std::vector<SetBitRun> runs;
510+
ASSERT_OK(VisitTwoSetBitRuns(left->data(), /*left_offset=*/1, right->data(),
511+
/*right_offset=*/2, /*length=*/12,
512+
[&](int64_t position, int64_t length) {
513+
runs.push_back({position, length});
514+
return Status::OK();
515+
}));
516+
ASSERT_THAT(runs, ElementsAreArray({SetBitRun{2, 1}, SetBitRun{8, 4}}));
517+
}
518+
519+
TEST_F(TestSetBitRunReader, VisitTwoSetBitRunsWithNullBitmap) {
520+
auto bitmap = BitmapFromString("01101101");
521+
522+
std::vector<SetBitRun> runs;
523+
ASSERT_OK(VisitTwoSetBitRuns(nullptr, /*left_offset=*/0, bitmap->data(),
524+
/*right_offset=*/1, /*length=*/6,
525+
[&](int64_t position, int64_t length) {
526+
runs.push_back({position, length});
527+
return Status::OK();
528+
}));
529+
ASSERT_THAT(runs, ElementsAreArray({SetBitRun{0, 2}, SetBitRun{3, 2}}));
530+
}
531+
487532
// Tests for BitRunReader.
488533

489534
TEST(BitRunReader, ZeroLength) {

0 commit comments

Comments
 (0)