add sort bench

polars/Cargo.toml

@@ -205,6 +205,10 @@ harness = false
 name = "take"
 harness = false
 
+[[bench]]
+name = "sort"
+harness = false
+
 [package.metadata.docs.rs]
 # not all because arrow 4.3 does not compile with simd
 # all-features = true

polars/benches/collect.rs

@@ -22,7 +22,7 @@ fn bench_collect_optional_str(v: &[Option<String>]) {
     criterion::black_box(f());
 }
 
-fn create_array(size: i32, null_percentage: f32) -> Vec<Option<i32>> {
+pub fn create_array(size: i32, null_percentage: f32) -> Vec<Option<i32>> {
     let mut rng = StdRng::seed_from_u64(0);
     (0..size)
         .map(|i| {

polars/benches/sort.rs

@@ -0,0 +1,29 @@
+#[macro_use]
+extern crate criterion;
+use criterion::Criterion;
+
+use polars::prelude::*;
+
+fn bench_sort(s: &Int32Chunked) {
+    criterion::black_box(s.sort(false));
+}
+
+fn add_benchmark(c: &mut Criterion) {
+    (10..=20).step_by(2).for_each(|log2_size| {
+        let size = 2usize.pow(log2_size);
+
+        let ca = Int32Chunked::init_rand(size, 0.0, 10);
+
+        c.bench_function(&format!("sort 2^{} i32", log2_size), |b| {
+            b.iter(|| bench_sort(&ca))
+        });
+
+        let ca = Int32Chunked::init_rand(size, 0.1, 10);
+        c.bench_function(&format!("sort null 2^{} i32", log2_size), |b| {
+            b.iter(|| bench_sort(&ca))
+        });
+    });
+}
+
+criterion_group!(benches, add_benchmark);
+criterion_main!(benches);

polars/polars-core/src/chunked_array/random.rs

@@ -4,7 +4,7 @@ use num::{Float, NumCast};
 use rand::distributions::Bernoulli;
 use rand::prelude::*;
 use rand::seq::IteratorRandom;
-use rand_distr::{Distribution, Normal, StandardNormal, Uniform};
+use rand_distr::{Distribution, Normal, Standard, StandardNormal, Uniform};
 
 fn create_rand_index_with_replacement(n: usize, len: usize) -> UInt32Chunked {
     let mut rng = rand::thread_rng();
@@ -24,6 +24,25 @@ fn create_rand_index_no_replacement(n: usize, len: usize) -> UInt32Chunked {
     UInt32Chunked::new_from_aligned_vec("", buf)
 }
 
+impl<T> ChunkedArray<T>
+where
+    T: PolarsNumericType,
+    Standard: Distribution<T::Native>,
+{
+    pub fn init_rand(size: usize, null_density: f32, seed: u64) -> Self {
+        let mut rng = StdRng::seed_from_u64(seed);
+        (0..size)
+            .map(|_| {
+                if rng.gen::<f32>() < null_density {
+                    None
+                } else {
+                    Some(rng.gen())
+                }
+            })
+            .collect()
+    }
+}
+
 impl<T> ChunkedArray<T>
 where
     ChunkedArray<T>: ChunkTake,