This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-nanoarrow.git
The following commit(s) were added to refs/heads/main by this push:
new ad83497 feat: Add `ArrowBitsUnpackInt32()` (#278)
ad83497 is described below
commit ad83497bc8daaec17005fa92b8690c97195bef52
Author: Dewey Dunnington <[email protected]>
AuthorDate: Thu Aug 17 15:38:28 2023 -0300
feat: Add `ArrowBitsUnpackInt32()` (#278)
As a follow-up to #276. The `int32` version is useful because R uses
32-bit integers to represent boolean (i.e., logical) arrays. This
results in a significant speedup in boolean conversion!
@WillAyd: I updated a few things that you *just* added (Sorry! 😬 ):
- I changed `Bitmap` -> `Bits` and removed `Unsafe` to make it more
consistent with the other functions that accept `const uint8_t* bits`
- I updated the test function so that it tests both the int32 and int8
types at once
Before this PR:
``` r
library(nanoarrow)
lgls <- nanoarrow:::vec_gen(logical(), 1e6)
bool_array <- as_nanoarrow_array(lgls)
bool_array_arrow <- arrow::as_arrow_array(bool_array)
bench::mark(
convert_array(bool_array, logical()),
as.vector(bool_array_arrow),
as.logical(lgls)
)
#> # A tibble: 3 × 6
#> expression min median `itr/sec` mem_alloc
`gc/sec`
#> <bch:expr> <bch> <bch:> <dbl> <bch:byt>
<dbl>
#> 1 convert_array(bool_array, logical()) 556µs 749µs 1.33e3 3.82MB
156.
#> 2 as.vector(bool_array_arrow) 558µs 780µs 1.30e3 3.82MB
144.
#> 3 as.logical(lgls) 0 1ns 2.28e8 0B
0
bench::mark(
convert_array(bool_array, integer()),
as.integer(lgls)
)
#> # A tibble: 2 × 6
#> expression min median `itr/sec` mem_alloc
`gc/sec`
#> <bch:expr> <bch> <bch:> <dbl> <bch:byt>
<dbl>
#> 1 convert_array(bool_array, integer()) 733µs 912µs 1093. 3.81MB
167.
#> 2 as.integer(lgls) 615µs 788µs 1273. 3.81MB
182.
```
After this PR:
``` r
library(nanoarrow)
lgls <- nanoarrow:::vec_gen(logical(), 1e6)
bool_array <- as_nanoarrow_array(lgls)
bool_array_arrow <- arrow::as_arrow_array(bool_array)
bench::mark(
convert_array(bool_array, logical()),
as.vector(bool_array_arrow),
as.logical(lgls)
)
#> # A tibble: 3 × 6
#> expression min median `itr/sec` mem_alloc
`gc/sec`
#> <bch:expr> <bch> <bch:> <dbl> <bch:byt>
<dbl>
#> 1 convert_array(bool_array, logical()) 105µs 308µs 3.21e3 3.83MB
367.
#> 2 as.vector(bool_array_arrow) 559µs 772µs 1.30e3 3.82MB
143.
#> 3 as.logical(lgls) 0 0 5.87e8 0B
0
bench::mark(
convert_array(bool_array, integer()),
as.integer(lgls)
)
#> # A tibble: 2 × 6
#> expression min median `itr/sec` mem_alloc
`gc/sec`
#> <bch:expr> <bch> <bch:> <dbl> <bch:byt>
<dbl>
#> 1 convert_array(bool_array, integer()) 104µs 310µs 3181. 3.81MB
423.
#> 2 as.integer(lgls) 615µs 784µs 1278. 3.81MB
142.
```
<sup>Created on 2023-08-17 with [reprex
v2.0.2](https://reprex.tidyverse.org)</sup>
---
r/src/materialize_int.h | 13 +++++++
r/src/materialize_lgl.h | 5 ++-
src/nanoarrow/buffer_inline.h | 59 ++++++++++++++++++++++++++++---
src/nanoarrow/buffer_test.cc | 82 ++++++++++++++++++++++++++-----------------
src/nanoarrow/nanoarrow.h | 12 ++++---
5 files changed, 127 insertions(+), 44 deletions(-)
diff --git a/r/src/materialize_int.h b/r/src/materialize_int.h
index 3d8a392..ffd39dd 100644
--- a/r/src/materialize_int.h
+++ b/r/src/materialize_int.h
@@ -56,6 +56,19 @@ static inline int nanoarrow_materialize_int(struct
ArrayViewSlice* src,
}
break;
case NANOARROW_TYPE_BOOL:
+ ArrowBitsUnpackInt32(
+ src->array_view->buffer_views[1].data.as_uint8 + raw_src_offset,
raw_src_offset,
+ dst->length, result + dst->offset);
+
+ // Set any nulls to NA_LOGICAL
+ if (is_valid != NULL && src->array_view->array->null_count != 0) {
+ for (R_xlen_t i = 0; i < dst->length; i++) {
+ if (!ArrowBitGet(is_valid, raw_src_offset + i)) {
+ result[dst->offset + i] = NA_LOGICAL;
+ }
+ }
+ }
+ break;
case NANOARROW_TYPE_INT8:
case NANOARROW_TYPE_UINT8:
case NANOARROW_TYPE_INT16:
diff --git a/r/src/materialize_lgl.h b/r/src/materialize_lgl.h
index cf9e149..a23802c 100644
--- a/r/src/materialize_lgl.h
+++ b/r/src/materialize_lgl.h
@@ -40,9 +40,8 @@ static int nanoarrow_materialize_lgl(struct ArrayViewSlice*
src, struct VectorSl
}
break;
case NANOARROW_TYPE_BOOL:
- for (R_xlen_t i = 0; i < dst->length; i++) {
- result[dst->offset + i] = ArrowBitGet(data_buffer, src->offset + i);
- }
+ ArrowBitsUnpackInt32(data_buffer, raw_src_offset, dst->length,
+ result + dst->offset);
// Set any nulls to NA_LOGICAL
if (is_valid != NULL && src->array_view->array->null_count != 0) {
diff --git a/src/nanoarrow/buffer_inline.h b/src/nanoarrow/buffer_inline.h
index efb5ca7..232dd42 100644
--- a/src/nanoarrow/buffer_inline.h
+++ b/src/nanoarrow/buffer_inline.h
@@ -222,7 +222,18 @@ static inline int64_t _ArrowBytesForBits(int64_t bits) {
return (bits >> 3) + ((bits & 7) != 0);
}
-static inline void _ArrowBitmapUnpackInt8(const uint8_t word, int8_t* out) {
+static inline void _ArrowBitsUnpackInt8(const uint8_t word, int8_t* out) {
+ out[0] = (word >> 0) & 1;
+ out[1] = (word >> 1) & 1;
+ out[2] = (word >> 2) & 1;
+ out[3] = (word >> 3) & 1;
+ out[4] = (word >> 4) & 1;
+ out[5] = (word >> 5) & 1;
+ out[6] = (word >> 6) & 1;
+ out[7] = (word >> 7) & 1;
+}
+
+static inline void _ArrowBitsUnpackInt32(const uint8_t word, int32_t* out) {
out[0] = (word >> 0) & 1;
out[1] = (word >> 1) & 1;
out[2] = (word >> 2) & 1;
@@ -247,8 +258,47 @@ static inline int8_t ArrowBitGet(const uint8_t* bits,
int64_t i) {
return (bits[i >> 3] >> (i & 0x07)) & 1;
}
-static inline void ArrowBitmapUnpackInt8Unsafe(const uint8_t* bits, int64_t
start_offset,
- int64_t length, int8_t* out) {
+static inline void ArrowBitsUnpackInt8(const uint8_t* bits, int64_t
start_offset,
+ int64_t length, int8_t* out) {
+ if (length == 0) {
+ return;
+ }
+
+ const int64_t i_begin = start_offset;
+ const int64_t i_end = start_offset + length;
+ const int64_t i_last_valid = i_end - 1;
+
+ const int64_t bytes_begin = i_begin / 8;
+ const int64_t bytes_last_valid = i_last_valid / 8;
+
+ if (bytes_begin == bytes_last_valid) {
+ for (int i = 0; i < length; i++) {
+ out[i] = ArrowBitGet(&bits[bytes_begin], i + i_begin % 8);
+ }
+
+ return;
+ }
+
+ // first byte
+ for (int i = 0; i < 8 - (i_begin % 8); i++) {
+ *out++ = ArrowBitGet(&bits[bytes_begin], i + i_begin % 8);
+ }
+
+ // middle bytes
+ for (int64_t i = bytes_begin + 1; i < bytes_last_valid; i++) {
+ _ArrowBitsUnpackInt8(bits[i], out);
+ out += 8;
+ }
+
+ // last byte
+ const int bits_remaining = i_end % 8 == 0 ? 8 : i_end % 8;
+ for (int i = 0; i < bits_remaining; i++) {
+ *out++ = ArrowBitGet(&bits[bytes_last_valid], i);
+ }
+}
+
+static inline void ArrowBitsUnpackInt32(const uint8_t* bits, int64_t
start_offset,
+ int64_t length, int32_t* out) {
if (length == 0) {
return;
}
@@ -261,7 +311,6 @@ static inline void ArrowBitmapUnpackInt8Unsafe(const
uint8_t* bits, int64_t star
const int64_t bytes_last_valid = i_last_valid / 8;
if (bytes_begin == bytes_last_valid) {
- // count bits within a single byte
for (int i = 0; i < length; i++) {
out[i] = ArrowBitGet(&bits[bytes_begin], i + i_begin % 8);
}
@@ -276,7 +325,7 @@ static inline void ArrowBitmapUnpackInt8Unsafe(const
uint8_t* bits, int64_t star
// middle bytes
for (int64_t i = bytes_begin + 1; i < bytes_last_valid; i++) {
- _ArrowBitmapUnpackInt8(bits[i], out);
+ _ArrowBitsUnpackInt32(bits[i], out);
out += 8;
}
diff --git a/src/nanoarrow/buffer_test.cc b/src/nanoarrow/buffer_test.cc
index 1907d37..7ca4874 100644
--- a/src/nanoarrow/buffer_test.cc
+++ b/src/nanoarrow/buffer_test.cc
@@ -272,83 +272,101 @@ TEST(BitmapTest, BitmapTestElement) {
}
template <int offset, int length>
-void TestArrowBitmapUnpackInt8Unsafe(const uint8_t* bitmap, int8_t* out,
- std::vector<uint8_t> expected) {
- ArrowBitmapUnpackInt8Unsafe(bitmap, offset, length, out);
- for (int i = 0; i < expected.size(); i++) {
+void TestArrowBitmapUnpackUnsafe(const uint8_t* bitmap, std::vector<int8_t>
expected) {
+ int8_t out[length];
+ int32_t out32[length];
+ memset(out, 0, sizeof(out));
+ memset(out32, 0, sizeof(out32));
+
+ ASSERT_EQ(length, expected.size());
+
+ ArrowBitsUnpackInt8(bitmap, offset, length, out);
+ for (int i = 0; i < length; i++) {
EXPECT_EQ(out[i], expected[i]);
}
+
+ ArrowBitsUnpackInt32(bitmap, offset, length, out32);
+ for (int i = 0; i < length; i++) {
+ EXPECT_EQ(out32[i], expected[i]);
+ }
}
-TEST(BitmapTest, BitmapTestBitmapUnpackInt8Unsafe) {
+TEST(BitmapTest, BitmapTestBitmapUnpack) {
uint8_t bitmap[3];
- int8_t result[sizeof(bitmap) * 8];
+ int64_t n_values = sizeof(bitmap) * 8;
+ int8_t result[n_values];
+ int32_t result32[n_values];
+ // Basic test of a validity buffer that is all true
memset(bitmap, 0xff, sizeof(bitmap));
- ArrowBitmapUnpackInt8Unsafe(bitmap, 0, sizeof(result), result);
- for (int i = 0; i < sizeof(result); i++) {
+ memset(result, 0, sizeof(result));
+ memset(result32, 0, sizeof(result32));
+
+ ArrowBitsUnpackInt8(bitmap, 0, sizeof(result), result);
+ for (int i = 0; i < n_values; i++) {
EXPECT_EQ(result[i], 1);
}
+ ArrowBitsUnpackInt32(bitmap, 0, sizeof(result), result32);
+ for (int i = 0; i < n_values; i++) {
+ EXPECT_EQ(result32[i], 1);
+ }
+
+ // Ensure that the first byte/middle byte/last byte logic is correct
+ // Note that TestArrowBitmapUnpack tests both the int8 and int32 version
bitmap[0] = 0x93; // 10010011
bitmap[1] = 0x55; // 01010101
bitmap[2] = 0xaa; // 10101010
// offset 0, length boundary, one byte
- TestArrowBitmapUnpackInt8Unsafe<0, 8>(bitmap, result, {1, 1, 0, 0, 1, 0, 0,
1});
+ TestArrowBitmapUnpackUnsafe<0, 8>(bitmap, {1, 1, 0, 0, 1, 0, 0, 1});
// offset 0, length boundary, different bytes
- TestArrowBitmapUnpackInt8Unsafe<0, 16>(
- bitmap, result, {1, 1, 0, 0, 1, 0, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0});
+ TestArrowBitmapUnpackUnsafe<0, 16>(bitmap,
+ {1, 1, 0, 0, 1, 0, 0, 1, 1, 0, 1, 0, 1,
0, 1, 0});
// offset 0, length non-boundary, one byte
- TestArrowBitmapUnpackInt8Unsafe<0, 5>(bitmap, result, {1, 1, 0, 0, 1});
+ TestArrowBitmapUnpackUnsafe<0, 5>(bitmap, {1, 1, 0, 0, 1});
// offset boundary, length boundary, one byte
- TestArrowBitmapUnpackInt8Unsafe<8, 8>(bitmap, result, {1, 0, 1, 0, 1, 0, 1,
0});
+ TestArrowBitmapUnpackUnsafe<8, 8>(bitmap, {1, 0, 1, 0, 1, 0, 1, 0});
// offset boundary, length boundary, different bytes
- TestArrowBitmapUnpackInt8Unsafe<8, 16>(
- bitmap, result, {1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1});
+ TestArrowBitmapUnpackUnsafe<8, 16>(bitmap,
+ {1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0,
1, 0, 1});
// offset boundary, length non-boundary, one byte
- TestArrowBitmapUnpackInt8Unsafe<8, 5>(bitmap, result, {1, 0, 1, 0, 1});
+ TestArrowBitmapUnpackUnsafe<8, 5>(bitmap, {1, 0, 1, 0, 1});
// offset boundary, length non-boundary, different bytes
- TestArrowBitmapUnpackInt8Unsafe<8, 13>(bitmap, result,
- {1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 0, 1,
0});
+ TestArrowBitmapUnpackUnsafe<8, 13>(bitmap, {1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 0,
1, 0});
// offset non-boundary, length boundary, one byte
- TestArrowBitmapUnpackInt8Unsafe<3, 5>(bitmap, result, {0, 1, 0, 0, 1});
+ TestArrowBitmapUnpackUnsafe<3, 5>(bitmap, {0, 1, 0, 0, 1});
// offset non-boundary, length boundary, different bytes
- TestArrowBitmapUnpackInt8Unsafe<3, 13>(bitmap, result,
- {0, 1, 0, 0, 1, 1, 0, 1, 0, 1, 0, 1,
0});
+ TestArrowBitmapUnpackUnsafe<3, 13>(bitmap, {0, 1, 0, 0, 1, 1, 0, 1, 0, 1, 0,
1, 0});
// offset non-boundary, length non-boundary, one byte
- TestArrowBitmapUnpackInt8Unsafe<3, 3>(bitmap, result, {0, 1, 0});
+ TestArrowBitmapUnpackUnsafe<3, 3>(bitmap, {0, 1, 0});
// offset non-boundary, length non-boundary, different bytes
- TestArrowBitmapUnpackInt8Unsafe<3, 11>(bitmap, result,
- {0, 1, 0, 0, 1, 1, 0, 1, 0, 1, 0});
+ TestArrowBitmapUnpackUnsafe<3, 11>(bitmap, {0, 1, 0, 0, 1, 1, 0, 1, 0, 1,
0});
// offset non-boundary non-first byte, length boundary, one byte
- TestArrowBitmapUnpackInt8Unsafe<11, 5>(bitmap, result, {0, 1, 0, 1, 0});
+ TestArrowBitmapUnpackUnsafe<11, 5>(bitmap, {0, 1, 0, 1, 0});
// offset non-boundary non-first byte, length boundary, different bytes
- TestArrowBitmapUnpackInt8Unsafe<11, 13>(bitmap, result,
- {0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 1, 0,
1});
+ TestArrowBitmapUnpackUnsafe<11, 13>(bitmap, {0, 1, 0, 1, 0, 0, 1, 0, 1, 0,
1, 0, 1});
// offset non-boundary non-first byte, length non-boundary, one byte
- TestArrowBitmapUnpackInt8Unsafe<11, 3>(bitmap, result, {0, 1, 0});
+ TestArrowBitmapUnpackUnsafe<11, 3>(bitmap, {0, 1, 0});
// offset non-boundary non-first byte, length non-boundary, different bytes
- TestArrowBitmapUnpackInt8Unsafe<11, 11>(bitmap, result,
- {0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 1, 0,
1});
+ TestArrowBitmapUnpackUnsafe<11, 11>(bitmap, {0, 1, 0, 1, 0, 0, 1, 0, 1, 0,
1});
// non-boundary, three byte span
- TestArrowBitmapUnpackInt8Unsafe<7, 11>(bitmap, result,
- {1, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1});
+ TestArrowBitmapUnpackUnsafe<7, 11>(bitmap, {1, 1, 0, 1, 0, 1, 0, 1, 0, 0,
1});
}
TEST(BitmapTest, BitmapTestSetTo) {
diff --git a/src/nanoarrow/nanoarrow.h b/src/nanoarrow/nanoarrow.h
index 09a80a8..6f4d905 100644
--- a/src/nanoarrow/nanoarrow.h
+++ b/src/nanoarrow/nanoarrow.h
@@ -680,6 +680,14 @@ static inline void ArrowBitsSetTo(uint8_t* bits, int64_t
start_offset, int64_t l
/// \brief Count true values in a bitmap
static inline int64_t ArrowBitCountSet(const uint8_t* bits, int64_t i_from,
int64_t i_to);
+/// \brief Extract int8 boolean values from a range in a bitmap
+static inline void ArrowBitsUnpackInt8(const uint8_t* bits, int64_t
start_offset,
+ int64_t length, int8_t* out);
+
+/// \brief Extract int32 boolean values from a range in a bitmap
+static inline void ArrowBitsUnpackInt32(const uint8_t* bits, int64_t
start_offset,
+ int64_t length, int32_t* out);
+
/// \brief Initialize an ArrowBitmap
///
/// Initialize the builder's buffer, empty its cache, and reset the size to
zero
@@ -716,10 +724,6 @@ static inline ArrowErrorCode ArrowBitmapAppend(struct
ArrowBitmap* bitmap,
static inline void ArrowBitmapAppendUnsafe(struct ArrowBitmap* bitmap,
uint8_t bits_are_set, int64_t
length);
-/// \brief Extract boolean values from a range in a bitmap
-static inline void ArrowBitmapUnpackInt8Unsafe(const uint8_t* bits, int64_t
start_offset,
- int64_t length, int8_t* out);
-
/// \brief Append boolean values encoded as int8_t to a bitmap
///
/// The values must all be 0 or 1.
