module ArraysOfArrays

Efficient storage and handling of nested arrays.

ArraysOfArrays provides two different types of nested arrays: ArrayOfSimilarArrays and VectorOfArrays.

AbstractArrayOfSimilarArrays{T,M,N,ET<:AbstractArray{T,M}} <: AbstractSlices{ET,N}

An array that contains arrays that have the same size/axes. The array is internally stored in flattened form as some kind of array of dimension M + N, in memory order. The flattened form can be accessed via flatview(A).

Implementation

Subtypes must implement (in addition to typical array operations)

ArraysOfArrays.fused(A::SomeArrayOfSimilarArrays)::AbstractArray{T,M+N}

which must return the underlying flat array. All split-mode operations (getsplitmode, stacked, flatview, innersize, getslicemap, etc.) are then provided automatically.

The following type aliases are defined:

• AbstractVectorOfSimilarArrays{T,M,ET} = AbstractArrayOfSimilarArrays{T,M,1,ET}
• AbstractArrayOfSimilarVectors{T,N,ET} = AbstractArrayOfSimilarArrays{T,1,N,ET}
• AbstractVectorOfSimilarVectors{T,ET} = AbstractArrayOfSimilarArrays{T,1,1,ET}

abstract type ArraysOfArrays.AbstractNestedArrayStyle{N} <: Broadcast.AbstractArrayStyle{N}

Supertype of the broadcast styles of nested array types like VectorOfArrays and ArrayOfSimilarArrays.

Packages that define array types with their own broadcast style can resolve style combination by specializing Base.Broadcast.BroadcastStyle(::AbstractNestedArrayStyle{N}, ::TheirStyle). Dispatch on this supertype, not on its subtypes.

abstract type AbstractPartMode{M,N} <: AbstractSplitMode

Abstract supertype for array partition modes with M inner dimensions and N outer dimensions.

The mode need not represent a true partition: partitions that discard part of the original array are allowed. Depending on the mode, the parts may also be reshaped.

Use getsplitmode to get the split mode of a split array.

Implementation

In addition to the requirements of AbstractSplitMode, subtypes must specialize ArraysOfArrays._bcast_expand to support outer-value arguments in bcastat.

abstract type AbstractSlicingMode{M} <: AbstractSplitMode

Abstract supertype for array slicing modes with M inner dimensions. The number of outer dimensions follows from the context, e.g. the dimensionality of the array that is split.

Use getsplitmode to get the split mode of a split array.

Implementation

In addition to the requirements of AbstractSplitMode, subtypes must specialize ArraysOfArrays.getinnerdims and ArraysOfArrays.getouterdims, which the slicing-specific operations rely on.

abstract type AbstractSplitMode <: Function

Abstract supertype for array split modes.

Use getsplitmode to get the split mode of an array.

Use splitup or call smode::AbstractSplitMode as a function to split an array:

splitup(A, smode) === smode(A)

See also fused, which undoes splitting.

Subtypes of AbstractSplitMode (except UnknownSplitMode) support InverseFunctions.inverse(smode) (typically returning an ArraysOfArrays.FuseArrays instance).

Implementation

Subtypes of AbstractSplitMode should specialize

• splitup(A, smode::SomeSplitMode) for arrays A
• is_memordered_splitmode(smode::SomeSplitMode)
• ArraysOfArrays.getinnerdims and ArraysOfArrays.getouterdims, which slicing modes require
• == and hash, if the mode carries state

(smode::SomeSplitMode)(A) calls splitup(A, smode) by default and should not be specialized.

ArrayOfSimilarArrays{T,M,N,P,ET} <: AbstractArrayOfSimilarArrays{T,M,N,ET}

Represents a view of an array of dimension M + N as an N-dimensional array with elements that are M-dimensional arrays. All element arrays implicitly have equal size/axes.

User code should typically not instantiate ArrayOfSimilarArrays directly, but use splitup with a SplitSlices mode.

Implementation

Constructors:

ArrayOfSimilarArrays{T,M,N}(data::AbstractArray)
ArrayOfSimilarArrays{T,M}(data::AbstractArray)

The following type aliases are defined:

