Add options to extend compute backends

.github/workflows/ci.yml

@@ -1,11 +1,13 @@
 name: CI
 on:
+  # All pull requests, so that stacked PRs are tested as well
   pull_request:
-    branches:
-      - main
   push:
     branches:
       - main
+      # TODO remove: GitHub creates no merge ref for this stack while #308 conflicts
+      # with main, so `pull_request` runs never start
+      - as/backend_fallbacks
     tags: '*'
 jobs:
   test:

Project.toml

@@ -8,13 +8,21 @@ ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 Mmap = "a63ad114-7e13-5084-954f-fe012c677804"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
+
+[weakdeps]
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+
+[extensions]
+StatisticsExt = "Statistics"
 
 [compat]
 Aqua = "0.8"
 ConstructionBase = "1"
 LRUCache = "1"
 Mmap = "1"
 OffsetArrays = "1"
+Preferences = "1.5.2"
 Statistics = "1.9"
 Suppressor = "0.2.8"
 Test = "1.9"
@@ -24,8 +32,8 @@ julia = "1.9"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TraceFuns = "626f57eb-f6ca-464c-8f04-384fee2a352e"
 
 [targets]

README.md

@@ -204,7 +204,46 @@ the desired range.
 
 ## Arrays that do not implement eachchunk
 
-There are arrays that live on disk but which are not split into rectangular chunks, so that the `haschunks` trait returns `Unchunked()`. In order to still enable broadcasting and reductions for these arrays, a chunk size will be estimated in a way that a certain memory limit per chunk is not exceeded. This memory limit defaults to 100MB and can be modified by changing `DiskArrays.default_chunk_size[]`. Then a chunk size is computed based on the element size of the array. However, there are cases where the size of the element type is undefined, e.g. for Strings or variable-length vectors. In these cases one can overload the `DiskArrays.element_size` function for certain container types which returns an approximate element size (in bytes). Otherwise the size of an element will simply be assumed to equal the value stored in `DiskArrays.fallback_element_size` which defaults to 100 bytes. 
+There are arrays that live on disk but which are not split into rectangular chunks, so that the `haschunks` trait returns `Unchunked()`. In order to still enable broadcasting and reductions for these arrays, a chunk size will be estimated in a way that a certain memory limit per chunk is not exceeded. This memory limit defaults to 100MB and can be modified by changing `DiskArrays.default_chunk_size[]`. Then a chunk size is computed based on the element size of the array. However, there are cases where the size of the element type is undefined, e.g. for Strings or variable-length vectors. In these cases one can overload the `DiskArrays.element_size` function for certain container types which returns an approximate element size (in bytes). Otherwise the size of an element will simply be assumed to equal the value stored in `DiskArrays.fallback_element_size` which defaults to 100 bytes.
+
+
+## Compute backend
+
+Computations on `AbstractDiskArray` objects (e.g. `sum`, `mean`, `mapreduce`, `minimum`) use a
+backend dispatch mechanism. The default is a single-threaded chunked iterator; packages like
+DiskArrayEngine.jl provide threaded or distributed backends.
+
+### Selecting the global backend
+
+The backend active when DiskArrays loads is determined by the `"backend"` preference (defaults to
+`"default"`). It can be changed at any time during a session:
+
+```julia
+# Switch to DiskArrayEngine and save the preference (persists across sessions)
+DiskArrays.set_backend("DiskArrayEngine")
+
+# Or switch session-only, without saving a preference
+DiskArrays.set_dynamic_backend!(DiskArrays.DefaultBackend())
+```
+
+A Julia session restart is required for `set_backend` to take effect, as the preference is read
+during module loading. `set_dynamic_backend!` takes effect immediately.
+
+Valid backend names are `"default"` and `"DiskArrayEngine"`.
+
+### Per-array backend with `withbackend`
+
+`withbackend` wraps an array so that all computations use a specified backend, independent of the
+global setting:
+
+```julia
+# Use the DiskArrayEngine backend only for this array
+result = sum(withbackend(my_diskarray, DiskArrayEngineBackend()))
+```
+
+This is mainly useful for testing or for mixing arrays backed by different computation engines
+within the same expression. Note that `BroadcastStyle` uses the global backend — it only sees the
+array type, not the `withbackend` wrapper.
 
 
 [ci-img]: https://github.com/JuliaIO/DiskArrays.jl/workflows/CI/badge.svg

