Parts of sorting broken due to `scratch` kwarg

[palday]

I noticed that sorting in AxisKeys.jl is broken. I'm posting here because it seems related to the changes in sorting that happened in Julia 1.9 and the relevant piece of the error (keyword argument scratch, see below) does not occur in AxisKeys.jl, but rather from Base.Sort.__sort!.

MWE:

using Pkg
Pkg.activate(; temp=true)
@info "Installing packages"
Pkg.add(["AxisKeys"]; io=devnull)
using AxisKeys

data = rand(2, 10, 3)
A = KeyedArray(data; channel=[:left, :right], time=range(13, step=2.5, length=10), iter=31:33)
sort!(A; dims=1)

Error:

julia> sort!(A; dims=1)
ERROR: TypeError: in keyword argument scratch, expected Union{Nothing, Vector{<:Integer}}, got a value of type Vector{Float64}
Stacktrace:
 [1] sort!(A::KeyedArray{…}; kw::@Kwargs{…})
   @ AxisKeys ~/.julia/packages/AxisKeys/QsNqy/src/functions.jl:274
 [2] __sort!(A::KeyedArray{…}, ::Val{…}, alg::Base.Sort.MissingOptimization{…}, order::Base.Order.ForwardOrdering, scratch::Vector{…})
   @ Base.Sort ./sort.jl:1879
 [3] sort!(A::KeyedArray{…}; dims::Int64, alg::Base.Sort.MissingOptimization{…}, lt::Function, by::Function, rev::Nothing, order::Base.Order.ForwardOrdering, scratch::Vector{…})
   @ Base.Sort ./sort.jl:1866
 [4] top-level scope
   @ REPL[8]:1
Some type information was truncated. Use `show(err)` to see complete types.

Notably, the error does not happen if the eltype of the array is Int.

I bisected to this:

commit cee0a0494c70208b6cd5a32ccdf75d954a429870
Author: Lilith Orion Hafner <lilithhafner@gmail.com>
Date:   Sat Dec 3 06:19:20 2022 +0600

    Refactor and document sorting dispatch (#47383)
    
    * create an internal `_sort!` function and use it (rename the existing `_sort!` to `__sort!`)
    
    * test for several of bugs that slipped through test suite
    
    * Give each sorting pass and DEFAULT_STABLE a docstring
    
    * add pretty printing for the new algorithms that are much more flexible than the old ones
    
    * fix unexpected allocations in Radix Sort
    
    fixes #47474
    in this PR rather than separate to avoid dealing with the merge
    
    * support and test backwards compatibility with packages that depend in sorting internals
    
    * support 3-, 5-, and 6-argument sort! for backwards compatibility
    
    * overhall scratch space handling
    
    make _sort! return scratch space rather than sorted vector
    so that things like IEEEFloatOptimization can re-use the
    scratch space allocated on their first recursive call
    
    * test handling -0.0 in IEEEFloatOptimization
    
    * fix and test bug where countsort's correct overflow behavior triggers error due to unexpected promotion to UInt

 base/sort.jl    | 1287 +++++++++++++++++++++++++++++++++++--------------------
 test/sorting.jl |  171 +++++++-
 2 files changed, 968 insertions(+), 490 deletions(-)

cc @LilithHafner @mcabbott

[ericphanson]

scratch is internal (not documented) so this seems like an AxisKeys issue

[barucden]

AxisKeys does not create a scratch though. It is initially created here in Base:

julia/base/sort.jl

Line 2096 in 75c38e3

scratch::Union{Vector{T}, Nothing}=size(A, dims) < 10 ? nothing : Vector{T}(undef, size(A, dims))) where T

as Vector{Float64}(undef, size(A, 1)) (notice Float64). A bunch of sorting logic then follows until it gets to sorting a vector that is a view of the original array in a particular dimension. It hits this line of AxisKeys:

perm = sortperm(parent(A); kw...)

which simply passes all keyword arguments it received from a function in Base (including scratch from above).

It ends up calling this method

julia/base/sort.jl

Lines 1868 to 1875 in 75c38e3

	function sortperm(A::AbstractArray;
	alg::Algorithm=DEFAULT_UNSTABLE,
	lt=isless,
	by=identity,
	rev::Union{Bool,Nothing}=nothing,
	order::Ordering=Forward,
	scratch::Union{Vector{<:Integer}, Nothing}=nothing,
	dims...) #to optionally specify dims argument

(notice Integer in scratch::Union{Vector{<:Integer}, Nothing}), and fails.

I am not sure what to do about it, but hopefully this might help someone with debugging.

[ericphanson]