• VectorOfSimilarArrays{T,M} = ArrayOfSimilarArrays{T,M,1}
• ArrayOfSimilarVectors{T,N} = ArrayOfSimilarArrays{T,1,N}
• VectorOfSimilarVectors{T} = ArrayOfSimilarArrays{T,1,1}

VectorOfSimilarArrays supports push!(), etc., provided the underlying array supports resizing of its last dimension (e.g. an ElasticArray).

The nested array can also be created using the function sliced and the wrapped flat array can be accessed using flatview afterwards:

A_flat = rand(2,3,4,5,6)
A_nested = sliced(A_flat, Val(2))
A_nested isa AbstractArray{<:AbstractArray{T,2},3} where T
flatview(A_nested) === A_flat

struct BaseSlicing{M,TPL<:Tuple{Vararg{Union{Colon,Int}}}} <: AbstractSlicingMode{M}

The split mode of Base.Slices (as returned by eachslice, eachcol and eachrow).

Constructor:

BaseSlicing{M,TPL}(slicemap::TPL)

slicemap equals the slicemap property of Base.Slices objects.

See also AbstractSlicingMode.

struct ArraysOfArrays.FuseArrays{S<:AbstractSplitMode} <: Function

Represents the inverse of a split mode: FuseArrays(smode) maps arrays split with smode back to their unsplit form.

FuseArrays(smode)(A) acts like fused(A) (resp. stacked(A) for SplitSlices modes), but FuseArrays(smode) has an InverseFunctions.inverse(FuseArrays(smode)) == smode while fused itself has no inverse. FuseArrays(smode)(A) also checks that the structure of A is compatible with smode.

Users should not instantiate FuseArrays directly, but use InverseFunctions.inverse(smode).

ArraysOfArrays.NestedArrayStyle{N}()

Broadcast style of nested array types like VectorOfArrays and ArrayOfSimilarArrays.

Broadcasts at this level apply f to whole element arrays, as in (x -> 2 .* x).(A). Such a broadcast returns a VectorOfArrays instead of a Vector of arrays if it runs over a single outer dimension with Base.OneTo axes and its result type is inferred as a concrete Array type with at least one dimension. The elements may be ragged, even if A is an ArrayOfSimilarArrays. All other broadcasts behave like the default broadcast machinery. Use bcastat to broadcast over the contents of the element arrays instead.

See ArraysOfArrays.AbstractNestedArrayStyle for resolving broadcast style combination with foreign array styles.

struct NonSplitMode{N} <: AbstractSplitMode

The split mode of unsplit collections that have N dimensions.

Constructor: NonSplitMode{N}()

PartsView{T,...} = VectorOfArrays{T,1,0,...}

A vector of vectors (that may differ in length), stored in contiguous, partitioned form. See VectorOfArrays for details.

User code should typically not construct a PartsView directly, but use splitup with a SplitParts mode or use consgroupedview instead.

Implementation

Constructors:

PartsView(A::AbstractVector{<:AbstractVector})
PartsView{T}(A::AbstractVector{<:AbstractVector}) where {T}

PartsView(
    data::AbstractVector, elem_ptr::AbstractVector{<:Integer},
    checks::Function = full_consistency_checks
)

struct SplitParts{M,VI,VD} <: AbstractPartMode{M,1}

The split mode of VectorOfArrays: a partition of a vector into consecutive parts of possibly different size, viewed as a vector of M-dimensional arrays.

Constructor:

SplitParts(
    elem_ptr::AbstractVector{<:Integer},
    kernel_size::AbstractVector{Dims{M-1}}
)

elem_ptr and kernel_size equal the equivalent properties of VectorOfArrays.

See also AbstractPartMode.

Implementation

getsplitmode(A::VectorOfArrays) copies the shape information of A where required (typically an O(length) operation), so the resulting mode is not affected if A is resized afterwards. Shape vectors that cannot be resized may be shared instead. A VectorOfArrays created via splitup, on the other hand, shares the vectors of the mode it was created from, like the VectorOfArrays inner constructor.

struct SplitSlices{M} <: AbstractSlicingMode{M}

The split mode of ArrayOfSimilarArrays: memory-ordered slicing with M inner dimensions.

Constructor:

SplitSlices{M}()

See also AbstractSlicingMode.

