diff --git a/arrow-buffer/Cargo.toml b/arrow-buffer/Cargo.toml index d1651abb795b..ae5e121164ff 100644 --- a/arrow-buffer/Cargo.toml +++ b/arrow-buffer/Cargo.toml @@ -59,3 +59,8 @@ harness = false [[bench]] name = "offset" harness = false + +[[bench]] +name = "mutable_buffer_repeat_slice" +harness = false + diff --git a/arrow-buffer/benches/mutable_buffer_repeat_slice.rs b/arrow-buffer/benches/mutable_buffer_repeat_slice.rs new file mode 100644 index 000000000000..3a3bc63e776e --- /dev/null +++ b/arrow-buffer/benches/mutable_buffer_repeat_slice.rs @@ -0,0 +1,80 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_buffer::Buffer; +use criterion::*; +use rand::distr::Alphanumeric; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use std::hint; + +fn criterion_benchmark(c: &mut Criterion) { + let mut group = c.benchmark_group("MutableBuffer repeat slice"); + let mut rng = StdRng::seed_from_u64(42); + + for slice_length in [3, 20, 100] { + let slice_to_repeat: Vec = hint::black_box( + (&mut rng) + .sample_iter(&Alphanumeric) + .take(slice_length) + .collect(), + ); + let slice_to_repeat: &[u8] = slice_to_repeat.as_ref(); + + for repeat_count in [3, 64, 1024, 8192] { + let parameter_string = format!("slice_len={slice_length} n={repeat_count}"); + + group.bench_with_input( + BenchmarkId::new("repeat_slice_n_times", ¶meter_string), + &(repeat_count), + |b, &repeat_count| { + b.iter(|| { + let mut mutable_buffer = arrow_buffer::MutableBuffer::with_capacity(0); + + mutable_buffer.repeat_slice_n_times(slice_to_repeat, repeat_count); + + let buffer: Buffer = mutable_buffer.into(); + + buffer + }) + }, + ); + group.bench_with_input( + BenchmarkId::new("extend_from_slice loop", ¶meter_string), + &(repeat_count), + |b, &repeat_count| { + b.iter(|| { + let mut mutable_buffer = arrow_buffer::MutableBuffer::with_capacity( + size_of_val(slice_to_repeat) * repeat_count, + ); + + for _ in 0..repeat_count { + mutable_buffer.extend_from_slice(slice_to_repeat); + } + + let buffer: Buffer = mutable_buffer.into(); + + buffer + }) + }, + ); + } + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/arrow-buffer/src/buffer/mutable.rs b/arrow-buffer/src/buffer/mutable.rs index 93d9d6b9ad84..a12037c6433f 100644 --- a/arrow-buffer/src/buffer/mutable.rs +++ b/arrow-buffer/src/buffer/mutable.rs @@ -222,6 +222,75 @@ impl MutableBuffer { } } + /// Adding to this mutable buffer `slice_to_repeat` repeated `repeat_count` times. + /// + /// # Example + /// + /// ## Repeat the same string bytes multiple times + /// ``` + /// # use arrow_buffer::buffer::MutableBuffer; + /// let mut buffer = MutableBuffer::new(0); + /// let bytes_to_repeat = b"ab"; + /// buffer.repeat_slice_n_times(bytes_to_repeat, 3); + /// assert_eq!(buffer.as_slice(), b"ababab"); + /// ``` + pub fn repeat_slice_n_times( + &mut self, + slice_to_repeat: &[T], + repeat_count: usize, + ) { + if repeat_count == 0 || slice_to_repeat.is_empty() { + return; + } + + let bytes_to_repeat = size_of_val(slice_to_repeat); + + // Ensure capacity + self.reserve(repeat_count * bytes_to_repeat); + + // Save the length before we do all the copies to know where to start from + let length_before = self.len; + + // Copy the initial slice once so we can use doubling strategy on it + self.extend_from_slice(slice_to_repeat); + + // This tracks how much bytes we have added by repeating so far + let added_repeats_length = bytes_to_repeat; + assert_eq!( + self.len - length_before, + added_repeats_length, + "should copy exactly the same number of bytes" + ); + + // Number of times the slice was repeated + let mut already_repeated_times = 1; + + // We will use doubling strategy to fill the buffer in log(repeat_count) steps + while already_repeated_times < repeat_count { + // How many slices can we copy in this iteration + // (either double what we have, or just the remaining ones) + let number_of_slices_to_copy = + already_repeated_times.min(repeat_count - already_repeated_times); + let number_of_bytes_to_copy = number_of_slices_to_copy * bytes_to_repeat; + + unsafe { + // Get to the start of the data before we started copying anything + let src = self.data.as_ptr().add(length_before) as *const u8; + + // Go to the current location to copy to (end of current data) + let dst = self.data.as_ptr().add(self.len); + + // SAFETY: the pointers are not overlapping as there is `number_of_bytes_to_copy` or less between them + std::ptr::copy_nonoverlapping(src, dst, number_of_bytes_to_copy) + } + + // Advance the length by the amount of data we just copied (doubled) + self.len += number_of_bytes_to_copy; + + already_repeated_times += number_of_slices_to_copy; + } + } + #[cold] fn reallocate(&mut self, capacity: usize) { let new_layout = Layout::from_size_align(capacity, self.layout.align()).unwrap(); @@ -1184,4 +1253,125 @@ mod tests { assert_eq!(pool.used(), 0); } } + + fn create_expected_repeated_slice( + slice_to_repeat: &[T], + repeat_count: usize, + ) -> Buffer { + let mut expected = MutableBuffer::new(size_of_val(slice_to_repeat) * repeat_count); + for _ in 0..repeat_count { + // Not using push_slice_repeated as this is the function under test + expected.extend_from_slice(slice_to_repeat); + } + expected.into() + } + + // Helper to test a specific repeat count with various slice sizes + fn test_repeat_count( + repeat_count: usize, + test_data: &[T], + ) { + let mut buffer = MutableBuffer::new(0); + buffer.repeat_slice_n_times(test_data, repeat_count); + + let expected = create_expected_repeated_slice(test_data, repeat_count); + let result: Buffer = buffer.into(); + + assert_eq!( + result, + expected, + "Failed for repeat_count={}, slice_len={}", + repeat_count, + test_data.len() + ); + } + + #[test] + fn test_repeat_slice_count_edge_cases() { + // Empty slice + test_repeat_count(100, &[] as &[i32]); + + // Zero repeats + test_repeat_count(0, &[1i32, 2, 3]); + } + + #[test] + fn test_small_repeats_counts() { + // test any special implementation for small repeat counts + let data = &[1u8, 2, 3, 4, 5]; + + for _ in 1..=10 { + test_repeat_count(2, data); + } + } + + #[test] + fn test_different_size_of_i32_repeat_slice() { + let data: &[i32] = &[1, 2, 3]; + let data_with_single_item: &[i32] = &[42]; + + for data in &[data, data_with_single_item] { + for item in 1..=9 { + let base_repeat_count = 2_usize.pow(item); + test_repeat_count(base_repeat_count - 1, data); + test_repeat_count(base_repeat_count, data); + test_repeat_count(base_repeat_count + 1, data); + } + } + } + + #[test] + fn test_different_size_of_u8_repeat_slice() { + let data: &[u8] = &[1, 2, 3]; + let data_with_single_item: &[u8] = &[10]; + + for data in &[data, data_with_single_item] { + for item in 1..=9 { + let base_repeat_count = 2_usize.pow(item); + test_repeat_count(base_repeat_count - 1, data); + test_repeat_count(base_repeat_count, data); + test_repeat_count(base_repeat_count + 1, data); + } + } + } + + #[test] + fn test_different_size_of_u16_repeat_slice() { + let data: &[u16] = &[1, 2, 3]; + let data_with_single_item: &[u16] = &[10]; + + for data in &[data, data_with_single_item] { + for item in 1..=9 { + let base_repeat_count = 2_usize.pow(item); + test_repeat_count(base_repeat_count - 1, data); + test_repeat_count(base_repeat_count, data); + test_repeat_count(base_repeat_count + 1, data); + } + } + } + + #[test] + fn test_various_slice_lengths() { + // Test different slice lengths with same repeat pattern + let repeat_count = 37; // Arbitrary non-power-of-2 + + // Single element + test_repeat_count(repeat_count, &[42i32]); + + // Small slices + test_repeat_count(repeat_count, &[1i32, 2]); + test_repeat_count(repeat_count, &[1i32, 2, 3]); + test_repeat_count(repeat_count, &[1i32, 2, 3, 4]); + test_repeat_count(repeat_count, &[1i32, 2, 3, 4, 5]); + + // Larger slices + let data_10: Vec = (0..10).collect(); + test_repeat_count(repeat_count, &data_10); + + let data_100: Vec = (0..100).collect(); + test_repeat_count(repeat_count, &data_100); + + let data_1000: Vec = (0..1000).collect(); + test_repeat_count(repeat_count, &data_1000); + } }