Implement unique and set operations by sorting

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUArrays"
 uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
-version = "12.2.0"
+version = "12.3.0"
 
 [workspace]
 projects = ["lib/GPUArraysCore", "lib/JLArrays", "test", "docs"]
@@ -29,7 +29,7 @@ JLD2Ext = "JLD2"
 
 [compat]
 AcceleratedKernels = "0.5.8"
-Adapt = "4.7.3"
+Adapt = "4.8"
 GPUArraysCore = "= 0.2.2"
 JLD2 = "0.4, 0.5, 0.6"
 KernelAbstractions = "0.10"

src/GPUArrays.jl

@@ -31,6 +31,7 @@ include("host/construction.jl")
 include("host/base.jl")
 include("host/indexing.jl")
 include("host/sorting.jl")
+include("host/set.jl")
 include("host/accumulate.jl")
 include("host/broadcast.jl")
 include("host/mapreduce.jl")

src/host/set.jl

@@ -0,0 +1,193 @@
+# unique elements and set operations, by sorting
+#
+# Sorting puts elements that are `isequal` next to each other, as long as the order agrees with
+# `isequal`. `isless` does, and so does `isless` on the real and imaginary parts of complex
+# numbers. To return elements in the order Base does, the first element of each group is marked
+# at its original position, and the marked elements are selected with a mask.
+
+# Mixed with complex numbers, a real number sorts as a complex number with zero imaginary part.
+complex_key(x::Complex) = reim(x)
+complex_key(x) = (x, zero(x))
+complex_isless(x, y) = isless(complex_key(x), complex_key(y))
+
+has_complex(T) = T isa Union ? any(has_complex, Base.uniontypes(T)) : T <: Complex
+
+# Plain `isless` where possible, because AcceleratedKernels can only use its faster sorting
+# algorithms with the default order.
+set_isless(::Type{T}) where {T} = has_complex(T) ? complex_isless : isless
+
+# Returns the sorting permutation of `keys` and the sorted keys. The sort is stable, so the
+# first element of each group of equal keys is the one that occurs first in `keys`.
+function sorted_keys(keys::AbstractVector)
+    lt = set_isless(eltype(keys))
+    p = sortperm(keys; lt)
+    return p, keys[p]
+end
+
+# For each element of the sorted `s`, whether it differs from the one before it
+function group_starts(s::AbstractVector)
+    n = length(s)
+    head = similar(s, Bool, n)
+    n == 0 && return head
+    head[1:1] .= true
+    head[2:n] .= .!isequal.(view(s, 2:n), view(s, 1:n-1))
+    return head
+end
+
+# Moves the flags `heads`, which are in sorted order, back to the original positions
+function scatter_mask(p::AbstractVector, heads::AbstractVector{Bool})
+    mask = similar(heads)
+    mask[p] = heads
+    return mask
+end
+
+function Base.unique(f, A::AnyGPUArray)
+    v = vec(A)
+    length(v) <= 1 && return copy(v)
+    p, s = sorted_keys(f === identity ? v : map(f, v))
+    return v[scatter_mask(p, group_starts(s))]
+end
+
+function Base.unique(A::AnyGPUArray; dims::Union{Colon, Integer}=:)
+    dims isa Colon && return unique(identity, A)
+    # Base returns a copy when `dims` is out of range
+    1 <= dims <= ndims(A) || return copy(A)
+    return unique_dims(A, Int(dims))
+end
+
+# The slices along `d` are compared in a 3-d reshape of `A`, whose middle dimension is `d`
+function unique_dims(A::AbstractArray, d::Int)
+    n = size(A, d)
+    n <= 1 && return copy(A)
+    B = reshape(A, prod(size(A)[1:d-1]), n, prod(size(A)[d+1:end]))
+    ltel = set_isless(eltype(A))
+    lt = let B = B, ltel = ltel
+        (i, j) -> begin
+            @inbounds for k in axes(B, 3), l in axes(B, 1)
+                a, b = B[l, i, k], B[l, j, k]
+                ltel(a, b) && return true
+                ltel(b, a) && return false
+            end
+            return false
+        end
+    end
+    eq = let B = B
+        (i, j) -> begin
+            @inbounds for k in axes(B, 3), l in axes(B, 1)
+                isequal(B[l, i, k], B[l, j, k]) || return false
+            end
+            return true
+        end
+    end
+    p = similar(A, Int, n)
+    p .= 1:n
+    sort!(p; lt)
+    head = similar(A, Bool, n)
+    head[1:1] .= true
+    head[2:n] .= .!eq.(view(p, 1:n-1), view(p, 2:n))
+    mask = scatter_mask(p, head)
+    return A[ntuple(i -> i == d ? mask : Colon(), ndims(A))...]
+end
+
+function Base.unique!(f, v::AnyGPUVector)
+    u = unique(f, v)
+    copyto!(v, firstindex(v), u, 1, length(u))
+    return resize!(v, length(u))
+end
+Base.unique!(v::AnyGPUVector) = unique!(identity, v)
+
+function Base.allunique(A::AnyGPUArray)
+    n = length(A)
+    n <= 1 && return true
+    lt = set_isless(eltype(A))
+    s = sort(vec(A); lt)
+    return !any(isequal.(view(s, 2:n), view(s, 1:n-1)))
+end
+# Base has `allunique(f, A)` since Julia 1.11
+if hasmethod(allunique, Tuple{Any, Any})
+    Base.allunique(f, A::AnyGPUArray) = allunique(map(f, vec(A)))
+end
+
+
+## set operations
+
+# Converts `x` to a vector with element type `T`. Like Base, which builds a set of type `T`,
+# this throws if an element can't be represented exactly.
+function as_eltype(::Type{T}, x::AbstractArray) where {T}
+    v = vec(x)
+    eltype(v) === T && return v
+    y = similar(v, T)
+    y .= v
+    inexact = .!isequal.(y, v)
+    if any(inexact)
+        # throws the same error as Base
+        push!(Set{T}(), @allowscalar v[findall(inexact)[1]])
+    end
+    return y
+end
+
+# For each element of `x`, whether it is `isequal` to an element of `y`. Base's sets compare
+# this way. Base's `issubset` and `isdisjoint` use `==` for arrays of fewer than 70 elements.
+function isin(x::AbstractVector, y::AbstractArray)
+    found = similar(x, Bool)
+    if isempty(y)
+        fill!(found, false)
+        return found
+    end
+    # Comparing elements of different types without converting them, as `isequal` does
+    lt = set_isless(promote_type(eltype(x), eltype(y)))
+    s = sort(vec(y); lt)
+    insorted = let s = s, lt = lt
+        x -> begin
+            i = searchsortedfirst(s, x; lt)
+            return i <= length(s) && isequal(@inbounds(s[i]), x)
+        end
+    end
+    found .= insorted.(x)
+    return found
+end
+
+function Base.union(a::AnyGPUArray, bs::AnyGPUArray...)
+    T = Base.promote_eltype(a, bs...)
+    return unique(vcat(as_eltype(T, a), map(b -> as_eltype(T, b), bs)...))
+end
+
+function Base.intersect(a::AnyGPUArray, bs::AnyGPUArray...)
+    u = unique(as_eltype(Base.promote_eltype(a, bs...), a))
+    keep = fill!(similar(u, Bool), true)
+    for b in bs
+        keep .&= isin(u, b)
+    end
+    return u[keep]
+end
+
+function Base.setdiff(a::AnyGPUArray, bs::AnyGPUArray...)
+    # Base promotes the element type before Julia 1.13 (JuliaLang/julia#59439)
+    T = VERSION >= v"1.13-" ? eltype(a) : Base.promote_eltype(a, bs...)
+    u = unique(as_eltype(T, a))
+    keep = fill!(similar(u, Bool), true)
+    for b in bs
+        keep .&= .!isin(u, b)
+    end
+    return u[keep]
+end
+
+# Keeps the elements that occur in an odd number of the arguments, counting each argument once
+function Base.symdiff(a::AnyGPUArray, bs::AnyGPUArray...)
+    T = Base.promote_eltype(a, bs...)
+    c = vcat(unique(as_eltype(T, a)), map(b -> unique(as_eltype(T, b)), bs)...)
+    length(c) <= 1 && return c
+    p, s = sorted_keys(c)
+    head = group_starts(s)
+    starts = findall(head)
+    k = length(starts)
+    lens = similar(starts)
+    lens[1:k-1] .= view(starts, 2:k) .- view(starts, 1:k-1)
+    lens[k:k] .= length(s) + 1 .- view(starts, k:k)
+    head[starts] .= isodd.(lens)
+    return c[scatter_mask(p, head)]
+end
+
+Base.issubset(a::AnyGPUArray, b::AnyGPUArray) = all(isin(vec(a), b))
+Base.issetequal(a::AnyGPUArray, b::AnyGPUArray) = issubset(a, b) && issubset(b, a)
+Base.isdisjoint(a::AnyGPUArray, b::AnyGPUArray) = !any(isin(vec(a), b))

