Fixes for low hanging performance "bugs" (AI driven)

benchmark/append_rows.jl

@@ -0,0 +1,19 @@
+using BenchmarkTools
+using StructArrays
+
+n = 1_000_000
+rows = [(a=i, b=2i, c=3.0i, d=4.0i) for i in 1:n]
+base = StructArray((a=Int[], b=Int[], c=Float64[], d=Float64[]))
+
+probe = copy(base)
+append!(probe, rows)
+@assert probe == rows
+
+trial = @benchmark append!(dest, $rows) setup=(dest=copy($base)) samples=10 evals=1 seconds=60
+estimate = minimum(trial)
+println(
+    "append! rows: ",
+    BenchmarkTools.prettytime(estimate.time), ", ",
+    BenchmarkTools.prettymemory(estimate.memory), ", ",
+    estimate.allocs, " allocations",
+)

benchmark/group_wide.jl

@@ -0,0 +1,26 @@
+using BenchmarkTools
+using Random
+using StructArrays
+
+function countgroups(keys, permutation)
+    count = 0
+    for _ in StructArrays.GroupPerm(keys, permutation)
+        count += 1
+    end
+    return count
+end
+
+rng = MersenneTwister(42)
+columns = ntuple(_ -> rand(rng, 1:16, 100_000), 64)
+table = StructArray(columns)
+permutation = sortperm(table)
+@assert countgroups(table, permutation) == length(unique(table))
+
+trial = @benchmark countgroups($table, $permutation) samples=10 evals=1 seconds=60
+estimate = minimum(trial)
+println(
+    "64-column grouping: ",
+    BenchmarkTools.prettytime(estimate.time), ", ",
+    BenchmarkTools.prettymemory(estimate.memory), ", ",
+    estimate.allocs, " allocations",
+)

benchmark/sortperm_recursive.jl

@@ -0,0 +1,23 @@
+using BenchmarkTools
+using Random
+using StructArrays
+
+function nested_table(::Type{T}; n=100_000, width=6, cardinality=16, seed=1234) where {T}
+    rng = MersenneTwister(seed)
+    columns = ntuple(_ -> T.(rand(rng, 1:cardinality, n)), width)
+    return StructArray(columns)
+end
+
+for T in (Float64, Int)
+    table = nested_table(T)
+    @assert issorted(table[sortperm(table)])
+
+    trial = @benchmark sortperm($table) samples=20 evals=1 seconds=60
+    estimate = minimum(trial)
+    println(
+        T, " keys: ",
+        BenchmarkTools.prettytime(estimate.time), ", ",
+        BenchmarkTools.prettymemory(estimate.memory), ", ",
+        estimate.allocs, " allocations",
+    )
+end

benchmark/wide_constructor.jl

@@ -0,0 +1,18 @@
+using BenchmarkTools
+using StructArrays
+
+construct_wide(columns) = StructArray(columns)
+
+names = ntuple(i -> Symbol(:x, i), 128)
+columns = NamedTuple{names}(ntuple(_ -> rand(1), 128))
+probe = construct_wide(columns)
+@assert size(probe) == (1,)
+
+trial = @benchmark construct_wide($columns)
+estimate = minimum(trial)
+println(
+    "128-column constructor: ",
+    BenchmarkTools.prettytime(estimate.time), ", ",
+    BenchmarkTools.prettymemory(estimate.memory), ", ",
+    estimate.allocs, " allocations",
+)

benchmark/wide_getindex.jl

@@ -0,0 +1,15 @@
+using BenchmarkTools
+using StructArrays
+
+table = StructArray(ntuple(_ -> rand(100), 128))
+probe = table[50]
+@assert probe == ntuple(i -> components(table)[i][50], 128)
+
+trial = @benchmark $table[50]
+estimate = minimum(trial)
+println(
+    "128-column getindex: ",
+    BenchmarkTools.prettytime(estimate.time), ", ",
+    BenchmarkTools.prettymemory(estimate.memory), ", ",
+    estimate.allocs, " allocations",
+)

src/sort.jl

@@ -35,8 +35,24 @@ Base.eltype(::Type{<:GroupPerm}) = UnitRange{Int}
     return eq
 end
 