docs/make.jl

@@ -6,7 +6,7 @@ makedocs(;
     format=Documenter.HTML(; prettyurls=get(ENV, "CI", nothing) == "true"),
     authors="Fabian Gans",
     sitename="DiskArrays.jl",
-    pages=Any["index.md"],
+    pages=Any["index.md", "Backend System" => "backend.md"],
     # strict = true,
     # clean = true,
     # checkdocs = :exports,

docs/src/backend.md

@@ -0,0 +1,70 @@
+# Backend System
+
+## Overview
+
+Computations on `AbstractDiskArray` objects use a backend dispatch mechanism. Each reduction
+function (e.g. `sum`, `mean`, `mapreduce`) calls a `diskarrays_<func>_impl(f, a, backend)`
+dispatch function, which defaults to a single-threaded chunked iterator but can be overridden
+for any `ComputeBackend` subtype.
+
+The global backend is selected at module load time via the `"backend"` preference and stored
+in `const compute_backend`. It is wrapped in a `DynamicBackend` which allows runtime switching
+via `set_dynamic_backend!` or `set_backend`. All entry points call `get_backend(a)`, which walks
+up the parent chain of `a` looking for a `WithBackendDiskArray` wrapper (created by `withbackend`),
+then falls back to the global backend.
+
+## Built-in Types
+
+| Type | Description |
+|------|-------------|
+| `ComputeBackend` | Abstract supertype for all backends |
+| `DefaultBackend` | Single-threaded chunked iterator |
+| `DiskArrayEngineBackend` | Used when DiskArrayEngine.jl is loaded |
+
+## Extendable Functions
+
+The following `diskarrays_*_impl` functions are the extension points. Each has a default
+implementation for `::ComputeBackend` and can be specialized for any `ComputeBackend` subtype.
+
+All functions accept a function `f` as the first argument, the disk array as the second, and a
+`ComputeBackend` as the third. Additional keyword arguments are passed through.
+
+| Function | Signature |
+|----------|-----------|
+| `diskarrays_sum_impl` | `diskarrays_sum_impl(f, a::AbstractDiskArray, ::ComputeBackend; kwargs...)` |
+| `diskarrays_prod_impl` | `diskarrays_prod_impl(f, a::AbstractDiskArray, ::ComputeBackend; kwargs...)` |
+| `diskarrays_all_impl` | `diskarrays_all_impl(f, a::AbstractDiskArray, ::ComputeBackend; kwargs...)` |
+| `diskarrays_any_impl` | `diskarrays_any_impl(f, a::AbstractDiskArray, ::ComputeBackend; kwargs...)` |
+| `diskarrays_minimum_impl` | `diskarrays_minimum_impl(f, a::AbstractDiskArray, ::ComputeBackend; kwargs...)` |
+| `diskarrays_maximum_impl` | `diskarrays_maximum_impl(f, a::AbstractDiskArray, ::ComputeBackend; kwargs...)` |
+| `diskarrays_mapreduce_impl` | `diskarrays_mapreduce_impl(f, op, a, dims, init, backend::ComputeBackend)` |
+| `diskarrays_mapreducedim_impl` | `diskarrays_mapreducedim_impl(f, op, R, a::AbstractDiskArray, backend::ComputeBackend)` |
+| `diskarrays_count_impl` | `diskarrays_count_impl(f, v::AbstractDiskArray, ::ComputeBackend)` |
+| `diskarrays_unique_impl` | `diskarrays_unique_impl(f, v::AbstractDiskArray, ::ComputeBackend)` |
+| `diskarrays_extrema_impl` | `diskarrays_extrema_impl(f, a::AbstractDiskArray, ::ComputeBackend; kwargs...)` |
+| `diskarrays_mean_impl` | `diskarrays_mean_impl(f, a::AbstractDiskArray, ::ComputeBackend; kwargs...)` |
+| `diskarrays_median_impl` | `diskarrays_median_impl(f, a::AbstractDiskArray, ::ComputeBackend; kwargs...)` |
+
+## Per-array Backend Override
+
+`withbackend(a, backend)` wraps an array so that `get_backend(a)` returns the specified backend
+instead of the global one. This is mainly useful for testing or for mixing arrays with
+different backends in the same expression.
+
+```julia
+a = withbackend(my_diskarray, DiskArrayEngineBackend())
+sum(a)  # uses DiskArrayEngineBackend regardless of the global setting
+```
+
+Note that `BroadcastStyle` uses the global backend — it only sees the array's type, not the
+`withbackend` wrapper.
+
+## Error Hints
+
+DiskArrays registers a `MethodError` hint for `DiskArrayEngineBackend`. When a method is not
+found, it prints:
+
+- `DiskArrayEngine.jl does not implement this operation for DiskArrayEngineBackend.` (if
+  DiskArrayEngine is already loaded)
+- `DiskArrayEngineBackend methods are implemented in DiskArrayEngine.jl, make sure it is loaded
+  with using DiskArrayEngine.` (if DiskArrayEngine is not yet loaded)

docs/src/index.md

@@ -7,3 +7,7 @@ CurrentModule = DiskArrays
 ```@autodocs
 Modules = [DiskArrays, DiskArrays.TestTypes]
 ```
+
+## User Guide
+
+- [Compute Backend](backend.md) — extending the compute backend mechanism

ext/StatisticsExt.jl

@@ -0,0 +1,17 @@
+module StatisticsExt
+import Statistics
+import DiskArrays: AbstractDiskArray, ComputeBackend, get_backend, diskarrays_mean_impl
+
+function Statistics.mean(f::F, a::AbstractDiskArray; kwargs...) where {F<:Function}
+    diskarrays_mean_impl(f, a, get_backend(a); kwargs...)
+end
+Statistics.mean(a::AbstractDiskArray; kwargs...) = Statistics.mean(identity, a; kwargs...)
+
+diskarrays_mean_impl(f, a::AbstractDiskArray, ::ComputeBackend;kwargs...) =
+    invoke(Statistics.mean,Tuple{typeof(f),AbstractArray{eltype(a),ndims(a)}},f,a;kwargs...)
+
+
+
+
+
+end

src/DiskArrays.jl

@@ -14,10 +14,11 @@ using Base: tail
     read(path, String)
 end DiskArrays
 
-export AbstractDiskArray, eachchunk, chunkexists, ChunkIndex, ChunkIndices, MissingTile
+export AbstractDiskArray, eachchunk, chunkexists, ChunkIndex, ChunkIndices, backend, withbackend,MissingTile
 
 include("scalar.jl")
 include("chunks.jl")
+include("compute.jl")
 include("diskarray.jl")
 include("batchgetindex.jl")
 include("diskindex.jl")
@@ -38,6 +39,7 @@ include("show.jl")
 include("chunktiledarray.jl")
 include("cached.jl")
 include("pad.jl")
+include("withbackend.jl")
 
 # The all-in-one macro
 
@@ -49,7 +51,6 @@ macro implement_diskarray(t)
         @implement_setindex $t
         @implement_broadcast $t
         @implement_iteration $t
-        @implement_mapreduce $t
         @implement_reshape $t
         @implement_array_methods $t
         @implement_permutedims $t
@@ -70,7 +71,6 @@ macro implement_diskarray_skip_zip(t)
         @implement_setindex $t
         @implement_broadcast $t
         @implement_iteration $t
-        @implement_mapreduce $t
         @implement_reshape $t
         @implement_array_methods $t
         @implement_permutedims $t
@@ -87,5 +87,9 @@ end
 # And we define the test types
 include("util/testtypes.jl")
 
+function __init__()
+    Base.Experimental.register_error_hint(_backend_error_hint, MethodError)
+end
+
 
 end # module

src/broadcast.jl

@@ -150,33 +150,36 @@ repsingle(size, range) = size == 1 ? (1:1) : range
 
 # Implementation macro
 
+function diskarrays_coptyo!(dest, bc, ::ComputeBackend)
+    foreach(eachchunk(dest)) do c
+        ar = [bc[i] for i in CartesianIndices(to_ranges(c))]
+        dest[to_ranges(c)...] = ar
+    end
+    return dest
+end
+# This is a heavily allocating implementation, but working for all cases.
+# As performance optimization one might:
+# Allocate the array only once if all chunks have the same size
+# Use FillArrays, if the DiskArray accepts these
+function diskarrays_fill!(dest, value, ::ComputeBackend)
+    foreach(eachchunk(dest)) do c
+        ar = fill(value, length.(to_ranges(c)))
+        dest[to_ranges(c)...] = ar
+    end
+    return dest
+end
+diskarrays_broadcaststyle(T, ::ComputeBackend) = ChunkStyle{ndims(T)}()
+
 macro implement_broadcast(t)
     t = esc(t)
     quote
         # Broadcasting with a DiskArray on LHS
-        function Base.copyto!(dest::$t, bc::Broadcasted{Nothing})
-            foreach(eachchunk(dest)) do c
-                ar = [bc[i] for i in CartesianIndices(to_ranges(c))]
-                dest[to_ranges(c)...] = ar
-            end
-            return dest
-        end
-        Base.BroadcastStyle(T::Type{<:$t}) = ChunkStyle{ndims(T)}()
+        Base.copyto!(dest::$t, bc::Broadcasted{Nothing}) = diskarrays_coptyo!(dest, bc, get_backend(dest))
+        Base.BroadcastStyle(T::Type{<:$t}) = diskarrays_broadcaststyle(T, get_backend(compute_backend))
         function DiskArrays.subsetarg(arg::$t, ranges)
             ashort = maybeonerange(size(arg), ranges)
             return arg[ashort...]
         end
-
-        # This is a heavily allocating implementation, but working for all cases.
-        # As performance optimization one might:
-        # Allocate the array only once if all chunks have the same size
-        # Use FillArrays, if the DiskArray accepts these
-        function Base.fill!(dest::$t, value)
-            foreach(eachchunk(dest)) do c
-                ar = fill(value, length.(to_ranges(c)))
-                dest[to_ranges(c)...] = ar
-            end
-            return dest
-        end
+        Base.fill!(dest::$t, value) = diskarrays_fill!(dest, value, get_backend(dest))
     end
 end

src/compute.jl

@@ -0,0 +1,63 @@
+using Preferences
+
+abstract type ComputeBackend end
+
+#Fallback: get_backend returns the backend itself
+get_backend(b::ComputeBackend) = b
+
+# Default Backend, simple nonthreaded, non-distributed computations
+struct DefaultBackend <: ComputeBackend end
+
+# Backend type for DiskArrayEngine.jl, which implements the `diskarrays_*_impl` methods for it
+struct DiskArrayEngineBackend <: ComputeBackend end
+
+# The global backend is always wrapped in a `DynamicBackend` so it can be switched at runtime.
+# The field is a small Union of concrete types, so dispatch on `get_backend` is
+# union-split into a branch instead of a dynamic dispatch. Backends that are not
+# known to DiskArrays can be attached to an array via `withbackend`.
+mutable struct DynamicBackend <: ComputeBackend
+    current_backend::Union{DefaultBackend,DiskArrayEngineBackend}
+end
+get_backend(b::DynamicBackend) = b.current_backend
+
+# The preference only selects the backend that is active when DiskArrays is loaded
+const backend = @load_preference("backend", "default")
+
+function _backend_from_name(name::String)
+    name == "default" && return DefaultBackend()
+    name == "DiskArrayEngine" && return DiskArrayEngineBackend()
+    throw(ArgumentError("Invalid backend: \"$(name)\""))
+end
+
+const compute_backend = DynamicBackend(_backend_from_name(backend))
+
+"""
+    set_dynamic_backend!(backend::Union{DefaultBackend,DiskArrayEngineBackend})
+
+Switch the active backend for the current session.
+"""
+set_dynamic_backend!(b::Union{DefaultBackend,DiskArrayEngineBackend}) =
+    compute_backend.current_backend = b
+
+"""
+    set_backend(new_backend::String)
+
+Switch the active backend to `"default"` or `"DiskArrayEngine"` and save it as
+a preference, so it is also the active backend in future sessions.
+"""
+function set_backend(new_backend::String)
+    b = _backend_from_name(new_backend)
+    @set_preferences!("backend" => new_backend)
+    set_dynamic_backend!(b)
+end
+
+const _DAE_PKGID = Base.PkgId(Base.UUID("2d4b2e14-ccd6-4284-b8b0-2378ace7c126"), "DiskArrayEngine")
+
+function _backend_error_hint(io, exc, argtypes, kwargs)
+    any(T -> T <: DiskArrayEngineBackend, argtypes) || return
+    if haskey(Base.loaded_modules, _DAE_PKGID)
+        print(io, "\nDiskArrayEngine.jl does not implement this operation for `DiskArrayEngineBackend`.")
+    else
+        print(io, "\n`DiskArrayEngineBackend` methods are implemented in DiskArrayEngine.jl, make sure it is loaded with `using DiskArrayEngine`.")
+    end
+end