Function take_fixed_size_binary

arrow-select/src/take.rs:798–864

The take kernel implementation for `FixedSizeBinaryArray`. The computation is done in two steps: - Compute the values buffer - Compute the null buffer

(
    values: &FixedSizeBinaryArray,
    indices: &PrimitiveArray<IndexType>,
    size: i32,
)

/// - Compute the values buffer
/// - Compute the null buffer
fn take_fixed_size_binary<IndexType: ArrowPrimitiveType>(
    values: &FixedSizeBinaryArray,
    indices: &PrimitiveArray<IndexType>,
    size: i32,
) -> Result<FixedSizeBinaryArray, ArrowError> {
    let size_usize = usize::try_from(size).map_err(|_| {
        ArrowError::InvalidArgumentError(format!("Cannot convert size '{}' to usize", size))
    })?;

    let result_buffer = match size_usize {
        1 => take_fixed_size::<IndexType, 1>(values.values(), indices),
        2 => take_fixed_size::<IndexType, 2>(values.values(), indices),
        4 => take_fixed_size::<IndexType, 4>(values.values(), indices),
        8 => take_fixed_size::<IndexType, 8>(values.values(), indices),
        16 => take_fixed_size::<IndexType, 16>(values.values(), indices),
        _ => take_fixed_size_binary_buffer_dynamic_length(values, indices, size_usize),
    };

    let value_nulls = take_nulls(values.nulls(), indices);
    let final_nulls = NullBuffer::union(value_nulls.as_ref(), indices.nulls());
    let array_data = ArrayDataBuilder::new(DataType::FixedSizeBinary(size))
        .len(indices.len())
        .nulls(final_nulls)
        .offset(0)
        .add_buffer(result_buffer)
        .build()?;

    return Ok(FixedSizeBinaryArray::from(array_data));

    /// Implementation of the take kernel for fixed size binary arrays.
    #[inline(never)]
    fn take_fixed_size_binary_buffer_dynamic_length<IndexType: ArrowPrimitiveType>(
        values: &FixedSizeBinaryArray,
        indices: &PrimitiveArray<IndexType>,
        size_usize: usize,
    ) -> Buffer {
        let values_buffer = values.values().as_slice();
        let mut values_buffer_builder = BufferBuilder::new(indices.len() * size_usize);

        if indices.null_count() == 0 {
            let array_iter = indices.values().iter().map(|idx| {
                let offset = idx.as_usize() * size_usize;
                &values_buffer[offset..offset + size_usize]
            });
            for slice in array_iter {
                values_buffer_builder.append_slice(slice);
            }
        } else {
            // The indices nullability cannot be ignored here because the values buffer may contain
            // nulls which should not cause a panic.
            let array_iter = indices.iter().map(|idx| {
                idx.map(|idx| {
                    let offset = idx.as_usize() * size_usize;
                    &values_buffer[offset..offset + size_usize]
                })
            });
            for slice in array_iter {
                match slice {