Function repeat_impl

datafusion/functions/src/string/repeat.rs:257–324

(
    string_array: &S,
    number_array: &Int64Array,
    max_item_capacity: usize,
    mut builder: B,
)

}

fn repeat_impl<'a, S, B>(
    string_array: &S,
    number_array: &Int64Array,
    max_item_capacity: usize,
    mut builder: B,
) -> Result<ArrayRef>
where
    S: StringArrayType<'a> + 'a,
    B: BulkNullStringArrayBuilder,
{
    // Reusable buffer to avoid allocations in string.repeat()
    let mut buffer = Vec::<u8>::with_capacity(max_item_capacity);

    // Helper function to repeat a string into a buffer using doubling strategy
    // count must be > 0
    #[inline]
    fn repeat_to_buffer(buffer: &mut Vec<u8>, string: &str, count: usize) {
        buffer.clear();
        if !string.is_empty() {
            let src = string.as_bytes();
            // Initial copy
            buffer.extend_from_slice(src);
            // Doubling strategy: copy what we have so far until we reach the target
            while buffer.len() < src.len() * count {
                let copy_len = buffer.len().min(src.len() * count - buffer.len());
                // SAFETY: we're copying valid UTF-8 bytes that we already verified
                buffer.extend_from_within(..copy_len);
            }
        }
    }

    // Output is null IFF either input is null
    let nulls = NullBuffer::union(string_array.nulls(), number_array.nulls());

    if let Some(ref n) = nulls {
        for i in 0..string_array.len() {
            if n.is_null(i) {
                builder.append_placeholder();
                continue;
            }
            // SAFETY: index `i` in both arrays is valid
            let string = unsafe { string_array.value_unchecked(i) };
            let count = unsafe { number_array.value_unchecked(i) };
            if count > 0 {
                repeat_to_buffer(&mut buffer, string, count as usize);
                // SAFETY: buffer contains valid UTF-8 since we only copy from a valid &str
                builder.append_value(unsafe { std::str::from_utf8_unchecked(&buffer) });
            } else {
                builder.append_value("");
            }
        }
    } else {
        for i in 0..string_array.len() {
            // SAFETY: no nulls, so every index in both arrays is valid
            let string = unsafe { string_array.value_unchecked(i) };
            let count = unsafe { number_array.value_unchecked(i) };
            if count > 0 {
                repeat_to_buffer(&mut buffer, string, count as usize);