Function initcap_ascii_array

datafusion/functions/src/unicode/initcap.rs:191–243

Fast path for `Utf8` or `LargeUtf8` arrays that are ASCII-only. We can use a single pass over the buffer and operate directly on bytes.

(
    string_array: &GenericStringArray<T>,
)

/// Fast path for `Utf8` or `LargeUtf8` arrays that are ASCII-only. We can use a
/// single pass over the buffer and operate directly on bytes.
fn initcap_ascii_array<T: OffsetSizeTrait>(
    string_array: &GenericStringArray<T>,
) -> ArrayRef {
    let offsets = string_array.offsets();
    let src = string_array.value_data();
    let first_offset = offsets.first().unwrap().as_usize();
    let last_offset = offsets.last().unwrap().as_usize();

    // For sliced arrays, only convert the visible bytes, not the entire input
    // buffer.
    let mut out = Vec::with_capacity(last_offset - first_offset);

    for window in offsets.windows(2) {
        let start = window[0].as_usize();
        let end = window[1].as_usize();

        let mut prev_is_alnum = false;
        for &b in &src[start..end] {
            let converted = if prev_is_alnum {
                b.to_ascii_lowercase()
            } else {
                b.to_ascii_uppercase()
            };
            out.push(converted);
            prev_is_alnum = b.is_ascii_alphanumeric();
        }
    }

    let values = Buffer::from_vec(out);
    let out_offsets = if first_offset == 0 {
        offsets.clone()
    } else {
        // For sliced arrays, we need to rebase the offsets to reflect that the
        // output only contains the bytes in the visible slice.
        let rebased_offsets = offsets
            .iter()
            .map(|offset| T::usize_as(offset.as_usize() - first_offset))
            .collect::<Vec<_>>();
        OffsetBuffer::<T>::new(rebased_offsets.into())
    };

    // SAFETY: ASCII case conversion preserves byte length, so the original
    // string boundaries are preserved. `out_offsets` is either identical to
    // the input offsets or a rebased version relative to the compacted values
    // buffer.
    Arc::new(unsafe {
        GenericStringArray::<T>::new_unchecked(
            out_offsets,
            values,
            string_array.nulls().cloned(),
        )
    })
}

fn initcap_utf8view(args: &[ArrayRef]) -> Result<ArrayRef> {
    let string_view_array = as_string_view_array(&args[0])?;