ffi: create bool arrays

lang/cpp/include/vortex/array.hpp

@@ -12,6 +12,7 @@
 #include <vortex.h>
 
 #include <cstddef>
+#include <cstdint>
 #include <initializer_list>
 #include <memory>
 #include <span>
@@ -31,6 +32,45 @@ class Array;
 class StringView;
 class BytesView;
 
+/**
+ * Readonly view over bitpacked booleans.
+ *
+ * Bits are laid out LSB-first.
+ * "bit_offset" is in [0; 8) and lets a view start at a non-byte-aligned bit.
+ *
+ * Example:
+ * "view" holds 6 boolean elements starting at bit 2, first 5 set to "true",
+ * last is "false".
+ *
+ * uint8_t word = 0b01111100;
+ * BoolView view = {&word, 6, 2};
+ */
+struct BoolView {
+    const uint8_t *ptr = nullptr;
+
+    /*
+     * Number of elements (bits) in the view. This is not the number of uint8_t
+     * words in "ptr", use words() for that.
+     */
+    size_t elements = 0;
+
+    /*
+     * Bit offset of first element in "bytes".
+     * Example: if bit_offset is 2, first bit is at bytes[0] & (1 << 2).
+     */
+    size_t bit_offset = 0;
+
+    // Get bit at "index". Index is in [0; elements()).
+    inline bool operator[](size_t index) const {
+        return vx_bool_view_nth(*this, index);
+    }
+
+    // Number of uint8_t words storing elements.
+    inline size_t words() const {
+        return vx_bool_view_words(*this);
+    }
+};
+
 /*
  * Validity type tells us whether there are null/invalid values in an Array.
  */
@@ -102,8 +142,7 @@ class ValidityBits {
     ValidityBits(const Session &session, const vx_array *canonical);
 
     const vx_array *owner_ = nullptr;
-    const uint8_t *bits_ = nullptr;
-    size_t bit_offset_ = 0;
+    BoolView view_;
     bool all_invalid_ = false;
 };
 } // namespace detail
@@ -132,6 +171,17 @@ class Array {
         return primitive_raw(detail::to_ptype<T>(), data.data(), data.size(), validity);
     }
 
