I am implementing the Poisson bootstrap in Rust and wanted to benchmark my repeat function against numpy's. Briefly, repeat takes in two arguments, data and weight, and repeats each element of data by the weight, e.g. [1, 2, 3], [1, 2, 0] -> [1, 2, 2]. My naive version was around 4.5x slower than np.repeat. pub fn repeat_by(arr: &[f64], repeats: &[u64]) -> Vec<f64> { // Use flat_map to create a single iterator of all repeated elements let result: Vec<f64> = arr .iter() .zip(repeats.iter()) .flat_map(|(&value, &count)| std::iter::repeat_n(value, count as usize)) .collect(); result } I also tried a couple of more versions, e.g. one where I pre-allocated a vector with the necessary capacity, but all performed similarly. While doing more investigating though, I found that np.repeat is actually way faster than other numpy functions that I expected to perform similarly. For example, we can build a list of indices and use numpy slicing / take to perform the same operation as np.repeat. However, doing this (and even removing the list construction from the timings), np.repeat is around 3x faster than numpy slicing / take. import timeit import numpy as np N_ROWS = 100_000 x = np.random.rand(N_ROWS) weight = np.random.poisson(1, len(data)) # pre-compute the indices so slow python looping doesn't affect the timing indices = [] for w in weight: for i in range(w): indices.append(i) print(timeit.timeit(lambda: np.repeat(x, weight), number=1_000)) # 0.8337333500003297 print(timeit.timeit(lambda: np.take(x, indices), number=1_000)) # 3.1320624930012855 My C is not so good, but it seems like the relevant implementation is here: https://github.com/numpy/numpy/blob/main/numpy/_core/src/multiarray/item_selection.c#L785. It would be amazing if someone could help me understand at a high level what this code is doing--on the surface, it doesn't look like anything particularly special (SIMD, etc.), and looks pretty similar to my naive Rust version (memcpy vs repeat_n). In addition, I am struggling to understand why it performs so much better than even numpy slicing. Continue reading...
