Method gather_nulls

arrow-array/src/array/union_array.rs:608–667

Computes the logical nulls for a sparse or dense union, by gathering individual bits from the null buffer of the selected field

(&self, nulls: Vec<(i8, NullBuffer)>)

    /// Computes the logical nulls for a sparse or dense union, by gathering individual bits from the null buffer of the selected field
    fn gather_nulls(&self, nulls: Vec<(i8, NullBuffer)>) -> BooleanBuffer {
        let one_null = NullBuffer::new_null(1);
        let one_valid = NullBuffer::new_valid(1);

        // Unsafe code below depend on it:
        // To remove one branch from the loop, if the a type_id is not utilized, or it's logical_nulls is None/all set,
        // we use a null buffer of len 1 and a index_mask of 0, or the true null buffer and usize::MAX otherwise.
        // We then unconditionally access the null buffer with index & index_mask,
        // which always return 0 for the 1-len buffer, or the true index unchanged otherwise
        // We also use a 256 array, so llvm knows that `type_id as u8 as usize` is always in bounds
        let mut logical_nulls_array = [(&one_valid, Mask::Zero); 256];

        for (type_id, nulls) in &nulls {
            if nulls.null_count() == nulls.len() {
                // Similarly, if all values are null, use a 1-null null-buffer to reduce cache pressure a bit
                logical_nulls_array[*type_id as u8 as usize] = (&one_null, Mask::Zero);
            } else {
                logical_nulls_array[*type_id as u8 as usize] = (nulls, Mask::Max);
            }
        }

        match &self.offsets {
            Some(offsets) => {
                assert_eq!(self.type_ids.len(), offsets.len());

                BooleanBuffer::collect_bool(self.type_ids.len(), |i| unsafe {
                    // SAFETY: BooleanBuffer::collect_bool calls us 0..self.type_ids.len()
                    let type_id = *self.type_ids.get_unchecked(i);
                    // SAFETY: We asserted that offsets len and self.type_ids len are equal
                    let offset = *offsets.get_unchecked(i);

                    let (nulls, offset_mask) = &logical_nulls_array[type_id as u8 as usize];

                    // SAFETY:
                    // If offset_mask is Max
                    // 1. Offset validity is checked at union creation
                    // 2. If the null buffer len equals it's array len is checked at array creation
                    // If offset_mask is Zero, the null buffer len is 1
                    nulls
                        .inner()
                        .value_unchecked(offset as usize & *offset_mask as usize)
                })
            }
            None => {
                BooleanBuffer::collect_bool(self.type_ids.len(), |index| unsafe {
                    // SAFETY: BooleanBuffer::collect_bool calls us 0..self.type_ids.len()
                    let type_id = *self.type_ids.get_unchecked(index);

                    let (nulls, index_mask) = &logical_nulls_array[type_id as u8 as usize];

                    // SAFETY:
                    // If index_mask is Max
                    // 1. On sparse union, every child len match it's parent, this is checked at union creation
                    // 2. If the null buffer len equals it's array len is checked at array creation
                    // If index_mask is Zero, the null buffer len is 1
                    nulls.inner().value_unchecked(index & *index_mask as usize)
                })
            }