struct StaticSlices{AT<:AbstractArray,M} <: AbstractSlicingMode{M}

Memory-ordered slicing with inner arrays of type T: splitup reinterprets a flat array as an array with element type T and fused reinterprets it back, both without copying data.

Constructor: StaticSlices(AT::Type{<:AbstractArray})

A primary use case is slicing a flat array into an array of static arrays. For example:

using StaticArrays
A = rand(2, 3, 1000)
smode = StaticSlices(SMatrix{2,3})
slicedA = splitup(A, smode)

AT may be a UnionAll, e.g. SVector{3} instead of SVector{3, Float64}, but must be specific enough to determine the number of inner dimensions. StaticSlices is about shape, not type: if AT specifies an element type it is ignored when the mode is applied, the element type of the result always follows from the element type of the flat array.

Splitting works by reinterpret and so requires an isbits element type. In Reactant-traced code (with the Reactant package loaded), splitup returns a lazy view with static-array elements instead, since traced arrays cannot be reinterpreted.

struct UnknownSplitMode{AT} <: AbstractSplitMode

Split mode of generic split objects of type AT that have been split in an unknown way, e.g. nested arrays of type Array{<:Array}.

Since the split parts may be non-contiguous in memory (and typically are), this split mode does not support fused or flatview. Nor can it be inferred whether the split object should be interpreted as a sliced array, a ragged array, or something else.

Constructor: UnknownSplitMode{T}()

VectorOfArrays{T,N,M,VT,VI,VD,ET<:AbstractArray{T,N}} <: AbstractVector{ET}

A VectorOfArrays represents a vector of N-dimensional arrays (that may differ in size). Internally, VectorOfArrays stores all elements of all arrays in a single flat vector. M must equal N - 1.

The VectorOfArrays itself supports push!, append!, etc., but the size of each individual array in the vector is fixed. resize! can be used to shrink, but not to grow, as the size of the additional element arrays in the vector would be unknown. However, memory space for up to n arrays with a maximum size s can be reserved via sizehint!(A::VectorOfArrays, n, s::Dims{N}).

User code should typically not instantiate VectorOfArrays directly, but use splitup with a SplitParts mode or use consgroupedview instead.

Implementation

Elements are stored via their length and the size of their leading (kernel) dimensions. If the kernel has zero length the size of the last dimension cannot be reconstructed, so e.g. an element of size (0, 3) reads back with size (0, 0).

A VectorOfArrays created via splitup or view shares its data and structural vectors with the split mode resp. the parent array. Mutating operations like push!, resize! and empty! write through to everything that shares them. They are rejected before any mutation unless the data and structural vectors are resizable.

Constructors:

VectorOfArrays{T,N}()

VectorOfArrays(A::AbstractVector{<:AbstractArray})
VectorOfArrays{T}(A::AbstractVector{<:AbstractArray})
VectorOfArrays{T,N}(A::AbstractVector{<:AbstractArray})

VectorOfArrays(
    data::AbstractVector,
    elem_ptr::AbstractVector{<:Integer},
    kernel_size::AbstractVector{<:Dims}
    checks::Function = ArraysOfArrays.full_consistency_checks
)

Other suitable values for checks are ArraysOfArrays.simple_consistency_checks and ArraysOfArrays.no_consistency_checks.

PartsView is defined as a type alias:

PartsView{T,VT,VI,VD,ET} = VectorOfArrays{T,1,0,VT,VI,VD,ET}

VectorOfVectors{T,...} = PartsView{T,...}

Alias for PartsView.

bcastat(f, ::Val{depth}, args...)
bcastat(f, depth::Integer, args...)

Broadcast f over the contents of nested arrays at nesting depth depth, with AwkwardArrays-like alignment, but with array-of-arrays nesting semantics:

• Nested (split) arrays with equal split modes are aligned at depth depth.
• Arrays that match the outer structure of a shallower nesting level contribute one value per element of that level, broadcast over everything below it.
• Scalars and Refs broadcast over everything.

bcastat(f, Val(1), args...) is equivalent to broadcast(f, args...), and as with mapat, a depth that exceeds the nesting depth of the arguments applies f at the innermost level. The Integer form relies on constant propagation for type stability; use the Val form when passing a non-constant depth.

