Add a keyword argument to disable multithreading

docs/src/lib/functions.md

@@ -6,8 +6,15 @@ CurrentModule = DataFrames
 
 ## Multi-threading support
 
-Selected operations in DataFrames.jl automatically use multiple threads when available.
-It is task-based and implemented using the `@spawn` macro from Julia Base.
+By default, selected operations in DataFrames.jl automatically use multiple threads
+when available. It is task-based and implemented using the `@spawn` macro from Julia Base.
+Multithreading can be disabled when running a particular block of code using the
+[`DataFrames.singlethreading() do... end`(@ref) syntax. It can also be disabled
+globally by calling [`DataFrames.setmultithreading(true)`](@ref).
+This is useful in particular to run functions which are not thread-safe, or when
+distribution of work across threads is managed separately.
+These functions are considered as experimental and may change or be removed once
+a cross-package mechanism for multithreading configuration is developed.
 
 This is a list of operations that currently make use of multi-threading:
 - `DataFrame` constructor with `copycols=true`; also recursively all functions

docs/src/lib/internals.md

@@ -15,7 +15,11 @@ gennames
 getmaxwidths
 ourshow
 ourstrwidth
+setmultithreading
+singlethreaded
 @spawn_for_chunks
+@spawn_or_async
+@spawn_or_run
 default_table_transformation
 isreadonly
 ```

src/dataframe/dataframe.jl

@@ -199,7 +199,8 @@ struct DataFrame <: AbstractDataFrame
         # we write into columns as we know that it is guaranteed
         # that it was freshly allocated in the outer constructor
         @static if VERSION >= v"1.4"
-            if copycols && len >= 1_000_000 && length(columns) > 1 && Threads.nthreads() > 1
+            if copycols && ismultithreaded() &&
+                len >= 1_000_000 && length(columns) > 1 && Threads.nthreads() > 1
                 @sync for i in eachindex(columns)
                     Threads.@spawn columns[i] = _preprocess_column(columns[i], len, copycols)
                 end
@@ -533,7 +534,7 @@ function _threaded_getindex(selected_rows::AbstractVector,
                             df_columns::AbstractVector,
                             idx::AbstractIndex)
     @static if VERSION >= v"1.4"
-        if length(selected_rows) >= 1_000_000 && Threads.nthreads() > 1
+        if length(selected_rows) >= 1_000_000 && ismultithreaded() && Threads.nthreads() > 1
             new_columns = Vector{AbstractVector}(undef, length(selected_columns))
             @sync for i in eachindex(new_columns)
                 Threads.@spawn new_columns[i] = df_columns[selected_columns[i]][selected_rows]

src/groupeddataframe/complextransforms.jl

@@ -261,7 +261,7 @@ function _combine_rows_with_first!((firstrow,)::Ref{Any},
     # Create up to one task per thread
     # This has lower overhead than creating one task per group,
     # but is optimal only if operations take roughly the same time for all groups
-    if VERSION >= v"1.4" && isthreadsafe(outcols, incols)
+    if VERSION >= v"1.4" && ismultithreaded() && isthreadsafe(outcols, incols)
         basesize = max(1, cld(len - 1, Threads.nthreads()))
         partitions = Iterators.partition(2:len, basesize)
     else
@@ -273,11 +273,11 @@ function _combine_rows_with_first!((firstrow,)::Ref{Any},
     tasks = Vector{Task}(undef, length(partitions))
     for (tid, idx) in enumerate(partitions)
         tasks[tid] =
-            @spawn _combine_rows_with_first_task!(tid, first(idx), last(idx), first(idx),
-                                                  outcols, outcolsref,
-                                                  type_widened, widen_type_lock,
-                                                  f, gd, starts, ends, incols, colnames,
-                                                  firstcoltype(firstmulticol))
+            @spawn_or_async _combine_rows_with_first_task!(tid, first(idx), last(idx), first(idx),
+                                                           outcols, outcolsref,
+                                                           type_widened, widen_type_lock,
+                                                           f, gd, starts, ends, incols, colnames,
+                                                           firstcoltype(firstmulticol))
     end
 
     # Workaround JuliaLang/julia#38931:

src/groupeddataframe/splitapplycombine.jl

@@ -666,8 +666,7 @@ function _combine(gd::GroupedDataFrame,
     for i in eachindex(cs_norm, optional_transform, tasks)
         cs_i = cs_norm[i]
         optional_i = optional_transform[i]
-
-        tasks[i] = @spawn if length(gd) > 0 && isagg(cs_i, gd)
+        tasks[i] = @spawn_or_async if length(gd) > 0 && isagg(cs_i, gd)
             _combine_process_agg(Ref{Any}(cs_i), optional_i, parentdf, gd,
                                  seen_cols, trans_res, idx_agg[])
         elseif keeprows && cs_i isa Pair && first(last(cs_i)) === identity &&
@@ -761,8 +760,8 @@ function _combine(gd::GroupedDataFrame,
 
     @sync for i in eachindex(trans_res)
         let i=i
-            @spawn reorder_cols!(trans_res, i, trans_res[i].col, trans_res[i].col_idx,
-                                 keeprows, idx_keeprows, gd)
+            @spawn_or_run reorder_cols!(trans_res, i, trans_res[i].col, trans_res[i].col_idx,
+                                        keeprows, idx_keeprows, gd)
         end
     end
 

src/groupeddataframe/utils.jl

@@ -354,7 +354,7 @@ function row_group_slots(cols::NTuple{N, AbstractVector},
             refmap
         end
         @sync for (seeni, range_chunk) in zip(seen_vec, range_chunks)
-            @spawn for i in range_chunk
+            @spawn_or_run for i in range_chunk
                 @inbounds begin
                     local refs_i
                     let i=i # Workaround for julia#15276
@@ -378,7 +378,7 @@ function row_group_slots(cols::NTuple{N, AbstractVector},
         end
     else
         @sync for (seeni, range_chunk) in zip(seen_vec, range_chunks)
-            @spawn for i in range_chunk
+            @spawn_or_run for i in range_chunk
                 @inbounds begin
                     local refs_i
                     let i=i # Workaround for julia#15276
@@ -414,10 +414,10 @@ function row_group_slots(cols::NTuple{N, AbstractVector},
         else
             xl = view(x, 1:len ÷ 2)
             xr = view(x, len ÷ 2 + 1:len)
-            t1 = @spawn reduce_or!(xl)
-            t2 = @spawn reduce_or!(xr)
-            fetch(t1)
-            fetch(t2)
+            @sync begin
+                @spawn_or_run reduce_or!(xl)
+                @spawn_or_run reduce_or!(xr)
+            end
             xl[1] .|= xr[1]
         end
         return

src/join/composer.jl

@@ -106,7 +106,7 @@ function compose_inner_table(joiner::DataFrameJoiner,
     left_ixs, right_ixs = find_inner_rows(joiner)
 
     @static if VERSION >= v"1.4"
-        if Threads.nthreads() > 1 && length(left_ixs) >= 1_000_000
+        if ismultithreaded() && Threads.nthreads() > 1 && length(left_ixs) >= 1_000_000
             dfl_task = Threads.@spawn joiner.dfl[left_ixs, :]
             dfr_noon_task = Threads.@spawn joiner.dfr[right_ixs, Not(joiner.right_on)]
             dfl = fetch(dfl_task)
@@ -238,7 +238,8 @@ function _compose_joined_table(joiner::DataFrameJoiner, kind::Symbol, makeunique
     @assert col_idx == ncol(joiner.dfl_on) + 1
 
     @static if VERSION >= v"1.4"
-        if Threads.nthreads() > 1 && target_nrow >= 1_000_000 && length(cols) > col_idx
+        if ismultithreaded() && Threads.nthreads() > 1 &&
+            target_nrow >= 1_000_000 && length(cols) > col_idx
             @sync begin
                 for col in eachcol(dfl_noon)
                     cols_i = left_idxs[col_idx]

src/other/utils.jl

@@ -131,6 +131,115 @@ end
 
 funname(c::ComposedFunction) = Symbol(funname(c.outer), :_, funname(c.inner))
 
+
+const SINGLETHREADING = Threads.Atomic{Bool}(false)
+const SINGLETHREADING_DEPTH = Threads.Atomic{Int}(0)
+
+ismultithreaded() = SINGLETHREADING_DEPTH[] == 0
+
+"""
+    DataFrames.singlethreaded(f)
+
+Run function `f` while disabling multithreading in all DataFrames.jl operations.
+This is useful in particular to run functions which are not thread-safe, or when
+distribution of work across threads is managed separately.
+
+*See also*: [`DataFrames.setmultithreading`](@ref) to disable multithreading globally
+
+!!! note
+
+    This function is considered as experimental and may change or be removed once
+    a cross-package mechanism for multithreading configuration is developed.
+
+    Currently, it disables multithreading for any DataFrames.jl
+    operations which may be run while `f` is running (e.g. if tasks using data
+    frames have been spawned on multiple threads).
+    This may change in the future.
+
+# Examples
+```jldoctest
+julia> df = DataFrame(x=repeat(1:5, inner=2), y=1:10);
+
+julia> gd = groupby(df, :x);
+
+julia> counter = 0;
+
+julia> f(x) = (sleep(0.1); global counter += 1); # Thread-unsafe function
+
+julia> DataFrames.singlethreaded() do
+           combine(gd, :y => f)
+       end
+5×2 DataFrame
+ Row │ x      y_f   
+     │ Int64  Int64 
+─────┼──────────────
+   1 │     1      1
+   2 │     2      2
+   3 │     3      3
+   4 │     4      4
+   5 │     5      5
+```
+"""
+function singlethreaded(f)
+    Threads.atomic_add!(SINGLETHREADING_DEPTH, 1)
+    try
+        return f()
+    finally
+        Threads.atomic_sub!(SINGLETHREADING_DEPTH, 1)
+    end
+end
+
+"""
+    DataFrames.setmultithreading(enable::Bool)
+
+Enable or disable multithreading permanently in all DataFrames.jl operations.
+This is useful in particular to run functions which are not thread-safe, or when
+distribution of work across threads is managed separately.
+
+*See also*: [`DataFrames.singlethreaded`](@ref) to disable multithreading only
+for a specific code block
+
+!!! note
+
+    This function is considered as experimental and may change or be removed once
+    a cross-package mechanism for multithreading configuration is developed.
+
+# Examples
+```jldoctest
+julia> df = DataFrame(x=repeat(1:5, inner=2), y=1:10);
+
+julia> gd = groupby(df, :x);
+
+julia> counter = 0;
+
+julia> f(x) = (sleep(0.1); global counter += 1); # Thread-unsafe function
+
+julia> DataFrames.setmultithreading(false);
+
+julia> combine(gd, :y => f)
+5×2 DataFrame
+ Row │ x      y_f   
+     │ Int64  Int64 
+─────┼──────────────
+   1 │     1      1
+   2 │     2      2
+   3 │     3      3
+   4 │     4      4
+   5 │     5      5
+
+julia> DataFrames.setmultithreading(true);
+```
+"""
+function setmultithreading(enable::Bool)
+    if !enable && !SINGLETHREADING[]
+        Threads.atomic_add!(SINGLETHREADING_DEPTH, 1)
+    elseif enable && SINGLETHREADING[]
+        Threads.atomic_sub!(SINGLETHREADING_DEPTH, 1)
+    end
+    SINGLETHREADING[] = !enable
+    return enable
+end
+
 # Compute chunks of indices, each with at least `basesize` entries
 # This method ensures balanced sizes by avoiding a small last chunk
 function split_indices(len::Integer, basesize::Integer)
@@ -159,7 +268,7 @@ if VERSION >= v"1.4"
 
                 nt = Threads.nthreads()
                 len = length(x)
-                if nt > 1 && len > basesize
+                if ismultithreaded() && nt > 1 && len > basesize
                     tasks = [Threads.@spawn begin
                                  for i in p
                                      local $(esc(lidx)) = @inbounds x[i]
@@ -215,6 +324,89 @@ macro spawn_for_chunks(basesize, ex)
     return _spawn_for_chunks_helper(ex.args[1], ex.args[2], basesize)
 end
 
+"""
+    @spawn_or_async expr
+
+Equivalent to `Threads.@spawn` if [`DataFrames.ismultithreaded() === true`](@ref)
+and to `@async` otherwise.
+"""
+spawn_or_async
+
+"""
+    @spawn_or_run expr
+
+Equivalent to `Threads.@spawn` if [`DataFrames.ismultithreaded() === true`](@ref),
+otherwise simply runs `expr`.
+"""
+spawn_or_run
+
+if VERSION >= v"1.4"
+    macro spawn_or_async(expr)
+        letargs = Base._lift_one_interp!(expr)
+
+        thunk = esc(:(()->($expr)))
+        var = esc(Base.sync_varname)
+        quote
+            let $(letargs...)
+                local task = Task($thunk)
+                task.sticky = !DataFrames.ismultithreaded()
+                if $(Expr(:islocal, var))
+                    @static if VERSION >= v"1.5.0"
+                        put!($var, task)
+                    else
+                        push!($var, task)
+                    end
+                end
+                schedule(task)
+                task
+            end
+        end
+    end
+
+    macro spawn_or_run(expr)
+        letargs = Base._lift_one_interp!(expr)
+
+        thunk = esc(:(()->($expr)))
+        var = esc(Base.sync_varname)
+        quote
+            let $(letargs...)
+                if DataFrames.ismultithreaded()
+                    local task = Task($thunk)
+                    task.sticky = false
+                    if $(Expr(:islocal, var))
+                        @static if VERSION >= v"1.5.0"
+                            put!($var, task)
+                        else
+                            push!($var, task)
+                        end
+                    end
+                    schedule(task)
+                else
+                    $thunk()
+                end
+                nothing
+            end
+        end
+    end
+else
+    # This is the definition of @async in Base
+    macro spawn_or_async(expr)
+        thunk = esc(:(()->($expr)))
+        var = esc(Base.sync_varname)
+        quote
+            local task = Task($thunk)
+            if $(Expr(:isdefined, var))
+                push!($var, task)
+            end
+            schedule(task)
+        end
+    end
+
+    macro spawn_or_run(expr)
+        esc(:($expr; nothing))
+    end
+end
+
 function _nt_like_hash(v, h::UInt)
     length(v) == 0 && return hash(NamedTuple(), h)