test/testsuite.jl

@@ -129,6 +129,7 @@ include("testsuite/aliasing.jl")
 include("testsuite/vector.jl")
 include("testsuite/reductions.jl")
 include("testsuite/sorting.jl")
+include("testsuite/set.jl")
 include("testsuite/accumulate.jl")
 include("testsuite/findall.jl")
 include("testsuite/broadcasting.jl")

test/testsuite/set.jl

@@ -0,0 +1,117 @@
+# random elements with many duplicates
+dups(ET, dims...) = ET <: Complex ? ET.(rand(1:3, dims...), rand(1:3, dims...)) :
+                                     ET.(rand(1:5, dims...))
+
+@testsuite "set/unique" (AT, eltypes)->begin
+    @testset "$ET" for ET in eltypes
+        # Base returns elements in the order they first occur
+        for n in (0, 1, 2, 10, 10000)
+            @test compare_exact(unique, AT, dups(ET, n))
+            @test compare_exact(allunique, AT, dups(ET, n))
+        end
+        @test compare_exact(unique, AT, dups(ET, 10, 20))
+        @test compare_exact(allunique, AT, ET.(collect(1:100)))
+        @test compare_exact(A -> unique(abs2, A), AT, dups(ET, 1000))
+        @test compare_exact(A -> unique(abs2, A), AT, ET[-2, 2, 1, -1])
+        @test compare_exact(A -> unique!(A), AT, dups(ET, 1000))
+        @test compare_exact(A -> unique!(abs2, A), AT, dups(ET, 1000))
+        if hasmethod(allunique, Tuple{Any, Any})
+            @test compare_exact(A -> allunique(abs2, A), AT, dups(ET, 3))
+        end
+
+        # unique slices
+        for dims in (1, 2, 3)
+            @test compare_exact(A -> unique(A; dims), AT, dups(ET, 100, 2))
+            @test compare_exact(A -> unique(A; dims), AT, dups(ET, 2, 3, 100))
+        end
+        @test compare_exact(A -> unique(A; dims=2), AT, dups(ET, 2, 0))
+        @test compare_exact(A -> unique(A; dims=1), AT, dups(ET, 1, 3))
+        # Base returns a copy when `dims` is out of range
+        for dims in (-1, 0, 4)
+            @test compare_exact(A -> unique(A; dims), AT, dups(ET, 2, 3, 4))
+        end
+
+        # views and reshaped arrays
+        @test compare_exact(A -> unique(view(A, 11:90)), AT, dups(ET, 100))
+        @test compare_exact(A -> unique(reshape(A, 10, 10); dims=2), AT, dups(ET, 100))
+    end
+
+    # with `isequal`, signed zeros differ and NaNs are equal
+    for ET in filter(isrealfloattype, eltypes)
+        x = ET[0, -0.0, NaN, 1, -NaN, NaN, -0.0, 1]
+        @test compare_exact(unique, AT, x)
+        @test compare_exact(allunique, AT, x)
+        @test compare_exact(A -> unique(A; dims=2), AT, [x x x]')
+    end
+
+    @test compare_exact(unique, AT, rand(Bool, 100))
+
+    # isbits unions, including ones with `missing`
+    for x in (Union{Int32, Float32}[Int32(1), 1.0f0, Int32(2), 2.0f0, Int32(1)],
+              Union{Missing, Float32}[1, missing, 2, missing, 1],
+              Union{Missing, ComplexF32}[1, missing, 1im, missing, 1])
+        @test compare_exact(unique, AT, x)
+        @test compare_exact(allunique, AT, x)
+        @test compare_exact(A -> unique(A; dims=1), AT, x)
+    end
+
+    # elements without `isless` are not supported (Base handles them by hashing)
+    if AT <: AbstractGPUArray
+        @test_throws Exception unique(AT([Some(Int32(1)), Some(Int32(2))]))
+    end
+end
+
+@testsuite "set/operations" (AT, eltypes)->begin
+    @testset "$ET" for ET in eltypes
+        for (m, n) in ((0, 0), (0, 10), (10, 0), (1, 1), (100, 200))
+            a, b, c = dups(ET, m), dups(ET, n), dups(ET, 50)
+            for f in (union, intersect, setdiff, symdiff, issubset, issetequal, isdisjoint)
+                @test compare_exact(f, AT, a, b)
+            end
+            for f in (union, intersect, setdiff, symdiff)
+                @test compare_exact(f, AT, a, b, c)
+                @test compare_exact(f, AT, a)
+            end
+        end
+        @test compare_exact(intersect, AT, dups(ET, 4, 5), dups(ET, 3))
+        @test compare_exact(issetequal, AT, ET[1, 2, 2], ET[2, 1])
+    end
+
+    # element types are promoted as in Base, but elements are compared without converting them
+    @test compare_exact(union, AT, Int32[3, 1], Float32[2, 3])
+    @test compare_exact(intersect, AT, Int32[3, 1, 2, 1], Float32[2, 1])
+    @test compare_exact(setdiff, AT, Int32[3, 1, 2, 1], Float32[2.5, 1])
+    @test compare_exact(symdiff, AT, Int16[1, 2], Float32[2, 4])
+    big, rounded = Int32[2^24 + 1, 1], Float32[2^24, 1]
+    for f in (union, intersect, setdiff, symdiff, issubset, isdisjoint)
+        # Base throws when it would convert an element inexactly, which depends on the order
+        @test compare_exact(f, AT, big, rounded)
+        @test compare_exact(f, AT, rounded, big)
+    end
+    @test compare_exact(issubset, AT, [-1], UInt[1])
+    @test compare_exact(intersect, AT, ComplexF32[1, 2im], Float32[1, 2])
+    # isbits unions
+    m = Union{Missing, Float32}[1, missing, 3]
+    mc = Union{Missing, ComplexF32}[1, missing, 2im]
+    for (x, y) in ((Float32[1, 2], m), (m, Float32[1, 2]), (ComplexF32[1, 2im], mc), (mc, Float32[1]))
+        for f in (union, intersect, setdiff, symdiff)
+            @test compare_exact(f, AT, x, y)
+        end
+    end
+
+    # with `isequal`, signed zeros differ and NaNs are equal
+    x, y = Float32[0, NaN, 1, -0.0], Float32[-0.0, NaN, 2]
+    for f in (union, intersect, setdiff, symdiff)
+        @test compare_exact(f, AT, x, y)
+    end
+    # The predicates use `isequal` too. Base's `issubset` and `isdisjoint` use `==` for arrays
+    # of fewer than 70 elements, so these are longer.
+    r = Float32.(1:100)
+    @test compare_exact(issubset, AT, [NaN32; r], [NaN32; r; 200])
+    @test compare_exact(issubset, AT, [0.0f0; r], [-0.0f0; r])
+    @test compare_exact(isdisjoint, AT, [NaN32; r], [NaN32; r .+ 200])
+    @test compare_exact(isdisjoint, AT, [0.0f0; r], [-0.0f0; r .+ 200])
+    mr = Union{Missing, Float32}[missing; r]
+    @test compare_exact(issubset, AT, mr, [mr; 200])
+    @test compare_exact(isdisjoint, AT, mr, [missing; r .+ 200])
+end

test/testsuite/sorting.jl

@@ -57,6 +57,19 @@ end
     @test compare(A -> sort(A; by=first, alg=MergeSort), AT, tagged)
     @test compare(A -> sortperm(A; by=first), AT, tagged)
 
+    # comparators that capture arrays, e.g. to sort the columns of a matrix
+    function sortcols(A; kwargs...)
+        lt = let A = A
+            (i, j) -> isless((A[1, i], A[2, i]), (A[1, j], A[2, j]))
+        end
+        idx = similar(A, Int, size(A, 2))
+        idx .= 1:size(A, 2)
+        return sort(idx; lt, kwargs...), sortperm(idx; lt, kwargs...)
+    end
+    @test compare(sortcols, AT, rand(1:3, 2, 1000))
+    @test compare(A -> sortcols(A; rev=true), AT, rand(1:3, 2, 1000))
+    @test compare(A -> sortcols(A; by=i -> 1001 - i), AT, rand(1:3, 2, 1000))
+
     # Arrays other than vectors need `dims`, and vectors take none, as in Base
     @test_throws MethodError sort(AT(x); dims=1)
     @test_throws MethodError sort!(AT(x); dims=1)