diff --git a/cpp/src/arrow/util/bitmap_ops.cc b/cpp/src/arrow/util/bitmap_ops.cc index 33a95150d53..c8571061d19 100644 --- a/cpp/src/arrow/util/bitmap_ops.cc +++ b/cpp/src/arrow/util/bitmap_ops.cc @@ -17,11 +17,14 @@ #include "arrow/util/bitmap_ops.h" +#include #include #include #include #include #include +#include +#include #include "arrow/buffer.h" #include "arrow/result.h" @@ -32,8 +35,7 @@ #include "arrow/util/bitmap_writer.h" #include "arrow/util/logging_internal.h" -namespace arrow { -namespace internal { +namespace arrow::internal { int64_t CountSetBits(const uint8_t* data, int64_t bit_offset, int64_t length) { constexpr int64_t pop_len = sizeof(uint64_t) * 8; @@ -106,8 +108,6 @@ int64_t CountAndSetBits(const uint8_t* left_bitmap, int64_t left_offset, namespace { -enum class TransferMode : bool { Copy, Invert }; - // Reverse all bits from entire byte(uint8) uint8_t ReverseUint8(uint8_t num) { num = ((num & 0xf0) >> 4) | ((num & 0x0f) << 4); @@ -123,29 +123,97 @@ uint8_t GetReversedBlock(uint8_t block_left, uint8_t block_right, uint8_t length return ReverseUint8(((block_right << 8) + block_left) >> length); } -template -void TransferBitmap(const uint8_t* data, int64_t offset, int64_t length, +/// Map output from readers and save it with the writer. +/// +/// All readers and writer must span over the same number of values. +/// +/// @tparam Op a function of as many input as there are readers. +template +void MapReadersWriter(auto&& writer, auto&& reader, auto&&... readers) { + constexpr auto kReaderCount = sizeof...(readers) + 1; + constexpr auto op = Op{}; + + // Need a real function so that the fold expression remains valid in release + [[maybe_unused]] const auto check_eq = [](auto a, auto b) { ARROW_DCHECK_EQ(a, b); }; + + auto nwords = reader.words(); + ((check_eq(readers.words(), nwords)), ...); + while (nwords--) { + writer.PutNextWord(op(reader.NextWord(), readers.NextWord()...)); + } + + auto nbytes = reader.trailing_bytes(); + ((check_eq(readers.trailing_bytes(), nbytes)), ...); + while (nbytes--) { + int valid_bits = 0; + std::array bytes = {}; + { + auto b = bytes.begin(); + *b++ = reader.NextTrailingByte(valid_bits); + [[maybe_unused]] auto read = [&](auto& r) { + int vb = 0; + *b++ = r.NextTrailingByte(vb); + check_eq(vb, valid_bits); + }; + (read(readers), ...); + } + writer.PutNextTrailingByte(std::apply(op, bytes), valid_bits); + } +} + +template +struct BitmapPtr { + Byte* data; + int64_t offset; + + BitmapPtr operator+(int64_t extra) { return {.data = data, .offset = offset + extra}; } +}; + +using BitmapConstPtr = BitmapPtr; +using BitmapMutPtr = BitmapPtr; + +/// Map inputs with a given operation and sace to output. +/// +/// This function assumes general non bit-aligned input and outputs. +/// It will first process less than a byte in order to bit-align the writer, and then +/// keep on going with an aligned writer. +/// Aligning the writer is what delivers significant speedup. +/// +/// @tparam Op a function of as many input as there are readers. +template +void FastMapReadersWriter(BitmapMutPtr out, int64_t length, auto&&... in) { + const int64_t out_bit_offset = out.offset % 8; + + if (length == 0) { + return; + } else if (out_bit_offset) { + using Reader = internal::BitmapWordReader; + using Writer = internal::BitmapWordWriter; + + const auto count = std::min(8 - out_bit_offset, length); + auto writer = Writer(out.data, out.offset, count); + MapReadersWriter(writer, Reader(in.data, in.offset, count)...); + FastMapReadersWriter(out + count, length - count, in + count...); + } else { + using Reader = internal::BitmapWordReader; + using Writer = internal::BitmapWordWriter; + + auto writer = Writer(out.data, out.offset, length); + MapReadersWriter(writer, Reader(in.data, in.offset, length)...); + } +} + +template +void MapBitmapUnary(const uint8_t* data, int64_t offset, int64_t length, int64_t dest_offset, uint8_t* dest) { - int64_t bit_offset = offset % 8; - int64_t dest_bit_offset = dest_offset % 8; + const int64_t bit_offset = offset % 8; + const int64_t dest_bit_offset = dest_offset % 8; if (bit_offset || dest_bit_offset) { - auto reader = internal::BitmapWordReader(data, offset, length); - auto writer = internal::BitmapWordWriter(dest, dest_offset, length); - - auto nwords = reader.words(); - while (nwords--) { - auto word = reader.NextWord(); - writer.PutNextWord(mode == TransferMode::Invert ? ~word : word); - } - auto nbytes = reader.trailing_bytes(); - while (nbytes--) { - int valid_bits; - auto byte = reader.NextTrailingByte(valid_bits); - writer.PutNextTrailingByte(mode == TransferMode::Invert ? ~byte : byte, valid_bits); - } - } else if (length) { - int64_t num_bytes = bit_util::BytesForBits(length); + FastMapReadersWriter({.data = dest, .offset = dest_offset}, length, + BitmapConstPtr{.data = data, .offset = offset}); + } else if (length > 0) { + const int64_t num_bytes = bit_util::BytesForBits(length); // Shift by its byte offset data += offset / 8; @@ -155,18 +223,19 @@ void TransferBitmap(const uint8_t* data, int64_t offset, int64_t length, // E.g., if trailing_bits = 5, last byte should be // - low 3 bits: new bits from last byte of data buffer // - high 5 bits: old bits from last byte of dest buffer - int64_t trailing_bits = num_bytes * 8 - length; - uint8_t trail_mask = (1U << (8 - trailing_bits)) - 1; + const int64_t trailing_bits = num_bytes * 8 - length; + const uint8_t trail_mask = (1U << (8 - trailing_bits)) - 1; uint8_t last_data; - if (mode == TransferMode::Invert) { - for (int64_t i = 0; i < num_bytes - 1; i++) { - dest[i] = static_cast(~(data[i])); - } - last_data = ~data[num_bytes - 1]; - } else { + if constexpr (std::is_same_v, std::identity>) { std::memcpy(dest, data, static_cast(num_bytes - 1)); last_data = data[num_bytes - 1]; + } else { + constexpr auto op = Op{}; + for (int64_t i = 0; i < num_bytes - 1; i++) { + dest[i] = static_cast(op(data[i])); + } + last_data = op(data[num_bytes - 1]); } // Set last byte @@ -218,29 +287,29 @@ void ReverseBlockOffsets(const uint8_t* data, int64_t offset, int64_t length, } } -} // namespace - -template -Result> TransferBitmap(MemoryPool* pool, const uint8_t* data, +template +Result> MapBitmapUnary(MemoryPool* pool, const uint8_t* data, int64_t offset, int64_t length, int64_t out_offset) { const int64_t phys_bits = length + out_offset; ARROW_ASSIGN_OR_RAISE(auto buffer, AllocateEmptyBitmap(phys_bits, pool)); uint8_t* dest = buffer->mutable_data(); - TransferBitmap(data, offset, length, out_offset, dest); + MapBitmapUnary(data, offset, length, out_offset, dest); return buffer; } +} // namespace + void CopyBitmap(const uint8_t* data, int64_t offset, int64_t length, uint8_t* dest, int64_t dest_offset) { - TransferBitmap(data, offset, length, dest_offset, dest); + MapBitmapUnary(data, offset, length, dest_offset, dest); } void InvertBitmap(const uint8_t* data, int64_t offset, int64_t length, uint8_t* dest, int64_t dest_offset) { - TransferBitmap(data, offset, length, dest_offset, dest); + MapBitmapUnary>(data, offset, length, dest_offset, dest); } void ReverseBitmap(const uint8_t* data, int64_t offset, int64_t length, uint8_t* dest, @@ -251,13 +320,12 @@ void ReverseBitmap(const uint8_t* data, int64_t offset, int64_t length, uint8_t* Result> CopyBitmap(MemoryPool* pool, const uint8_t* data, int64_t offset, int64_t length, int64_t out_offset) { - return TransferBitmap(pool, data, offset, length, out_offset); + return MapBitmapUnary(pool, data, offset, length, out_offset); } Result> InvertBitmap(MemoryPool* pool, const uint8_t* data, int64_t offset, int64_t length) { - return TransferBitmap(pool, data, offset, length, - /*out_offset=*/0); + return MapBitmapUnary>(pool, data, offset, length, /*out_offset=*/0); } Result> ReverseBitmap(MemoryPool* pool, const uint8_t* data, @@ -361,11 +429,11 @@ Result> OptionalBitmapAnd(MemoryPool* pool, namespace { -template