-roweq(t::Tuple{}, i, j) = true
-roweq(t::Tuple, i, j) = roweq(t[1], i, j) ? roweq(tail(t), i, j) : false
+_roweq(t::Tuple{}, i, j) = true
+_roweq(t::Tuple, i, j) = roweq(t[1], i, j) ? _roweq(tail(t), i, j) : false
+function roweq(t::T, i, j) where {T<:Tuple}
+    if @generated
+        types = fieldtypes(T)
+        if length(types) > 32 && all(==(types[1]), types) && isconcretetype(types[1])
+            return quote
+                for col in t
+                    roweq(col, i, j) || return false
+                end
+                return true
+            end
+        end
+        return :(_roweq(t, i, j))
+    else
+        return _roweq(t, i, j)
+    end
+end
 roweq(s::StructArray, i, j) = roweq(Tuple(components(s)), i, j)
 
 function uniquesorted(keys, perm=sortperm(keys))
@@ -72,30 +88,29 @@ forward_vec(::Ordering) = nothing
 
 # Methods from IndexedTables to refine sorting:
 # # assuming x[p] is sorted, sort by remaining columns where x[p] is constant
-function refine_perm!(p, cols, c, x, y′, lo, hi)
-    temp = similar(p, 0)
+function refine_perm!(p, cols, c, x, y′, lo, hi, temp=similar(p, 0), counts=Int[])
     order = Perm(Forward, y′)
     y = something(forward_vec(order), y′)
     nc = length(cols)
     for idxs in GroupPerm(x, p, lo:hi)
         i, i1 = extrema(idxs)
         if i1 > i
-            sort_sub_by!(p, i, i1, y, order, temp)
+            sort_sub_by!(p, i, i1, y, order, temp, counts)
             if c < nc-1
                 z = cols[c+2]
-                refine_perm!(p, cols, c+1, y, z, i, i1)
+                refine_perm!(p, cols, c+1, y, z, i, i1, temp, counts)
             end
         end
     end
 end
 
 # sort the values in v[i0:i1] in place, by array `by`
-Base.@noinline function sort_sub_by!(v, i0, i1, by, order, temp)
+Base.@noinline function sort_sub_by!(v, i0, i1, by, order, temp, counts=Int[])
     empty!(temp)
     sort!(v, i0, i1, MergeSort, order, temp)
 end
 
-Base.@noinline function sort_sub_by!(v, i0, i1, by::AbstractVector{T}, order, temp) where T<:Integer
+Base.@noinline function sort_sub_by!(v, i0, i1, by::AbstractVector{T}, order, temp, counts=Int[]) where T<:Integer
     min = max = by[v[i0]]
     @inbounds for i = i0+1:i1
         val = by[v[i]]
