Fix benchmark.

2 files changed, 30 insertions, 33 deletions

diff --git a/benchmarks/benchmarks/bench_core.py b/benchmarks/benchmarks/bench_core.py
index ac28e5b81..26cffcab1 100644
--- a/benchmarks/benchmarks/bench_core.py
+++ b/benchmarks/benchmarks/bench_core.py
@@ -11,9 +11,6 @@ class Core(Benchmark):
         self.l50 = range(50)
         self.l = [np.arange(1000), np.arange(1000)]
         self.l10x10 = np.ones((10, 10))
-        random = np.random.RandomState(123)
-        self.arr = np.arange(10 ** 4)
-        random.shuffle(self.arr)
 
     def time_array_1(self):
         np.array(1)
@@ -78,17 +75,6 @@ class Core(Benchmark):
     def time_tril_l10x10(self):
         np.tril(self.l10x10)
 
-    def time_quicksort(self):
-        np.sort(self.arr, kind='quick')
-
-    def time_sorts(self, kind):
-        try:
-            np.sort(self.arr, kind=kind)
-        except TypeError:
-            raise NotImplementedError('Sort kind not implemented or invalid dtype')
-
-    time_sorts.params = ['quick', 'merge', 'heap', 'radix']
-
 
 class Temporaries(Benchmark):
     def setup(self):
diff --git a/benchmarks/benchmarks/bench_function_base.py b/benchmarks/benchmarks/bench_function_base.py
index 9ef03262b..a37656542 100644
--- a/benchmarks/benchmarks/bench_function_base.py
+++ b/benchmarks/benchmarks/bench_function_base.py
@@ -96,35 +96,46 @@ class Select(Benchmark):
 
 
 class Sort(Benchmark):
-    def setup(self):
-        self.e = np.arange(10000, dtype=np.float32)
-        self.o = np.arange(10001, dtype=np.float32)
+    params = [
+        ['quick', 'merge', 'heap', 'radix', None],
+        ['float32', 'int32', 'uint32']
+    ]
+    param_names = ['kind', 'dtype']
+
+    def setup(self, kind, dtype):
+        self.e = np.arange(10000, dtype=dtype)
+        self.o = np.arange(10001, dtype=dtype)
         np.random.seed(25)
         np.random.shuffle(self.o)
         # quicksort implementations can have issues with equal elements
-        self.equal = np.ones(10000)
-        self.many_equal = np.sort(np.arange(10000) % 10)
+        self.equal = np.ones(10000, dtype=dtype)
+        self.many_equal = np.sort(np.arange(10000) % 10).astype(dtype)
+
+        try:
+            np.sort(self.e, kind=kind)
+        except TypeError:
+            raise NotImplementedError()
 
-    def time_sort(self):
-        np.sort(self.e)
+    def time_sort(self, kind, dtype):
+        np.sort(self.e, kind=kind)
 
-    def time_sort_random(self):
-        np.sort(self.o)
+    def time_sort_random(self, kind, dtype):
+        np.sort(self.o, kind=kind)
 
-    def time_sort_inplace(self):
-        self.e.sort()
+    def time_sort_inplace(self, kind, dtype):
+        self.e.sort(kind=kind)
 
-    def time_sort_equal(self):
-        self.equal.sort()
+    def time_sort_equal(self, kind, dtype):
+        self.equal.sort(kind=kind)
 
-    def time_sort_many_equal(self):
-        self.many_equal.sort()
+    def time_sort_many_equal(self, kind, dtype):
+        self.many_equal.sort(kind=kind)
 
-    def time_argsort(self):
-        self.e.argsort()
+    def time_argsort(self, kind, dtype):
+        self.e.argsort(kind=kind)
 
-    def time_argsort_random(self):
-        self.o.argsort()
+    def time_argsort_random(self, kind, dtype):
+        self.o.argsort(kind=kind)
 
 
 class SortWorst(Benchmark):