diff --git a/benchmark/append_rows.jl b/benchmark/append_rows.jl new file mode 100644 index 0000000..a12dde4 --- /dev/null +++ b/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", +) diff --git a/benchmark/group_wide.jl b/benchmark/group_wide.jl new file mode 100644 index 0000000..7484061 --- /dev/null +++ b/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", +) diff --git a/benchmark/sortperm_recursive.jl b/benchmark/sortperm_recursive.jl new file mode 100644 index 0000000..395f544 --- /dev/null +++ b/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 diff --git a/benchmark/wide_constructor.jl b/benchmark/wide_constructor.jl new file mode 100644 index 0000000..3c477ae --- /dev/null +++ b/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", +) diff --git a/benchmark/wide_getindex.jl b/benchmark/wide_getindex.jl new file mode 100644 index 0000000..7d23924 --- /dev/null +++ b/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", +) diff --git a/src/sort.jl b/src/sort.jl index 64316d9..278b43e 100644 --- a/src/sort.jl +++ b/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 diff --git a/src/structarray.jl b/src/structarray.jl index eba837b..1da2c82 100644 --- a/src/structarray.jl +++ b/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))) diff --git a/src/tables.jl b/src/tables.jl index 432a75b..58c7927 100644 --- a/src/tables.jl +++ b/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 diff --git a/src/utils.jl b/src/utils.jl index 00458e1..e8865fd 100644 --- a/src/utils.jl +++ b/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