Fall back to non-recursive Base.partialsort! for quickselect! of …

…very large arrays
JuliaSIMD · Sep 12, 2023 · 20951d3 · 20951d3 · brenhinkeller · Sep 12, 2023
1 parent bf90f9f
commit 20951d3
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "VectorizedStatistics"
 uuid = "3b853605-1c98-4422-8364-4bd93ee0529e"
 authors = ["C. Brenhin Keller", "Chris Elrod"]
-version = "0.5.5"
+version = "0.5.6"
 
 [deps]
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"

diff --git a/src/quicksort.jl b/src/quicksort.jl
@@ -73,6 +73,11 @@ sortnans!(A::AbstractArray{<:Integer}, iₗ::Int=firstindex(A), iᵤ::Int=lastin
 
 # Partially sort `A` around the `k`th sorted element and return that element
 function quickselect!(A::AbstractArray, iₗ::Int=firstindex(A), iᵤ::Int=lastindex(A), k=(iₗ+iᵤ)÷2)
+    # Fall back to Base implementation for very large arrays
+    if iᵤ-iₗ > 20000
+        return Base.Sort.partialsort!(view(A, iₗ:iᵤ), k-(iₗ-1))
+    end
+
     # Pick a pivot for partitioning
     N = iᵤ - iₗ + 1
     A[iₗ], A[k] = A[k], A[iₗ]