Nested array arguments must be split arrays (like ArrayOfSimilarArrays and VectorOfArrays): bcastat operates on their flat data with split-mode-specific argument expansions, without per-element iteration, and so works on GPU arrays.

consgrouped_ptrs(A::AbstractVector)

Compute an element pointer vector, suitable for creation of a PartsView that implies grouping equal consecutive entries of A.

Example:

    A = [1, 1, 2, 3, 3, 2, 2, 2]
    elem_ptr = consgrouped_ptrs(A)
    first.(PartsView(A, elem_ptr)) == [1, 2, 3, 2]

Typically, elem_ptr will be used to apply the computed grouping to other data:

    B = [1, 2, 3, 4, 5, 6, 7, 8]
    PartsView(B, elem_ptr) == [[1, 2], [3], [4, 5], [6, 7, 8]]

consgroupedview(source::AbstractVector, target)

Compute a grouping of equal consecutive elements on source via consgrouped_ptrs and apply the grouping to target, resp. each element of target. target may be a vector or a named or unnamed tuple of vectors. The result is a PartsView, resp. a tuple of such.

Example:

A = [1, 1, 2, 3, 3, 2, 2, 2]
B = [1, 2, 3, 4, 5, 6, 7, 8]
consgroupedview(A, B) == [[1, 2], [3], [4, 5], [6, 7, 8]]

consgroupedview plays well with columnar tables, too:

    using Tables, TypedTables
    data = Table(
        a = [1, 1, 2, 3, 3, 2, 2, 2],
        b = [1, 2, 3, 4, 5, 6, 7, 8],
        c = [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]
    )

    result = Table(consgroupedview(data.a, Tables.columns(data)))

will return

     a          b          c
   ┌──────────────────────────────────────
 1 │ [1, 1]     [1, 2]     [1.1, 2.2]
 2 │ [2]        [3]        [3.3]
 3 │ [3, 3]     [4, 5]     [4.4, 5.5]
 4 │ [2, 2, 2]  [6, 7, 8]  [6.6, 7.7, 8.8]

without copying any data:

    flatview(result.a) === data.a
    flatview(result.b) === data.b
    flatview(result.c) === data.c

deepmap(f, A::AbstractArray)
deepmap(f, A::AbstractArray{<:AbstractArray{<:...}})

Applies map at the deepest layer of nested arrays. If A is not a nested array, deepmap behaves identically to Base.map.

ArraysOfArrays.element_ptr(A::VectorOfArrays)

Returns a copy of the internal element pointer vector of A. The pointers are absolute indices into fused(A) and need not start at one, e.g. for views and partial partitioned results.

See also getsplitmode, which returns the element pointers together with the kernel sizes.

flatview(A::AbstractArray)
flatview(A::AbstractArray{<:AbstractArray})

View array A in a flattened form, with inner dimensions first. The shape of the flattened form will depend on the type of A. If A is not a nested array, the return value is A itself. Only specific types of nested arrays are supported.

flatview is a zero-copy O(1) operation.

If is_memordered_splitmode(getsplitmode(A)) is true and the elements of A cover the underlying data completely, flatview(A) is equivalent to fused(A).

For sliced arrays the result of flatview(A) will equal stacked(A). For partitioned vectors it will equal vecflattened(A), provided that the parts cover the underlying data completely.

flatview(A::VectorOfArrays{T})::AbstractVector{T}

Returns the data of all element arrays of A as a single vector, without copying. If the elements of A cover its internal storage completely, this is the internal storage vector itself, otherwise a view of the covered region (e.g. for views of vectors of arrays and partial partitioned views). Do not change the length of the returned vector, as this would break the internal consistency of A.

ArraysOfArrays.full_consistency_checks(A::VectorOfArrays)

Check the internal consistency of A completely, including whether the length of each element is compatible with its kernel size. Takes O(length(A)) time, and synchronizes with the device for device-resident structural vectors.

This is the default value of the checks argument of the VectorOfArrays constructor.

fused(A::AbstractArray)
fused(A::AbstractArray{<:AbstractArray})

View array A in unsplit form.

If A is not a nested array, return A itself. If A is a split array, return the original unsplit array.