Ohh, I see. I think AxisKeys does not provide sort! for KeyedArray (only KeyedVector), so it falls back to Base, which is implemented by sorting vectors. So sort! for KeyedArray falls back to sort! for KeyedVectors, and it all works pre-1.9. Then sort!(::AbstractArray) has the new default scratch kwarg (of type Vector{Float64} in this case), which it forwards to the sort! for KeyedVectors, which does hit the AxisKeys method, and passes those kwargs to the sortperm call, which needs integer scratch space.

Hm, I think this still isn't quite Base's breakage, but it's not a super obvious use of internals by AxisKeys. It was essentially relying on the implementation of sort!(::AbstractArray) to lower to sorting vectors (so it didn't need to provide sort!(::KeyedArray) just sort!(::KeyedVector)), and while that remained true, the particular implementation didn't remain compatible with the assumptions AxisKeys was making (sortperm taking the same kwargs as sort).

[LilithHafner]

sort!(x::AbstractArray; dims) allocates a single scratch-space to reuse when sorting each slice. When working with KeyedArray, the slice sorting algorithm requires a scratchspace that is not compatible with Vector{eltype(x)}. The generic sort!(x::AbstractArray; dims) method can't be expected to allocate the correct scratchspace type.

Two possible fixes from the Julia side:

1. Use Vector{UInt8} as a scratch space when sorting isbitstypes of any type. Then Flaot64 and Int are compatible. That fixes this case, but not the case where the eltype is not isbitstype.
2. Silently discard mis-typed scratch spaces, possibly dispatching to an algorithm which does not require a scratchspace. Might come at a noticeable performance cost in this case.

The slice sorting algorithm

function Base.sort!(A::KeyedVector; kw...)
    perm = sortperm(parent(A); kw...)
    permute!(axiskeys(A,1), perm) # error if keys cannot be sorted, could treat like push!
    permute!(parent(A), perm)
    A
end

is already fairly inefficient because it first computes sortperm and then permutes the inputs rather than sorting all at once, so an additional allocation shouldn't be noticed here.

I think option 2 is a reasonable thing to do, at least until we settle on a public API for scratchspaces.

[ericphanson]

I think (2) is nice to do and adds flexibility, but the immediate fix is for AxisKeys to not depend on the details of how sort!(::AbstractArray) works (by adding an explicit method for it).

[palday]

I'm not sure AxisKeys really depends on the details of how it works, but rather just on it working. My understanding is that the contract is more or less:

• getindex and setindex! are defined (potentially with a requirement for linear indexing)
• isless is defined for the eltype.

Both of those conditions are met.

[ericphanson]

I believe it does, AxisKeys adds a method for sort!(::KeyedVector) but not sort!(::KeyedArray). It assumes that sort!(::AbstractArray) will call sort!(::KeyedVector) in a way that that method is happy with, but sort!(::KeyedVector) assumes it may pass all kwargs on to sortperm, which isn't true.

But Base doesn't need to call sort!(::KeyedVector) at all (in which case the special axis handling stuff would stop happening when you sorted a KeyedArray), or it could change how it calls that function (which is what happened here).

[palday]

Okay, maybe for KeyedArray, but there is still an error when calling sort!(::KeyedVector) directly with 1-dimensional array and the body of that function definitely only depends on public API.

[ericphanson]

there is still an error when calling sort!(::KeyedVector) directly with 1-dimensional array

can you provide a MWE? I don't see that. For example,

julia> data = rand(10);

julia> A = KeyedArray(data; time=collect(range(13, step=2.5, length=10)))
1-dimensional KeyedArray(NamedDimsArray(...)) with keys:
↓   time ∈ 10-element Vector{Float64}
And data, 10-element Vector{Float64}:
 (13.0)  0.7906376794689575
 (15.5)  0.223747395427907
 (18.0)  0.37445787840685807
 (20.5)  0.5604038403956254
 (23.0)  0.7965313358416087
 (25.5)  0.6041031656437967
 (28.0)  0.21207823790093805
 (30.5)  0.4116456128979581
 (33.0)  0.2747546678987417
 (35.5)  0.04284820974498316

julia> sort!(A)
1-dimensional KeyedArray(NamedDimsArray(...)) with keys:
↓   time ∈ 10-element Vector{Float64}
And data, 10-element Vector{Float64}:
 (35.5)  0.04284820974498316
 (28.0)  0.21207823790093805
 (15.5)  0.223747395427907
 (33.0)  0.2747546678987417
 (18.0)  0.37445787840685807
 (30.5)  0.4116456128979581
 (20.5)  0.5604038403956254
 (25.5)  0.6041031656437967
 (13.0)  0.7906376794689575
 (23.0)  0.7965313358416087

Note that if I don't include collect, it does error, but the same occurs on Julia 1.6, and is an unrelated issue.