tmap for enumerate(chunks(...))

JuliaFolds2 · Sep 18, 2024 · 66bc292 · 66bc292
1 parent 2702ac4
commit 66bc292
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/src/implementation.jl b/src/implementation.jl
@@ -320,7 +320,7 @@ function tmap(f, ::Type{T}, A::AbstractArray, _Arrs::AbstractArray...; kwargs...
 end
 
 function tmap(f,
-        A::Union{AbstractArray, ChunkSplitters.Chunk},
+        A::Union{AbstractArray, ChunkSplitters.Chunk, ChunkSplitters.Enumerate},
         _Arrs::AbstractArray...;
         scheduler::MaybeScheduler = NotGiven(),
         kwargs...)
@@ -333,7 +333,8 @@ function tmap(f,
        _scheduler.split != :batch
         error("Only `split == :batch` is supported because the parallel operation isn't commutative. (Scheduler: $_scheduler)")
     end
-    if A isa ChunkSplitters.Chunk && chunking_enabled(_scheduler)
+    if (A isa ChunkSplitters.Chunk || A isa ChunkSplitters.Enumerate) &&
+       chunking_enabled(_scheduler)
         auto_disable_chunking_warning()
         if _scheduler isa DynamicScheduler
             _scheduler = DynamicScheduler(;
@@ -377,7 +378,7 @@ end
 # w/o chunking (DynamicScheduler{NoChunking}): ChunkSplitters.Chunk
 function _tmap(scheduler::DynamicScheduler{NoChunking},
         f,
-        A::ChunkSplitters.Chunk,
+        A::Union{ChunkSplitters.Chunk, ChunkSplitters.Enumerate},
         _Arrs::AbstractArray...)
     (; threadpool) = scheduler
     tasks = map(A) do idcs

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -92,6 +92,9 @@ end;
     @test tmapreduce(+, enumerate(chunks(data; size=5)); chunking=false) do (i, idcs)
         [i, sum(@view(data[idcs]))]
     end == [sum(1:20), sum(data)]
+    @test tmap(enumerate(chunks(data; n=5)); chunking=false) do (i, idcs)
+        [i, idcs]
+    end == [[1, 1:20], [2, 21:40], [3, 41:60], [4, 61:80], [5, 81:100]]
 end;
 
 @testset "macro API" begin
@@ -267,6 +270,11 @@ end;
         @set reducer = +
         [i, sum(@view(data[idcs]))]
     end) == [sum(1:20), sum(data)]
+    @test @tasks(for (i, idcs) in enumerate(chunks(1:100; n=5))
+        @set chunking=false
+        @set collect=true
+        [i, idcs]
+    end) == [[1, 1:20], [2, 21:40], [3, 41:60], [4, 61:80], [5, 81:100]]
 end;
 
 @testset "WithTaskLocals" begin