Merge pull request #553 from JuliaFolds/fixed-size-collect

Add fast pathway for `copy`, `collect`, `tcollect`, and `tcopy` for size-stable operations

.github/workflows/test.yml

@@ -22,7 +22,7 @@ jobs:
           - julia-version: '1.8'
             envname: "main_v2"
           - julia-version: 'nightly'
-            envname: "main_v2"
+            envname: "main_v3"
       fail-fast: false
     name: Test Julia ${{ matrix.julia-version }}
     steps:

Project.toml

@@ -1,7 +1,7 @@
 name = "Transducers"
 uuid = "28d57a85-8fef-5791-bfe6-a80928e7c999"
 authors = ["Takafumi Arakaki <aka.tkf@gmail.com>"]
-version = "0.4.75"
+version = "0.4.76"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/Transducers.jl

@@ -101,7 +101,7 @@ using InitialValues:
     asmonoid,
     hasinitialvalue
 using Logging: @logmsg, LogLevel
-using MicroCollections: UndefVector
+using MicroCollections: UndefVector, UndefArray
 using Requires
 using Setfield: @lens, @set, set, setproperties
 using SplittablesBase: SplittablesBase, amount, halve

src/core.jl

@@ -887,6 +887,8 @@ identityof(::typeof(right), ::Any) = nothing
 abstract type Reducible end
 abstract type Foldable <: Reducible end
 
+Base.IteratorSize(::Type{<:Reducible}) = Base.SizeUnknown()
+
 """
    Transducers.asfoldable(x) -> foldable
    

src/dreduce.jl

@@ -75,13 +75,14 @@ function dtransduce(
     end
     @argcheck basesize > 0
     @argcheck threads_basesize > 0
-    futures = map(firstindex(coll):basesize:lastindex(coll)) do start
+    chunks = split_into_chunks(coll, basesize)
+    futures = map(chunks) do chunk
         Distributed.remotecall(
             _remote_reduce,
             pool,
             rf,
             init,
-            coll[start:min(end, start - 1 + basesize)],
+            chunk,
             threads_basesize,
         )
     end

src/processes.jl

@@ -803,7 +803,10 @@ julia> collect(Interpose(missing), 1:3)
  3
 ```
 """
-function Base.collect(xf::Transducer, coll)
+Base.collect(xf::XF, coll) where {XF <: Transducer} = _collect(xf, coll, OutputSize(XF), Base.IteratorSize(coll))
+Base.collect(foldable::Foldable) = collect(extract_transducer(foldable)...)
+
+function _collect(xf, coll, ::Any, ::Any)
     result = finish!(unreduced(transduce(
         Map(SingletonVector) ∘ xf,
         wheninit(collector, append!!),
@@ -816,9 +819,17 @@ function Base.collect(xf::Transducer, coll)
     end
     return result
 end
-# Base.collect(xf, coll) = append!([], xf, coll)
 
-Base.collect(foldable::Foldable) = collect(extract_transducer(foldable)...)
+function _collect(xf, arr::Array, ::SizeStable, ::Union{Base.HasLength, Base.HasShape})
+    rf(coll, (i, val)) = @inbounds setindex!!(coll, val, i)
+    dest = UndefArray(size(arr)...)
+    unreduced(transduce(
+        Enumerate() ∘ xf,
+        wheninit(() -> dest, rf),
+        dest,
+        arr
+    ))
+end
 
 """
     copy(xf::Transducer, T, foldable) :: Union{T, Empty{T}}
@@ -869,7 +880,9 @@ julia> @assert copy(
        ) == DataFrame(A = [2], B = [3])
 ```
 """
-Base.copy(xf::Transducer, ::Type{T}, foldable) where {T} = append!!(xf, Empty(T), foldable)
+function Base.copy(xf::XF, T, foldable::F) where {XF <: Transducer, F}
+    _copy(xf, T, foldable, OutputSize(XF), Base.IteratorSize(F))
+end
 Base.copy(xf::Transducer, foldable) = copy(xf, _materializer(foldable), foldable)
 
 function Base.copy(::Type{T}, ed::Foldable) where {T}
@@ -882,6 +895,40 @@ function Base.copy(ed::Foldable)
     return copy(xf, foldable)
 end
 
+_copy(xf, ::Type{T}, foldable, ::Any, ::Any) where {T} = append!!(xf, Empty(T), foldable)
+
+function _copy(xf, ::Type{Vector{<:Any}}, arr, ::SizeStable, ::Base.HasLength)
+    rf(coll, (i, val)) = @inbounds setindex!!(coll, val, i)
+    dest = UndefVector(length(arr))
+    unreduced(transduce(
+        Enumerate() ∘ xf,
+        wheninit(() -> dest, rf),
+        dest,
+        arr
+    ))
+end
+
+function _copy(xf, ::Type{Array{<:Any, N}}, arr, ::SizeStable, ::Base.HasShape) where {N}
+    sz_arr = size(arr)
+    M = length(sz_arr)
+    if N > M
+        sz = (sz_arr..., ntuple(_ -> 1, N - ndims(arr))...)
+    elseif N < M
+        l, r = sz_arr[1:(N-1)], sz_arr[N:end]
+        sz = (l..., prod(r))
+    else
+        sz = sz_arr
+    end
+    rf(coll, (i, val)) = @inbounds setindex!!(coll, val, i)
+    dest = UndefArray(sz)
+    unreduced(transduce(
+        Enumerate() ∘ xf,
+        wheninit(() -> dest, rf),
+        dest,
+        arr
+    ))
+end
+
 Base.Set(foldable::Foldable) = copy(Set, foldable)
 Base.Dict(foldable::Foldable) = copy(Dict, foldable)
 

src/reduce.jl

@@ -351,8 +351,19 @@ julia> @assert foldxt(
        ) == Table(a = [1, 2])
 ```
 """
-tcopy(xf, T, reducible; kwargs...) =
-    foldxt(append!!, Map(SingletonVector) ∘ xf, reducible; init = Empty(T), kwargs...)
+tcopy(xf::XF, T, reducible::R; kwargs...) where {XF, R} = _tcopy(xf, T, reducible, OutputSize(XF), Base.IteratorSize(R); kwargs...)
+_tcopy(xf, T, reducible, ::Any, ::Any; kwargs...) = foldxt(append!!, Map(SingletonVector) ∘ xf, reducible; init = Empty(T), kwargs...)
+function _tcopy(xf, ::Type{T}, reducible, ::SizeStable, ::Union{Base.HasLength, Base.HasShape};
+                basesize=max(amount(reducible) ÷ Threads.nthreads(), 1), kwargs...) where {T <: Array}
+    chunks = split_into_chunks(reducible, basesize)
+    foldxt(append!!, Map(x -> copy(xf, T, x)), chunks; init = Empty(T), kwargs...)
+end
+
+# This can't be collect(Partition(sz), col) because of https://github.com/JuliaFolds/Transducers.jl/issues/554
+function split_into_chunks(coll, sz)
+    collect(Iterators.partition(coll, sz))
+end
+
 tcopy(xf, reducible; kwargs...) = tcopy(xf, _materializer(reducible), reducible; kwargs...)
 
 function tcopy(::Type{T}, itr; kwargs...) where {T}