splitup(fused(A), getsplitmode(A)) must equal A, and should have the same type as A if at all possible, except if getsplitmode(A) is an UnknownSplitMode.

If is_memordered_splitmode(getsplitmode(A)) is true and the elements of A cover the underlying data completely, fused(A) is equivalent to flatview(A).

fused should be a zero-copy O(1) operation, if at all possible.

ArraysOfArrays.getinnerdims(tpl::Tuple, smode::AbstractSplitMode)

Get the entries of tpl corresponding to the inner dimensions of split mode smode, in the order specified by smode.

ArraysOfArrays.getouterdims(tpl::Tuple, smode::AbstractSplitMode)

Get the entries of tpl corresponding to the outer dimensions of split mode smode, in the order specified by smode.

getslicemap(A::AbstractArray{<:AbstractArray})

Return the slicemap of A with respect to B = fused(A): a tuple with one entry per dimension of B, Colon() for sliced (inner) dimensions and k for dimensions indexed by dimension k of A, so that

A[i...] == view(B, map(s -> s isa Colon ? (:) : i[s], getslicemap(A))...)

E.g. A = eachslice(B, dims = (3,1,5)) of a five-dimensional B has the slicemap (2, :, 1, :, 3), since A[i1, i2, i3] == view(B, i2, :, i1, :, i3).

Equals the slicemap field of Base.Slices objects.

getsplitmode(A::AbstractArray)::NonSplitMode
getsplitmode(A::AbstractArray{<:AbstractArray})::AbstractSplitMode

Get the split mode of A.

splitup(fused(A), getsplitmode(A)) must equal A, and should have the same type as A if at all possible, except if getsplitmode(A) is an UnknownSplitMode.

getsplitmode should be a zero-copy O(1) operation, if at all possible.

innerlengths(A::AbstractArray{<:AbstractArray})

Returns the lengths of the element arrays of A, as an array of Int shaped like A.

innermap(f, A::AbstractArray)
innermap(f, A::AbstractArray{<:AbstractArray})

Nested map at depth 2. Equivalent to map(X -> map(f, X), A) for arrays of arrays, otherwise equivalent to Base.map.

innermapreduce(f, op, A::AbstractArray{<:AbstractArray}; [init])

Per-element mapreduce over the contents of the element arrays of A: returns an array shaped like A that contains mapreduce(f, op, A[i]; [init]) for each element A[i].

For split arrays (like ArrayOfSimilarArrays and VectorOfArrays) this uses segmented reductions over the underlying flat data where possible, and works on GPU arrays.

innerreduce(op, A::AbstractArray{<:AbstractArray}; [init])

Per-element reduce over the contents of the element arrays of A, equivalent to innermapreduce(identity, op, A; [init]).

innersize(A::AbstractArray{<:AbstractArray}, [dim])

Returns the size of the element arrays of A. Fails if the element arrays are not of equal size.

innersizes(A::AbstractArray{<:AbstractArray{T,M}})

Returns the sizes of the element arrays of A, as an array of Dims{M} shaped like A. In contrast to innersize, the element arrays do not need to be of equal size.

innersum(A::AbstractArray{<:AbstractArray})

Per-element sum over the contents of the element arrays of A. For numerical element types, empty element arrays sum to zero.

ArraysOfArrays.internal_element_ptr(A::VectorOfArrays)

Returns the internal element pointer vector of A.

Do not modify the returned vector: this would break the internal consistency of A. See element_ptr for a safe alternative.

is_memordered_splitmode(smode::AbstractSplitMode)::Bool

Check whether smode splits in memory order.

If true, inner arrays are stored contiguously in memory in Julia-native dimension order, and the same is true for the outer dimensions (no dimension reordering).

If true, flatview and fused are equivalent, provided that the elements of the array cover the underlying data completely (see flatview).

mapat(f, ::Val{depth}, As::AbstractArray...)
mapat(f, depth::Integer, As::AbstractArray...)

Nested map at nesting depth depth: apply f elementwise to the objects at depth depth of the nested arrays As. Depth 1 refers to the elements of the arrays themselves, so mapat(f, Val(1), As...) is equivalent to map(f, As...) and mapat(f, Val(2), A) is equivalent to innermap(f, A). If depth exceeds the nesting depth of the arrays, f is applied to the innermost elements, like deepmap.

