[Request] standard library BitVector container

[sa-]

Edit: look at the comments for the updated version

---

To implement the Apache Arrow spec, one of the components that you need is a validity bitmap. This function is a stepping stone to implement one, and seeing as no one has posted here yet I thought I would do the honours :D

from Vector import DynamicVector

var a = DynamicVector[Bool]()
a.push_back(True)
a.push_back(True)
a.push_back(True)
a.push_back(True)
a.push_back(True)
a.push_back(True)
a.push_back(True)
a.push_back(True)

fn bitpack_bools(bools: DynamicVector[Bool]) -> DynamicVector[UI8]:
    let num_bytes = (len(bools) + 7) // 8
    
    var packed = DynamicVector[UI8](num_bytes)
    packed.resize(num_bytes)
    
    for i in range(len(bools)):
        let idx = (i // 8)
        let bit = 1 if (bools[i]) else 0  # would prefer a branchless way of doing this
        packed[idx] |= bit << (i % 8)
    
    return packed

bitpacked = bitpack_bools(a)

for i in range(len(bitpacked)):
    print(bitpacked[i])

This outputs 255, as one might expect.

Of course, it might be the case that DynamicVector[Bool] is already bitpacked, but this function could be used with other sorts of iterables.

Any feedback or comments are always welcome. Happy Friday!

[lattner]

Mojo should be really great at at this. I'd recommend implementing a new BitVector (you pick the name) type, instead of overloading DynamicVector (there are analogies to C++ vector<bool> specialization which has well known problems)

[sa-]

Followed by an implementation of Buffers for fixed-width primitive types:

struct ArrowFixedWidthBuffer[T: AnyType]:  
    # maybe use Dtype for T instead of AnyType, but DynamicVector uses AnyType
    var data: Pointer[UI8]
    var length: Int
    var mem_use: Int
    
    fn __init__(inout self, values: DynamicVector[T]):
        let byte_width = sizeof[T]()
        let num_bytes = len(values) * byte_width
        let num_bytes_with_padding = ((num_bytes + PADDING - 1) // PADDING) * PADDING
        let ui8_ptr = Pointer[UI8].aligned_alloc(PADDING, num_bytes_with_padding)
        let ptr = ui8_ptr.bitcast[T]()
        memset_zero(ptr, num_bytes_with_padding)
        
        for i in range(values.size):
            ptr.store(i, values[i])
        
        self.data = ui8_ptr
        self.length = len(values)
        self.mem_use = num_bytes_with_padding
        
    fn __getitem__(self, index: Int) -> T:
        # TODO: bounds check
        let ptr = self.data.bitcast[T]()
        return ptr.load(index)
    
    fn __len__(self) -> Int:
        return self.length

[sa-]

The above code is hereby licensed as MIT for what it's worth! Feel free to use it in/modify it for the std lib

[guidorice]

@sa- thanks for your codes which you open-sourced here. I am interested in arrow as well- geoarrow in particular. Question: do you think it would might be possible to simply use a mojo stdlib container like Tensor instead of your ArrowFixedWidthBuffer? The alignment and padding constraints of arrow buffers seem to add a bit of complexity.

[guidorice]

Also, both of those structs would need a __del__(): free() I think . 😄

[sa-]

It should be possible to use the Tensor struct if the alignment & padding constraints can be met. I don't know the reason behind those requirements but I do know that it's part of the arrow spec.

This was also written 6 months ago, about a week or two after Mojo's public launch so there was no Tensor type back then