@@ -108,7 +123,7 @@ Base.@noinline function sort_sub_by!(v, i0, i1, by::AbstractVector{T}, order, te
     rangelen = max-min+1
     n = i1-i0+1
     if rangelen <= n
-        sort_int_range_sub_by!(v, i0-1, n, by, rangelen, min, temp)
+        sort_int_range_sub_by!(v, i0-1, n, by, rangelen, min, temp, counts)
     else
         empty!(temp)
         sort!(v, i0, i1, MergeSort, order, temp)
@@ -117,10 +132,11 @@ Base.@noinline function sort_sub_by!(v, i0, i1, by::AbstractVector{T}, order, te
 end
 
 # in-place counting sort of x[ioffs+1:ioffs+n] by values in `by`
-function sort_int_range_sub_by!(x, ioffs, n, by, rangelen, minval, temp)
+function sort_int_range_sub_by!(x, ioffs, n, by, rangelen, minval, temp, where=Int[])
     offs = 1 - minval
 
-    where = fill(0, rangelen+1)
+    resize!(where, rangelen+1)
+    fill!(where, 0)
     where[1] = 1
     @inbounds for i = 1:n
         where[by[x[i+ioffs]] + offs + 1] += 1

src/structarray.jl

@@ -350,6 +350,12 @@ Base.@propagate_inbounds function _getindex(x::StructArray{T}, I::Vararg{Int}) w
     return createinstance(T, get_ith(cols, I...)...)
 end
 
+Base.@propagate_inbounds function _getindex(x::StructArray{T}, I::Vararg{Int}) where {T<:Tup}
+    cols = components(x)
+    @boundscheck checkbounds(x, I...)
+    return T(get_ith(cols, I...))
+end
+
 @inline function _getindex(s::StructArray{T}, I...) where {T}
     @boundscheck checkbounds(s, I...)
     StructArray{T}(map(v -> @inbounds(getindex(v, I...)), components(s)))

src/tables.jl

@@ -31,6 +31,16 @@ end
 try_compatible_columns(rows::StructArray{T}, s::StructArray{T}) where {T} = Tables.columntable(rows)
 try_compatible_columns(rows::StructArray{R}, s::StructArray{S}) where {R,S} = nothing
 
+function _prepare_rows!(s, rows, ::typeof(push!))
+    _sizehint_rows!(s, rows, Base.IteratorSize(rows))
+end
+_prepare_rows!(s, rows, ::typeof(pushfirst!)) = s
+
+function _sizehint_rows!(s, rows, ::Union{Base.HasLength, Base.HasShape})
+    sizehint!(s, length(s) + length(rows))
+end
+_sizehint_rows!(s, rows, ::Any) = s
+
 for (f, g) in zip((:append!, :prepend!), (:push!, :pushfirst!))
     @eval function Base.$f(s::StructVector, rows)
         table = try_compatible_columns(rows, s)
@@ -42,6 +52,7 @@ for (f, g) in zip((:append!, :prepend!), (:push!, :pushfirst!))
         else
             # Otherwise, fallback to a generic implementation expecting
             # that `rows` is an iterator:
+            _prepare_rows!(s, rows, $g)
             return foldl($g, rows; init = s)
         end
     end

src/utils.jl

@@ -3,7 +3,18 @@ argtail(_, args...) = args
 split_tuple_type(T) = fieldtype(T, 1), Tuple{argtail(T.parameters...)...}
 
 eltypes(nt::NamedTuple{names}) where {names} = NamedTuple{names, eltypes(values(nt))}
-eltypes(t::Tuple) = Tuple{map(eltype, t)...}
+_eltypes(t::Tuple) = Tuple{map(eltype, t)...}
+function eltypes(t::T) where {T<:Tuple}
+    if @generated
+        types = fieldtypes(T)
+        if !isempty(types) && all(==(types[1]), types) && isconcretetype(types[1]) && types[1] <: AbstractArray
+            return :(NTuple{$(length(types)), $(eltype(types[1]))})
+        end
+        return :(_eltypes(t))
+    else
+        return _eltypes(t)
+    end
+end
 
 alwaysfalse(t) = false
 
@@ -199,8 +210,44 @@ maybe_convert_elt(::Type{T}, vals::NamedTuple) where T = T<:NamedTuple ? convert
 Compute the unique value that `f` takes on each `component ∈ components`.
 If not all values are equal, return `nothing`. Otherwise, return the unique value.
 """
-function findconsistentvalue(f::F, cols::Tup) where F
+function _findconsistentvalue(f, cols)
     val = f(first(cols))
     isconsistent = all(map(isequal(val) ∘ f, values(cols)))
     return ifelse(isconsistent, val, nothing)
 end
+
+function findconsistentvalue(f::F, cols::T) where {F, T<:Tuple}
+    if @generated
+        types = fieldtypes(T)
+        if length(types) > 32 && all(==(types[1]), types) && isconcretetype(types[1])
+            return quote
+                val = f(first(cols))
+                for col in cols
+                    isequal(val, f(col)) || return nothing
+                end
+                return val
+            end
+        end
+        return :(_findconsistentvalue(f, cols))
+    else
+        return _findconsistentvalue(f, cols)
+    end
+end
+
+function findconsistentvalue(f::F, cols::T) where {F, T<:NamedTuple}
+    if @generated
+        types = fieldtypes(T)
+        if length(types) > 32 && all(==(types[1]), types) && isconcretetype(types[1])
+            return quote
+                val = f(first(cols))
+                for col in cols
+                    isequal(val, f(col)) || return nothing
+                end
+                return val
+            end
+        end
+        return :(_findconsistentvalue(f, cols))
+    else
+        return _findconsistentvalue(f, cols)
+    end
+end