+    /**
+     * A Bool array copied from bitpacked storage.
+     *
+     * Example:
+     *
+     * std::vector<uint8_t> bits(2, 1);
+     * BoolView view = {bits, 0};
+     * auto array = Array::bool_array(view);
+     */
+    static Array bool_array(const BoolView &view, const Validity &validity = ValidityType::NonNullable);
+
     /**
      * Import an Arrow array. Consumes both "array" and "schema", do not use
      * or release them afterwards. For a record batch pass nullable = false.
@@ -296,17 +346,17 @@ class PrimitiveView {
 };
 
 /**
- * Read-only view over a Bool array. As Bool values are bit-packed, there's no
- * span. Read individual values with value(i).
+ * Read-only view over a Bool array.
  */
 template <>
 class PrimitiveView<bool> {
 public:
     /*
-     * Get raw value from this view. Values at null/invalid positions are
-     * unspecified.
+     * Get bitpacked storage for this view's values. Values at null/invalid
+     * positions are unspecified
      */
-    bool value(size_t index) const;
+    BoolView values() const;
+
     bool is_null(size_t index) const {
         return validity_.is_null(index);
     }

lang/cpp/src/array.cpp

@@ -78,11 +78,10 @@ bool ValidityBits::is_null(size_t index) const {
     if (all_invalid_) {
         return true;
     }
-    if (bits_ == nullptr) {
+    if (owner_ == nullptr) {
         return false;
     }
-    const size_t bit = bit_offset_ + index;
-    return (bits_[bit / 8] >> (bit % 8) & 1) == 0;
+    return !view_[index];
 }
 
 ValidityBits::ValidityBits(const Session &session, const vx_array *canonical) {
@@ -106,29 +105,28 @@ ValidityBits::ValidityBits(const Session &session, const vx_array *canonical) {
     vx_array_free(raw.array);
     throw_on_error(error);
 
-    bits_ = static_cast<const uint8_t *>(vx_array_data_ptr_bool(owner_, &bit_offset_, &error));
+    vx_bool_view view = vx_array_data_ptr_bool(owner_, &error);
     if (error != nullptr) {
         vx_array_free(owner_);
     }
     throw_on_error(error);
+    view_.bit_offset = view.bit_offset;
+    view_.ptr = view.ptr;
+    view_.elements = view.elements;
 }
 
 ValidityBits::ValidityBits(ValidityBits &&other) noexcept
-    : owner_(other.owner_), bits_(other.bits_), bit_offset_(other.bit_offset_),
-      all_invalid_(other.all_invalid_) {
+    : owner_(other.owner_), view_(std::move(other.view_)), all_invalid_(other.all_invalid_) {
     other.owner_ = nullptr;
-    other.bits_ = nullptr;
 }
 
 ValidityBits &ValidityBits::operator=(ValidityBits &&other) noexcept {
     if (this != &other) {
         vx_array_free(owner_);
         owner_ = other.owner_;
-        bits_ = other.bits_;
-        bit_offset_ = other.bit_offset_;
+        view_ = std::move(other.view_);
         all_invalid_ = other.all_invalid_;
         other.owner_ = nullptr;
-        other.bits_ = nullptr;
     }
     return *this;
 }
@@ -179,6 +177,22 @@ Array Array::primitive_raw(vx_ptype ptype, const void *data, size_t len, const V
     return Access::adopt<Array>(out);
 }
 
+Array Array::bool_array(const BoolView &view, const Validity &validity) {
+    std::optional<Array> keep_alive;
+    vx_validity raw {};
+    raw.type = static_cast<vx_validity_type>(validity.type());
+    if (validity.type() == ValidityType::FromArray) {
+        keep_alive = validity.array();
+        raw.array = Access::c_ptr(*keep_alive);
+    }
+
+    vx_error *error = nullptr;
+    vx_bool_view bool_view {view.ptr, view.elements, view.bit_offset};
+    const vx_array *out = vx_array_new_bool(&bool_view, &raw, &error);
+    throw_on_error(error);
+    return Access::adopt<Array>(out);
+}
+
 Array Array::from_arrow(const Session &session, ArrowArray *array, ArrowSchema *schema, bool nullable) {
     vx_error *error = nullptr;
     const vx_array *out = vx_array_from_arrow(Access::c_ptr(session), array, schema, nullable, &error);
@@ -333,13 +347,11 @@ Scalar Array::scalar_at(const Session &session, size_t index) const {
     return Access::adopt<Scalar>(scalar);
 }
 
-bool PrimitiveView<bool>::value(size_t i) const {
-    if (i >= size_) {
-        throw VortexException("index " + std::to_string(i) + " out of bounds for view of size " +
-                                  std::to_string(size_),
-                              ErrorCode::OutOfBounds);
-    }
-    return vx_array_get_bool(Access::c_ptr(canonical_), i);
+BoolView PrimitiveView<bool>::values() const {
+    vx_error *error = nullptr;
+    const vx_bool_view view = vx_array_data_ptr_bool(Access::c_ptr(canonical_), &error);
+    throw_on_error(error);
+    return {view.ptr, view.elements, view.bit_offset};
 }
 
 std::string_view StringView::operator[](size_t i) const {

lang/cpp/tests/array.cpp

@@ -170,6 +170,46 @@ TEST_CASE("Slice", "[array]") {
     REQUIRE_THROWS_AS(a.slice(2, 100), VortexException);
 }
 
+TEST_CASE("Bool array", "[array]") {
+    Session session;
+
+    constexpr size_t ELEMENTS = 9;
+    std::array<bool, ELEMENTS> data = {true, false, true, true, false, true, true, false, true};
+    std::vector<uint8_t> bitpacked(2);
+    for (size_t i = 0; i < ELEMENTS; ++i) {
+        if (data[i]) {
+            bitpacked[i / 8] |= 1 << (i % 8);
+        }
+    }
+
+    BoolView bool_view = {bitpacked.data(), ELEMENTS, 0};
+    Array array = Array::bool_array(bool_view);
+    REQUIRE(array.size() == data.size());
+    REQUIRE(array.has_dtype(DataTypeVariant::Bool));
+    REQUIRE_FALSE(array.nullable());
+
+    auto view = array.bools(session);
+    BoolView values = view.values();
+    for (size_t i = 0; i < values.elements; ++i) {
+        REQUIRE(values[i] == data[i]);
+    }
+
+    array = array.slice(3, array.size());
+    view = array.bools(session);
+    values = view.values();
+
+    Array roundtrip = Array::bool_array(values);
+    REQUIRE(roundtrip.size() == view.size());
+    REQUIRE_FALSE(roundtrip.nullable());
+
+    auto roundtrip_view = roundtrip.bools(session);
+    BoolView roundtrip_values = roundtrip_view.values();
+    for (size_t i = 0; i < view.size(); ++i) {
+        REQUIRE(roundtrip_values[i] == values[i]);
+        REQUIRE(roundtrip_values[i] == data[i + 3]);
+    }
+}
+
 TEST_CASE("Error with a code", "[array]") {
     std::vector<int16_t> data = {0};
     Array a = Array::primitive<int16_t>(data);

lang/cpp/tests/expression.cpp

@@ -79,13 +79,13 @@ TEST_CASE("Operator overloading", "[expr]") {
     Array array = Array::primitive<uint32_t>(data);
 
     Array applied = array.apply(expr::root() == expr::lit<uint32_t>(2));
-    auto bits = applied.bools(session);
-    REQUIRE(bits.size() == data.size());
-    REQUIRE_FALSE(bits.value(0));
-    REQUIRE(bits.value(1));
-    REQUIRE(bits.value(2));
-    REQUIRE_FALSE(bits.value(3));
-    REQUIRE_THROWS_AS(bits.value(data.size()), VortexException);
+    PrimitiveView<bool> view = applied.bools(session);
+    REQUIRE(view.size() == data.size());
+    BoolView values = view.values();
+    REQUIRE_FALSE(values[0]);
+    REQUIRE(values[1]);
+    REQUIRE(values[2]);
+    REQUIRE_FALSE(values[3]);
 }
 
 TEST_CASE("Apply error", "[expr]") {

vortex-ffi/cbindgen.toml

@@ -67,6 +67,12 @@ typedef struct ArrowSchema FFI_ArrowSchema;
 typedef struct ArrowArray FFI_ArrowArray;
 typedef struct ArrowArrayStream FFI_ArrowArrayStream;
 #endif
+
+// Number of uint8_t words holding view's elements
+#define vx_bool_view_words(X) (((X).elements + (X).bit_offset + 7) / 8)
+// I'th element (bit) in view, LSB-first. Already accounts for bit_offset
+#define vx_bool_view_nth(X, I) \
+    ((((X).ptr[((I) + (X).bit_offset) / 8] & (1 << (((I) + (X).bit_offset) % 8))) != 0))
 """
 
 trailer = """

vortex-ffi/cinclude/vortex.h

@@ -60,6 +60,12 @@ typedef struct ArrowArray FFI_ArrowArray;
 typedef struct ArrowArrayStream FFI_ArrowArrayStream;
 #endif
 
+// Number of uint8_t words holding view's elements
+#define vx_bool_view_words(X) (((X).elements + (X).bit_offset + 7) / 8)
+// I'th element (bit) in view, LSB-first. Already accounts for bit_offset
+#define vx_bool_view_nth(X, I)                                                                               \
+    ((((X).ptr[((I) + (X).bit_offset) / 8] & (1 << (((I) + (X).bit_offset) % 8))) != 0))
+
 #include <stdarg.h>
 #include <stdbool.h>
 #include <stddef.h>
@@ -475,6 +481,38 @@ typedef struct {
     const vx_array *array;
 } vx_validity;
 
+/**
+ * Readonly view over bitpacked booleans.
+ *
+ * "elements" is the number of bits/elements. Use vx_bool_view_words(view) to
+ * get the number of uint8_t words.
+ * Bits are laid out LSB-first.
+ *
+ * "bit_offset" is in [0; 8) and lets a view start at a non-byte-aligned bit.
+ * Use vx_bool_view_nth(view, index) macro to read a single element.
+ *
+ * Example:
+ * "view" holds 6 boolean elements, bit_offset=2, first 5 elements are "true",
+ * last is "false".
+ *
+ * uint8_t word = 0b01111100;
+ * vx_bool_view view = {&word, 6, 2};
+ */
+typedef struct {
+    /**
+     * Element 0 is bit "bit_offset" of "ptr".
+     */
+    const uint8_t *ptr;
+    /**
+     * Number of elements represented by "ptr".
+     */
+    size_t elements;
+    /**
+     * Bit offset of element 0 within the first byte of "ptr".
+     */
+    size_t bit_offset;
+} vx_bool_view;
+
 /**
  * A non owning view over a byte range.
  */
@@ -690,6 +728,25 @@ const vx_array *vx_array_new_primitive(vx_ptype ptype,
                                        const vx_validity *validity,
                                        vx_error **error);
 
+/**
+ * Create a new Bool array from vx_bool_view.
+ *
+ * Example:
+ *
+ * a Bool array with 9 elements, first 8 are "true", last is "false".
+ *
+ * const vx_error* error = NULL;
+ * vx_validity validity = {};
+ * validity.type = VX_VALIDITY_NON_NULLABLE;
+ *
+ * uint8_t words\[2\] = {0xff, 0}; // 11111111 00000000
+ * vx_bool_view view = {words, 9, 0};
+ *
+ * const vx_array* array = vx_array_new_bool(&view, &validity, &error);
+ * vx_array_free(array);
+ */
+const vx_array *vx_array_new_bool(const vx_bool_view *view, const vx_validity *validity, vx_error **error);
+
 /**
  * Create a Vortex array by importing an Arrow array via the Arrow C Data Interface.
  *
@@ -778,15 +835,25 @@ const vx_array *vx_array_canonicalize(const vx_session *session, const vx_array
 const void *vx_array_data_ptr_primitive(const vx_array *array, vx_error **error_out);
 
 /**
- * Return a pointer to the bitpacked buffer of a canonical Bool array.
- * Pointer is valid as long as "array" is valid.
- *
- * Writes bit offset of the first element into "bit_offset_out".
- * "bit_offset_out" must not be NULL.
+ * Return vx_bool_view for a canonical Bool array.
+ * View is valid as long as "array" is valid.
  *
  * Errors if array is not a canonical Bool.
+ *
+ * Example:
+ *
+ * vx_validity validity = {};
+ * validity.type = VX_VALIDITY_NON_NULLABLE;
+ *
+ * uint8_t words\[2\] = {0xff, 0}; // 11111111 00000000
+ * vx_bool_view view = {words, 9, 0};
+ *
+ * const vx_array* array = vx_array_new_bool(&view, &validity, &error);
+ * vx_bool_view other = vx_array_data_ptr_bool(array, &error);
+ *
+ * vx_array_free(array);
  */
-const void *vx_array_data_ptr_bool(const vx_array *array, size_t *bit_offset_out, vx_error **error_out);
+vx_bool_view vx_array_data_ptr_bool(const vx_array *array, vx_error **error_out);
 
 /**
  * Apply the expression to the array, wrapping it with a ScalarFnArray.