The Integer form relies on constant propagation for type stability; use the Val form when passing a non-constant depth.

All of As must have the same nesting structure down to depth. Split arrays must have equal split modes (see getsplitmode) on each nesting level; generic nested arrays are combined elementwise like map and so may differ in type as long as their shapes match.

For split arrays (like ArrayOfSimilarArrays and VectorOfArrays) mapat operates directly on the flat data covered by the elements, without per-element iteration, and so works on GPU arrays.

See also bcastat for broadcast semantics with arguments of different nesting depth.

nestedview(A::AbstractArray{T,M+N}, M::Integer)
nestedview(A::AbstractArray{T,2})

Deprecated, use sliced instead.

ArraysOfArrays.no_consistency_checks(A::VectorOfArrays)

Don't check the internal consistency of A at all.

Suitable as the checks argument of the VectorOfArrays constructor for data and structural vectors that are known to be consistent with each other.

partitioned(A::AbstractVector, lengths::AbstractVector{<:Integer})
partitioned(A::AbstractVector, shapes::AbstractVector{<:Dims})

Return a partitioned view of A, as a vector of arrays.

The parts are consecutive, non-overlapping views of A, with sizes given by lengths (resulting in a vector of vectors) or shapes (resulting in a vector of arrays).

ArraysOfArrays.simple_consistency_checks(A::VectorOfArrays)

Check the internal consistency of A without scanning its structural vectors: only the first and last element pointer are read, so the check takes O(1) time. On device-resident structural vectors those two reads require a device synchronization each.

Suitable as the checks argument of the VectorOfArrays constructor.

sliced(A::AbstractArray{T,2})
sliced(A::AbstractArray{T,M+N}, Val(M))
sliced(A::AbstractArray{T,M+N}, M::Integer)

Return a sliced view of A, using the columns or the first M dimensions as inner dimensions.

With StaticArrays loaded, sliced(A, SA) with a static array type SA (e.g. SVector{3} or SMatrix{2,3}) returns a reinterpreted array with SA elements instead of an array of views (see StaticSlices).

splitup(A::AbstractArray, smode::AbstractSplitMode)

View array A in split form, as an array of arrays.

splitup should be a zero-copy operation, if at all possible. Splitting with a SplitParts mode validates the mode against the data (an O(n) check), since an inconsistent mode would result in silently corrupt elements.

See also fused and getsplitmode.

stacked(A::AbstractArray{T,N})::AbstractArray{T,N}
stacked(A::AbstractArray{<:AbstractArray{T,M},N})::AbstractArray{T,M+N}

Join the element arrays of a nested array into a single array along one or more new dimensions, return non-nested arrays unchanged.

Similar to Base.stack, but can return the original underlying array of sliced arrays in more cases. Empty arrays of arrays, which Base.stack rejects, yield an empty array with the innersize of A as inner dimensions.

unstackmode(A::AbstractArray)
unstackmode(A::AbstractArray{<:AbstractArray})

Get the split mode required to restore A from stacked(A), so that splitup(stacked(A), unstackmode(A)) == A.

The result of splitup(stacked(A), unstackmode(A)) may have a different type and underlying memory layout than A.

vecflattened(A::AbstractArray{T})::AbstractVector{T}
vecflattened(A::AbstractArray{<:AbstractArray})::AbstractVector{T}

Concatenate nested arrays into a single vector, return non-nested vectors unchanged.

If A is a nested view of a vector, vecflattened(A) should return the underlying vector in a zero-copy O(1) fashion. So in contrast to reduce(vcat, A) and mapreduce(vec, vcat, A), the result may share memory with A.

Implementation

The default implementations are

vecflattened(A::AbstractVector) = A
vecflattened(A::AbstractArray) = vec(A)
vecflattened(A::AbstractVector{<:AbstractVector}) = reduce(vcat, A)
vecflattened(A::AbstractArray{<:AbstractArray}) = mapreduce(vec, vcat, A)

Memory-ordered slicings (see is_memordered_splitmode) return vec(fused(A)) without copying data.

Specialize vecflattened for custom nested array types that can provide a zero-copy implementation.